Update chunk size and add named columns

SOAP/compression/compress_soap_catalogue.py

@@ -19,7 +19,8 @@
     # Load empty dictionary if wrong_compression.yml is empty
     compression_fixes = yaml.safe_load(cfile) or {}
 
-chunksize = 1000
+# 4096 gives a good balance between full and masked reads
+chunksize = 4096
 compression_opts = {"compression": "gzip", "compression_opts": 4}
 
 
@@ -85,14 +86,29 @@ def __call__(self, name, h5obj):
                 print(name)
 
 
-def create_lossy_dataset(file, name, shape, filter):
+def get_chunk_shape(shape, named_columns):
+    """
+    Work out the chunk shape for a dataset of the given (uncompressed) shape.
+
+    For an ordinary multi-column dataset we chunk as (chunksize, Ncol) so the
+    whole row lives in one chunk, since these are read/used as a whole.
+
+    For a named columns datasets we instead chunk as (chunksize * Ncol, 1) so
+    there is one column per chunk. This lets a reader read a single column.
+    """
+    if len(shape) == 1:
+        return (min(shape[0], chunksize),)
+    elif named_columns:
+        return (min(shape[0], chunksize * shape[1]), 1)
+    else:
+        return (min(shape[0], chunksize), shape[1])
+
+
+def create_lossy_dataset(file, name, shape, filter, named_columns=False):
     fprops = filterdict[filter]
     type = h5py.h5t.decode(fprops["type"])
     new_plist = h5py.h5p.create(h5py.h5p.DATASET_CREATE)
-    if len(shape) == 1:
-        chunk = (min(shape[0], chunksize),)
-    else:
-        chunk = (min(shape[0], chunksize), shape[1])
+    chunk = get_chunk_shape(shape, named_columns)
     new_plist.set_chunk(chunk)
     for f in fprops["filters"]:
         new_plist.set_filter(f[0], f[1], tuple(f[2]))
@@ -112,25 +128,26 @@ def compress_dataset(input_name, output_name, dset):
 
     with h5py.File(input_name, "r") as ifile, h5py.File(output_name, "r+") as ofile:
         group_name = dset.split("/")[0]
-        if group_name == "Cells":
+        if group_name in ("Cells", "SubgridScheme"):
             filter = "None"
         else:
             filter = ifile[dset].attrs["Lossy compression filter"]
         dset_name = dset.split("/")[-1]
         if dset_name in compression_fixes:
             filter = compression_fixes[dset_name]
+        named_columns = (
+            "SubgridScheme/NamedColumns" in ifile
+            and dset_name in ifile["SubgridScheme/NamedColumns"]
+        )
         data = ifile[dset][:]
         if filter == "None":
             if len(data.shape) == 1:
                 compression_opts["chunks"] = min(chunksize, data.shape[0])
             else:
-                compression_opts["chunks"] = (
-                    min(chunksize, data.shape[0]),
-                    data.shape[1],
-                )
+                compression_opts["chunks"] = get_chunk_shape(data.shape, named_columns)
             ofile.create_dataset("data", data=data, **compression_opts)
         else:
-            create_lossy_dataset(ofile, "data", data.shape, filter)
+            create_lossy_dataset(ofile, "data", data.shape, filter, named_columns)
             ofile["data"][:] = data
         for attr in ifile[dset].attrs:
             if attr == "Is Compressed":

SOAP/core/combine_args.py

@@ -15,6 +15,7 @@
         "chunks",
         "dmo",
         "centrals_only",
+        "skip_named_columns",
         "record_halo_timings",
         "record_property_timings",
         "max_halos",

SOAP/core/combine_chunks.py

@@ -61,6 +61,68 @@ def spatial_sort(halo_cofp, halo_index, cellgrid, comm):
     return order, cell_counts
 
 
