From 85310d848ba07cc62794b3942d628161a502fbb5 Mon Sep 17 00:00:00 2001
From: Noe Brucy <noe.brucy@cea.fr>
Date: Thu, 24 Jun 2021 10:54:08 +0200
Subject: [PATCH] [ism] update ism auto file

---
 ism.py | 27 +++++++++++----------------
 1 file changed, 11 insertions(+), 16 deletions(-)

diff --git a/ism.py b/ism.py
index bfa6c31..2fa0ddf 100644
--- a/ism.py
+++ b/ism.py
@@ -36,16 +36,18 @@ def convert_coldens_s(n0):
 convert_coldens = np.vectorize(convert_coldens_s)
 
 
-def get_dispersion(dset, i):
+def get_dispersion(dset, name):
     """
-    Compute dispersion from dset["name"]
+    Compute dispersion from dset[name]
     """
-    vel = dset["vel"][:, i]
+    vel = dset[name]
     mass = dset["mass"]
     mass_tot = np.sum(mass)
-    mean = np.sum(mass * vel) / mass_tot
-    return np.sqrt(np.sum(mass * (vel - mean) ** 2) / mass_tot)
-
+    if mass_tot > 0:
+        mean = np.sum(mass * vel) / mass_tot
+        return np.sqrt(np.sum(mass * (vel - mean) ** 2) / mass_tot)
+    else:
+        return 0
 
 def get_sinks(pp):
     csv_name = f"{pp.path}/output_{pp.num:05}/sink_{pp.num:05}.csv"
@@ -54,20 +56,18 @@ def get_sinks(pp):
 
 def analyze_box(pp):
     pp.cells["mass"] = snapshotprocessor.mass_func(pp.cells)
-    pp.coldens("z")
     coldens = pp.get_value("/maps/coldens_z", unit=U.coldens)
     sinks = get_sinks(pp)
-    sinks["vel"] = np.array([sinks.vx, sinks.vy, sinks.vz])
-    sinks["vel"] *= pp.info["unit_velocity"].express(U.km_s)
     sinks["mass"] = sinks.M
     res = {}
     dirs = ["x", "y", "z"]
     res["time"] = pp.info["time"] * pp.info["unit_time"].express(U.Myr)
     for i, dir in enumerate(dirs):
-        res[f"sigma_{dir}"] = get_dispersion(pp.cells, i) * pp.info[
+        pp.cells[f"v{dir}"] = pp.cells["vel"][:, i]
+        res[f"sigma_{dir}"] = get_dispersion(pp.cells, f"v{dir}") * pp.info[
             "unit_velocity"
         ].express(U.km_s)
-        res[f"sigma_sinks_{dir}"] = get_dispersion(sinks, i) * pp.info[
+        res[f"sigma_sinks_{dir}"] = get_dispersion(sinks, f"v{dir}") * pp.info[
             "unit_velocity"
         ].express(U.km_s)
     res["coldens_mean"] = np.mean(coldens)
@@ -92,13 +92,8 @@ def load_wrapper(pp, fun):
     """
     pp.load_cells(keys=["pos", "vel", "dx", "rho"])
     pp.coldens("z")
-
     res = fun(pp)
     pp.unload_cells()
-    pp.unload_parts()
-    del pp.dm
-    del pp.gas
-    del pp.stars
     return res