Remove redundant particle sorting

.github/workflows/main.yml

@@ -9,6 +9,7 @@ on:
     branches: [ master ]
   pull_request:
     branches: [ master ]
+    types: [ opened, synchronize, reopened, ready_for_review ]
 
   # Allows you to run this workflow manually from the Actions tab
   workflow_dispatch:

SOAP/compute_halo_properties.py

@@ -202,21 +202,30 @@ def compute_halo_properties():
     )
 
     filters = parameter_file.get_filters()
+    for filter_name, filter_info in filters.items():
+        for prop in filter_info.get("properties", []):
+            assert prop.split("/")[0] == "BoundSubhalo", (
+                f'Filter "{filter_name}" uses "{prop}", but filters can only '
+                "use BoundSubhalo properties."
+            )
     category_filter = CategoryFilter(filters, dmo=args.dmo)
 
     # Get the full list of property calculations we can do
-    # Note that the order matters: we need to do the BoundSubhalo first,
-    # since quantities are filtered based on the particle numbers in there
-    # Similarly, things like SO 5xR500_crit can only be done after
-    # SO 500_crit for obvious reasons
-    halo_prop_list = []
+    # Each kind of calculation is collected separately so that the final list
+    # can be built in a deliberate order (see where it is assembled below),
+    # rather than relying on the order things happen to be created in.
+    subhalo_props = []
+    so_props = []
+    exclusive_apertures = []
+    inclusive_apertures = []
+    projected_apertures = []
 
     # We require BoundSubhalo since it's used for filters
     if comm_world_rank == 0:
         if "SubhaloProperties" not in parameter_file.parameters:
             print("SubhaloProperties must be in the parameter file", flush=True)
             comm_world.Abort(1)