+def write_named_columns(outfile, args, cellgrid, all_metadata):
+    """
+    Write SubgridScheme/NamedColumns metadata for properties that declare
+    a `columns_from_snapshot` source (see property_table.Property)
+    """
+    # Skip if we want the old behaviour
+    if args.skip_named_columns:
+        return
+
+    # Properties actually present in this output, keyed by basename, with
+    # their output shape excluding the halo axis (e.g. (9,) or ()).
+    present_props = {}
+    for metadata in all_metadata:
+        name, size = metadata[0], metadata[1]
+        present_props[name.split("/")[-1]] = size
+
+    # Of those, which declare a snapshot field to source column names from.
+    wanted = {
+        prop.name: prop.columns_from_snapshot
+        for prop in PropertyTable.full_property_list.values()
+        if prop.columns_from_snapshot is not None and prop.name in present_props
+    }
+    if not wanted:
+        return
+
+    snap_named_columns = {}
+    snap_filename = cellgrid.snap_filename.format(file_nr=0)
+    with h5py.File(snap_filename, "r") as snap_file:
+        try:
+            group = snap_file["SubgridScheme/NamedColumns"]
+            snap_named_columns = {
+                key: [x.decode("utf-8") for x in group[key][:]] for key in group.keys()
+            }
+        except KeyError:
+            pass
+
+    named_columns_group = outfile.require_group("SubgridScheme/NamedColumns")
+    for prop_name, snap_field in wanted.items():
+        columns = snap_named_columns.get(snap_field)
+        expected_shape = present_props[prop_name]
+        expected_len = expected_shape[0] if len(expected_shape) else 1
+        if columns is None:
+            print(
+                f"named columns requested for {prop_name}, but "
+                f"SubgridScheme/NamedColumns/{snap_field} was not found in "
+                f"the input snapshot ({snap_filename}); skipping.",
+                flush=True,
+            )
+            continue
+        if len(columns) != expected_len:
+            print(
+                f"named columns for {prop_name} (from snapshot field "
+                f"{snap_field}) have length {len(columns)}, but the property "
+                f"has shape {expected_len}; skipping.",
+                flush=True,
+            )
+            continue
+        named_columns_group.create_dataset(
+            prop_name, data=[c.encode("utf-8") for c in columns]
+        )
+
+
 def combine_chunks(
     args,
     cellgrid,
@@ -348,6 +410,11 @@ def combine_chunks(
                 for attr_name, attr_value in attrs.items():
                     dataset.attrs[attr_name] = attr_value
 
+            # Write named-column metadata for properties that support it
+            write_named_columns(
+                outfile, args, cellgrid, ref_metadata + soap_metadata + fof_metadata
+            )
+
             # Save the names of the groups containing the data
             subhalo_types = set()
             for metadata in ref_metadata + soap_metadata + fof_metadata:

SOAP/core/soap_args.py

@@ -106,6 +106,12 @@ def get_soap_args(comm):
         action="store_true",
         help="Only process central halos, discarding satellites",
     )
+    parser.add_argument(
+        "--skip-named-columns",
+        action="store_true",
+        help="Skip writing SubgridScheme/NamedColumns metadata for properties that "
+        "support it (e.g. StellarLuminosity)",
+    )
     parser.add_argument(
         "--record-halo-timings",
         action="store_true",
@@ -220,6 +226,7 @@ def get_soap_args(comm):
     args.snapshot_nr = all_args["Parameters"]["snap_nr"]
     args.chunks = all_args["Parameters"]["chunks"]
     args.centrals_only = all_args["Parameters"]["centrals_only"]
+    args.skip_named_columns = all_args["Parameters"]["skip_named_columns"]
     args.record_halo_timings = all_args["Parameters"]["record_halo_timings"]
     args.record_property_timings = all_args["Parameters"]["record_property_timings"]
     args.dmo = all_args["Parameters"]["dmo"]

SOAP/property_table.py

@@ -90,6 +90,9 @@ class Property:
     particle_properties: list
     output_physical: bool
     a_scale_exponent: int
+    # Name of the snapshot particle dataset (e.g. "Luminosities")
+    # whose NamedColumns entry should be copied for this property
+    columns_from_snapshot: str = None
 
 
 class PropertyTable:
@@ -2536,6 +2539,7 @@ class PropertyTable:
             particle_properties=["PartType4/Luminosities"],
             output_physical=True,
             a_scale_exponent=None,
+            columns_from_snapshot="Luminosities",
         ),
         "Tgas": Property(
             name="GasTemperature",

scripts/COLIBRE/halo_properties_hybrid.sh

@@ -39,6 +39,6 @@ dmo_flag=""
 #TODO: Set nodes and chunks
 mpirun -- python3 -u -m mpi4py SOAP/compute_halo_properties.py \
        parameter_files/COLIBRE_HYBRID.yml \
-       --sim-name=${sim} --snap-nr=${snapnum} --chunks=1 ${dmo_flag}
+       --sim-name=${sim} --snap-nr=${snapnum} --chunks=1 --skip-named-columns ${dmo_flag}
 
 echo "Job complete!"

scripts/COLIBRE/halo_properties_thermal.sh

@@ -40,6 +40,6 @@ dmo_flag=""
 #TODO: Set nodes and chunks
 mpirun -- python3 -u -m mpi4py SOAP/compute_halo_properties.py \
        parameter_files/COLIBRE_THERMAL.yml \
-       --sim-name=${sim} --snap-nr=${snapnum} --chunks=1 ${dmo_flag}
+       --sim-name=${sim} --snap-nr=${snapnum} --chunks=1 --skip-named-columns ${dmo_flag}
 
 echo "Job complete!"