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

147
galsec.py
View File

@@ -298,7 +298,6 @@ def regroup(
class GalsecAnalysis(ABC): class GalsecAnalysis(ABC):
def __init__(self): def __init__(self):
raise NotImplementedError("This method should be implemented in a subclass") raise NotImplementedError("This method should be implemented in a subclass")
@@ -358,7 +357,7 @@ class GalsecAnalysis(ABC):
""" """
bin_deltas = {"x": delta_x, "y": delta_y, "z": delta_z} bin_deltas = {"x": delta_x, "y": delta_y, "z": delta_z}
filter_bounds = {"x": [xmin, xmax], "y": [ymin, ymax], "z": [zmin, zmax]} filter_bounds = {"x": [xmin, xmax], "y": [ymin, ymax], "z": [zmin, zmax]}
self.grid = self.analysis(bin_deltas, filter_bounds) self.grid = self.analysis(bin_deltas, filter_bounds)
return self.grid return self.grid
@@ -411,15 +410,16 @@ class GalsecAnalysis(ABC):
"delta_y": np.array([0.01, 0.1, 1, 10, 20]) * u.kpc, "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, "delta_z": np.array([0.01, 0.1, 1, 10, 20]) * u.kpc,
}, },
**kwargs): **kwargs,
):
nb_scales = len(list(scales.values())[0]) nb_scales = len(list(scales.values())[0])
methods = { methods = {
"cartesian_analysis" : self.cartesian_analysis, "cartesian_analysis": self.cartesian_analysis,
"sector_analysis" : self.sector_analysis, "sector_analysis": self.sector_analysis,
"ring_analysis" : self.ring_analysis, "ring_analysis": self.ring_analysis,
"vertical_ring_analysis" : self.vertical_ring_analysis, "vertical_ring_analysis": self.vertical_ring_analysis,
} }
for i in range(nb_scales): for i in range(nb_scales):
@@ -429,12 +429,16 @@ class GalsecAnalysis(ABC):
print(f"Doing {sub_analysis_method} for scale {i}/{nb_scales}: {scale_i}") print(f"Doing {sub_analysis_method} for scale {i}/{nb_scales}: {scale_i}")
grid = methods[sub_analysis_method](**scale_i, **kwargs) grid = methods[sub_analysis_method](**scale_i, **kwargs)
if i==0: if i == 0:
all_grids = grid all_grids = grid
else: 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 return all_grids
class Galsec(GalsecAnalysis): class Galsec(GalsecAnalysis):
""" """
Galactic sector extractor Galactic sector extractor
@@ -545,7 +549,7 @@ class Galsec(GalsecAnalysis):
dset["vely"] = dset["velocity"][:, 1] dset["vely"] = dset["velocity"][:, 1]
dset["velz"] = dset["velocity"][:, 2] dset["velz"] = dset["velocity"][:, 2]
def binning(self, bins, filter_bounds, binning_mode = "delta", dataset=None): def binning(self, bins, filter_bounds, binning_mode="delta", dataset=None):
""" """
Bin the data Bin the data
""" """
@@ -561,8 +565,8 @@ class Galsec(GalsecAnalysis):
if binning_mode == "delta": if binning_mode == "delta":
delta = bins[name] delta = bins[name]
elif binning_mode == "number": elif binning_mode == "number":
bin_min = filter_bounds[name][0].value bin_min = filter_bounds[name][0].value
bin_max = filter_bounds[name][1].value bin_max = filter_bounds[name][1].value
delta = (bin_max - bin_min) / bins[name] delta = (bin_max - bin_min) / bins[name]
if delta is not None: if delta is not None:
@@ -579,7 +583,9 @@ class Galsec(GalsecAnalysis):
elif name == "phi": elif name == "phi":
delta_phi = (delta / data["r_bin"]) * u.rad delta_phi = (delta / data["r_bin"]) * u.rad
phi_bin = (np.trunc(data["phi"] / delta_phi) + 0.5) * delta_phi 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 data["phi_bin"] = phi_bin
else: else:
print(f"Unsupported binning variable {name}") print(f"Unsupported binning variable {name}")
@@ -592,7 +598,6 @@ class Galsec(GalsecAnalysis):
bin = np.digitize(data[name], bin_array, right=False) bin = np.digitize(data[name], bin_array, right=False)
bin_array = np.append(bin_array, np.max(data[name])) bin_array = np.append(bin_array, np.max(data[name]))
# Store the middle value # Store the middle value
data[f"{name}_bin"] = (bin_array[bin - 1] + bin_array[bin]) / 2 data[f"{name}_bin"] = (bin_array[bin - 1] + bin_array[bin]) / 2
@@ -600,8 +605,7 @@ class Galsec(GalsecAnalysis):
keys_binning.append(name) keys_binning.append(name)
return keys_binning return keys_binning
def densify(self, processed_data: dict, bin_deltas: dict, filter_bounds: dict):
def densify(self, processed_data:dict, bin_deltas:dict, filter_bounds:dict):
""" """
Densify the data by adding void bins where no data is present Densify the data by adding void bins where no data is present
""" """
@@ -612,25 +616,49 @@ class Galsec(GalsecAnalysis):
if "r" in filter_bounds: if "r" in filter_bounds:
rmax = filter_bounds["r"][1].value rmax = filter_bounds["r"][1].value
else: 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 delta_phi = bin_deltas["phi"].value / rmax
bins_array[bin_name] = np.arange(0, 2 * np.pi, delta_phi*0.9) bins_array[bin_name] = np.arange(0, 2 * np.pi, delta_phi * 0.9)
else: else:
if bin_name in filter_bounds: if bin_name in filter_bounds:
bin_min = filter_bounds[bin_name][0].value + bin_deltas[bin_name].value / 2 bin_min = (
bin_max = filter_bounds[bin_name][1].value filter_bounds[bin_name][0].value
+ bin_deltas[bin_name].value / 2
)
bin_max = filter_bounds[bin_name][1].value
else: else:
bin_min = np.min([np.min(processed_data[fluid][bin_name].value) for fluid in processed_data]) bin_min = np.min(
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) 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())) 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: for fluid in processed_data:
keys = processed_data[fluid].keys() keys = processed_data[fluid].keys()
units = {key: processed_data[fluid][key].unit for key in 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: for bin_name in bin_deltas:
void_array[bin_name] = bins_array[bin_name] void_array[bin_name] = bins_array[bin_name]
@@ -639,18 +667,23 @@ class Galsec(GalsecAnalysis):
dense_array = vstack([processed_data[fluid], void_array]) dense_array = vstack([processed_data[fluid], void_array])
processed_data[fluid] = dense_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: for fluid in processed_data:
bin_names_with_bin = [bin_name + "_bin" for bin_name in bin_deltas] 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: for bin_name in bin_deltas:
processed_data[fluid].remove_column(bin_name) processed_data[fluid].remove_column(bin_name)
processed_data[fluid].rename_column(bin_name + "_bin", bin_name) processed_data[fluid].rename_column(bin_name + "_bin", bin_name)
return processed_data return processed_data
def analysis(self, bins: dict, filter_bounds: dict, binning_mode="delta"): def analysis(self, bins: dict, filter_bounds: dict, binning_mode="delta"):
result = {} result = {}
@@ -692,51 +725,63 @@ class Galsec(GalsecAnalysis):
return result return result
class GalsecTimeSeries(GalsecAnalysis): class GalsecTimeSeries(GalsecAnalysis):
def __init__(self, galsecs: list, loader=None): def __init__(self, galsecs: list, loader=None):
self.galsecs = galsecs self.galsecs = galsecs
self.loader = loader self.loader = loader
self.loaded = loader is None self.loaded = loader is None
def _analysis_single(self, galsec, bins, filter_bounds, binning_mode="delta"): def _analysis_single(self, galsec, bins, filter_bounds, binning_mode="delta"):
if not self.loaded: if not self.loaded:
galsec = self.loader(galsec) 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: for fluid in analysis_result:
analysis_result[fluid]["time"] = galsec.time analysis_result[fluid]["time"] = galsec.time
return analysis_result return analysis_result
def analysis(
def analysis(self, bins: dict, filter_bounds: dict, self,
binning_mode="delta", bins: dict,
aggregate=True, std=True, percentiles=[], filter_bounds: dict,
num_processes=1): binning_mode="delta",
aggregate=True,
std=True,
percentiles=[],
num_processes=1,
):
timed_data = [] timed_data = []
if num_processes == 1: if num_processes == 1:
for galsec in self.galsecs: 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: else:
from multiprocessing import Pool from multiprocessing import Pool
with Pool(num_processes) as p: with Pool(num_processes) as p:
timed_data = p.starmap( timed_data = p.starmap(
self._analysis_single, 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 = {} averaged_data = {}
for fluid in timed_data[0]: for fluid in timed_data[0]:
averaged_data[fluid] = vstack([result[fluid] for result in timed_data]) averaged_data[fluid] = vstack([result[fluid] for result in timed_data])
if aggregate: if aggregate:
grouped_data = averaged_data[fluid].group_by(list(bins.keys())).groups grouped_data = averaged_data[fluid].group_by(list(bins.keys())).groups
if std: if std:
averaged_data[fluid + "_std"] = grouped_data.aggregate(lambda x: np.std(x.value)) averaged_data[fluid + "_std"] = grouped_data.aggregate(
for q in percentiles: lambda x: np.std(x.value)
averaged_data[fluid + f"_q{q}"] = grouped_data.aggregate(lambda x: np.percentile(x.value, q)) )
averaged_data[fluid] = grouped_data.aggregate(np.mean) for q in percentiles:
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 return averaged_data