From 84666b694ea5520bedf821baefb240c0e41fd921 Mon Sep 17 00:00:00 2001 From: Noe Brucy Date: Thu, 26 Jan 2023 16:03:27 +0100 Subject: [PATCH] added galsec --- .pre-commit-config.yaml | 15 +- galturb/galsec.py | 355 +++++++++++++++++++++++++++ galturb/galsec_plot.py | 53 ++++ galturb/params_gal.yml | 96 ++++++++ galturb/pymsesrc | 13 + galturb/sectors_extraction_ramses.py | 108 ++++++++ 6 files changed, 632 insertions(+), 8 deletions(-) create mode 100644 galturb/galsec.py create mode 100644 galturb/galsec_plot.py create mode 100644 galturb/params_gal.yml create mode 100644 galturb/pymsesrc create mode 100644 galturb/sectors_extraction_ramses.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6fbbb70..45eeeec 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -20,15 +20,14 @@ repos: hooks: - id: remove-crlf - repo: https://github.com/psf/black - rev: 20.8b1 + rev: 22.6.0 hooks: - id: black - - repo: https://github.com/asottile/blacken-docs - rev: v1.8.0 - hooks: - - id: blacken-docs - additional_dependencies: [ black==20.8b1 ] - - repo: https://gitlab.com/PyCQA/flake8 - rev: 3.8.4 +# - repo: https://github.com/asottile/blacken-docs +# rev: v1.8.0 +# hooks: +# - id: blacken-docs + - repo: https://github.com/pycqa/flake8/ + rev: 6.0.0 hooks: - id: flake8 diff --git a/galturb/galsec.py b/galturb/galsec.py new file mode 100644 index 0000000..5e0e538 --- /dev/null +++ b/galturb/galsec.py @@ -0,0 +1,355 @@ +# coding: utf-8 +""" + Galaxy kiloparsec extractor + +""" +import numpy as np +from astropy.table import QTable, hstack +from astropy import units as u +from astropy.units.quantity import Quantity + + +def vect_r(position: np.array, vector: np.array) -> np.array: + """Radial component of a vector + + Parameters + ---------- + position : np.array (N, 3) + (only reads x and y) + vector : np.array (N, 3) + (only reads x and y components) + + Returns + ------- + np.array (N) + Radial component + """ + r = position[:, :2] + ur = np.transpose((np.transpose(r) / np.sqrt(np.sum(r**2, axis=1)))) + return np.einsum("ij, ij -> i", vector[:, :2], ur) + + +def vect_phi(position: np.array, vector: np.array): + """Azimuthal component of a vecto + + Parameters + ---------- + position : np.array (N, 3) + (only reads x and y) + vector : np.array (N, 3) + (only reads x and y components) + + Returns + ------- + np.array (N) + Azimuthal component + """ + r = position[:, :2] + r_norm = np.sqrt(np.sum(r**2, axis=1)) + rot = np.array([[0, -1], [1, 0]]) + uphi = np.transpose(np.einsum("ij, kj -> ik", rot, r) / r_norm) + return np.einsum("ij,ij -> i", vector[:, :2], uphi) + + +def get_sfr( + stars: QTable, time: Quantity[u.Myr], average_on: Quantity[u.Myr] = 30 * u.Myr +) -> Quantity[u.Msun / u.year]: + """_summary_ + + Parameters + ---------- + stars : QTable + _description_ + time : Quantity[u.Myr], + time at wich the SFR should be computed + average_on : Quantity, optional + compute the sfr on the last `average_on` years, by default 30*u.Myr + + Returns + ------- + float + SFR in Msun/yr + """ + try: + mask = (stars["birth_time"] > max(time - average_on, 0 * u.Myr)) & ( + stars["birth_time"] < time + ) + dtime = min(average_on, time).to(u.year) + if dtime == 0: + sfr = 0.0 * u.Msun / u.year + else: + sfr = np.sum(stars["mass"][mask]) / dtime + except KeyError: + sfr = 0.0 * u.Msun / u.year + return sfr + + +def aggregate( + grouped_data: QTable, + weight_field: str = "mass", + extensive_fields: str = ["mass", "ek"], + weighted_fields: list = ["velr", "velphi", "velz"], +) -> QTable: + """Aggregate wisely from grouped data + + Parameters + ---------- + grouped_data : QTable + should already have group information + weight_field : str, optional + the field used for the weighting of the non extensive quantities, by default "mass" + extensive_fields : str, optional + these will be summed. Should include weight_field, by default ["mass", "ek"] + weighted_fields : list, optional + the field that will be weighted, by default ["velr", "velphi", "velz"] + + Returns + ------- + QTable + a tabble with the aggregated value for each group + """ + assert weight_field in extensive_fields + + for field in weighted_fields: + grouped_data[field] *= grouped_data[weight_field] + + binned_data = grouped_data[extensive_fields + weighted_fields].groups.aggregate( + np.add + ) + + for field in weighted_fields: + binned_data[field] /= binned_data[weight_field] # weighted mean + + # Veocity dispersion + grouped_data[field] /= grouped_data[weight_field] + for i in range(len(grouped_data.groups)): + slice = grouped_data.groups.indices[i], grouped_data.groups.indices[i + 1] + grouped_data[slice[0] : slice[1]][field] -= binned_data[i][field] + grouped_data[field] = grouped_data[weight_field] * grouped_data[field] ** 2 + binned_data[f"sigma_{field}"] = grouped_data[field,].groups.aggregate( + np.add + )[field] + binned_data[f"sigma_{field}"] = np.sqrt( + binned_data[f"sigma_{field}"] / binned_data[weight_field] + ) + + return binned_data + + +class Galsec: + """ + Galactic sector extractor + """ + + fluids = ["gas", "stars", "dm"] + units = { + "gas": { + "position": u.kpc, + "volume": u.kpc**3, + "velocity": u.km / u.s, + "mass": u.Msun, + }, + "stars": { + "position": u.kpc, + "velocity": u.km / u.s, + "mass": u.Msun, + "birth_time": u.Myr, + }, + "dm": {"position": u.kpc, "velocity": u.km / u.s, "mass": u.Msun}, + } + + def __init__(self, data: dict, copy=True) -> None: + """Initiazise a Galaxasec object + + Parameters + ---------- + cell_data : dict + A dataset of cells in the following format: + header: + time float [Myr] time since the start of the simulation + gas: + position float array (Ngas, 3) [kpc], centered + volume float array (Ngas) [pc^3] + velocity float array (Ngas, 3) [km/s] + mass float array (Ngas) [Msun] + stars: + position float array (Nstar, 3) [kpc], centered + velocity float array (Nstar, 3) [km/s] + mass float array (Nstar) [Msun] + birth_time float array (Nstar) [Myr] + dm: + position float array (Ngas, 3) [kpc], centered + velocity float array (Ngas, 3) [km/s] + mass float array (Ngas) [Msun] + maps: + extent float list [xmin, xmax, ymin, ymax] Coordinates of the edges of the map, centered + gas_coldens float array (Nx, Ny) [Msun/pc^2] Map of column density + copy : bool, default=True + Wheter the data should be copied. + """ + + self.data = {} + + for fluid in self.fluids: + self.data[fluid] = QTable(data[fluid], units=self.units[fluid], copy=copy) + + self.time = data["header"]["time"] * u.Myr + + self.compute_derived_values() + + def compute_derived_values(self) -> None: + """ + Helper function to computed values derivated from input data + """ + for fluid in self.fluids: + dset = self.data[fluid] + dset["r"] = np.sqrt( + np.sum(dset["position"][:, :2] ** 2, axis=1) + ) # galactic radius + dset["phi"] = np.angle( + dset["position"][:, 0] + dset["position"][:, 1] * 1j + ) # galactic longitude + dset["phi"][dset["phi"] < 0] += ( + 2 * np.pi * u.rad + ) # rescaling to get only positive values + dset["velphi"] = vect_phi( + dset["position"], dset["velocity"] + ) # azimuthal velocity + dset["velr"] = vect_r(dset["position"], dset["velocity"]) # radial velocity + dset["velz"] = dset["velocity"][:, 2] # vertical velocity + dset["ek"] = (dset["mass"] * np.sum(dset["velocity"] ** 2, axis=1)).to( + u.erg + ) # kinetic energy + + def ring_binning(self, delta_r: Quantity[u.kpc]): + """Add radial bin informations to the dataset + + Parameters + ---------- + delta_r : Quantity[u.kpc] + spacing between two radial bins + """ + for fluid in self.fluids: + r_bin = np.trunc(self.data[fluid]["r"] / delta_r) + r_bin = (r_bin + 0.5) * delta_r # Store the middle value + self.data[fluid]["r_bin"] = r_bin + + def sector_binning(self, delta_r: Quantity[u.kpc], delta_l: Quantity[u.kpc]): + """Add sector bin information to the dataset + + Parameters + ---------- + delta_r : Quantity[u.kpc] + spacing between two radial bins + delta_l : Quantity[u.kpc] + spacing (in spatial units) between two azimuthal bins + """ + + self.ring_binning(delta_r) + + for fluid in self.fluids: + delta_phi = (delta_l / self.data[fluid]["r_bin"]) * u.rad + phi_bin = (np.trunc(self.data[fluid]["phi"] / delta_phi) + 0.5) * delta_phi + self.data[fluid]["phi_bin"] = phi_bin + + def sector_analysis( + self, + delta_r: Quantity[u.kpc] = u.kpc, + delta_l: Quantity[u.kpc] = u.kpc, + rmin: Quantity[u.kpc] = 1 * u.kpc, + rmax: Quantity[u.kpc] = 12 * u.kpc, + ): + """Compute the aggration of quantities in sectors bins + + Parameters + ---------- + delta_r : Quantity[u.kpc], optional + spacing between two radial bins, by default u.kpc + delta_l : Quantity[u.kpc], optional + spacing (in spatial units) between two azimuthal bins, by default u.kpc + rmin : Quantity[u.kpc], optional + filter out bin below that radius, by default 1*u.kpc + rmax : Quantity[u.kpc], optional + filter out bin beyond that radius, by default 12*u.kpc + """ + + self.sector_binning(delta_r, delta_l) + grouped_data = {} + self.sectors = {} + + for fluid in self.fluids: + if fluid == "gas": + extensive_fields = ["mass", "ek", "volume"] + else: + extensive_fields = ["mass", "ek"] + filtered_data = self.data[fluid][ + np.logical_and( + self.data[fluid]["r"] > rmin, self.data[fluid]["r"] < rmax + ) + ] + grouped_data[fluid] = filtered_data.group_by(["r_bin", "phi_bin"]) + self.sectors[fluid] = hstack( + [ + grouped_data[fluid]["r_bin", "phi_bin"].groups.aggregate(np.fmin), + aggregate(grouped_data[fluid], extensive_fields=extensive_fields), + ] + ) + self.sectors[fluid].rename_column("r_bin", "r") + self.sectors[fluid].rename_column("phi_bin", "phi") + + self.sectors["stars"]["sfr"] = ( + np.zeros(len(self.sectors["stars"]["mass"])) * u.Msun / u.year + ) + for i, group in enumerate(grouped_data["stars"].groups): + self.sectors["stars"]["sfr"][i] = get_sfr(group, self.time) + + self.sectors["stars"]["sfr"][i] = get_sfr(group, self.time) + + def ring_analysis( + self, + delta_r: Quantity[u.kpc] = u.kpc, + rmin: Quantity[u.kpc] = 1 * u.kpc, + rmax: Quantity[u.kpc] = 12 * u.kpc, + ): + """Compute the aggration of quantities in radial bins + + Parameters + ---------- + delta_r : Quantity[u.kpc], optional + spacing between two radial bins, by default u.kpc + rmin : Quantity[u.kpc], optional + filter out bin below that radius, by default 1*u.kpc + rmax : Quantity[u.kpc], optional + filter out bin beyond that radius, by default 12*u.kpc + """ + + self.ring_binning(delta_r) + grouped_data = {} + self.rings = {} + + for fluid in self.fluids: + if fluid == "gas": + extensive_fields = ["mass", "ek", "volume"] + else: + extensive_fields = ["mass", "ek"] + filtered_data = self.data[fluid][ + np.logical_and( + self.data[fluid]["r"] > rmin, self.data[fluid]["r"] < rmax + ) + ] + grouped_data[fluid] = filtered_data.group_by(["r_bin"]) + self.rings[fluid] = hstack( + [ + grouped_data[fluid][ + "r_bin", + ].groups.aggregate(np.fmin), + aggregate(grouped_data[fluid], extensive_fields=extensive_fields), + ] + ) + self.rings[fluid].rename_column("r_bin", "r") + + self.rings["stars"]["sfr"] = ( + np.zeros(len(self.rings["stars"]["mass"])) * u.Msun / u.year + ) + for i, group in enumerate(grouped_data["stars"].groups): + self.rings["stars"]["sfr"][i] = get_sfr(group, self.time) diff --git a/galturb/galsec_plot.py b/galturb/galsec_plot.py new file mode 100644 index 0000000..0229ad9 --- /dev/null +++ b/galturb/galsec_plot.py @@ -0,0 +1,53 @@ +# coding: utf-8 +""" + Galaxy kiloparsec plotter +""" +from astropy.table import QTable +import matplotlib.pyplot as plt +import numpy as np + + +def plot_radial(table: QTable): + + fig = plt.figure(figsize=(6, 5)) + # setting the axis limits in [left, bottom, width, height] + rect = [0.1, 0.1, 0.8, 0.8] + + # the polar axis: + ax_polar = fig.add_axes(rect, polar=True, frameon=False) + rmax = 10 + ax_polar.set_rmax(rmax) + + ax_polar.set_xticklabels([]) + + sc = ax_polar.scatter(table["phi"], table["r"], c=table["mass"], s=100, alpha=1) + plt.colorbar(sc) + + fig, ax = plt.subplots(subplot_kw=dict(projection="polar"), figsize=(6, 5)) + + rbins = np.unique(table["r"]) + delta_r = rbins[1] - rbins[0] + + for r in rbins: + + mask = table["r"] == r + phibins = table["phi"][mask].value + + C = np.log10(table["mass"][mask].value) + N = len(C) + C = C.reshape(1, N) + P = np.zeros(shape=(2, N + 1)) + + R = np.ones(shape=(2, N + 1)) * (r + delta_r / 2.0) + R[0, :] -= delta_r + R = R.value + + deltaphi = phibins[1] - phibins[0] + P[0, :-1] = phibins - deltaphi / 2 + P[1, :-1] = phibins - deltaphi / 2 + P[0, -1] = P[0, 0] + P[1, -1] = P[1, 0] + + pc = ax.pcolormesh(P, R, C, cmap="magma_r", vmin=6, vmax=9) + print(P, R, C) + fig.colorbar(pc) diff --git a/galturb/params_gal.yml b/galturb/params_gal.yml new file mode 100644 index 0000000..8bb03ca --- /dev/null +++ b/galturb/params_gal.yml @@ -0,0 +1,96 @@ +plot : # Plot parameters + put_title : False # Add a title to plot + + # Maps + ntick : 6 # Number of ticks for maps + + # Overlays + vel_red : 40 # Take point each vel_red for velocities + + time_fmt : "time = {:.3g} {}" # Time format string, 1st field is time and 2nd is unit + +disk: # Disk specific parameters + enable : True # Enable specific disk analysis + center : [0.5, 0.5, 0.5] # Position of the center + binning : "log" # Kind of binning (lin : linear, log : logarithmic) + mass_star : 1. # Mass of the central star + nb_bin : 256 # Number of bins for averaged quantities + bin_in : 1e-3 # Outer radius of the inner bin + bin_out : 18 # Inner radius of the outer bin + rmin_pdf : 1 # Inner radius for PDF computation + rmax_pdf : 18 # Outer radius for PDF computation + + +pdf: # parameters for probability density functions + nb_bin : 100 # Number of bins for the PDF + range : [-1.5, 2.5] # Range of the PDF (log of fluctuation) + xmin_fit : 0.3 # Lower boundary of the fit (log of fluctuation) + xmax_fit : 1.5 # Upper boundary of the fit (log of fluctuation) + fit_cut : 1e-4 # Exclude value that are < fit_cut * maximum + + +pymses: # Parameters for Pymses reader + + # Source settings + variables : ["rho", "vel", "P"] # Read these variables + part_variables : ["vel","mass","id","level","epoch"] # Read these variables + order : '<' # Bit order + + + # Processing options + levelmax : 20 # Maximal AMR level visited when looking levels + fft : False # Quick and dirty rendering using FFT + verbose : True # Let pymses write on standart output + multiprocessing : True # Whether to use multiprocessing + + # Camera settings + center : [0.5, 0.5, 0.5] # Center of the image + zoom : 4. # Zoom of the image + map_size : 2048 # Size of the computed maps in pixel + + # Filter parameters + filter : False # Enable filtering + min_coords : [0.35, 0.35, 0.45] + max_coords : [0.65, 0.65, 0.55] + +input: # Parameters on how to look for input files (= output from Ramses) + + log_prefix : "run.log" # Prefix of the log file + label_filename : "label.txt" # Name of the label file + nml_filename : "galaxy.nml" # name of the namelist file + ramses_ism : False # If ramses-ism is used + +out: # Parameters for post processing + tag : "" # Tag for the image + interactive : True # Interactive mode (keep figures open) + save : True # Save the plots on the disk + ext : '.jpeg' # extension for plots + fmt : "" # Format of the output filename for plots + # The following keys are accepted + # {out} : The output directory (where hdf5 files are also stored) + # {run} : Name of the relevant run + # {num} : Name of the input file (from Ramses) + # {ext} : Extension defined above + # {name} : Name of the rule + # {tag} : Tag defined above + # {nml[nml_key]} : The value of nml_key in the namelist (ex: {nml[amr_params/levelmin]}) + + +process: # General setting of the post-processor module + verbose : True # Give more infos on what is going on + num_process : 1 # Number of forks + save_cells : True # Save cells structure on disk + save_parts : True # Save particles on disk + unload_cells : True # Save memory usage + +rules: # Specific rules parameters + turb_energy_threshold : -1 # Remove invalid data (<0 = no threshold) + + +astrophysix: # Parameters for astrophysix and galactica + simu_fmt : "{tag}_{run}" # Format of the name of simulation + descr_fmt : "{tag}_{run}" # Format of the default description + # The following keys are accepted + # {run} : Name of the relevant run + # {tag} : Tag defined above + # {nml[nml_key]} : The value of nml_key in the namelist (ex: {nml[amr_params/levelmin]}) diff --git a/galturb/pymsesrc b/galturb/pymsesrc new file mode 100644 index 0000000..bae6e2e --- /dev/null +++ b/galturb/pymsesrc @@ -0,0 +1,13 @@ +{ + "Version": 1, + "RAMSES": { + "ndimensions": 3, + "amr_field_descr": [ + {"__type__": "scalar_field", "__file_type__": "hydro", "name": "rho", "ivar": 0}, + {"__type__": "vector_field", "__file_type__": "hydro", "name": "vel", "ivars": [1, 2, 3]}, + {"__type__": "scalar_field", "__file_type__": "hydro", "name": "P", "ivar": 4}, + {"__type__": "scalar_field", "__file_type__": "grav", "name": "phi", "ivar": 0}, + {"__type__": "vector_field", "__file_type__": "grav", "name": "g", "ivars": [1, 2, 3]} + ] + } +} diff --git a/galturb/sectors_extraction_ramses.py b/galturb/sectors_extraction_ramses.py new file mode 100644 index 0000000..1bf1b0b --- /dev/null +++ b/galturb/sectors_extraction_ramses.py @@ -0,0 +1,108 @@ +# coding: utf-8 + +import numpy as np +from snapshotprocessor import U +from tables import NoSuchNodeError + + +def load_fields(pp): + """ + Parameters + ---------- + path : _type_ + _description_ + output : _type_ + _description_ + + Returns + ------- + dict + A dataset of cells in the following format: + gas: + position (Ngas, 3) [kpc], centered + volume (Ngas) [pc^3] + velocity (Ngas, 3) [km/s] + mass (Ngas) [Msun] + stars: + position (Nstar, 3) [kpc], centered + velocity (Nstar, 3) [km/s] + mass (Nstar) [Msun] + birth_time (Nstar) [Myr] + dm: + position (Ngas, 3) [kpc], centered + velocity (Ngas, 3) [km/s] + mass (Ngas) [Msun] + maps: + extent (xmin, xmax, ymin, ymax) Coordinates of the edges of the map, centered + gas_coldens (Nx, Ny) [Msun/pc^2], map of column density + """ + + # Load arrays + pp.load_cells(keys=["pos", "vel", "dx", "rho"]) + + try: + pp.load_parts(keys=["pos", "vel", "mass", "epoch"]) + except (KeyError, NoSuchNodeError): + pp.load_parts(keys=["pos", "vel", "mass"]) + + cells = pp.cells + parts = pp.parts + + if "epoch" not in parts: + parts["epoch"] = np.zeros(len(pp.parts["mass"])) + + # Compute extra fields and convert units + for dset in (cells, parts): + dset["position"] = dset["pos"] - np.array([0.5, 0.5, 0.5]) + + cells["mass"] = ( + cells["rho"] + * cells["dx"] ** 3 + * (pp.info["unit_density"] * pp.info["unit_length"] ** 3).express(U.Msun) + ) + + cells["volume"] = cells["dx"] ** 3 * (pp.info["unit_length"] ** 3).express( + U.kpc**3 + ) + + # Separate DM from stars + mass_dm = np.max(parts["mass"]) + mask_dm = parts["mass"] == mass_dm + mask_star = parts["mass"] < mass_dm + + # Create dataset + data = { + "header": {"time": pp.time * pp.info["unit_time"].express(U.Myr)}, + "gas": { + "position": cells["position"] * pp.info["unit_length"].express(U.kpc), + "volume": cells["volume"], + "velocity": cells["vel"] * pp.info["unit_velocity"].express(U.km_s), + "mass": cells["mass"], + }, + "stars": { + "position": parts["position"][mask_star] + * pp.info["unit_length"].express(U.kpc), + "velocity": parts["vel"][mask_star] + * pp.info["unit_velocity"].express(U.km_s), + "mass": parts["mass"][mask_star] * pp.info["unit_mass"].express(U.Msun), + "birth_time": parts["epoch"][mask_star] + * pp.info["unit_time"].express(U.Myr), + }, + "dm": { + "position": parts["position"][mask_dm] + * pp.info["unit_length"].express(U.kpc), + "velocity": parts["vel"][mask_dm] + * pp.info["unit_velocity"].express(U.km_s), + "mass": parts["mass"][mask_dm] * pp.info["unit_mass"].express(U.Msun), + }, + } + return data + + +if __name__ == "__main__": + from snapshotprocessor import SnapshotProcessor + + pp = SnapshotProcessor( + "/home/nbrucy/simus/F20H_alfven_frig", num=1, params="params_gal.yml" + ) + data = load_fields(pp)