Opt-in properties, halo finder docs, and install fixes

README.md

@@ -14,19 +14,32 @@ Please cite SOAP using the
 
 ## Installation
 
-The code is written in python and uses mpi4py for parallelism.
-IO is also intended to run in parallel, and so
-[parallel h5py](https://docs.h5py.org/en/stable/mpi.html) is recommended.
-SOAP and its dependencies can be
-installed directly using the command
-`pip install git+https://github.com/SWIFTSIM/SOAP.git`
-but this may install a serial version of h5py. Therefore the following
-steps are recommended for install
+The code is written in python and uses mpi4py for parallelism. Whichever
+install method you use, you will need MPI available so that mpi4py can be
+installed.
+
+### Quick install (serial HDF5)
+
+SOAP and its dependencies can be installed directly using the command
+```
+pip install git+https://github.com/SWIFTSIM/SOAP.git
+```
+This will usually install a serial version of h5py. SOAP will run correctly,
+but for large simulations the I/O will be slower.
+
+### Recommended install (parallel HDF5)
+
+For large runs, [parallel h5py](https://docs.h5py.org/en/stable/mpi.html) is
+recommended so that I/O is carried out in parallel. This requires an HDF5
+library which was itself built with MPI support. h5py must then be built from
+source against that library before installing SOAP:
 ```
 pip install mpi4py
 export HDF5_MPI="ON"; export CC=mpicc; pip install --no-binary=h5py h5py
 pip install git+https://github.com/SWIFTSIM/SOAP.git
 ```
+If SOAP (and therefore serial h5py) is already installed, then add the flags
+`--no-cache-dir` and `--force-reinstall` when reinstalling h5py.
 
 ### Installation on COSMA
 
@@ -38,7 +51,7 @@ you can install an SOAP virtual environment by running
 
 The command `./tests/run_small_volume.sh` will download a small example
 simulation, run the group membership and halo properties scripts on it.
-This uses the parameter file at `./tests/run_small_volume.yml`, and the
+This uses the parameter file at `./tests/small_volume.yml`, and the
 resulting catalogue is placed in the `output` directory. It also generates the
 pdf documentation to describe the output file (which is written to
 `documentation/SOAP.pdf`). 
@@ -56,7 +69,7 @@ the snapshot number, and a parameter file. For example:
 ```
 snapnum=0077
 sim=L1000N0900/DMO_FIDUCIAL
-mpirun python python SOAP/group_membership.py \
+mpirun python SOAP/group_membership.py \
     --sim-name=${sim} --snap-nr=${snapnum} parameter_files/FLAMINGO.yml
 ```
 
@@ -167,7 +180,8 @@ the halo finder to use, which halo definitions to use, and
 which properties to calculate for each halo definition. A description
 of all possible fields can be found in
 [`parameter_files/README.md`](parameter_files/README.md), alongside a number
-of example parameter files.
+of example parameter files. How to specify each of the supported halo finders is
+described in [`parameter_files/halo_finders.md`](parameter_files/halo_finders.md).
 
 ### Compression
 
@@ -197,12 +211,13 @@ the job name with the slurm sbatch -J flag.
 
 ## Modifying the code
 
-You can install an editable version of SOAP by cloning this repository and running:
+You can install an editable version of SOAP by cloning this repository and
+following the [installation](#installation) steps above, but replacing the
+final command with
 ```
-pip install mpi4py
-export HDF5_MPI="ON"; export CC=mpicc; pip install --no-binary=h5py h5py
-pip install -e .
+pip install -e ".[test]"
 ```
+This also installs the optional dependencies required to run the tests.
 
 The property calculations are defined in the following files in the `SOAP/particle_selection` directory:
 
@@ -217,7 +232,13 @@ Adding new quantities to already defined SOAP apertures is relatively easy. Ther
   * Next you have to add the quantity to the type of aperture you want it to be calculated for (`aperture_properties.py`, `SO_properties.py`, `subhalo_properties.py`, or `projected_aperture_properties.py`). In all these files there is a class named `property_list` which defines the subset of all properties that are calculated for this specific aperture.
   * To calculate your quantity you have to define a `@lazy_property` with the same name in the `XXParticleData` class in the same file. There should be a lot of examples of different quantities that are already calculated. An important thing to note is that fields that are used for multiple calculations should have their own `@lazy_property` to avoid loading things multiple times, so check if the things that you need are already there.
   * Add the property to the parameter file.
-  * At this point everything should now work. To test the newly added quantities you can run a unit test using `pytest -W error -m pytest tests/test_{NAME_OF_FILE}`. This checks whether the code crashes, and whether there are problems with units and overflows. This should make sure that SOAP never crashes while calculating the new properties.
+  * At this point everything should now work. To test the newly added quantities you can run a unit test using `pytest -W error tests/test_{NAME_OF_FILE}.py`. This checks whether the code crashes, and whether there are problems with units and overflows. This should make sure that SOAP never crashes while calculating the new properties.
+
+If your property is expensive to compute, or requires additional dependencies,
+set `opt_in_reason` in its `SOAP/property_table.py` entry to a short (20
+characters or less) reason. It will then only be calculated if it is explicitly
+enabled in the parameter file. Any additional dependencies must be imported
+within the `@lazy_property`, not at the top of the file.
 
 If SOAP does crash while evaluating your new property it will try to
 output the ID of the halo it was processing when it crashed. Then you

SOAP/catalogue_readers/read_hbtplus.py

@@ -14,13 +14,20 @@ def hbt_filename(hbt_basename, file_nr):
     return f"{hbt_basename}.{file_nr}.hdf5"
 
 
-def read_hbtplus_groupnr(basename, read_potential_energies=False, registry=None):
+def read_hbtplus_groupnr(
+    basename, read_potential_energies=False, registry=None, index_by_track_id=False
+):
     """
     Read HBTplus output and return group number for each particle ID
 
     Potential energies will not be returned by default. To return the potential
     energies a unit registry must be passed.
 
+    If index_by_track_id is True then the group number of each particle is the
+    TrackId of its subhalo, rather than the position of the subhalo in the
+    catalogue. This only has an effect for unsorted catalogues, since for
+    sorted catalogues the position is already equal to the TrackId.
+
     """
 
     from mpi4py import MPI
@@ -106,6 +113,7 @@ def read_hbtplus_groupnr(basename, read_potential_energies=False, registry=None)
 
         # Number of particles in each subhalo
         halo_size = halos["Nbound"]
+        halo_track_id = halos["TrackId"]
         del halos
 
         # Apply same combination process to potential energies
@@ -152,6 +160,8 @@ def read_hbtplus_groupnr(basename, read_potential_energies=False, registry=None)
     total_nr_halos = comm.allreduce(nr_local_halos)
     halo_offset = comm.scan(len(halo_size), op=MPI.SUM) - len(halo_size)
     halo_index = np.arange(nr_local_halos, dtype=int) + halo_offset
+    if index_by_track_id and not sorted_file:
+        halo_index = halo_track_id
     grnr_bound = np.repeat(halo_index, halo_size)
 
     # Assign ranking by binding energy to the particles
@@ -182,7 +192,13 @@ def read_hbtplus_groupnr(basename, read_potential_energies=False, registry=None)
 
 
 def read_hbtplus_catalogue(
-    comm, basename, a_unit, registry, boxsize, keep_orphans=False
+    comm,
+    basename,
+    a_unit,
+    registry,
+    boxsize,
+    keep_orphans=False,
+    index_by_track_id=False,
 ):
     """
     Read in the HBTplus halo catalogue, distributed over communicator comm.
@@ -192,6 +208,9 @@ def read_hbtplus_catalogue(
     a_unit   - unyt a factor
     registry - unyt unit registry
     boxsize  - box size as a unyt quantity
+    index_by_track_id - use the TrackId as the index of each halo. This only
+                        has an effect for unsorted catalogues, since for sorted
+                        catalogues the index is already equal to the TrackId
 
     Returns a dict of unyt arrays with the halo properies.
     Arrays which must always be returned:
@@ -304,6 +323,8 @@ def read_hbtplus_catalogue(
     nr_local_halos = len(keep)
     local_offset = comm.scan(nr_local_halos) - nr_local_halos
     index = np.arange(nr_local_halos, dtype=int) + local_offset
+    if index_by_track_id and not sorted_file:
+        index = subhalo["TrackId"]
     index = index[keep]
     index = unyt.unyt_array(
         index, units=unyt.dimensionless, dtype=int, registry=registry

SOAP/compute_halo_properties.py

@@ -523,6 +523,7 @@ def compute_halo_properties():
             print("Storing processing time for each property")
         parameter_file.print_unregistered_properties(halo_prop_list, dmo=args.dmo)
         parameter_file.print_skipped_properties(halo_prop_list, dmo=args.dmo)
+        parameter_file.print_optin_skipped_properties(halo_prop_list, dmo=args.dmo)
         parameter_file.print_invalid_properties(halo_prop_list)
         parameter_file.print_variation_warnings()
         if not parameter_file.renclose_enabled():

SOAP/core/combine_chunks.py

@@ -773,6 +773,7 @@ def combine_chunks(
                     cellgrid.a_unit,
                     cellgrid.snap_unit_registry,
                     cellgrid.boxsize,
+                    index_by_track_id=args.index_by_track_id,
                 )
                 prev_order, _ = spatial_sort(
                     prev_data["cofp"],

SOAP/core/halo_centres.py

@@ -78,7 +78,12 @@ def __init__(
             )
         elif args.halo_format == "HBTplus":
             halo_data = read_hbtplus.read_hbtplus_catalogue(
-                comm, halo_basename, a_unit, registry, boxsize
+                comm,
+                halo_basename,
+                a_unit,
+                registry,
+                boxsize,
+                index_by_track_id=args.index_by_track_id,
             )
         elif args.halo_format == "Subfind":
             halo_data = read_subfind.read_gadget4_catalogue(

SOAP/core/parameter_file.py

@@ -37,21 +37,32 @@ def _property_by_name(name: str):
     return _PROPERTY_BY_NAME.get(name)
 
 
+# Keys in the HaloFinder section (other than type and filename) which each halo
+# finder supports. A warning is printed if a key is set for a halo finder which
+# does not support it. Add a new halo finder here, and document it in
+# parameter_files/halo_finders.md.
+_HALO_FINDER_KEYS = {
+    "HBTplus": {
+        "fof_filename",
+        "fof_radius_filename",
+        "read_potential_energies",
+        "index_by_track_id",
+    },
+    "VR": set(),
+    "Subfind": set(),
+    "SubfindEagle": set(),
+    "Rockstar": set(),
+}
+
 # Known parameter file structure, used by check_schema to flag typos. A value
 # of None means the keys directly under that section are user-named or free-form
 # and are not checked; a set lists the only keys allowed directly under that
 # section. Add a key here when a new option is introduced.
 _ALLOWED_KEYS = {
     "Parameters": None,
-    "Snapshots": {"filename", "fof_filename"},
-    "HaloFinder": {
-        "type",
-        "filename",
-        "fof_filename",
-        "fof_radius_filename",
-        "read_potential_energies",
-    },
-    "GroupMembership": {"filename"},
+    "Snapshots": {"filename"},
+    "HaloFinder": {"type", "filename"}.union(*_HALO_FINDER_KEYS.values()),
+    "GroupMembership": {"filename", "fof_ids_filename"},
     "ExtraInput": None,
     "HaloProperties": {"filename", "chunk_dir"},
     "SubhaloProperties": {"properties"},
@@ -73,6 +84,21 @@ def _property_by_name(name: str):
 }
 
 
+def halo_finder_warnings(halo_finder: Dict) -> List[str]:
+    """
+    Return a warning for each key in the HaloFinder section which is not
+    supported by the chosen halo finder type, and so will be ignored.
+    """
+    supported = _HALO_FINDER_KEYS.get(halo_finder.get("type"), set())
+    optional_keys = _ALLOWED_KEYS["HaloFinder"] - {"type", "filename"}
+    return [
+        f'Warning: "HaloFinder/{key}" is not supported for halo finder '
+        f'"{halo_finder.get("type")}" and will be ignored'
+        for key in halo_finder
+        if key in optional_keys and key not in supported
+    ]
+
+
 class ParameterFile:
     """
     Internal representation of the parameter file.
@@ -136,6 +162,12 @@ def __init__(
         # Only used when calculate_missing_properties is True.
         self.skipped_properties = set()
 
+        # Properties which are not in the parameter file, and which are not
+        # calculated because they are flagged as opt-in in the property table.
+        # Stored as {property name: reason}. Only used when
+        # calculate_missing_properties is True.
+        self.optin_skipped_properties = {}
+
         # Properties which are enabled in the parameter file, but which cannot
         # be calculated because the input files lack the datasets they need.
         self.uncomputable_properties = {}
@@ -228,6 +260,16 @@ def missing_datasets(self, property_name: str) -> List[str]:
                 missing.append(dataset_name)
         return missing
 
+    def _opt_in_reason(self, property_name: str):
+        """
+        Get the reason a property is only calculated when explicitly enabled
+        in the parameter file, or None if it is not an opt-in property.
+        """
+        prop = _property_by_name(property_name)
+        if prop is None:
+            return None
+        return prop.opt_in_reason
+
     def _filter_property_names(self) -> set:
         """
         Get the names of the properties used by the filters defined in the
@@ -304,6 +346,18 @@ def get_property_filters(self, base_halo_type: str, full_list: List[str]) -> Dic
             # Property is not listed in the parameter file for this base_halo_type
             elif not self.calculate_missing_properties():
                 filters[property] = False
+            elif self._opt_in_reason(property) is not None:
+                # Opt-in properties must be explicitly enabled in the parameter file
+                if property in self._filter_property_names():
+                    raise ValueError(
+                        f"{property} is used by a filter, but is an opt-in property "
+                        f'("{self._opt_in_reason(property)}") and so is not calculated '
+                        f"unless it is explicitly enabled. Please enable it in the "
+                        f'"{base_halo_type}" section of the parameter file.'
+                    )
+                filters[property] = False
+                listed[property] = False
+                self.optin_skipped_properties[property] = self._opt_in_reason(property)
             elif missing and property not in self._filter_property_names():
                 # The property was not asked for explicitly and cannot be
                 # computed, so it is skipped. Properties used by a filter are
@@ -409,6 +463,30 @@ def print_skipped_properties(self, halo_prop_list=None, dmo: bool = False) -> No
             for property in sorted(skipped):
                 print(f"  {property}")
 
+    def print_optin_skipped_properties(
+        self, halo_prop_list=None, dmo: bool = False
+    ) -> None:
+        """
+        Print a list of the properties which are not in the parameter file, and
+        which are not calculated because they are flagged as opt-in in the
+        property table, along with the reason for each one.
+        """
+        skipped = dict(self.optin_skipped_properties)
+
+        # In a DMO run, drop properties that would be skipped anyway
+        # because they are not DMO properties
+        if dmo and halo_prop_list is not None:
+            for name in self._non_dmo_property_names(halo_prop_list):
+                skipped.pop(name, None)
+
+        if len(skipped):
+            print(
+                "Not computing the following properties for the reason given, "
+                "they must be explicitly enabled in the parameter file:"
+            )
+            for property in sorted(skipped):
+                print(f"  {property.ljust(40)}{skipped[property]}")
+
     def print_uncomputable_properties(self) -> None:
         """
         Print a list of the properties which are enabled in the parameter file,
@@ -557,12 +635,22 @@ def print_variation_warnings(self) -> None:
 
     def check_schema(self) -> None:
         """
-        Abort if the parameter file has an unrecognised section, or a mistyped
+        Abort if the parameter file has an unrecognised section, a mistyped
         key directly under a section which has a fixed set of keys (see
-        _ALLOWED_KEYS). This catches typos which would otherwise be silently
-        ignored. It does not check value types, or keys nested more deeply.
+        _ALLOWED_KEYS), or an unknown halo finder type. This catches typos which
+        would otherwise be silently ignored. It does not check value types, or
+        keys nested more deeply. Also warns about HaloFinder keys which the
+        chosen halo finder does not support.
         """
         errors = []
+        halo_finder = self.parameters.get("HaloFinder", {})
+        if "type" in halo_finder and halo_finder["type"] not in _HALO_FINDER_KEYS:
+            errors.append(
+                f'unknown halo finder type "{halo_finder["type"]}", the supported '
+                f"types are {', '.join(_HALO_FINDER_KEYS)}"
+            )
+        for warning in halo_finder_warnings(halo_finder):
+            print(warning)
         for section, block in self.parameters.items():
             if section not in _ALLOWED_KEYS:
                 errors.append(f'unknown section "{section}"')

SOAP/core/soap_args.py

@@ -220,6 +220,7 @@ def get_soap_args(comm):
     args.read_potential_energies = all_args["HaloFinder"].get(
         "read_potential_energies", False
     )
+    args.index_by_track_id = all_args["HaloFinder"].get("index_by_track_id", False)
     args.fof_group_filename = all_args["HaloFinder"].get("fof_filename", "")
     args.fof_radius_filename = all_args["HaloFinder"].get("fof_radius_filename", "")
     args.output_file = all_args["HaloProperties"]["filename"]