Improved pruning

examples/martini_fire.ipynb

@@ -512,9 +512,9 @@
     "            directory_download(url_base, urllib.parse.urljoin(directory, link[\"href\"]))\n",
     "\n",
     "\n",
-    "# if os.path.isdir(simulation_directory):\n",
-    "#     print(\"Found directory matching simulation_directory locally, skipping download.\")\n",
-    "# else:\n",
+    "if os.path.isdir(simulation_directory):\n",
+    "    print(\"Found directory matching simulation_directory locally, skipping download.\")\n",
+    "else:\n",
     "directory_download(data_url, simulation_directory)"
    ]
   },

martini/martini.py

@@ -143,7 +143,9 @@ def __init__(
 
         return
 
-    def _prune_particles(self, spatial=True, spectral=True, obj_type_str="data cube"):
+    def _prune_particles(
+        self, spatial=True, spectral=True, mass=True, obj_type_str="data cube"
+    ):
         """
         Determines which particles cannot contribute to the DataCube and
         removes them to speed up calculation. Assumes the kernel is 0 at
@@ -158,6 +160,8 @@ def _prune_particles(self, spatial=True, spectral=True, obj_type_str="data cube"
         spectral : bool
             If ``True``, prune particles that fall outside the spectral bandwidth.
             (Default: ``True``)
+        mass : bool
+            If ``True``, prune particles that have zero HI mass. (Default: ``True``)
         obj_type_str : str
             String describing the object to be pruned for messages.
             (Default: ``"data cube"``)
@@ -194,12 +198,16 @@ def _prune_particles(self, spatial=True, spectral=True, obj_type_str="data cube"
             if spectral
             else tuple()
         )
-        reject_mask = np.zeros(self.source.pixcoords[0].shape)
-        # this could be a logical_or.reduce?:
-        for condition in spatial_reject_conditions + spectral_reject_conditions:
-            reject_mask = np.logical_or(reject_mask, condition)
-        self.source.apply_mask(np.logical_not(reject_mask))
-        self.sph_kernel._apply_mask(np.logical_not(reject_mask))
+        mass_reject_conditions = (self.source.mHI_g == 0,) if mass else tuple()
+        accept_mask = np.logical_not(
+            np.logical_or.reduce(
+                spatial_reject_conditions
+                + spectral_reject_conditions
+                + mass_reject_conditions
+            )
+        )
+        self.source.apply_mask(accept_mask)
+        self.sph_kernel._apply_mask(accept_mask)
         if not self.quiet:
             print(
                 f"Pruned particles that will not contribute to {obj_type_str}, "
@@ -1551,6 +1559,7 @@ def __init__(
             _prune_kwargs=dict(
                 spatial=False,
                 spectral=True,
+                mass=True,
                 obj_type_str="spectrum",
             ),
             quiet=quiet,

martini/martini.pyi

@@ -33,7 +33,11 @@ class _BaseMartini:
         _prune_kwargs: T.Dict[str, T.Union[bool, str]] = ...,
     ) -> None: ...
     def _prune_particles(
-        self, spatial: bool = ..., spectral: bool = ..., obj_type_str: str = ...
+        self,
+        spatial: bool = ...,
+        spectral: bool = ...,
+        mass: bool = ...,
+        obj_type_str: str = ...,
     ) -> None: ...
     def _evaluate_pixel_spectrum(
         self,

martini/sources/sph_source.py

@@ -312,7 +312,7 @@ def apply_mask(self, mask):
             raise ValueError("Mask must have same length as particle arrays.")
         mask_sum = np.sum(mask)
         if mask_sum == 0:
-            raise RuntimeError("No source particles in target region.")
+            raise RuntimeError("No non-zero mHI source particles in target region.")
         self.npart = mask_sum
         if not self.T_g.isscalar:
             self.T_g = self.T_g[mask]

tests/test_martini.py

@@ -306,8 +306,10 @@ def test_add_noise(self, m_init):
             (-7 * U.km / U.s, False),
         ),
     )
+    @pytest.mark.parametrize(("mass_off", "mass_in"), ((0, False), (1, True)))
     @pytest.mark.parametrize("spatial", (True, False))
     @pytest.mark.parametrize("spectral", (True, False))
+    @pytest.mark.parametrize("mass", (True, False))
     def test_prune_particles(
         self,
         ra_off,
@@ -316,26 +318,41 @@ def test_prune_particles(
         dec_in,
         v_off,
         v_in,
+        mass_off,
+        mass_in,
         single_particle_source,
         spatial,
         spectral,
+        mass,
     ):
         """
         Check that a particle offset by a specific set of (RA, Dec, v) is inside/outside
         the cube as expected.
         """
-        if spatial and spectral:
+        if spatial and spectral and mass:
+            expect_particle = all((ra_in, dec_in, v_in, mass_in))
+        elif spatial and spectral and not mass:
             expect_particle = all((ra_in, dec_in, v_in))
-        elif spatial and not spectral:
+        elif spatial and not spectral and mass:
+            expect_particle = all((ra_in, dec_in, mass_in))
+        elif spatial and not spectral and not mass:
             expect_particle = all((ra_in, dec_in))
-        elif spectral and not spatial:
+        elif not spatial and spectral and mass:
+            expect_particle = all((v_in, mass_in))
+        elif not spatial and spectral and not mass:
             expect_particle = v_in
-        elif not spectral and not spatial:
+        elif not spectral and not spatial and mass:
+            expect_particle = mass_in
+        elif not spectral and not spatial and not mass:
             expect_particle = True
         # set distance so that 1kpc = 1arcsec
         distance = (1 * U.kpc / 1 / U.arcsec).to(U.Mpc, U.dimensionless_angles())
         source = single_particle_source(
-            distance=distance, ra=ra_off, dec=dec_off, vpeculiar=v_off
+            distance=distance,
+            ra=ra_off,
+            dec=dec_off,
+            vpeculiar=v_off,
+            mHI_g=mass_off * np.ones(1) * 1.0e4 * U.Msun,
         )
         datacube = DataCube(
             n_px_x=2,
@@ -359,13 +376,14 @@ def test_prune_particles(
             noise=None,
             sph_kernel=sph_kernel,
             spectral_model=spectral_model,
-            _prune_kwargs=dict(spatial=spatial, spectral=spectral),
+            _prune_kwargs=dict(spatial=spatial, spectral=spectral, mass=mass),
         )
         # if more than 1px (datacube) + 5px (pad) + 2px (sm_range) then expect to prune
         # if more than 1px (datacube) + 4px (4*spectrum_half_width) then expect to prune
         if not expect_particle:
             with pytest.raises(
-                RuntimeError, match="No source particles in target region."
+                RuntimeError,
+                match="No non-zero mHI source particles in target region.",
             ):
                 _BaseMartini(**kwargs)
         else:
@@ -480,7 +498,10 @@ def centre_pixels_slice(m):
             dec=source.dec,
             coordinate_frame=FK5(equinox="J1950"),
         )
-        with pytest.raises(RuntimeError, match="No source particles in target region."):
+        with pytest.raises(
+            RuntimeError,
+            match="No non-zero mHI source particles in target region.",
+        ):
             Martini(
                 source=source,
                 datacube=datacube_fk5_J1950,
@@ -542,7 +563,10 @@ def centre_channels_slice(m):
             specsys="galactocentric",
         )
         assert datacube_galactocentric.wcs.wcs.specsys == "galactocentric"
-        with pytest.raises(RuntimeError, match="No source particles in target region."):
+        with pytest.raises(
+            RuntimeError,
+            match="No non-zero mHI source particles in target region.",
+        ):
             Martini(
                 source=source,
                 datacube=datacube_galactocentric,
@@ -767,7 +791,8 @@ def test_prune_particles(
         # if more than 1px (datacube) + 4px (4*spectrum_half_width) then expect to prune
         if not expect_particle:
             with pytest.raises(
-                RuntimeError, match="No source particles in target region."
+                RuntimeError,
+                match="No non-zero mHI source particles in target region.",
             ):
                 GlobalProfile(**kwargs)
         else:

tests/test_sources.py

@@ -339,7 +339,9 @@ def test_apply_badmask(self, s):
             ValueError, match="Mask must have same length as particle arrays."
         ):
             s.apply_mask(np.array([1, 2, 3]))
-        with pytest.raises(RuntimeError, match="No source particles in target region."):
+        with pytest.raises(
+            RuntimeError, match="No non-zero mHI source particles in target region."
+        ):
             s.apply_mask(np.zeros(s.npart, dtype=int))
 
     def test_rotate_axis_angle(self, s):