-    halo_prop_list.append(
+    subhalo_props.append(
         subhalo_properties.SubhaloProperties(
             cellgrid,
             parameter_file,
@@ -236,7 +245,7 @@ def compute_halo_properties():
         ):
             continue
         if "core_excision_fraction" in SO_variations[variation]:
-            halo_prop_list.append(
+            so_props.append(
                 SO_properties.CoreExcisedSOProperties(
                     cellgrid,
                     parameter_file,
@@ -251,7 +260,7 @@ def compute_halo_properties():
                 )
             )
         else:
-            halo_prop_list.append(
+            so_props.append(
                 SO_properties.SOProperties(
                     cellgrid,
                     parameter_file,
@@ -268,7 +277,7 @@ def compute_halo_properties():
             "radius_multiple" in SO_variations[variation]
             and SO_variations[variation]["radius_multiple"] > 0.0
         ):
-            halo_prop_list.append(
+            so_props.append(
                 SO_properties.RadiusMultipleSOProperties(
                     cellgrid,
                     parameter_file,
@@ -302,7 +311,11 @@ def compute_halo_properties():
     assert inclusive_radii_kpc == sorted(inclusive_radii_kpc)
     assert exclusive_radii_kpc == sorted(exclusive_radii_kpc)
 
-    # Add the apertures defined with fixed physical radii
+    # Add the apertures defined with fixed physical radii, followed by those
+    # whose radius is defined by a SOAP property. Exclusive and inclusive
+    # apertures go into separate lists; within each, aperture_variations is
+    # sorted by radius so they stay in ascending order, which is what the
+    # skip_gt_enclose_radius logic requires.
     for variation in aperture_variations:
         if "radius_in_kpc" not in aperture_variations[variation]:
             continue
@@ -316,7 +329,7 @@ def compute_halo_properties():
             if aperture_variations[variation].get("skip_gt_enclose_radius", False):
                 radii_kpc = inclusive_radii_kpc
 
-            halo_prop_list.append(
+            inclusive_apertures.append(
                 aperture_properties.InclusiveSphereProperties(
                     cellgrid,
                     parameter_file,
@@ -331,7 +344,7 @@ def compute_halo_properties():
                 )
             )
         else:
-            halo_prop_list.append(
+            exclusive_apertures.append(
                 aperture_properties.ExclusiveSphereProperties(
                     cellgrid,
                     parameter_file,
@@ -346,7 +359,7 @@ def compute_halo_properties():
                 )
             )
 
-    # Add the apertures based on SOAP properties
+    # Apertures based on SOAP properties
     for variation in aperture_variations:
         if "radius_in_kpc" in aperture_variations[variation]:
             continue
@@ -356,7 +369,7 @@ def compute_halo_properties():
         # struggle to handle the group names
         assert int(radius_multiple) == radius_multiple
         if aperture_variations[variation]["inclusive"]:
-            halo_prop_list.append(
+            inclusive_apertures.append(
                 aperture_properties.InclusiveSphereProperties(
                     cellgrid,
                     parameter_file,
@@ -371,7 +384,7 @@ def compute_halo_properties():
                 )
             )
         else:
-            halo_prop_list.append(
+            exclusive_apertures.append(
                 aperture_properties.ExclusiveSphereProperties(
                     cellgrid,
                     parameter_file,
@@ -413,7 +426,7 @@ def compute_halo_properties():
             continue
         assert "property" not in projected_aperture_variations[variation]
         assert "radius_multiple" not in projected_aperture_variations[variation]
-        halo_prop_list.append(
+        projected_apertures.append(
             projected_aperture_properties.ProjectedApertureProperties(
                 cellgrid,
                 parameter_file,
@@ -429,11 +442,15 @@ def compute_halo_properties():
         if "radius_in_kpc" in projected_aperture_variations[variation]:
             continue
         assert "property" in projected_aperture_variations[variation]
+        assert (
+            projected_aperture_variations[variation]["property"].split("/")[0]
+            == "BoundSubhalo"
+        ), "Projected apertures can only be defined by a BoundSubhalo property"
         radius_multiple = projected_aperture_variations[variation].get(
             "radius_multiple", 1
         )
         assert int(radius_multiple) == radius_multiple
-        halo_prop_list.append(
+        projected_apertures.append(
             projected_aperture_properties.ProjectedApertureProperties(
                 cellgrid,
                 parameter_file,
@@ -454,6 +471,35 @@ def compute_halo_properties():
     if comm_world_rank == 0 and args.output_parameters:
         parameter_file.write_parameters(args.output_parameters)
 
+    # Assemble the calculations in the order they will be run for each halo.
+    # This order matters, for four separate reasons:
+    #
+    #  - BoundSubhalo must come first: its results are used by the category
+    #    filters and by the enclose radius check of every aperture.
+    #  - Within each group of apertures the radii must ascend, because an
+    #    aperture may copy its results from the previous (smaller) aperture of
+    #    the same type. aperture_variations is sorted by radius, so appending in
+    #    order gives this.
+    #  - Calculations which see the same particles are kept together, so that
+    #    the shared particle arrays can be dropped as soon as the last
+    #    calculation needing them has run. Everything using only the bound
+    #    particles comes first, then everything using every particle in the
+    #    search radius.
+    #  - The SO calculations come last, after the inclusive apertures they share
+    #    their particle arrays with. SO adds quantities to that shared object
+    #    which no aperture uses (the sorted mass profile, the masks flagging
+    #    particles bound to another halo), and for the largest halos those are
+    #    several GB. Running SO last means they only exist while the
+    #    calculations which need them are running.
+    #
+    halo_prop_list = (
+        subhalo_props
+        + exclusive_apertures
+        + projected_apertures
+        + inclusive_apertures
+        + so_props
+    )
+
     if len(halo_prop_list) < 1:
         raise Exception("Must select at least one halo property calculation!")
 

SOAP/core/category_filter.py

@@ -86,7 +86,6 @@ def get_do_calculation(
             precomputed_properties["BoundSubhalo/NumberOfGasParticles"] = 0
             precomputed_properties["BoundSubhalo/NumberOfStarParticles"] = 0
             precomputed_properties["BoundSubhalo/NumberOfBlackHoleParticles"] = 0
-            precomputed_properties["SO/200_crit/NumberOfGasParticles"] = 0
         for name, filter_info in self.filters.items():
             if (len(filter_info["properties"]) == 1) or (
                 filter_info["combine_properties"] == "sum"

SOAP/core/halo_tasks.py

@@ -6,6 +6,7 @@
 import unyt
 
 from SOAP.core import memory_use, shared_array
+from SOAP.core.shared_particle_data import ParticleDataCache
 from SOAP.core.dataset_names import mass_dataset, ptypes_for_so_masses
 from SOAP.particle_selection.halo_properties import SearchRadiusTooSmallError
 from SOAP.property_table import PropertyTable
@@ -30,6 +31,7 @@ def process_single_halo(
     boxsize,
     input_halo,
     target_density,
+    shared_keys=None,
 ):
     """
     This computes properties for one halo and runs on a single
@@ -116,6 +118,14 @@ def process_single_halo(
                 offset = input_halo["cofp"] - 0.5 * boxsize
                 pos[:, :] = ((pos - offset) % boxsize) + offset
 
+            # Cache for quantities derived from these particles which more than
+            # one property calculation needs.
+            shared_particle_data = ParticleDataCache()
+            # The key each calculation will look up, so that an entry can be
+            # dropped as soon as no calculation which is still to run needs it.
+            if shared_keys is None:
+                shared_keys = [hp.shared_key(particle_data) for hp in halo_prop_list]
+
             # Try to compute properties of this halo which haven't been done yet
             for prop_nr, halo_prop in enumerate(halo_prop_list):
                 if halo_prop_done[prop_nr]:
@@ -124,7 +134,11 @@ def process_single_halo(
                 try:
                     t0_halo_prop = time.time()
                     halo_prop.calculate(
-                        input_halo, current_radius, particle_data, halo_result
+                        input_halo,
+                        current_radius,
+                        particle_data,
+                        halo_result,
+                        shared_particle_data,
                     )
                 except SearchRadiusTooSmallError:
                     # Search radius was too small, so will need to try again with a larger radius.
@@ -161,6 +175,14 @@ def process_single_halo(
                             time.time() - t0_halo_prop
                         )
 
+                # Drop any shared particle arrays which none of the calculations
+                # still to do for this halo will ask for
+                shared_particle_data.keep_only(
+                    key
+                    for nr, key in enumerate(shared_keys)
+                    if nr > prop_nr and not halo_prop_done[nr]
+                )
+
             # If we computed all of the properties, we're done with this halo
             if np.all(halo_prop_done):
                 break
@@ -318,6 +340,9 @@ def process_halos(
             if target_density is None or density < target_density:
                 target_density = density
 
+    # The shared particle data key each calculation uses
+    shared_keys = [hp.shared_key(data) for hp in halo_prop_list]
+
     # Allocate shared storage for a single integer and initialize to zero
     if comm.Get_rank() == 0:
         local_shape = (1,)
@@ -381,6 +406,7 @@ def process_halos(
                     boxsize,
                     input_halo,
                     target_density if input_halo["is_central"] == 1 else None,
+                    shared_keys,
                 )
                 if halo_result is not None:
                     # Store results and flag this halo as done

SOAP/core/shared_particle_data.py

@@ -0,0 +1,52 @@
+#!/bin/env python
+
+"""
+shared_particle_data.py
+
+Cache of particle quantities that are shared between the property
+calculations of a single halo.
+"""
+
+from typing import Any, Callable, Hashable, Iterable
+
+
+class ParticleDataCache:
+    """
+    Entries are created on first use, so nothing is computed for a halo unless
+    a property calculation actually asks for it, and discarded once the
+    calculations which need them have all run.
+    """
+
+    def __init__(self):
+        """
+        Constructor. Creates an empty cache.
+        """
+        self.cache = {}
+
+    def get(self, key: Hashable, factory: Callable[[], Any]) -> Any:
+        """
+        Return the cached entry for key, creating it with factory() if this is
+        the first time it has been requested.
+
+        Parameters:
+         - key: Hashable
+           Identifies the quantity being requested. See HaloProperty.shared_key().
+         - factory: Callable
+           Function taking no arguments which computes the entry. It is only
+           called if the key is not already in the cache.
+        """
+        if key not in self.cache:
+            self.cache[key] = factory()
+        return self.cache[key]
+
+    def keep_only(self, keys: Iterable[Hashable]):
+        """
+        Drop every entry whose key is not in keys.
+
+        Parameters:
+         - keys: Iterable
+           Keys to keep. Anything else is discarded.
+        """
+        keys = set(keys)
+        for key in [k for k in self.cache if k not in keys]:
+            del self.cache[key]