From 155813e3918f0e5579f3b8b977266e51141b53b3 Mon Sep 17 00:00:00 2001 From: robjmcgibbon Date: Mon, 14 Sep 2026 11:07:39 +0100 Subject: [PATCH 1/2] Update chunk size and add named columns --- SOAP/compression/compress_soap_catalogue.py | 41 +++++++++---- SOAP/core/combine_args.py | 1 + SOAP/core/combine_chunks.py | 67 +++++++++++++++++++++ SOAP/core/soap_args.py | 7 +++ SOAP/property_table.py | 4 ++ scripts/COLIBRE/halo_properties_hybrid.sh | 2 +- scripts/COLIBRE/halo_properties_thermal.sh | 2 +- 7 files changed, 111 insertions(+), 13 deletions(-) diff --git a/SOAP/compression/compress_soap_catalogue.py b/SOAP/compression/compress_soap_catalogue.py index 6114f419..14c2b641 100644 --- a/SOAP/compression/compress_soap_catalogue.py +++ b/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,28 @@ 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": diff --git a/SOAP/core/combine_args.py b/SOAP/core/combine_args.py index 2f71e2ee..a2d85969 100644 --- a/SOAP/core/combine_args.py +++ b/SOAP/core/combine_args.py @@ -15,6 +15,7 @@ "chunks", "dmo", "centrals_only", + "skip_named_columns", "record_halo_timings", "record_property_timings", "max_halos", diff --git a/SOAP/core/combine_chunks.py b/SOAP/core/combine_chunks.py index 90d5478f..a65b6850 100644 --- a/SOAP/core/combine_chunks.py +++ b/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: diff --git a/SOAP/core/soap_args.py b/SOAP/core/soap_args.py index b15fe788..4f1015fa 100644 --- a/SOAP/core/soap_args.py +++ b/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"] diff --git a/SOAP/property_table.py b/SOAP/property_table.py index 3fab6be7..3c5d4bb8 100644 --- a/SOAP/property_table.py +++ b/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", diff --git a/scripts/COLIBRE/halo_properties_hybrid.sh b/scripts/COLIBRE/halo_properties_hybrid.sh index 486dbf0e..c3bc09a3 100644 --- a/scripts/COLIBRE/halo_properties_hybrid.sh +++ b/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!" diff --git a/scripts/COLIBRE/halo_properties_thermal.sh b/scripts/COLIBRE/halo_properties_thermal.sh index 28120f98..546b6f07 100644 --- a/scripts/COLIBRE/halo_properties_thermal.sh +++ b/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!" From 15cfcd31673dabc92903968364e022585c463c8a Mon Sep 17 00:00:00 2001 From: robjmcgibbon Date: Mon, 14 Sep 2026 13:01:43 +0100 Subject: [PATCH 2/2] Format --- SOAP/compression/compress_soap_catalogue.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/SOAP/compression/compress_soap_catalogue.py b/SOAP/compression/compress_soap_catalogue.py index 14c2b641..fd5233f9 100644 --- a/SOAP/compression/compress_soap_catalogue.py +++ b/SOAP/compression/compress_soap_catalogue.py @@ -144,9 +144,7 @@ def compress_dataset(input_name, output_name, dset): if len(data.shape) == 1: compression_opts["chunks"] = min(chunksize, data.shape[0]) else: - compression_opts["chunks"] = get_chunk_shape( - data.shape, named_columns - ) + 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, named_columns)