noe/galaxy_ninja
1
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Activity

black galsec

Browse Source
This commit is contained in:
Noe Brucy
2024-11-28 11:51:25 +01:00
parent b76d4fc688
commit ffdc0db867
1 changed files with 120 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

105
galsec.py
View File

@@ -298,7 +298,6 @@ def regroup(
class GalsecAnalysis(ABC):
def __init__(self):
raise NotImplementedError("This method should be implemented in a subclass")
@@ -411,7 +410,8 @@ class GalsecAnalysis(ABC):
"delta_y": np.array([0.01, 0.1, 1, 10, 20]) * u.kpc,
"delta_z": np.array([0.01, 0.1, 1, 10, 20]) * u.kpc,
},
**kwargs):
**kwargs,
):
nb_scales = len(list(scales.values())[0])
@@ -432,9 +432,13 @@ class GalsecAnalysis(ABC):
if i == 0:
all_grids = grid
else:
all_grids = {fluid: vstack([all_grids[fluid], grid[fluid]]) for fluid in self.fluids}
all_grids = {
fluid: vstack([all_grids[fluid], grid[fluid]])
for fluid in self.fluids
}
return all_grids
class Galsec(GalsecAnalysis):
"""
Galactic sector extractor
@@ -579,7 +583,9 @@ class Galsec(GalsecAnalysis):
elif name == "phi":
delta_phi = (delta / data["r_bin"]) * u.rad
phi_bin = (np.trunc(data["phi"] / delta_phi) + 0.5) * delta_phi
phi_bin[phi_bin >= 2 * np.pi * u.rad] = 0.5 * delta_phi[phi_bin >= 2 * np.pi * u.rad]
phi_bin[phi_bin >= 2 * np.pi * u.rad] = (
0.5 * delta_phi[phi_bin >= 2 * np.pi * u.rad]
)
data["phi_bin"] = phi_bin
else:
print(f"Unsupported binning variable {name}")
@@ -592,7 +598,6 @@ class Galsec(GalsecAnalysis):
bin = np.digitize(data[name], bin_array, right=False)
bin_array = np.append(bin_array, np.max(data[name]))
# Store the middle value
data[f"{name}_bin"] = (bin_array[bin - 1] + bin_array[bin]) / 2
@@ -600,7 +605,6 @@ class Galsec(GalsecAnalysis):
keys_binning.append(name)
return keys_binning
def densify(self, processed_data: dict, bin_deltas: dict, filter_bounds: dict):
"""
Densify the data by adding void bins where no data is present
@@ -612,25 +616,49 @@ class Galsec(GalsecAnalysis):
if "r" in filter_bounds:
rmax = filter_bounds["r"][1].value
else:
rmax = np.max([np.max(processed_data[fluid]["r"].value) for fluid in processed_data])
rmax = np.max(
[
np.max(processed_data[fluid]["r"].value)
for fluid in processed_data
]
)
delta_phi = bin_deltas["phi"].value / rmax
bins_array[bin_name] = np.arange(0, 2 * np.pi, delta_phi * 0.9)
else:
if bin_name in filter_bounds:
bin_min = filter_bounds[bin_name][0].value + bin_deltas[bin_name].value / 2
bin_min = (
filter_bounds[bin_name][0].value
+ bin_deltas[bin_name].value / 2
)
bin_max = filter_bounds[bin_name][1].value
else:
bin_min = np.min([np.min(processed_data[fluid][bin_name].value) for fluid in processed_data])
bin_max = np.max([np.max(processed_data[fluid][bin_name].value) for fluid in processed_data])
bins_array[bin_name] = np.arange(bin_min, bin_max, bin_deltas[bin_name].value*0.9)
bin_min = np.min(
[
np.min(processed_data[fluid][bin_name].value)
for fluid in processed_data
]
)
bin_max = np.max(
[
np.max(processed_data[fluid][bin_name].value)
for fluid in processed_data
]
)
bins_array[bin_name] = np.arange(
bin_min, bin_max, bin_deltas[bin_name].value * 0.9
)
grids = np.meshgrid(*list(bins_array.values()))
bins_array = {key: grid.flatten() for key, grid in zip(bins_array.keys(), grids)}
bins_array = {
key: grid.flatten() for key, grid in zip(bins_array.keys(), grids)
}
for fluid in processed_data:
keys = processed_data[fluid].keys()
units = {key: processed_data[fluid][key].unit for key in keys}
void_array = {key: np.zeros_like(list(bins_array.values())[0]) for key in keys}
void_array = {
key: np.zeros_like(list(bins_array.values())[0]) for key in keys
}
for bin_name in bin_deltas:
void_array[bin_name] = bins_array[bin_name]
@@ -639,18 +667,23 @@ class Galsec(GalsecAnalysis):
dense_array = vstack([processed_data[fluid], void_array])
processed_data[fluid] = dense_array
self.binning(dataset=processed_data, bins=bin_deltas, filter_bounds=filter_bounds)
self.binning(
dataset=processed_data, bins=bin_deltas, filter_bounds=filter_bounds
)
for fluid in processed_data:
bin_names_with_bin = [bin_name + "_bin" for bin_name in bin_deltas]
processed_data[fluid] = processed_data[fluid].group_by(bin_names_with_bin).groups.aggregate(np.add)
processed_data[fluid] = (
processed_data[fluid]
.group_by(bin_names_with_bin)
.groups.aggregate(np.add)
)
for bin_name in bin_deltas:
processed_data[fluid].remove_column(bin_name)
processed_data[fluid].rename_column(bin_name + "_bin", bin_name)
return processed_data
def analysis(self, bins: dict, filter_bounds: dict, binning_mode="delta"):
result = {}
@@ -692,38 +725,48 @@ class Galsec(GalsecAnalysis):
return result
class GalsecTimeSeries(GalsecAnalysis):
def __init__(self, galsecs: list, loader=None):
self.galsecs = galsecs
self.loader = loader
self.loaded = loader is None
def _analysis_single(self, galsec, bins, filter_bounds, binning_mode="delta"):
if not self.loaded:
galsec = self.loader(galsec)
analysis_result = galsec.analysis(bins, filter_bounds, binning_mode=binning_mode)
analysis_result = galsec.analysis(
bins, filter_bounds, binning_mode=binning_mode
)
for fluid in analysis_result:
analysis_result[fluid]["time"] = galsec.time
return analysis_result
def analysis(self, bins: dict, filter_bounds: dict,
def analysis(
self,
bins: dict,
filter_bounds: dict,
binning_mode="delta",
aggregate=True, std=True, percentiles=[],
num_processes=1):
aggregate=True,
std=True,
percentiles=[],
num_processes=1,
):
timed_data = []
if num_processes == 1:
for galsec in self.galsecs:
timed_data.append(self._analysis_single(galsec, bins, filter_bounds, binning_mode))
timed_data.append(
self._analysis_single(galsec, bins, filter_bounds, binning_mode)
)
else:
from multiprocessing import Pool
with Pool(num_processes) as p:
timed_data = p.starmap(
self._analysis_single,
[(galsec, bins, filter_bounds, binning_mode) for galsec in self.galsecs],
[
(galsec, bins, filter_bounds, binning_mode)
for galsec in self.galsecs
],
)
averaged_data = {}
@@ -732,11 +775,13 @@ class GalsecTimeSeries(GalsecAnalysis):
if aggregate:
grouped_data = averaged_data[fluid].group_by(list(bins.keys())).groups
if std:
averaged_data[fluid + "_std"] = grouped_data.aggregate(lambda x: np.std(x.value))
averaged_data[fluid + "_std"] = grouped_data.aggregate(
lambda x: np.std(x.value)
)
for q in percentiles:
averaged_data[fluid + f"_q{q}"] = grouped_data.aggregate(lambda x: np.percentile(x.value, q))
averaged_data[fluid + f"_q{q}"] = grouped_data.aggregate(
lambda x: np.percentile(x.value, q)
)
averaged_data[fluid] = grouped_data.aggregate(np.mean)
return averaged_data