Fixing contamination_fraction in server mode

tangos/input_handlers/pynbody.py

@@ -246,16 +246,17 @@ def get_properties(self):
         if len(timesteps)>0:
             f = self.load_timestep_without_caching(sorted(timesteps)[-1])
             if self.quicker:
-                res = self._estimate_resolution_quicker(f)
+                res_kpc = self._estimate_spatial_resolution_quicker(f)
+                res_msol = self._estimate_mass_resolution_quicker(f)
             else:
-                res = self._estimate_resolution(f)
-
-            return {'approx_resolution_kpc': res}
+                res_kpc = self._estimate_spatial_resolution(f)
+                res_msol = self._estimate_mass_resolution(f)
+            return {'approx_resolution_kpc': res_kpc, 'approx_resolution_Msol': res_msol}
 
         else:
             return {}
 
-    def _estimate_resolution(self, f):
+    def _estimate_spatial_resolution(self, f):
         f.physical_units()
         if "eps" in f.dm.loadable_keys():
             # Interpret the eps array as a softening, and assume that it is not a comoving softening (as the
@@ -273,14 +274,29 @@ def _estimate_resolution(self, f):
             estimated_eps = 0.01 * frac_length * f.properties['boxsize'].in_units('kpc a', **f.conversion_context())
             return float(estimated_eps)
 
-    def _estimate_resolution_quicker(self, f):
+    def _estimate_spatial_resolution_quicker(self, f):
         interparticle_distance = float(f.properties['boxsize'].in_units("kpc a",**f.conversion_context()))/(float(len(f))**(1./3))
         res = 0.01*interparticle_distance
         logger.warn("Because 'quicker' flag is set, estimating res %.2g kpc from the file size; this could be inaccurate",
                     res)
         logger.warn(" -- it will certainly be wrong for e.g. zoom simulations")
         return res
 
+    def _estimate_mass_resolution(self, f):
+        f.physical_units()
+        if "mass" in f.dm.loadable_keys():
+            min_simulation_mass = f.dm['mass'].min()
+            return float(min_simulation_mass)
+
+    def _estimate_mass_resolution_quicker(self, f):
+        import pynbody.analysis.cosmology as cosmo
+        rho_m = cosmo.rho_M(f, z=0, unit="Msol kpc**-3 a**-3")
+        volume_box = float(f.properties['boxsize'].in_units("kpc a",**f.conversion_context()) ** 3)
+        estimated_part_mass = rho_m * volume_box / len(f)
+        logger.warn("Because 'quicker' flag is set, estimating mass res %.2g msol from the file size; this could be inaccurate",
+                    estimated_part_mass)
+        logger.warn(" -- it will certainly be wrong for e.g. zoom simulations")
+        return estimated_part_mass
 
 class RamsesHOPInputHandler(PynbodyInputHandler):
     patterns = ["output_0????"]

tangos/properties/pynbody/zoom.py

@@ -1,12 +1,28 @@
 from . import PynbodyPropertyCalculation
 
+
 class Contamination(PynbodyPropertyCalculation):
+    """ Calculates the contamination of a zoomed halo by heavier, unzoomed particles.
+        The current behaviour returns:
+            1.0 if the halo contains only heavy particles
+            X.X if the halo contains a mixture of heavy and light particles
+            0.0 if the halo contains exclusively light particles
+        Heavies are defined as any particle heavier than the deepest zoom level if multiple zooms are performed
+    """
+
+    names = "contamination_fraction"
 
     def calculate(self, halo, exist):
-        n_heavy = (halo.dm['mass'] > self.min_dm_mass).sum()
-        return float(n_heavy) / len(halo.dm)
+        loaded_data_min_dm_mass = halo.dm['mass'].min()
 
-    def preloop(self, sim, db_timestep):
-        self.min_dm_mass = sim.dm['mass'].min()
+        # If mass resolution not present as a property, ensure backward compatibility by setting the min mass to infinity
+        import numpy as np
+        simulation_min_dm_mass = self._simulation.get("approx_resolution_Msol", default=np.inf)
 
-    names = "contamination_fraction"
+        if loaded_data_min_dm_mass > simulation_min_dm_mass:
+            # Loaded data contains only "heavy" particles, e.g. an unzoomed halo in server mode
+                return 1.0
+        else:
+            # Loaded data contains heavy and/or light particles, e.g
+            n_heavy = (halo.dm['mass'] > loaded_data_min_dm_mass).sum()
+            return float(n_heavy) / len(halo.dm)

tests/test_simulation_outputs.py

@@ -15,7 +15,6 @@ def setup():
     db.config.base = os.path.join(os.path.dirname(__file__), "test_simulations")
     output_manager = pynbody_outputs.ChangaInputHandler("test_tipsy")
 
-
 def test_get_handler():
     assert db.input_handlers.get_named_handler_class('pynbody.ChangaInputHandler') == pynbody_outputs.ChangaInputHandler
 
@@ -25,6 +24,20 @@ def test_get_deprecated_handler():
 def test_handler_name():
     assert pynbody_outputs.ChangaInputHandler.handler_class_name()=="pynbody.ChangaInputHandler"
 
+def test_handler_properties():
+    prop = output_manager.get_properties()
+    assert len(prop) == 10
+    assert 'approx_resolution_kpc' in prop
+    assert 'approx_resolution_Msol' in prop
+    npt.assert_allclose(prop['approx_resolution_kpc'], 0.3499348849)
+    npt.assert_allclose(prop['approx_resolution_Msol'], 144411.17640)
+
+def test_handler_properties_quicker_flag():
+    output_manager.quicker = True
+    prop = output_manager.get_properties()
+    npt.assert_allclose(prop['approx_resolution_kpc'], 33.4360203909648)
+    npt.assert_allclose(prop['approx_resolution_Msol'], 40370039199.44858)
+
 def test_enumerate():
     assert set(output_manager.enumerate_timestep_extensions())==set(["tiny.000640","tiny.000832"])