From 8c3db9b7cbd946d31f097889651c45174ca6b3d8 Mon Sep 17 00:00:00 2001 From: Noe Brucy Date: Mon, 30 Jan 2023 11:32:30 +0100 Subject: [PATCH] Moved galsec --- galturb/galsec.py | 355 ---------------------------------- galturb/galsec_plot.py | 53 ----- galturb/sectors_extraction.py | 265 ------------------------- 3 files changed, 673 deletions(-) delete mode 100644 galturb/galsec.py delete mode 100644 galturb/galsec_plot.py delete mode 100644 galturb/sectors_extraction.py diff --git a/galturb/galsec.py b/galturb/galsec.py deleted file mode 100644 index 5e0e538..0000000 --- a/galturb/galsec.py +++ /dev/null @@ -1,355 +0,0 @@ -# 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 deleted file mode 100644 index 0229ad9..0000000 --- a/galturb/galsec_plot.py +++ /dev/null @@ -1,53 +0,0 @@ -# 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/sectors_extraction.py b/galturb/sectors_extraction.py deleted file mode 100644 index 28a24a1..0000000 --- a/galturb/sectors_extraction.py +++ /dev/null @@ -1,265 +0,0 @@ -# coding: utf-8 - -import numpy as np -import pandas as pd -from plotter import U - - -def get_gas_dm_stars(pp): - # Load arrays - try: - pp.load_parts(keys=["pos", "vel", "mass", "epoch"]) - except KeyError: - pp.load_parts(keys=["pos", "vel", "mass"]) - - pp.load_cells(keys=["pos", "vel", "dx", "rho"]) - cells = pp.cells - parts = pp.parts - - # Compute extra fields and convert units - for dset in (cells, parts): - dset["pos_kpc"] = dset["pos"] - np.array([0.5, 0.5, 0.5]) - dset["pos_kpc"] *= pp.info["unit_length"].express(U.kpc) - dset["r"] = np.sqrt(np.sum(dset["pos_kpc"][:, :2] ** 2, axis=1)) - dset["phi"] = np.angle(dset["pos_kpc"][:, 0] + dset["pos_kpc"][:, 1] * 1j) - dset["velphi"] = pp.getter_vect_phi(dset, "vel") * pp.info[ - "unit_velocity" - ].express(U.km_s) - dset["velr"] = pp.getter_vect_r(dset, "vel") * pp.info["unit_velocity"].express( - U.km_s - ) - dset["velz"] = dset["vel"][:, 2] * pp.info["unit_velocity"].express(U.km_s) - - cells["mass_kg"] = ( - cells["rho"] - * cells["dx"] ** 3 - * (pp.info["unit_density"] * pp.info["unit_length"] ** 3).express(U.kg) - ) - - parts["mass_kg"] = parts["mass"] * pp.info["unit_mass"].express(U.kg) - - for dset in (cells, parts): - dset["ek"] = dset["mass_kg"] * np.sum(dset["vel"] ** 2, axis=1) - dset["ek"] *= (U.kg * pp.info["unit_velocity"] ** 2).express(U.J) - - # Separate DM from stars - mass_dm = np.max(parts["mass_kg"]) - mask_dm = parts["mass_kg"] == mass_dm - mask_star = parts["mass_kg"] < mass_dm - - # Create separated arrays - gas = cells - dm = {key: parts[key][mask_dm] for key in parts} - stars = {key: parts[key][mask_star] for key in parts} - - # Store arrays and return them - pp.gas = gas - pp.dm = dm - pp.stars = stars - return gas, dm, stars - - -def get_dispersion(dset, name): - """ - Compute dispersion from dset["name"] - """ - vel = dset[name] - mass = dset["mass_kg"] - mass_tot = np.sum(mass) - mean = np.sum(mass * vel) / mass_tot - return np.sqrt(np.sum(mass * (vel - mean) ** 2) / mass_tot) - - -def get_polar_sigma(dset): - """ - Get speed dispersion in polar coordinates - """ - return { - velname: get_dispersion(dset, velname) for velname in ["velphi", "velr", "velz"] - } - - -def get_sfr(pp, stars): - try: - epoch = stars["epoch"].copy() - epoch *= pp.info["unit_time"].express(U.year) - mass = stars["mass"].copy() - mass *= pp.info["unit_mass"].express(U.Msun) - mask = epoch > 0 - masstot, time = np.histogram(epoch[mask], weights=mass[mask], bins=200) - dtime = np.diff(time) - sfr = masstot[-1] / dtime[-1] - except KeyError: - sfr = 0.0 - return sfr - - -def get_coldens(pp): - """ - Get mean column density in a sector - """ - pp.coldens("z") - pp.coldens_map = pp.get_value("/maps/coldens_z", unit=U.coldens) - - im_extent = np.array(pp.get_attribute("/maps", "im_extent")) - im_extent *= pp.info["unit_length"].express(U.kpc) - map_size = pp.params.pymses.map_size - center = np.array(pp.params.disk.center) - center *= pp.info["unit_length"].express(U.kpc) - - # Physical size of cells - dx = (im_extent[1] - im_extent[0]) / map_size - dy = (im_extent[3] - im_extent[2]) / map_size - - # Physical coordinates of the center of the cells - x = np.linspace(im_extent[0], im_extent[1], map_size) + 0.5 * dx - center[0] - y = np.linspace(im_extent[2], im_extent[3], map_size) + 0.5 * dy - center[1] - - xx, yy = np.meshgrid(x, y) - - # Physical radius - pp.rr_map = np.sqrt(xx ** 2 + yy ** 2) - pp.phi_map = np.angle(xx + yy * 1j) - - -def sector_analysis(pp, gds_ring, mask_ring, phi=0, dphi=0.125): - """ - Analyze box at given coordinates - """ - masks_box = [(np.abs(dset["phi"] - phi) < dphi) for i, dset in enumerate(gds_ring)] - gds_box = [ - {key: dset[key][mask] for key in dset if key in keys} - for dset, mask in zip(gds_ring, masks_box) - ] - - res = {} - - # Generic Info - res["phi"] = phi - res["dphi"] = dphi - - res["sfr"] = get_sfr(pp, gds_box[2]) - res["coldens"] = np.mean( - pp.coldens_map[mask_ring & (np.abs(pp.phi_map - phi) < dphi)] - ) - - for dset, fluid in zip(gds_box, ["gas", "dm", "stars"]): - res[f"ek_{fluid}"] = np.sum(dset["ek"]) # J - res[f"mass_{fluid}"] = np.sum(dset["mass_kg"]) # kg - res[f"ek_spe_{fluid}"] = res[f"ek_{fluid}"] / res[f"mass_{fluid}"] # J.kg^-1 - sigma = get_polar_sigma(dset) - for dir in sigma: - res[f"sigma_{dir}_{fluid}"] = sigma[dir] - return res - - -keys = ["epoch", "ek", "mass_kg", "pos_kpc", "velphi", "velr", "velz", "mass", "phi"] - - -def ring_analysis(pp, r=4, dr=0.5, dphi=0.125, z=0, dz=0.5): - """ - Compute the average at a given of quantities computed in polar sectors. - """ - gds = [pp.gas, pp.dm, pp.stars] - phi_sectors = np.arange(-np.pi + dphi, np.pi - dphi, 2 * dphi) - masks_ring = [ - (np.abs(dset["r"] - r) < dr) & (np.abs(dset["pos_kpc"][:, 2] - z) < dz) - for dset in gds - ] - gds_ring = [ - {key: dset[key][mask] for key in dset if key in keys} - for dset, mask in zip(gds, masks_ring) - ] - mask_ring = np.abs(pp.rr_map - r) < dr - data_sec = [ - sector_analysis(pp, gds_ring, mask_ring, phi, dphi) for phi in phi_sectors - ] - - res = {} - for key in data_sec[0]: - res[key] = [d[key] for d in data_sec] - - res["r"] = [r] * len(phi_sectors) - res["dr"] = [dr] * len(phi_sectors) - return res - - -def get_time_from_relax(pp): - try: - epoch = pp.parts["epoch"].copy() - epoch *= pp.info["unit_time"].express(U.Myr) - trelax = np.min(epoch[epoch > 0]) - tfromrelax = np.max(epoch - trelax) - except KeyError: - tfromrelax = 0.0 - return tfromrelax - - -def analyse_rings(pp, radius=[4]): - res = {} - for i, r in enumerate(radius): - ring = ring_analysis(pp, r, dphi=1.0 / (2 * r)) - - if i == 0: - res = ring - else: - for key in res: - res[key].extend(ring[key]) - - res["run"] = [pp.run] * len(res["r"]) - res["num"] = [pp.num] * len(res["r"]) - time = get_time_from_relax(pp) - res["time"] = [time] * len(res["r"]) - - return res - - -def analyse_disk(pp, rmax=12.0): - - res = {} - res["run"] = pp.run - res["num"] = pp.num - res["coldens"] = np.mean(pp.coldens_map[pp.rr_map < rmax]) - res["sfr"] = get_sfr(pp, pp.stars) - res["time"] = get_time_from_relax(pp) - - for dset, fluid in zip([pp.gas, pp.dm, pp.stars], ["gas", "dm", "stars"]): - res[f"ek_{fluid}"] = np.sum(dset["ek"]) # J - res[f"mass_{fluid}"] = np.sum(dset["mass_kg"]) # kg - res[f"ek_spe_{fluid}"] = res[f"ek_{fluid}"] / res[f"mass_{fluid}"] # J.kg^-1 - return res - - -def load_wrapper(pp, fun): - """ - Wrapper to load_unload data and map function - """ - get_gas_dm_stars(pp) - get_coldens(pp) - - res = fun(pp) - pp.unload_cells() - pp.unload_parts() - del pp.dm - del pp.gas - del pp.stars - return res - - -def allinone(pp, redo=False): - def fun(pp): - return analyse_disk(pp), analyse_rings(pp, [4, 5, 6, 7, 8]) - - try: - assert not redo - sectors = pd.read_csv("{pp.run}/disk_{pp.run}_{pp.num}.csv") - disk = pd.read_csv(f"{pp.run}/disk_{pp.run}_{pp.num}.csv") - - except (AssertionError, FileNotFoundError): - res = load_wrapper(pp, fun) - disk = pd.DataFrame({key: [res[0][key]] for key in res[0]}) - sectors = pd.DataFrame({key: res[1][key] for key in res[1]}) - sectors.to_csv(f"{pp.run}/sectors_{pp.run}_{pp.num}.csv") - disk.to_csv(f"{pp.run}/disk_{pp.run}_{pp.num}.csv") - - return disk, sectors