noe/pipeline
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Moved galsec

Browse Source
This commit is contained in:
Noe Brucy
2023-01-30 11:32:30 +01:00
parent 84666b694e
commit 8c3db9b7cb
3 changed files with 0 additions and 673 deletions
Download Patch File Download Diff File Expand all files Collapse all files
galturb/galsec.py
-355
View File
@@ -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)
galturb/galsec_plot.py
-53
View File
@@ -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)
galturb/sectors_extraction.py
-265
View File
@@ -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