Pass halo indices from a file

README.md

@@ -102,14 +102,66 @@ The optional `--max-ranks-reading` flag determines how many MPI ranks per node
 read the snapshot. This can be used to avoid overloading the file system. The
 default value is 32.
 
+### Selecting which subhalos to process
+
+By default SOAP calculates properties for every subhalo in the input catalogue.
+There are a number of flags which can be used to process only a subset of them.
+
+The `--centrals-only` flag discards satellites, so that only central subhalos are
+processed.
+
+Individual subhalos can be selected with the `--halo-indices` flag.
+This specifies the index of the required subhalos in the halo catalogue, which is
+the quantity written to `InputHalos/HaloCatalogueIndex`.
+
+For larger numbers of subhalos the indices can be listed in a text file, which is
+passed with the `--halo-indices-file` flag.
+The file must contain one index per line. Blank lines, and anything following a
+`#`, are ignored. Duplicate indices are discarded.
+
+### Command line arguments
+
+The arguments listed here are passed on the command line, and cannot be set in
+the parameter file.
+
+`SOAP/group_membership.py`:
+
+| Argument | Description |
+| --- | --- |
+| `config_file` | Name of the yaml parameter file. Required |
+| `--sim-name` | Name of the simulation to process |
+| `--snap-nr` | Snapshot number to process |
+
+`SOAP/compute_halo_properties.py`:
+
+| Argument | Default | Description |
+| --- | --- | --- |
+| `config_file` | | Name of the yaml parameter file. Required |
+| `--sim-name` | | Name of the simulation to process |
+| `--snap-nr` | | Snapshot number to process |
+| `--chunks` | 1 | Number of chunks to split the volume into. Should be at least the number of compute nodes |
+| `--dmo` | off | Run in dark matter only mode, skipping any hydro-only properties |
+| `--centrals-only` | off | Only process central halos, discarding satellites. See [Selecting which subhalos to process](#selecting-which-subhalos-to-process) |
+| `--halo-indices` | | Only process the listed halo indices. See [Selecting which subhalos to process](#selecting-which-subhalos-to-process) |
+| `--halo-indices-file` | | Only process the halo indices listed in the given file. See [Selecting which subhalos to process](#selecting-which-subhalos-to-process) |
+| `--max-halos` | 0 (all) | Only process the first N halos in the catalogue. See [Debugging](#debugging) |
+| `--record-halo-timings` | off | Record the time taken to process each halo. See [Timing](#timing) |
+| `--record-property-timings` | off | Record the time taken to calculate each property. This doubles the size of the output catalogue. See [Timing](#timing) |
+| `--reference-snapshot` | | Number of a snapshot which contains all particle types. Used to determine the datasets and units of any particle types which are missing from the snapshot being processed, e.g. stars or black holes at high redshift |
+| `--snipshot` / `--snapshot` | auto | Force snipshot or snapshot mode. By default this is determined from the value of `SelectOutput` in the snapshot header |
+| `--profile` | 0 | Run with profiling. See [Profiling](#profiling) |
+| `--max-ranks-reading` | 32 | Number of MPI ranks per node which read snapshot data. Can be reduced to avoid overloading the file system |
+| `--output-parameters` | | Where to write the parameters used by this run, in yaml format |
+
 ### Parameter files
 
 To run either of the programs a parameters file must be passed. This
 contains information including the input and output directories,
 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, and a number of example parameter files
-can be found in the `parameters_files` directory.
+of all possible fields can be found in
+[`parameter_files/README.md`](parameter_files/README.md), alongside a number
+of example parameter files.
 
 ### Compression
 
@@ -194,8 +246,9 @@ mpirun.
 The `-Werror` flag is useful for making pdb stop on warnings. E.g. division by
 zero in the halo property calculations will be caught.
 
-It is also possible to select individual halos to process with the `--halo-indices`
-flag. This specifies the index of the required halos in the halo catalogue. E.g.
+If SOAP crashes while processing a particular halo it will try to report the
+index of that halo, which can then be re-run on its own with the
+`--halo-indices` flag, e.g.
 ```
 python3 -Werror -m pdb ./compute_halo_properties.py --halo-indices 1 2 3 ...
 ```

SOAP/catalogue_readers/read_hbtplus.py

@@ -41,7 +41,7 @@ def read_hbtplus_groupnr(basename, read_potential_energies=False, registry=None)
             nr_files = 1
             sorted_file = True
         else:
-            print(f"No HBT files found for basename {basename}")
+            print(f"No HBT files found for basename {basename}", flush=True)
             comm.Abort()
     else:
         nr_files = None

SOAP/compression/compress_soap_catalogue.py

@@ -206,7 +206,7 @@ def assign_datasets(nr_files, nr_ranks, comm_rank):
             datasets = h5copy.dsets.copy()
 
         except Exception as e:
-            print(f"Error: {e}")
+            print(f"Error: {e}", flush=True)
             comm.Abort(1)
     else:
         tmp_dir = None

SOAP/compression/create_empty_SOAP_catalogue.py

@@ -115,7 +115,6 @@ def __call__(self, name, h5obj):
             elif name == "Parameters":
                 for attr in self.ifile[name].attrs:
                     self.ofile[name].attrs[attr] = self.ifile[name].attrs[attr]
-                self.ofile[name].attrs["halo_indices"] = np.array([], dtype="int64")
                 self.ofile[name].attrs["snapshot_nr"] = self.snapnum
             else:
                 for attr in self.ifile[name].attrs:

SOAP/compute_halo_properties.py

@@ -126,7 +126,7 @@ def compute_halo_properties():
                 swift_filename, extra_input, swift_filename_ref, extra_input_ref
             )
         except Exception as err_msg:
-            print(err_msg)
+            print(err_msg, flush=True)
             # Thrown if there are issues with the input files
             comm_world.Abort(1)
         parsec_cgs = cellgrid.constants["parsec"]
@@ -234,7 +234,7 @@ def compute_halo_properties():
     # 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")
+            print("SubhaloProperties must be in the parameter file", flush=True)
             comm_world.Abort(1)
     halo_prop_list.append(
         subhalo_properties.SubhaloProperties(
@@ -544,7 +544,7 @@ def compute_halo_properties():
         try:
             os.makedirs(os.path.dirname(args.output_file), exist_ok=True)
         except OSError as e:
-            print(f"Error creating output directory: {e}")
+            print(f"Error creating output directory: {e}", flush=True)
             comm_world.Abort(1)
     comm_world.barrier()
 
@@ -586,7 +586,7 @@ def compute_halo_properties():
             try:
                 os.makedirs(scratch_file_dir, exist_ok=True)
             except OSError as e:
-                print(f"Error creating scratch directory: {e}")
+                print(f"Error creating scratch directory: {e}", flush=True)
                 comm_world.Abort(1)
     comm_world.barrier()
 

SOAP/core/chunk_tasks.py

@@ -244,7 +244,7 @@ def message(m):
                 try:
                     cellgrid.check_datasets_exist(properties, self.halo_prop_list)
                 except KeyError as err_msg:
-                    print(err_msg)
+                    print(err_msg, flush=True)
                     comm.Abort(1)
             else:
                 properties = None

SOAP/core/combine_args.py

@@ -4,6 +4,30 @@
 
 from virgo.util.partial_formatter import PartialFormatter
 
+# Arguments which must be passed on the command line. The Parameters section
+# of the config file is intended for values which are substituted into the
+# other sections, so we don't allow these to be set there. Note that most of
+# them would be silently ignored if they were, since command line arguments
+# which have a default value always take precedence over the config file.
+COMMAND_LINE_ONLY_PARAMETERS = frozenset(
+    (
+        "config_file",
+        "chunks",
+        "dmo",
+        "centrals_only",
+        "record_halo_timings",
+        "record_property_timings",
+        "max_halos",
+        "halo_indices",
+        "halo_indices_file",
+        "profile",
+        "max_ranks_reading",
+        "output_parameters",
+        "snipshot",
+        "snapshot",
+    )
+)
+
 
 def combine_arguments(command_line_args, config_file):
     """
@@ -21,6 +45,16 @@ def combine_arguments(command_line_args, config_file):
     with open(config_file, "r") as infile:
         config_file_args = yaml.safe_load(infile)
 
+    # Check the config file doesn't set arguments which must be passed
+    # on the command line
+    invalid = COMMAND_LINE_ONLY_PARAMETERS.intersection(config_file_args["Parameters"])
+    if invalid:
+        raise ValueError(
+            "The following cannot be set in the Parameters section of the "
+            f"config file, they must be passed on the command line: "
+            f"{', '.join(sorted(invalid))}"
+        )
+
     # Combine the two
     all_args = {"Parameters": {}}
     for name in config_file_args["Parameters"]:

SOAP/core/combine_chunks.py

@@ -224,11 +224,6 @@ def combine_chunks(
             params.attrs["centrals_only"] = 0 if args.centrals_only == False else 1
             calc_names = sorted([hp.name for hp in halo_prop_list])
             params.attrs["calculations"] = calc_names
-            params.attrs["halo_indices"] = (
-                args.halo_indices
-                if args.halo_indices is not None
-                else np.ndarray(0, dtype=int)
-            )
             if recently_heated_gas_filter.initialised:
                 recently_heated_gas_metadata = recently_heated_gas_filter.get_metadata()
                 recently_heated_gas_params = params.create_group(

SOAP/core/halo_centres.py

@@ -106,14 +106,36 @@ def __init__(
         del halo_data
 
         # Only keep halos in the supplied list of halo IDs.
-        if (args.halo_indices is not None) and (local_halo["index"].shape[0]):
+        if args.halo_indices is not None:
             halo_indices = np.asarray(args.halo_indices, dtype=np.int64)
-            keep = np.zeros_like(local_halo["index"], dtype=bool)
-            matching_index = virgo.util.match.match(halo_indices, local_halo["index"])
-            have_match = matching_index >= 0
-            keep[matching_index[have_match]] = True
-            for name in local_halo:
-                local_halo[name] = local_halo[name][keep, ...]
+            nr_requested = halo_indices.shape[0]
+            have_match = np.zeros(nr_requested, dtype=np.int8)
+            if local_halo["index"].shape[0]:
+                keep = np.zeros_like(local_halo["index"], dtype=bool)
+                matching_index = virgo.util.match.match(
+                    halo_indices, local_halo["index"]
+                )
+                have_match[:] = matching_index >= 0
+                keep[matching_index[have_match.astype(bool)]] = True
+                for name in local_halo:
+                    local_halo[name] = local_halo[name][keep, ...]
+
+            # Report any requested halos which are not in the catalogue. This
+            # is collective, so it must be done on all ranks.
+            comm.Allreduce(MPI.IN_PLACE, have_match, op=MPI.MAX)
+            have_match = have_match.astype(bool)
+            nr_matched = np.sum(have_match)
+            if comm_rank == 0:
+                print(f"Matched {nr_matched} of {nr_requested} requested halo indices")
+                if nr_matched < nr_requested:
+                    # Written to the directory SOAP is being run from, and
+                    # named after the output catalogue
+                    output_file = sub_snapnum(args.output_file, args.snapshot_nr)
+                    basename = os.path.basename(output_file)
+                    basename = os.path.splitext(basename)[0]
+                    filename = f"{basename}_unmatched_halo_indices.txt"
+                    np.savetxt(filename, halo_indices[~have_match], fmt="%d")
+                    print(f"WARNING: wrote unmatched halo indices to {filename}")
 
         # Discard satellites, if necessary
         if args.centrals_only:
@@ -144,7 +166,7 @@ def __init__(
 
         # Exit if we don't have any halos
         if (total_nr_halos == 0) and (comm_rank == 0):
-            print("No halos found, aborting run")
+            print("No halos found, aborting run", flush=True)
             comm.Abort(1)
 
         # Assign halos to chunk tasks: