118 lines
3.6 KiB
Python
118 lines
3.6 KiB
Python
import os
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib as mpl
|
|
import osyris
|
|
from scipy.io import FortranFile
|
|
|
|
from scipy.spatial import KDTree
|
|
|
|
def link_tracer(data):
|
|
mask_tracer = data["part"]["family"].values == 0
|
|
|
|
pos_cells = data["mesh"]["position"].to("pc").values
|
|
pos_tracers = data["part"]["position"].to("pc").values[mask_tracer]
|
|
kdtree = KDTree(pos_cells)
|
|
dist, ind_cell = kdtree.query(pos_tracers, k=1, p=np.inf)
|
|
return ind_cell
|
|
|
|
def get_cell_data(data, ind_cell, key, unit):
|
|
cell_data = data["mesh"][key].to(unit).values[ind_cell]
|
|
return cell_data
|
|
|
|
|
|
def get_tracer_info_osyris(path, snap):
|
|
data = osyris.RamsesDataset(snap, path=os.path.expanduser(path)).load()
|
|
mesh = data["mesh"]
|
|
|
|
mesh["temperature"] = ((mesh["pressure"] / mesh["density"] ) * (1.4 *osyris.units("proton_mass") / osyris.units("boltzmann_constant"))).to("K")
|
|
mesh["n"] = (mesh["density"] / (1.4 *osyris.units("proton_mass"))).to("cm**-3")
|
|
mesh["mass"] = (mesh["density"] * mesh["dx"]**3).to("M_sun")
|
|
|
|
ind_cell = link_tracer(data)
|
|
|
|
mask_tracer = data["part"]["family"].values == 0
|
|
id_tracers = data["part"]["identity"].values[mask_tracer]
|
|
|
|
sort = np.argsort(id_tracers)
|
|
pos_tracers = data["part"]["position"].to("pc").values[mask_tracer][sort]
|
|
T_tracers = get_cell_data(data, ind_cell, "temperature", "K")[sort]
|
|
n_tracers = get_cell_data(data, ind_cell, "n", "cm**-3")[sort]
|
|
|
|
tracer_properties = dict(
|
|
id=id_tracers[sort],
|
|
position=pos_tracers,
|
|
temperature=T_tracers,
|
|
density=n_tracers
|
|
)
|
|
|
|
time = data.meta["time"].to("Myr").magnitude
|
|
tracer_properties["time"] = time
|
|
|
|
del data, mesh
|
|
|
|
return tracer_properties
|
|
|
|
def get_tracer_fast(path, snap, reextract=False, delete=True):
|
|
print("Extracting tracers from snapshot ", snap)
|
|
UTILS = os.path.expanduser("~/src/ramses/utils/f90")
|
|
out = f"{path}/tracers_{snap:05d}.dat"
|
|
|
|
if reextract or not os.path.exists(out):
|
|
command = f"{UTILS}/part2list -inp ~/simus/mass_inflow/n0.66_rms00000_tracers_memory_70/output_{snap:05d} -out {out}"
|
|
os.system(command)
|
|
f = FortranFile(out, 'r')
|
|
n = f.read_ints(np.int32)[0]
|
|
t = f.read_reals(np.float64)[0]
|
|
id_tracers = f.read_ints(np.int32)
|
|
x = f.read_reals(np.float64)
|
|
y = f.read_reals(np.float64)
|
|
z = f.read_reals(np.float64)
|
|
f.close()
|
|
|
|
sort = np.argsort(id_tracers)
|
|
#pos_tracers = np.array([x[sort], y[sort], z[sort]]).T
|
|
|
|
|
|
|
|
tracer_properties = dict(
|
|
id=id_tracers[sort],
|
|
z=z[sort],
|
|
time = t
|
|
)
|
|
|
|
if delete:
|
|
os.remove(out)
|
|
return tracer_properties
|
|
|
|
|
|
|
|
path = os.path.expanduser("~/simus/mass_inflow/n0.66_rms00000_tracers_memory_70")
|
|
snaps = np.arange(173, 185, 1)
|
|
extractor = "fast"
|
|
|
|
extractors = {"fast": get_tracer_fast,
|
|
"osyris": get_tracer_info_osyris}
|
|
|
|
|
|
tracers_all = []
|
|
|
|
for i in snaps:
|
|
tracers_all.append(extractors[extractor](path, i))
|
|
|
|
|
|
traj_tracers = {}
|
|
for key in tracers_all[0]:
|
|
if key not in ["id", "time"]:
|
|
traj_tracers[key] = np.array([tracers_all[i][key] for i in range(len(tracers_all))])
|
|
if traj_tracers[key].ndim == 2:
|
|
traj_tracers[key] = traj_tracers[key].transpose(1, 0)
|
|
elif traj_tracers[key].ndim == 3:
|
|
traj_tracers[key] = traj_tracers[key].transpose(1, 0, 2)
|
|
|
|
traj_tracers["id"] = tracers_all[0]["id"]
|
|
traj_tracers["time"] = [tracers_all[i]["time"] for i in range(len(tracers_all))]
|
|
|
|
|
|
name = os.path.basename(path)
|
|
np.savez(f"{path}/{name}_tracers.npz", **traj_tracers) |