From eb69f672f469b94942fce758a6cab5b01b8a0ff6 Mon Sep 17 00:00:00 2001 From: robjmcgibbon Date: Wed, 9 Sep 2026 20:15:45 +0100 Subject: [PATCH 01/11] Shared particle data for computing SO radius --- SOAP/core/halo_tasks.py | 13 +- SOAP/core/shared_particle_data.py | 55 ++++ SOAP/particle_selection/SO_properties.py | 286 +++++++++++++----- .../particle_selection/aperture_properties.py | 5 + .../projected_aperture_properties.py | 5 + SOAP/particle_selection/subhalo_properties.py | 7 +- 6 files changed, 289 insertions(+), 82 deletions(-) create mode 100644 SOAP/core/shared_particle_data.py diff --git a/SOAP/core/halo_tasks.py b/SOAP/core/halo_tasks.py index 4da85e23..8df0cfe6 100644 --- a/SOAP/core/halo_tasks.py +++ b/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 SharedParticleData 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 @@ -116,6 +117,12 @@ 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. It is created here, inside the + # search radius loop, so that it is discarded as soon as the set of + # particles changes. + shared_particle_data = SharedParticleData() + # 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 +131,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. diff --git a/SOAP/core/shared_particle_data.py b/SOAP/core/shared_particle_data.py new file mode 100644 index 00000000..00baaa61 --- /dev/null +++ b/SOAP/core/shared_particle_data.py @@ -0,0 +1,55 @@ +#!/bin/env python + +""" +shared_particle_data.py + +Cache of particle quantities that are shared between the property +calculations of a single halo. + +process_single_halo() in halo_tasks.py hands the same set of particles to +every property calculation it runs for a halo. Several of those calculations +begin by deriving the same quantities from that set (concatenated masses and +radii, sorted radial profiles, ...), which is wasted work when it is repeated +once per calculation. + +A SharedParticleData object lets those calculations look up quantities that +have already been derived from the same particles. It is created inside the +search radius loop of process_single_halo(), so a new (empty) cache is used +whenever the set of particles changes. +""" + +from typing import Any, Callable, Hashable + + +class SharedParticleData: + """ + Cache of quantities derived from the particles of a single halo. + + Entries are created on first use, so nothing is computed for a halo unless + a property calculation actually asks for it. + """ + + 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. Calculations that want to + share an entry have to agree on the key, so it needs to include + everything the entry depends on (e.g. the particle types that were + used to compute it). + - 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] diff --git a/SOAP/particle_selection/SO_properties.py b/SOAP/particle_selection/SO_properties.py index b5478558..bcd86d60 100644 --- a/SOAP/particle_selection/SO_properties.py +++ b/SOAP/particle_selection/SO_properties.py @@ -45,6 +45,7 @@ from SOAP.core.parameter_file import ParameterFile from SOAP.core.snapshot_datasets import SnapshotDatasets from SOAP.core.swift_cells import SWIFTCellGrid +from SOAP.core.shared_particle_data import SharedParticleData def cumulative_mass_intersection(r: float, rho_dim: float, slope_dim: float) -> float: @@ -217,25 +218,25 @@ def find_SO_radius_and_mass( return SO_r, SO_mass, SO_volume -class SOParticleData: +class SOSharedParticleData: """ - Halo calculation class. - - All properties we want to compute are implemented as lazy methods of this - class. - - Note that unlike other halo properties that use apertures, SO calculations - only require a single mask, since they are always inclusive. - That said, we still require a types==PartTypeX mask - (see aperture_properties.py) to access some arrays that have been - precomputed for all particles. - - Note that SOs are the only halo types that can include neutrino particles - (these are never bound to a subhalo). They are however only included in - the spherical overdensity radius calculation and in the calculation of - neutrino specific properties (i.e. neutrino masses), and are not taken into - account for other properties, like the total particle mass, velocity - dispersion... + Particle quantities that are shared by all SO variations of a halo. + + Every SO variation of a halo (200_crit, 200_mean, 500_crit, ...) is handed + the same particles, and differs only in the threshold that determines the + SO radius. The quantities computed here are therefore identical for all of + them: the concatenated particle arrays, the cumulative mass profile that + the SO radius is read off, and the masks flagging particles that are bound + to another halo. Computing them once rather than once per variation avoids + repeatedly sorting every particle within the search radius. + + The quantities that do depend on the threshold (the SO radius itself, and + the particle selections derived from it) are computed by SOParticleData, + which holds a reference to one of these objects. + + Note that neutrinos are kept separate from the other particle types, since + they only contribute to the spherical overdensity radius calculation and to + neutrino specific properties. """ def __init__( @@ -243,13 +244,8 @@ def __init__( input_halo: Dict, data: Dict, types_present: List[str], - recently_heated_gas_filter: RecentlyHeatedGasFilter, - observer_position: unyt.unyt_array, snapshot_datasets: SnapshotDatasets, - core_excision_fraction: float, softening_of_parttype: unyt.unyt_array, - virial_definition: bool, - search_radius: unyt.unyt_quantity, cosmology: dict, boxsize: unyt.unyt_quantity, ): @@ -264,26 +260,11 @@ def __init__( - types_present: List List of all particle types (e.g. 'PartType0') that are present in the data dictionary. - - recently_heated_gas_filter: RecentlyHeatedGasFilter - Filter used to mask out gas particles that were recently heated by - AGN feedback. - - observer_position: unyt.unyt_array - Position of an observer, used to determine the observer direction for - Doppler B calculations. - snapshot_datasets: SnapshotDatasets Object containing metadata about the datasets in the snapshot, like appropriate aliases and column names. - - core_excision_fraction: float - Ignore particles within a sphere of core_excision_fraction * SORadius - when calculating CoreExcision properties - softening_of_parttype: unyt.unyt_array Softening length of each particle types - - virial_definition: bool - Whether to calculate the properties that are only valid for virial SO - definitions - - search_radius: unyt.unyt_quantity - Current search radius. Particles are guaranteed to be included up to - this radius. - cosmology: dict Cosmological parameters required for SO calculation - boxsize: unyt.unyt_quantity @@ -293,16 +274,12 @@ def __init__( self.data = data self.has_neutrinos = "PartType6" in data self.types_present = types_present - self.recently_heated_gas_filter = recently_heated_gas_filter - self.observer_position = observer_position self.snapshot_datasets = snapshot_datasets - self.core_excision_fraction = core_excision_fraction self.softening_of_parttype = softening_of_parttype - self.virial_definition = virial_definition - self.search_radius = search_radius self.cosmology = cosmology self.boxsize = boxsize self.compute_basics() + self.compute_mass_profile() def get_dataset(self, name: str) -> unyt.unyt_array: """ @@ -353,28 +330,18 @@ def compute_basics(self): self.fofid = np.concatenate(fofid) self.softening = np.concatenate(softening) - def compute_SO_radius_and_mass( - self, reference_density: unyt.unyt_quantity, physical_radius: unyt.unyt_quantity - ) -> bool: + def compute_mass_profile(self): """ - Compute the SO radius from the density profile of the particles. + Compute the cumulative mass profile used to determine the SO radius. Adds the contribution from neutrinos (if present) to the masses and - radii. Sorts the particles by radius and computes the cumulative mass - profile. Calls find_SO_radius_and_mass(), unless a radius multiple is - used as aperture radius. + radii, sorts the particles by radius, and computes the cumulative mass + profile and the mean density within the radius of each particle. Also + determines the FOF ID of this object from its central particle, and + uses that to flag the particles which are bound to another halo. - Parameters: - - reference_density: unyt.unyt_quantity - Threshold density value that determines the SO radius. - - physical_radius: unyt.unyt_quantity - Physical radius that determines the SO radius in case a radius - multiple is used (e.g. 5xR500_crit). - - Returns True if an SO radius was found, i.e. when both SO_radius and - SO_mass are non-zero. - - Rethrows any SearchRadiusTooSmallError thrown by find_SO_radius_and_mass(). + None of this depends on the density threshold of an individual SO + variation, so it is computed once and used by all of them. """ # add neutrinos if self.has_neutrinos: @@ -414,10 +381,145 @@ def compute_SO_radius_and_mass( # particle *should* be at r=0. We need to manually exclude it, in case round # off error places it at a very small non-zero radius. nskip = max(1, np.argmax(ordered_radius > 0.0 * ordered_radius.units)) - ordered_radius = ordered_radius[nskip:] - cumulative_mass = cumulative_mass[nskip:] - nr_parts = len(ordered_radius) - density = cumulative_mass / (4.0 / 3.0 * np.pi * ordered_radius**3) + self.ordered_radius = ordered_radius[nskip:] + self.cumulative_mass = cumulative_mass[nskip:] + self.nr_parts = len(self.ordered_radius) + self.density = self.cumulative_mass / ( + 4.0 / 3.0 * np.pi * self.ordered_radius**3 + ) + + # figure out which particles in the list are bound to a halo that is not the + # central halo + self.is_bound_to_satellite = ( + (self.groupnr >= 0) & (self.groupnr != self.index) & (self.fofid == fofid) + ) + self.is_bound_to_external = ( + (self.groupnr >= 0) & (self.groupnr != self.index) & (self.fofid != fofid) + ) + + +class SOParticleData: + """ + Halo calculation class. + + All properties we want to compute are implemented as lazy methods of this + class. + + Note that unlike other halo properties that use apertures, SO calculations + only require a single mask, since they are always inclusive. + That said, we still require a types==PartTypeX mask + (see aperture_properties.py) to access some arrays that have been + precomputed for all particles. + + Note that SOs are the only halo types that can include neutrino particles + (these are never bound to a subhalo). They are however only included in + the spherical overdensity radius calculation and in the calculation of + neutrino specific properties (i.e. neutrino masses), and are not taken into + account for other properties, like the total particle mass, velocity + dispersion... + """ + + def __init__( + self, + shared: "SOSharedParticleData", + recently_heated_gas_filter: RecentlyHeatedGasFilter, + observer_position: unyt.unyt_array, + core_excision_fraction: float, + virial_definition: bool, + search_radius: unyt.unyt_quantity, + ): + """ + Constructor. + + Parameters: + - shared: SOSharedParticleData + Object holding the particle quantities that are the same for every + SO variation of this halo. + - recently_heated_gas_filter: RecentlyHeatedGasFilter + Filter used to mask out gas particles that were recently heated by + AGN feedback. + - observer_position: unyt.unyt_array + Position of an observer, used to determine the observer direction for + Doppler B calculations. + - core_excision_fraction: float + Ignore particles within a sphere of core_excision_fraction * SORadius + when calculating CoreExcision properties + - virial_definition: bool + Whether to calculate the properties that are only valid for virial SO + definitions + - search_radius: unyt.unyt_quantity + Current search radius. Particles are guaranteed to be included up to + this radius. + """ + self.shared = shared + + # Quantities that are the same for every SO variation of this halo. + # Note that compute_SO_radius_and_mass() only ever rebinds these arrays + # (it never modifies them in place), so it is safe to share them. + self.input_halo = shared.input_halo + self.data = shared.data + self.has_neutrinos = shared.has_neutrinos + self.types_present = shared.types_present + self.snapshot_datasets = shared.snapshot_datasets + self.softening_of_parttype = shared.softening_of_parttype + self.cosmology = shared.cosmology + self.boxsize = shared.boxsize + self.centre = shared.centre + self.index = shared.index + self.mass = shared.mass + self.radius = shared.radius + self.position = shared.position + self.velocity = shared.velocity + self.types = shared.types + self.groupnr = shared.groupnr + self.fofid = shared.fofid + self.softening = shared.softening + if shared.has_neutrinos: + self.nu_mass = shared.nu_mass + self.nu_radius = shared.nu_radius + self.nu_softening = shared.nu_softening + + # Quantities that differ between the SO variations of this halo + self.recently_heated_gas_filter = recently_heated_gas_filter + self.observer_position = observer_position + self.core_excision_fraction = core_excision_fraction + self.virial_definition = virial_definition + self.search_radius = search_radius + + def get_dataset(self, name: str) -> unyt.unyt_array: + """ + Local wrapper for SnapshotDatasets.get_dataset(). + """ + return self.snapshot_datasets.get_dataset(name, self.data) + + def compute_SO_radius_and_mass( + self, reference_density: unyt.unyt_quantity, physical_radius: unyt.unyt_quantity + ) -> bool: + """ + Compute the SO radius from the density profile of the particles. + + Uses the cumulative mass profile computed once for this halo by + SOSharedParticleData, and calls find_SO_radius_and_mass(), unless a + radius multiple is used as aperture radius. Particles outside the SO + radius are then removed. + + Parameters: + - reference_density: unyt.unyt_quantity + Threshold density value that determines the SO radius. + - physical_radius: unyt.unyt_quantity + Physical radius that determines the SO radius in case a radius + multiple is used (e.g. 5xR500_crit). + + Returns True if an SO radius was found, i.e. when both SO_radius and + SO_mass are non-zero. + + Rethrows any SearchRadiusTooSmallError thrown by find_SO_radius_and_mass(). + """ + # The radial profile is the same for every SO variation of this halo + ordered_radius = self.shared.ordered_radius + cumulative_mass = self.shared.cumulative_mass + density = self.shared.density + nr_parts = self.shared.nr_parts # Check if we ever reach the density threshold if reference_density > 0: @@ -457,13 +559,9 @@ def compute_SO_radius_and_mass( SO_exists = self.SO_r > 0 and self.SO_mass > 0 # figure out which particles in the list are bound to a halo that is not the - # central halo - self.is_bound_to_satellite = ( - (self.groupnr >= 0) & (self.groupnr != self.index) & (self.fofid == fofid) - ) - self.is_bound_to_external = ( - (self.groupnr >= 0) & (self.groupnr != self.index) & (self.fofid != fofid) - ) + # central halo (also the same for every SO variation of this halo) + self.is_bound_to_satellite = self.shared.is_bound_to_satellite + self.is_bound_to_external = self.shared.is_bound_to_external if SO_exists: # Estimate DMO mass fraction found at SO_r @@ -3572,6 +3670,7 @@ def calculate( search_radius: unyt.unyt_quantity, data: Dict, halo_result: Dict, + shared_particle_data: SharedParticleData = None, ): """ Compute spherical masses and overdensities for a halo @@ -3583,6 +3682,10 @@ def calculate( has the particle coordinates for type 1 halo_result - dict with halo properties computed so far. Properties computed here should be added to halo_result. + shared_particle_data - cache of particle quantities shared with the other + property calculations for this halo. If None, the + quantities this calculation needs are computed for + its own use only. Input particle data arrays are unyt_arrays. """ @@ -3627,19 +3730,35 @@ def calculate( if input_halo["is_central"] and do_calculation[self.halo_filter]: types_present = [type for type in self.particle_properties if type in data] + # Quantities which are the same for every SO variation of this halo + # are computed once and reused by the other variations. The particle + # types are part of the cache key because they determine the order + # in which the particle arrays are concatenated. + def make_shared(): + return SOSharedParticleData( + input_halo, + data, + types_present, + self.snapshot_datasets, + self.softening_of_parttype, + self.cosmology, + self.boxsize, + ) + + if shared_particle_data is None: + shared = make_shared() + else: + shared = shared_particle_data.get( + ("SOSharedParticleData", tuple(types_present)), make_shared + ) + part_props = SOParticleData( - input_halo, - data, - types_present, + shared, self.filter, self.observer_position, - self.snapshot_datasets, self.core_excision_fraction, - self.softening_of_parttype, self.virial_definition, search_radius, - self.cosmology, - self.boxsize, ) # we need to make sure the physical radius uses the correct unit @@ -3890,6 +4009,7 @@ def calculate( search_radius: unyt.unyt_quantity, data: Dict, halo_result: Dict, + shared_particle_data: SharedParticleData = None, ): """ Calculate the properties of an SO of which the radius is the multiple of @@ -3907,6 +4027,10 @@ def calculate( Dictionary in which halo properties for this halo are stored. Should contain a valid result for the "parent" SO, i.e. the SO that determines the radius of this SO. + - shared_particle_data: SharedParticleData + Cache of particle quantities shared with the other property + calculations for this halo. If None, the quantities this calculation + needs are computed for its own use only. Throws a RuntimeError if the "parent" SO radius cannot be obtained from halo_result. @@ -3929,5 +4053,7 @@ def calculate( "SO radius multiple estimate was too small!" ) - super().calculate(input_halo, search_radius, data, halo_result) + super().calculate( + input_halo, search_radius, data, halo_result, shared_particle_data + ) return diff --git a/SOAP/particle_selection/aperture_properties.py b/SOAP/particle_selection/aperture_properties.py index 9c4fe74a..3e6441ed 100644 --- a/SOAP/particle_selection/aperture_properties.py +++ b/SOAP/particle_selection/aperture_properties.py @@ -172,6 +172,7 @@ def MetalFracStar(self): from SOAP.core.category_filter import CategoryFilter from SOAP.core.parameter_file import ParameterFile from SOAP.core.snapshot_datasets import SnapshotDatasets +from SOAP.core.shared_particle_data import SharedParticleData class ApertureParticleData: @@ -4027,6 +4028,7 @@ def calculate( search_radius: unyt.unyt_quantity, data: Dict, halo_result: Dict, + shared_particle_data: SharedParticleData = None, ): """ Compute centre of mass etc of bound particles @@ -4039,6 +4041,9 @@ def calculate( has the particle coordinates for type 1 halo_result - dict with halo properties computed so far. Properties computed here should be added to halo_result. + shared_particle_data - cache of particle quantities shared with the other + property calculations for this halo. Not used yet by + this calculation. Input particle data arrays are unyt_arrays. The halo_result dictionary is updated with the properties computed by this function. diff --git a/SOAP/particle_selection/projected_aperture_properties.py b/SOAP/particle_selection/projected_aperture_properties.py index 58e3d324..7a7ff2c3 100644 --- a/SOAP/particle_selection/projected_aperture_properties.py +++ b/SOAP/particle_selection/projected_aperture_properties.py @@ -32,6 +32,7 @@ from SOAP.core.category_filter import CategoryFilter from SOAP.core.parameter_file import ParameterFile from SOAP.core.snapshot_datasets import SnapshotDatasets +from SOAP.core.shared_particle_data import SharedParticleData from SOAP.core.dataset_names import mass_dataset from SOAP.property_calculation.half_mass_radius import ( get_half_mass_radius, @@ -1819,6 +1820,7 @@ def calculate( search_radius: unyt.unyt_quantity, data: Dict, halo_result: Dict, + shared_particle_data: SharedParticleData = None, ): """ Compute centre of mass etc of bound particles @@ -1831,6 +1833,9 @@ def calculate( has the particle coordinates for type 1 halo_result - dict with halo properties computed so far. Properties computed here should be added to halo_result. + shared_particle_data - cache of particle quantities shared with the other + property calculations for this halo. Not used yet by + this calculation. Input particle data arrays are unyt_arrays. The halo_result dictionary is updated with the properties computed by this function. diff --git a/SOAP/particle_selection/subhalo_properties.py b/SOAP/particle_selection/subhalo_properties.py index f5bb7e1f..3834ce62 100644 --- a/SOAP/particle_selection/subhalo_properties.py +++ b/SOAP/particle_selection/subhalo_properties.py @@ -2584,7 +2584,9 @@ def __init__( if not dset in self.particle_properties[pgroup]: self.particle_properties[pgroup].append(dset) - def calculate(self, input_halo, search_radius, data, halo_result): + def calculate( + self, input_halo, search_radius, data, halo_result, shared_particle_data=None + ): """ Compute centre of mass etc of bound particles @@ -2595,6 +2597,9 @@ def calculate(self, input_halo, search_radius, data, halo_result): has the particle coordinates for type 1 halo_result - dict with halo properties computed so far. Properties computed here should be added to halo_result. + shared_particle_data - cache of particle quantities shared with the other + property calculations for this halo. Not used yet by + this calculation. Input particle data arrays are unyt_arrays. """ From c46cd49900d07a6a121be27bd2ba387191043702 Mon Sep 17 00:00:00 2001 From: robjmcgibbon Date: Wed, 9 Sep 2026 21:26:27 +0100 Subject: [PATCH 02/11] Share particles across apertures --- SOAP/compute_halo_properties.py | 168 ++++++++------- .../particle_selection/aperture_properties.py | 162 ++++++--------- .../projected_aperture_properties.py | 123 +++++------ .../shared_halo_particle_data.py | 195 ++++++++++++++++++ SOAP/particle_selection/subhalo_properties.py | 122 +++++------ 5 files changed, 464 insertions(+), 306 deletions(-) create mode 100644 SOAP/particle_selection/shared_halo_particle_data.py diff --git a/SOAP/compute_halo_properties.py b/SOAP/compute_halo_properties.py index 1cf0279f..d5cd67b6 100644 --- a/SOAP/compute_halo_properties.py +++ b/SOAP/compute_halo_properties.py @@ -302,89 +302,101 @@ 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 - for variation in aperture_variations: - if "radius_in_kpc" not in aperture_variations[variation]: - continue - assert "property" not in aperture_variations[variation] - assert "radius_multiple" not in aperture_variations[variation] - if aperture_variations[variation]["inclusive"]: - # If skip_gt_enclose_radius is False (which is the default) then - # we always want to calculate its properties, regardless of the - # size of the next smallest aperture. - radii_kpc = [aperture_variations[variation]["radius_in_kpc"]] - if aperture_variations[variation].get("skip_gt_enclose_radius", False): - radii_kpc = inclusive_radii_kpc - - halo_prop_list.append( - aperture_properties.InclusiveSphereProperties( - cellgrid, - parameter_file, - aperture_variations[variation]["radius_in_kpc"], - None, - recently_heated_gas_filter, - stellar_age_calculator, - cold_dense_gas_filter, - category_filter, - aperture_variations[variation].get("filter", "basic"), - radii_kpc, + # Add the apertures defined with fixed physical radii, followed by those + # whose radius is defined by a SOAP property. Exclusive and inclusive + # apertures are added as two separate groups, rather than interleaved by + # radius, so that the calculations which see the same set of particles run + # consecutively. Radii are still in ascending order within each group, + # which is what the skip_gt_enclose_radius logic requires. + for inclusive in (False, True): + + # Apertures defined with fixed physical radii + for variation in aperture_variations: + if "radius_in_kpc" not in aperture_variations[variation]: + continue + if aperture_variations[variation]["inclusive"] != inclusive: + continue + assert "property" not in aperture_variations[variation] + assert "radius_multiple" not in aperture_variations[variation] + if inclusive: + # If skip_gt_enclose_radius is False (which is the default) then + # we always want to calculate its properties, regardless of the + # size of the next smallest aperture. + radii_kpc = [aperture_variations[variation]["radius_in_kpc"]] + if aperture_variations[variation].get("skip_gt_enclose_radius", False): + radii_kpc = inclusive_radii_kpc + + halo_prop_list.append( + aperture_properties.InclusiveSphereProperties( + cellgrid, + parameter_file, + aperture_variations[variation]["radius_in_kpc"], + None, + recently_heated_gas_filter, + stellar_age_calculator, + cold_dense_gas_filter, + category_filter, + aperture_variations[variation].get("filter", "basic"), + radii_kpc, + ) ) - ) - else: - halo_prop_list.append( - aperture_properties.ExclusiveSphereProperties( - cellgrid, - parameter_file, - aperture_variations[variation]["radius_in_kpc"], - None, - recently_heated_gas_filter, - stellar_age_calculator, - cold_dense_gas_filter, - category_filter, - aperture_variations[variation].get("filter", "basic"), - exclusive_radii_kpc, + else: + halo_prop_list.append( + aperture_properties.ExclusiveSphereProperties( + cellgrid, + parameter_file, + aperture_variations[variation]["radius_in_kpc"], + None, + recently_heated_gas_filter, + stellar_age_calculator, + cold_dense_gas_filter, + category_filter, + aperture_variations[variation].get("filter", "basic"), + exclusive_radii_kpc, + ) ) - ) - # Add the apertures based on SOAP properties - for variation in aperture_variations: - if "radius_in_kpc" in aperture_variations[variation]: - continue - assert "property" in aperture_variations[variation] - radius_multiple = aperture_variations[variation].get("radius_multiple", 1) - # Only allow integer radius mutiples, otherwise swiftsimio will - # struggle to handle the group names - assert int(radius_multiple) == radius_multiple - if aperture_variations[variation]["inclusive"]: - halo_prop_list.append( - aperture_properties.InclusiveSphereProperties( - cellgrid, - parameter_file, - None, - (aperture_variations[variation]["property"], radius_multiple), - recently_heated_gas_filter, - stellar_age_calculator, - cold_dense_gas_filter, - category_filter, - aperture_variations[variation].get("filter", "basic"), - [], + # Apertures based on SOAP properties + for variation in aperture_variations: + if "radius_in_kpc" in aperture_variations[variation]: + continue + if aperture_variations[variation]["inclusive"] != inclusive: + continue + assert "property" in aperture_variations[variation] + radius_multiple = aperture_variations[variation].get("radius_multiple", 1) + # Only allow integer radius mutiples, otherwise swiftsimio will + # struggle to handle the group names + assert int(radius_multiple) == radius_multiple + if inclusive: + halo_prop_list.append( + aperture_properties.InclusiveSphereProperties( + cellgrid, + parameter_file, + None, + (aperture_variations[variation]["property"], radius_multiple), + recently_heated_gas_filter, + stellar_age_calculator, + cold_dense_gas_filter, + category_filter, + aperture_variations[variation].get("filter", "basic"), + [], + ) ) - ) - else: - halo_prop_list.append( - aperture_properties.ExclusiveSphereProperties( - cellgrid, - parameter_file, - None, - (aperture_variations[variation]["property"], radius_multiple), - recently_heated_gas_filter, - stellar_age_calculator, - cold_dense_gas_filter, - category_filter, - aperture_variations[variation].get("filter", "basic"), - [], + else: + halo_prop_list.append( + aperture_properties.ExclusiveSphereProperties( + cellgrid, + parameter_file, + None, + (aperture_variations[variation]["property"], radius_multiple), + recently_heated_gas_filter, + stellar_age_calculator, + cold_dense_gas_filter, + category_filter, + aperture_variations[variation].get("filter", "basic"), + [], + ) ) - ) projected_aperture_variations = parameter_file.get_halo_type_variations( "ProjectedApertureProperties" diff --git a/SOAP/particle_selection/aperture_properties.py b/SOAP/particle_selection/aperture_properties.py index 3e6441ed..bc8fed8c 100644 --- a/SOAP/particle_selection/aperture_properties.py +++ b/SOAP/particle_selection/aperture_properties.py @@ -173,6 +173,9 @@ def MetalFracStar(self): from SOAP.core.parameter_file import ParameterFile from SOAP.core.snapshot_datasets import SnapshotDatasets from SOAP.core.shared_particle_data import SharedParticleData +from SOAP.particle_selection.shared_halo_particle_data import ( + SharedHaloParticleData, +) class ApertureParticleData: @@ -201,16 +204,11 @@ class ApertureParticleData: def __init__( self, - input_halo: Dict, - data: Dict, - types_present: List[str], - inclusive: bool, + shared: SharedHaloParticleData, aperture_radius: unyt.unyt_quantity, stellar_age_calculator: StellarAgeCalculator, recently_heated_gas_filter: RecentlyHeatedGasFilter, cold_dense_gas_filter: ColdDenseGasFilter, - snapshot_datasets: SnapshotDatasets, - softening_of_parttype: unyt.unyt_array, boxsize: unyt.unyt_quantity, cosmology: dict, ): @@ -218,16 +216,10 @@ def __init__( Constructor. Parameters: - - input_halo: Dict - Dictionary containing properties of the halo read from the VR catalogue. - - data: Dict - Dictionary containing particle data. - - types_present: List - List of all particle types (e.g. 'PartType0') that are present in the data - dictionary. - - inclusive: bool - Whether or not to include particles not gravitationally bound to the subhalo - in the property calculations. + - shared: SharedHaloParticleData + Object holding the concatenated particle arrays for this halo, shared + with the other aperture calculations that use the same particles + (i.e. that have the same value of "inclusive"). - aperture_radius: unyt.unyt_quantity Aperture radius. - stellar_age_calculator: StellarAgeCalculator @@ -238,26 +230,22 @@ def __init__( AGN feedback. - cold_dense_gas_filter: ColdDenseGasFilter Filter used to mask out gas particles containing cold, dense gas. - - snapshot_datasets: SnapshotDatasets - Object containing metadata about the datasets in the snapshot, like - appropriate aliases and column names. - - softening_of_parttype: unyt.unyt_array - Softening length of each particle types - boxsize: unyt.unyt_quantity Boxsize for correcting periodic boundary conditions - cosmology: dict Cosmological parameters required for SO calculation """ - self.input_halo = input_halo - self.data = data - self.types_present = types_present - self.inclusive = inclusive + self.shared = shared + self.input_halo = shared.input_halo + self.data = shared.data + self.types_present = shared.types_present + self.inclusive = shared.inclusive + self.snapshot_datasets = shared.snapshot_datasets + self.softening_of_parttype = shared.softening_of_parttype self.aperture_radius = aperture_radius self.stellar_age_calculator = stellar_age_calculator self.recently_heated_gas_filter = recently_heated_gas_filter self.cold_dense_gas_filter = cold_dense_gas_filter - self.snapshot_datasets = snapshot_datasets - self.softening_of_parttype = softening_of_parttype self.boxsize = boxsize self.cosmology = cosmology self.compute_basics() @@ -270,52 +258,26 @@ def get_dataset(self, name: str) -> unyt.unyt_array: def compute_basics(self): """ - Compute some properties that are always needed, regardless of which - properties we actually want to compute. - """ - self.centre = self.input_halo["cofp"] - self.index = self.input_halo["index"] - mass = [] - position = [] - radius = [] - velocity = [] - types = [] - softening = [] - for ptype in self.types_present: - grnr = self.get_dataset(f"{ptype}/GroupNr_bound") - if self.inclusive: - in_halo = np.ones(grnr.shape, dtype=bool) - else: - in_halo = grnr == self.index - mass.append(self.get_dataset(f"{ptype}/{mass_dataset(ptype)}")[in_halo]) - pos = ( - self.get_dataset(f"{ptype}/Coordinates")[in_halo, :] - - self.centre[None, :] - ) - position.append(pos) - r = np.sqrt(pos[:, 0] ** 2 + pos[:, 1] ** 2 + pos[:, 2] ** 2) - radius.append(r) - velocity.append(self.get_dataset(f"{ptype}/Velocities")[in_halo, :]) - typearr = int(ptype[-1]) * np.ones(r.shape, dtype=np.int32) - types.append(typearr) - s = np.ones(r.shape, dtype=np.float64) * self.softening_of_parttype[ptype] - softening.append(s) - - self.mass = np.concatenate(mass) - self.position = np.concatenate(position) - self.radius = np.concatenate(radius) - self.velocity = np.concatenate(velocity) - self.types = np.concatenate(types) - self.softening = np.concatenate(softening) - - self.mask = self.radius <= self.aperture_radius - - self.mass = self.mass[self.mask] - self.position = self.position[self.mask] - self.velocity = self.velocity[self.mask] - self.radius = self.radius[self.mask] - self.type = self.types[self.mask] - self.softening = self.softening[self.mask] + Select the particles that are inside the aperture radius. + + The concatenated arrays covering the whole halo are computed once by + SharedHaloParticleData; all that is left to do here is apply the + aperture mask. Note that self.types is deliberately left unmasked, + since it is what the *_mask_ap masks below are indexed with, whereas + self.type refers only to the particles inside the aperture. + """ + self.centre = self.shared.centre + self.index = self.shared.index + self.types = self.shared.types + + self.mask = self.shared.radius <= self.aperture_radius + + self.mass = self.shared.mass[self.mask] + self.position = self.shared.position[self.mask] + self.velocity = self.shared.velocity[self.mask] + self.radius = self.shared.radius[self.mask] + self.type = self.shared.types[self.mask] + self.softening = self.shared.softening[self.mask] @lazy_property def gas_mask_ap(self) -> NDArray[bool]: @@ -548,11 +510,7 @@ def star_mask_all(self) -> NDArray[bool]: """ if self.Nstar == 0: return None - groupnr_bound = self.get_dataset("PartType4/GroupNr_bound") - if self.inclusive: - return np.ones(groupnr_bound.shape, dtype=bool) - else: - return groupnr_bound == self.index + return self.shared.in_halo_mask("PartType4") @lazy_property def Mstar_init(self) -> unyt.unyt_quantity: @@ -788,11 +746,7 @@ def bh_mask_all(self) -> NDArray[bool]: """ if self.Nbh == 0: return None - groupnr_bound = self.get_dataset("PartType5/GroupNr_bound") - if self.inclusive: - return np.ones(groupnr_bound.shape, dtype=bool) - else: - return groupnr_bound == self.index + return self.shared.in_halo_mask("PartType5") @lazy_property def BH_subgrid_masses(self) -> unyt.unyt_array: @@ -1715,11 +1669,7 @@ def gas_mask_all(self) -> NDArray[bool]: """ if self.Ngas == 0: return None - groupnr_bound = self.get_dataset("PartType0/GroupNr_bound") - if self.inclusive: - return np.ones(groupnr_bound.shape, dtype=bool) - else: - return groupnr_bound == self.index + return self.shared.in_halo_mask("PartType0") @lazy_property def gas_SFR(self) -> unyt.unyt_array: @@ -4042,8 +3992,9 @@ def calculate( halo_result - dict with halo properties computed so far. Properties computed here should be added to halo_result. shared_particle_data - cache of particle quantities shared with the other - property calculations for this halo. Not used yet by - this calculation. + property calculations for this halo. If None, the + quantities this calculation needs are computed for + its own use only. Input particle data arrays are unyt_arrays. The halo_result dictionary is updated with the properties computed by this function. @@ -4148,17 +4099,36 @@ def calculate( ) types_present = [type for type in self.particle_properties if type in data] + + # Every aperture with the same value of "inclusive" sees the same + # particles, so the concatenated arrays are computed once and shared + # (with the bound subhalo and the projected apertures too, for the + # exclusive ones). The particle types are part of the cache key + # because they determine the order of the concatenated arrays. + def make_shared(): + return SharedHaloParticleData( + input_halo, + data, + types_present, + self.inclusive, + self.snapshot_datasets, + self.softening_of_parttype, + ) + + if shared_particle_data is None: + shared = make_shared() + else: + shared = shared_particle_data.get( + ("SharedHaloParticleData", self.inclusive, tuple(types_present)), + make_shared, + ) + part_props = ApertureParticleData( - input_halo, - data, - types_present, - self.inclusive, + shared, aperture_radius, self.stellar_ages, self.recently_heated_gas_filter, self.cold_dense_gas_filter, - self.snapshot_datasets, - self.softening_of_parttype, self.boxsize, self.cosmology, ) diff --git a/SOAP/particle_selection/projected_aperture_properties.py b/SOAP/particle_selection/projected_aperture_properties.py index 7a7ff2c3..f59a1c18 100644 --- a/SOAP/particle_selection/projected_aperture_properties.py +++ b/SOAP/particle_selection/projected_aperture_properties.py @@ -33,6 +33,9 @@ from SOAP.core.parameter_file import ParameterFile from SOAP.core.snapshot_datasets import SnapshotDatasets from SOAP.core.shared_particle_data import SharedParticleData +from SOAP.particle_selection.shared_halo_particle_data import ( + SharedHaloParticleData, +) from SOAP.core.dataset_names import mass_dataset from SOAP.property_calculation.half_mass_radius import ( get_half_mass_radius, @@ -56,37 +59,29 @@ class ProjectedApertureParticleData: def __init__( self, - input_halo: Dict, - data: Dict, - types_present: List[str], + shared: SharedHaloParticleData, aperture_radius: unyt.unyt_quantity, - snapshot_datasets: SnapshotDatasets, boxsize: unyt.unyt_quantity, ): """ Constructor. Parameters: - - input_halo: Dict - Dictionary containing properties of the halo read from the VR catalogue. - - data: Dict - Dictionary containing particle data. - - types_present: List - List of all particle types (e.g. 'PartType0') that are present in the data - dictionary. + - shared: SharedHaloParticleData + Object holding the concatenated particle arrays for the bound + particles of this halo, shared with the exclusive aperture + calculations and the bound subhalo. - aperture_radius: unyt.unyt_quantity Aperture radius. - - snapshot_datasets: SnapshotDatasets - Object containing metadata about the datasets in the snapshot, like - appropriate aliases and column names. - boxsize: unyt.unyt_quantity Boxsize for correcting periodic boundary conditions """ - self.input_halo = input_halo - self.data = data - self.types_present = types_present + self.shared = shared + self.input_halo = shared.input_halo + self.data = shared.data + self.types_present = shared.types_present + self.snapshot_datasets = shared.snapshot_datasets self.aperture_radius = aperture_radius - self.snapshot_datasets = snapshot_datasets self.boxsize = boxsize self.compute_basics() @@ -98,45 +93,21 @@ def get_dataset(self, name: str) -> unyt.unyt_array: def compute_basics(self): """ - Compute some properties that are always needed, regardless of which - properties and projection we actually want to compute. - """ - self.centre = self.input_halo["cofp"] - self.index = self.input_halo["index"] - - mass = [] - position = [] - radius_projx = [] - radius_projy = [] - radius_projz = [] - velocity = [] - types = [] - for ptype in self.types_present: - grnr = self.get_dataset(f"{ptype}/GroupNr_bound") - in_halo = grnr == self.index - mass.append(self.get_dataset(f"{ptype}/{mass_dataset(ptype)}")[in_halo]) - pos = ( - self.get_dataset(f"{ptype}/Coordinates")[in_halo, :] - - self.centre[None, :] - ) - position.append(pos) - rprojx = np.sqrt(pos[:, 1] ** 2 + pos[:, 2] ** 2) - radius_projx.append(rprojx) - rprojy = np.sqrt(pos[:, 0] ** 2 + pos[:, 2] ** 2) - radius_projy.append(rprojy) - rprojz = np.sqrt(pos[:, 0] ** 2 + pos[:, 1] ** 2) - radius_projz.append(rprojz) - velocity.append(self.get_dataset(f"{ptype}/Velocities")[in_halo, :]) - typearr = int(ptype[-1]) * np.ones(rprojx.shape, dtype=np.int32) - types.append(typearr) - - self.mass = np.concatenate(mass) - self.position = np.concatenate(position) - self.radius_projx = np.concatenate(radius_projx) - self.radius_projy = np.concatenate(radius_projy) - self.radius_projz = np.concatenate(radius_projz) - self.velocity = np.concatenate(velocity) - self.types = np.concatenate(types) + Take the particle arrays that are always needed from the shared object, + and mask out the particles outside the aperture in each projection. + + The projected radii are computed once per halo by the shared object, so + they are not recomputed for every projected aperture. + """ + self.centre = self.shared.centre + self.index = self.shared.index + self.mass = self.shared.mass + self.position = self.shared.position + self.velocity = self.shared.velocity + self.types = self.shared.types + self.radius_projx = self.shared.radius_projx + self.radius_projy = self.shared.radius_projy + self.radius_projz = self.shared.radius_projz self.mask_projx = self.radius_projx <= self.aperture_radius self.mask_projy = self.radius_projy <= self.aperture_radius @@ -386,7 +357,7 @@ def star_mask_all(self) -> NDArray[bool]: """ if self.Nstar == 0: return None - return self.part_props.get_dataset("PartType4/GroupNr_bound") == self.index + return self.part_props.shared.in_halo_mask("PartType4") @lazy_property def Mstar_init(self) -> unyt.unyt_quantity: @@ -430,7 +401,7 @@ def bh_mask_all(self) -> NDArray[bool]: """ if self.Nbh == 0: return None - return self.part_props.get_dataset("PartType5/GroupNr_bound") == self.index + return self.part_props.shared.in_halo_mask("PartType5") @lazy_property def BH_subgrid_masses(self) -> unyt.unyt_array: @@ -1138,7 +1109,7 @@ def gas_mask_all(self) -> NDArray[bool]: """ if self.Ngas == 0: return None - return self.part_props.get_dataset("PartType0/GroupNr_bound") == self.index + return self.part_props.shared.in_halo_mask("PartType0") @lazy_property def gas_total_dust_mass_fractions(self) -> unyt.unyt_array: @@ -1834,8 +1805,9 @@ def calculate( halo_result - dict with halo properties computed so far. Properties computed here should be added to halo_result. shared_particle_data - cache of particle quantities shared with the other - property calculations for this halo. Not used yet by - this calculation. + property calculations for this halo. If None, the + quantities this calculation needs are computed for + its own use only. Input particle data arrays are unyt_arrays. The halo_result dictionary is updated with the properties computed by this function. @@ -1925,12 +1897,31 @@ def calculate( ) types_present = [type for type in self.particle_properties if type in data] + + # The concatenated arrays for the bound particles of this halo are + # also used by the bound subhalo and the exclusive apertures, so they + # are computed once and shared. The particle types are part of the + # cache key because they determine the order of the arrays. + def make_shared(): + return SharedHaloParticleData( + input_halo, + data, + types_present, + False, + self.snapshot_datasets, + self.softening_of_parttype, + ) + + if shared_particle_data is None: + shared = make_shared() + else: + shared = shared_particle_data.get( + ("SharedHaloParticleData", False, tuple(types_present)), make_shared + ) + part_props = ProjectedApertureParticleData( - input_halo, - data, - types_present, + shared, aperture_radius, - self.snapshot_datasets, self.boxsize, ) for projname in ["projx", "projy", "projz"]: diff --git a/SOAP/particle_selection/shared_halo_particle_data.py b/SOAP/particle_selection/shared_halo_particle_data.py new file mode 100644 index 00000000..c8436388 --- /dev/null +++ b/SOAP/particle_selection/shared_halo_particle_data.py @@ -0,0 +1,195 @@ +#!/bin/env python + +""" +shared_halo_particle_data.py + +Particle arrays shared between the aperture, projected aperture and subhalo +property calculations of a single halo. + +All of those calculations start by concatenating the same per particle +quantities (masses, positions, velocities and particle types) for either the +gravitationally bound particles of the halo (exclusive apertures, projected +apertures and the bound subhalo) or for every particle in the search radius +(inclusive apertures). Doing that once per halo instead of once per +calculation avoids repeating the work for every aperture. + +Quantities derived from those arrays are implemented as lazy properties, so a +calculation which does not need one never pays for it: the projected radii are +only computed if a projected aperture asks for them, and the softening and the +(3D) radius only if an aperture or the bound subhalo does. +""" + +from typing import Dict, List + +import numpy as np +from numpy.typing import NDArray +import unyt + +from SOAP.core.dataset_names import mass_dataset +from SOAP.core.lazy_properties import lazy_property +from SOAP.core.snapshot_datasets import SnapshotDatasets + + +class SharedHaloParticleData: + """ + Concatenated particle arrays for a single halo. + + One of these objects covers every calculation which sees the same set of + particles, which is determined by "inclusive": exclusive apertures, + projected apertures and the bound subhalo all use the gravitationally + bound particles of the halo, while inclusive apertures use every particle + that was read in. + + Note that the order of the concatenated arrays is set by types_present, so + calculations may only share one of these objects if they use the same + particle types in the same order. That is what the cache key in + SharedParticleData accounts for. + """ + + def __init__( + self, + input_halo: Dict, + data: Dict, + types_present: List[str], + inclusive: bool, + snapshot_datasets: SnapshotDatasets, + softening_of_parttype: unyt.unyt_array, + ): + """ + Constructor. + + Parameters: + - input_halo: Dict + Dictionary containing properties of the halo read from the halo + catalogue. + - data: Dict + Dictionary containing particle data. + - types_present: List + List of all particle types (e.g. 'PartType0') that are present in the + data dictionary. + - inclusive: bool + Whether or not to include particles that are not gravitationally + bound to the subhalo. + - snapshot_datasets: SnapshotDatasets + Object containing metadata about the datasets in the snapshot, like + appropriate aliases and column names. + - softening_of_parttype: unyt.unyt_array + Softening length of each particle types + """ + self.input_halo = input_halo + self.data = data + self.types_present = types_present + self.inclusive = inclusive + self.snapshot_datasets = snapshot_datasets + self.softening_of_parttype = softening_of_parttype + self.centre = input_halo["cofp"] + self.index = input_halo["index"] + self.in_halo_masks = {} + self.compute_basics() + + def get_dataset(self, name: str) -> unyt.unyt_array: + """ + Local wrapper for SnapshotDatasets.get_dataset(). + """ + return self.snapshot_datasets.get_dataset(name, self.data) + + def in_halo_mask(self, ptype: str) -> NDArray[bool]: + """ + Mask which selects the particles of ptype that are included in the + calculations: only the particles bound to this halo for exclusive + calculations, all of them for inclusive ones. This mask needs to be + applied _first_ to raw "PartTypeX" datasets. + + The mask is computed once per particle type and then reused, since + every calculation sharing this object needs the same one. + + Parameters: + - ptype: str + Particle type, e.g. 'PartType0'. + """ + if ptype not in self.in_halo_masks: + groupnr_bound = self.get_dataset(f"{ptype}/GroupNr_bound") + if self.inclusive: + mask = np.ones(groupnr_bound.shape, dtype=bool) + else: + mask = groupnr_bound == self.index + self.in_halo_masks[ptype] = mask + return self.in_halo_masks[ptype] + + def compute_basics(self): + """ + Concatenate the quantities which every calculation sharing this object + needs, over all particle types that are present. + + Also records the number of particles of each type, which the lazy + softening below uses to rebuild a per type quantity in the same order. + """ + mass = [] + position = [] + velocity = [] + types = [] + self.nr_part_of_type = [] + for ptype in self.types_present: + in_halo = self.in_halo_mask(ptype) + mass.append(self.get_dataset(f"{ptype}/{mass_dataset(ptype)}")[in_halo]) + pos = ( + self.get_dataset(f"{ptype}/Coordinates")[in_halo, :] + - self.centre[None, :] + ) + position.append(pos) + velocity.append(self.get_dataset(f"{ptype}/Velocities")[in_halo, :]) + nr_part = pos.shape[0] + types.append(int(ptype[-1]) * np.ones(nr_part, dtype=np.int32)) + self.nr_part_of_type.append((ptype, nr_part)) + + self.mass = np.concatenate(mass) + self.position = np.concatenate(position) + self.velocity = np.concatenate(velocity) + self.types = np.concatenate(types) + + @lazy_property + def radius(self) -> unyt.unyt_array: + """ + Distance of each particle from the halo centre. + """ + pos = self.position + return np.sqrt(pos[:, 0] ** 2 + pos[:, 1] ** 2 + pos[:, 2] ** 2) + + @lazy_property + def softening(self) -> unyt.unyt_array: + """ + Softening length of each particle. + """ + softening = [] + for ptype, nr_part in self.nr_part_of_type: + softening.append( + np.ones(nr_part, dtype=np.float64) * self.softening_of_parttype[ptype] + ) + return np.concatenate(softening) + + @lazy_property + def radius_projx(self) -> unyt.unyt_array: + """ + Distance of each particle from the halo centre, projected along the + x axis. + """ + pos = self.position + return np.sqrt(pos[:, 1] ** 2 + pos[:, 2] ** 2) + + @lazy_property + def radius_projy(self) -> unyt.unyt_array: + """ + Distance of each particle from the halo centre, projected along the + y axis. + """ + pos = self.position + return np.sqrt(pos[:, 0] ** 2 + pos[:, 2] ** 2) + + @lazy_property + def radius_projz(self) -> unyt.unyt_array: + """ + Distance of each particle from the halo centre, projected along the + z axis. + """ + pos = self.position + return np.sqrt(pos[:, 0] ** 2 + pos[:, 1] ** 2) diff --git a/SOAP/particle_selection/subhalo_properties.py b/SOAP/particle_selection/subhalo_properties.py index 3834ce62..27ca5db1 100644 --- a/SOAP/particle_selection/subhalo_properties.py +++ b/SOAP/particle_selection/subhalo_properties.py @@ -53,6 +53,9 @@ from SOAP.core.category_filter import CategoryFilter from SOAP.core.parameter_file import ParameterFile from SOAP.core.snapshot_datasets import SnapshotDatasets +from SOAP.particle_selection.shared_halo_particle_data import ( + SharedHaloParticleData, +) from SOAP.core.swift_cells import SWIFTCellGrid @@ -73,13 +76,9 @@ class SubhaloParticleData: def __init__( self, - input_halo: Dict, - data: Dict, - types_present: List[str], + shared: SharedHaloParticleData, stellar_age_calculator: StellarAgeCalculator, recently_heated_gas_filter: RecentlyHeatedGasFilter, - snapshot_datasets: SnapshotDatasets, - softening_of_parttype: unyt.unyt_array, boxsize: unyt.unyt_quantity, cosmology: dict, ): @@ -87,34 +86,28 @@ def __init__( Constructor. Parameters: - - input_halo: Dict - Dictionary containing properties of the halo read from the VR catalogue. - - data: Dict - Dictionary containing particle data. - - types_present: List - List of all particle types (e.g. 'PartType0') that are present in the data - dictionary. + - shared: SharedHaloParticleData + Object holding the concatenated particle arrays for the bound + particles of this halo, shared with the aperture calculations. - stellar_age_calculator: StellarAgeCalculator Object used to compute stellar ages from the current cosmological scale factor and the birth scale factors of star particles. - recently_heated_gas_filter: RecentlyHeatedGasFilter Filter used to mask out gas particles that were recently heated by AGN feedback. - - snapshot_datasets: SnapshotDatasets - Object containing metadata about the datasets in the snapshot, like - appropriate aliases and column names. - boxsize: unyt.unyt_quantity Boxsize for correcting periodic boundary conditions - cosmology: dict Cosmological parameters required for SO calculation """ - self.input_halo = input_halo - self.data = data - self.types_present = types_present + self.shared = shared + self.input_halo = shared.input_halo + self.data = shared.data + self.types_present = shared.types_present + self.snapshot_datasets = shared.snapshot_datasets + self.softening_of_parttype = shared.softening_of_parttype self.stellar_age_calculator = stellar_age_calculator self.recently_heated_gas_filter = recently_heated_gas_filter - self.snapshot_datasets = snapshot_datasets - self.softening_of_parttype = softening_of_parttype self.boxsize = boxsize self.cosmology = cosmology self.compute_basics() @@ -127,41 +120,20 @@ def get_dataset(self, name: str) -> unyt.unyt_array: def compute_basics(self): """ - Compute some properties that are always needed, regardless of which - properties we actually want to compute. - """ - self.centre = self.input_halo["cofp"] - self.index = self.input_halo["index"] - - mass = [] - position = [] - radius = [] - velocity = [] - types = [] - softening = [] - for ptype in self.types_present: - grnr = self.get_dataset(f"{ptype}/GroupNr_bound") - in_halo = grnr == self.index - mass.append(self.get_dataset(f"{ptype}/{mass_dataset(ptype)}")[in_halo]) - pos = ( - self.get_dataset(f"{ptype}/Coordinates")[in_halo, :] - - self.centre[None, :] - ) - position.append(pos) - r = np.sqrt(pos[:, 0] ** 2 + pos[:, 1] ** 2 + pos[:, 2] ** 2) - radius.append(r) - velocity.append(self.get_dataset(f"{ptype}/Velocities")[in_halo, :]) - typearr = int(ptype[-1]) * np.ones(r.shape, dtype=np.int32) - types.append(typearr) - s = np.ones(r.shape, dtype=np.float64) * self.softening_of_parttype[ptype] - softening.append(s) - - self.mass = np.concatenate(mass) - self.position = np.concatenate(position) - self.radius = np.concatenate(radius) - self.velocity = np.concatenate(velocity) - self.types = np.concatenate(types) - self.softening = np.concatenate(softening) + Take the particle arrays that are always needed from the shared object. + + The bound subhalo uses every bound particle, so unlike the aperture + calculations it applies no further mask and can use these arrays as + they are. + """ + self.centre = self.shared.centre + self.index = self.shared.index + self.mass = self.shared.mass + self.position = self.shared.position + self.radius = self.shared.radius + self.velocity = self.shared.velocity + self.types = self.shared.types + self.softening = self.shared.softening @lazy_property def gas_mask_sh(self) -> NDArray[bool]: @@ -387,7 +359,7 @@ def star_mask_all(self) -> NDArray[bool]: """ if self.Nstar == 0: return None - return self.get_dataset(f"PartType4/GroupNr_bound") == self.index + return self.shared.in_halo_mask("PartType4") @lazy_property def mass_star_init(self) -> unyt.unyt_array: @@ -505,7 +477,7 @@ def bh_mask_all(self) -> NDArray[bool]: """ if self.Nbh == 0: return None - return self.get_dataset(f"PartType5/GroupNr_bound") == self.index + return self.shared.in_halo_mask("PartType5") @lazy_property def Mbh_subgrid(self) -> unyt.unyt_quantity: @@ -909,7 +881,7 @@ def dm_mask_all(self) -> NDArray[bool]: Mask that can be used to filter out DM particles that belong to this subhalo in raw particle arrays, like PartType0/Masses. """ - return self.get_dataset(f"PartType1/GroupNr_bound") == self.index + return self.shared.in_halo_mask("PartType1") @lazy_property def potential_energy_dm(self) -> unyt.unyt_array: @@ -1932,7 +1904,7 @@ def gas_mask_all(self) -> NDArray[bool]: Mask that can be used to filter out gas particles that belong to this subhalo in raw particle arrays, like PartType0/Masses. """ - return self.get_dataset(f"PartType0/GroupNr_bound") == self.index + return self.shared.in_halo_mask("PartType0") @lazy_property def gas_SFR(self) -> unyt.unyt_array: @@ -2598,22 +2570,40 @@ def calculate( halo_result - dict with halo properties computed so far. Properties computed here should be added to halo_result. shared_particle_data - cache of particle quantities shared with the other - property calculations for this halo. Not used yet by - this calculation. + property calculations for this halo. If None, the + quantities this calculation needs are computed for + its own use only. Input particle data arrays are unyt_arrays. """ types_present = [type for type in self.particle_properties if type in data] + # The concatenated arrays for the bound particles of this halo are also + # used by the exclusive and projected aperture calculations, so they are + # computed once and shared. The particle types are part of the cache key + # because they determine the order in which the arrays are concatenated. + def make_shared(): + return SharedHaloParticleData( + input_halo, + data, + types_present, + False, + self.snapshot_datasets, + self.softening_of_parttype, + ) + + if shared_particle_data is None: + shared = make_shared() + else: + shared = shared_particle_data.get( + ("SharedHaloParticleData", False, tuple(types_present)), make_shared + ) + part_props = SubhaloParticleData( - input_halo, - data, - types_present, + shared, self.stellar_ages, self.filter, - self.snapshot_datasets, - self.softening_of_parttype, self.boxsize, self.cosmology, ) From d0f86dd21e9df1f1dd43fdd9d5680a3339d8f135 Mon Sep 17 00:00:00 2001 From: robjmcgibbon Date: Thu, 10 Sep 2026 16:53:07 +0100 Subject: [PATCH 03/11] Combine SO shared data with other apertures --- SOAP/compute_halo_properties.py | 223 ++++++++++------- SOAP/core/halo_tasks.py | 66 ++++- SOAP/core/memory_use.py | 19 ++ SOAP/core/shared_particle_data.py | 42 +++- SOAP/particle_selection/SO_properties.py | 231 +++--------------- .../particle_selection/aperture_properties.py | 14 +- SOAP/particle_selection/halo_properties.py | 30 +++ .../projected_aperture_properties.py | 14 +- .../shared_halo_particle_data.py | 114 ++++++++- SOAP/particle_selection/subhalo_properties.py | 10 +- 10 files changed, 438 insertions(+), 325 deletions(-) diff --git a/SOAP/compute_halo_properties.py b/SOAP/compute_halo_properties.py index d5cd67b6..ecab86a6 100644 --- a/SOAP/compute_halo_properties.py +++ b/SOAP/compute_halo_properties.py @@ -209,14 +209,21 @@ def compute_halo_properties(): # 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 +243,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 +258,7 @@ def compute_halo_properties(): ) ) else: - halo_prop_list.append( + so_props.append( SO_properties.SOProperties( cellgrid, parameter_file, @@ -268,7 +275,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, @@ -304,99 +311,96 @@ def compute_halo_properties(): # Add the apertures defined with fixed physical radii, followed by those # whose radius is defined by a SOAP property. Exclusive and inclusive - # apertures are added as two separate groups, rather than interleaved by - # radius, so that the calculations which see the same set of particles run - # consecutively. Radii are still in ascending order within each group, - # which is what the skip_gt_enclose_radius logic requires. - for inclusive in (False, True): - - # Apertures defined with fixed physical radii - for variation in aperture_variations: - if "radius_in_kpc" not in aperture_variations[variation]: - continue - if aperture_variations[variation]["inclusive"] != inclusive: - continue - assert "property" not in aperture_variations[variation] - assert "radius_multiple" not in aperture_variations[variation] - if inclusive: - # If skip_gt_enclose_radius is False (which is the default) then - # we always want to calculate its properties, regardless of the - # size of the next smallest aperture. - radii_kpc = [aperture_variations[variation]["radius_in_kpc"]] - if aperture_variations[variation].get("skip_gt_enclose_radius", False): - radii_kpc = inclusive_radii_kpc - - halo_prop_list.append( - aperture_properties.InclusiveSphereProperties( - cellgrid, - parameter_file, - aperture_variations[variation]["radius_in_kpc"], - None, - recently_heated_gas_filter, - stellar_age_calculator, - cold_dense_gas_filter, - category_filter, - aperture_variations[variation].get("filter", "basic"), - radii_kpc, - ) + # 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 + assert "property" not in aperture_variations[variation] + assert "radius_multiple" not in aperture_variations[variation] + if aperture_variations[variation]["inclusive"]: + # If skip_gt_enclose_radius is False (which is the default) then + # we always want to calculate its properties, regardless of the + # size of the next smallest aperture. + radii_kpc = [aperture_variations[variation]["radius_in_kpc"]] + if aperture_variations[variation].get("skip_gt_enclose_radius", False): + radii_kpc = inclusive_radii_kpc + + inclusive_apertures.append( + aperture_properties.InclusiveSphereProperties( + cellgrid, + parameter_file, + aperture_variations[variation]["radius_in_kpc"], + None, + recently_heated_gas_filter, + stellar_age_calculator, + cold_dense_gas_filter, + category_filter, + aperture_variations[variation].get("filter", "basic"), + radii_kpc, ) - else: - halo_prop_list.append( - aperture_properties.ExclusiveSphereProperties( - cellgrid, - parameter_file, - aperture_variations[variation]["radius_in_kpc"], - None, - recently_heated_gas_filter, - stellar_age_calculator, - cold_dense_gas_filter, - category_filter, - aperture_variations[variation].get("filter", "basic"), - exclusive_radii_kpc, - ) + ) + else: + exclusive_apertures.append( + aperture_properties.ExclusiveSphereProperties( + cellgrid, + parameter_file, + aperture_variations[variation]["radius_in_kpc"], + None, + recently_heated_gas_filter, + stellar_age_calculator, + cold_dense_gas_filter, + category_filter, + aperture_variations[variation].get("filter", "basic"), + exclusive_radii_kpc, ) + ) - # Apertures based on SOAP properties - for variation in aperture_variations: - if "radius_in_kpc" in aperture_variations[variation]: - continue - if aperture_variations[variation]["inclusive"] != inclusive: - continue - assert "property" in aperture_variations[variation] - radius_multiple = aperture_variations[variation].get("radius_multiple", 1) - # Only allow integer radius mutiples, otherwise swiftsimio will - # struggle to handle the group names - assert int(radius_multiple) == radius_multiple - if inclusive: - halo_prop_list.append( - aperture_properties.InclusiveSphereProperties( - cellgrid, - parameter_file, - None, - (aperture_variations[variation]["property"], radius_multiple), - recently_heated_gas_filter, - stellar_age_calculator, - cold_dense_gas_filter, - category_filter, - aperture_variations[variation].get("filter", "basic"), - [], - ) + # Apertures based on SOAP properties + for variation in aperture_variations: + if "radius_in_kpc" in aperture_variations[variation]: + continue + assert "property" in aperture_variations[variation] + # Apertures are computed before the SO calculations, so they cannot + # be defined in terms of an SO property + assert not aperture_variations[variation]["property"].startswith( + "SO/" + ), "Apertures cannot be defined by an SO property" + radius_multiple = aperture_variations[variation].get("radius_multiple", 1) + # Only allow integer radius mutiples, otherwise swiftsimio will + # struggle to handle the group names + assert int(radius_multiple) == radius_multiple + if aperture_variations[variation]["inclusive"]: + inclusive_apertures.append( + aperture_properties.InclusiveSphereProperties( + cellgrid, + parameter_file, + None, + (aperture_variations[variation]["property"], radius_multiple), + recently_heated_gas_filter, + stellar_age_calculator, + cold_dense_gas_filter, + category_filter, + aperture_variations[variation].get("filter", "basic"), + [], ) - else: - halo_prop_list.append( - aperture_properties.ExclusiveSphereProperties( - cellgrid, - parameter_file, - None, - (aperture_variations[variation]["property"], radius_multiple), - recently_heated_gas_filter, - stellar_age_calculator, - cold_dense_gas_filter, - category_filter, - aperture_variations[variation].get("filter", "basic"), - [], - ) + ) + else: + exclusive_apertures.append( + aperture_properties.ExclusiveSphereProperties( + cellgrid, + parameter_file, + None, + (aperture_variations[variation]["property"], radius_multiple), + recently_heated_gas_filter, + stellar_age_calculator, + cold_dense_gas_filter, + category_filter, + aperture_variations[variation].get("filter", "basic"), + [], ) + ) projected_aperture_variations = parameter_file.get_halo_type_variations( "ProjectedApertureProperties" @@ -425,7 +429,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, @@ -441,11 +445,16 @@ def compute_halo_properties(): if "radius_in_kpc" in projected_aperture_variations[variation]: continue assert "property" in projected_aperture_variations[variation] + # Projected apertures are computed before the SO calculations, so they + # cannot be defined in terms of an SO property + assert not projected_aperture_variations[variation]["property"].startswith( + "SO/" + ), "Projected apertures cannot be defined by an SO 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, @@ -466,6 +475,30 @@ 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 three 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. Nothing outside the SO calculations reads an SO result, + # which is what lets them move after the apertures; the assertions above + # keep that true. + halo_prop_list = ( + subhalo_props + + exclusive_apertures + + projected_apertures + + so_props + + inclusive_apertures + ) + if len(halo_prop_list) < 1: raise Exception("Must select at least one halo property calculation!") diff --git a/SOAP/core/halo_tasks.py b/SOAP/core/halo_tasks.py index 8df0cfe6..b33cbea3 100644 --- a/SOAP/core/halo_tasks.py +++ b/SOAP/core/halo_tasks.py @@ -5,8 +5,10 @@ import numpy as np import unyt +from mpi4py import MPI + from SOAP.core import memory_use, shared_array -from SOAP.core.shared_particle_data import SharedParticleData +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 @@ -20,6 +22,10 @@ # Radius in Mpc at which we report halos which have a large search radius REPORT_RADIUS = 20.0 +# TEMPORARY (issue 64): number of shared particle data objects created for each +# halo processed by this rank. Strip this out after testing. +shared_data_created = [] + def process_single_halo( mesh, @@ -31,6 +37,7 @@ def process_single_halo( boxsize, input_halo, target_density, + shared_keys=None, ): """ This computes properties for one halo and runs on a single @@ -121,7 +128,14 @@ def process_single_halo( # one property calculation needs. It is created here, inside the # search radius loop, so that it is discarded as soon as the set of # particles changes. - shared_particle_data = SharedParticleData() + 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. + # The keys depend only on which particle types are present, so they + # are the same for every halo in this chunk. + 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): @@ -172,8 +186,18 @@ 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): + # TEMPORARY (issue 64) + shared_data_created.append(shared_particle_data.nr_created) break # Either the density is still too high or the property calculation failed. @@ -329,6 +353,11 @@ def process_halos( if target_density is None or density < target_density: target_density = density + # The shared particle data key each calculation uses depends only on which + # particle types were read in, so it is the same for every halo in this + # chunk and can be worked out once here. + 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,) @@ -392,6 +421,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 @@ -430,6 +460,38 @@ def process_halos( comm.barrier() nr_halos_left = comm.allreduce(np.sum(halo_arrays["done"].local.value == 0)) + # TEMPORARY (issue 64): how many shared particle data objects each halo + # needed. One per distinct set of particles is expected (two for a central, + # which also does the SO calculations, one for a satellite); anything more + # means an entry was evicted while a later calculation still wanted it. + # Strip this out after testing. + local_hist = np.zeros(5, dtype=np.int64) + for nr in shared_data_created: + local_hist[min(nr, 4)] += 1 + hist = comm.allreduce(local_hist, op=MPI.SUM) + if comm.Get_rank() == 0 and hist.sum() > 0: + print( + "SHARED_DATA_CREATED_PER_HALO " + + " ".join( + f"{n if n < 4 else '4+'}={hist[n]}" for n in range(5) if hist[n] + ), + flush=True, + ) + + # TEMPORARY (issue 64): report peak per-rank memory, to check the effect of + # sharing and evicting the particle arrays. Strip this out after testing. + peak_gb = memory_use.get_peak_rss_gb() + if peak_gb is not None: + peak_max = comm.allreduce(peak_gb, op=MPI.MAX) + peak_sum = comm.allreduce(peak_gb, op=MPI.SUM) + if comm.Get_rank() == 0: + print( + f"PEAK_RSS_PER_RANK max={peak_max:.3f}GB " + f"mean={peak_sum / comm.Get_size():.3f}GB " + f"over {comm.Get_size()} ranks", + flush=True, + ) + # Stop the clock comm.barrier() t1_all = time.time() diff --git a/SOAP/core/memory_use.py b/SOAP/core/memory_use.py index 22943f53..c667d219 100644 --- a/SOAP/core/memory_use.py +++ b/SOAP/core/memory_use.py @@ -22,3 +22,22 @@ def get_memory_use(): free_mem_gb = mem.available / GB return total_mem_gb, free_mem_gb + + +def get_peak_rss_gb(): + """ + Peak resident set size of this process, in GB. + + TEMPORARY (issue 64): used to check the effect of sharing and evicting the + particle arrays on per rank memory. Strip this out after testing. + + Returns None if VmHWM cannot be read. + """ + try: + with open("/proc/self/status") as f: + for line in f: + if line.startswith("VmHWM:"): + return float(line.split()[1]) / 1024**2 + except OSError: + pass + return None diff --git a/SOAP/core/shared_particle_data.py b/SOAP/core/shared_particle_data.py index 00baaa61..d45f3f75 100644 --- a/SOAP/core/shared_particle_data.py +++ b/SOAP/core/shared_particle_data.py @@ -12,21 +12,23 @@ radii, sorted radial profiles, ...), which is wasted work when it is repeated once per calculation. -A SharedParticleData object lets those calculations look up quantities that -have already been derived from the same particles. It is created inside the -search radius loop of process_single_halo(), so a new (empty) cache is used -whenever the set of particles changes. +A ParticleDataCache lets those calculations look up quantities that have +already been derived from the same particles. It is created inside the search +radius loop of process_single_halo(), so a new (empty) cache is used whenever +the set of particles changes, and entries are dropped as soon as no remaining +calculation needs them. """ -from typing import Any, Callable, Hashable +from typing import Any, Callable, Hashable, Iterable -class SharedParticleData: +class ParticleDataCache: """ Cache of quantities derived from the particles of a single halo. Entries are created on first use, so nothing is computed for a halo unless - a property calculation actually asks for it. + a property calculation actually asks for it, and discarded once the + calculations which need them have all run. """ def __init__(self): @@ -34,6 +36,11 @@ def __init__(self): Constructor. Creates an empty cache. """ self.cache = {} + # TEMPORARY (issue 64): how many entries this cache has had to create. + # One per distinct set of particles is expected; more means an entry was + # dropped while a later calculation still needed it, which is a waste of + # time rather than a correctness problem. Strip this out after testing. + self.nr_created = 0 def get(self, key: Hashable, factory: Callable[[], Any]) -> Any: """ @@ -45,11 +52,30 @@ def get(self, key: Hashable, factory: Callable[[], Any]) -> Any: Identifies the quantity being requested. Calculations that want to share an entry have to agree on the key, so it needs to include everything the entry depends on (e.g. the particle types that were - used to compute it). + used to compute it). 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() + self.nr_created += 1 # TEMPORARY (issue 64) return self.cache[key] + + def keep_only(self, keys: Iterable[Hashable]): + """ + Drop every entry whose key is not in keys. + + Called after each calculation with the keys the remaining calculations + for this halo still need, so that particle arrays are not kept alive + for longer than they are used. Dropping an entry too early is a + performance problem rather than a correctness one: the next calculation + that wants it simply rebuilds it. + + 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] diff --git a/SOAP/particle_selection/SO_properties.py b/SOAP/particle_selection/SO_properties.py index bcd86d60..3a70a948 100644 --- a/SOAP/particle_selection/SO_properties.py +++ b/SOAP/particle_selection/SO_properties.py @@ -45,7 +45,10 @@ from SOAP.core.parameter_file import ParameterFile from SOAP.core.snapshot_datasets import SnapshotDatasets from SOAP.core.swift_cells import SWIFTCellGrid -from SOAP.core.shared_particle_data import SharedParticleData +from SOAP.core.shared_particle_data import ParticleDataCache +from SOAP.particle_selection.shared_halo_particle_data import ( + SharedHaloParticleData, +) def cumulative_mass_intersection(r: float, rho_dim: float, slope_dim: float) -> float: @@ -218,186 +221,6 @@ def find_SO_radius_and_mass( return SO_r, SO_mass, SO_volume -class SOSharedParticleData: - """ - Particle quantities that are shared by all SO variations of a halo. - - Every SO variation of a halo (200_crit, 200_mean, 500_crit, ...) is handed - the same particles, and differs only in the threshold that determines the - SO radius. The quantities computed here are therefore identical for all of - them: the concatenated particle arrays, the cumulative mass profile that - the SO radius is read off, and the masks flagging particles that are bound - to another halo. Computing them once rather than once per variation avoids - repeatedly sorting every particle within the search radius. - - The quantities that do depend on the threshold (the SO radius itself, and - the particle selections derived from it) are computed by SOParticleData, - which holds a reference to one of these objects. - - Note that neutrinos are kept separate from the other particle types, since - they only contribute to the spherical overdensity radius calculation and to - neutrino specific properties. - """ - - def __init__( - self, - input_halo: Dict, - data: Dict, - types_present: List[str], - snapshot_datasets: SnapshotDatasets, - softening_of_parttype: unyt.unyt_array, - cosmology: dict, - boxsize: unyt.unyt_quantity, - ): - """ - Constructor. - - Parameters: - - input_halo: Dict - Dictionary containing properties of the halo read from the VR catalogue. - - data: Dict - Dictionary containing particle data. - - types_present: List - List of all particle types (e.g. 'PartType0') that are present in the data - dictionary. - - snapshot_datasets: SnapshotDatasets - Object containing metadata about the datasets in the snapshot, like - appropriate aliases and column names. - - softening_of_parttype: unyt.unyt_array - Softening length of each particle types - - cosmology: dict - Cosmological parameters required for SO calculation - - boxsize: unyt.unyt_quantity - Boxsize for correcting periodic boundary conditions - """ - self.input_halo = input_halo - self.data = data - self.has_neutrinos = "PartType6" in data - self.types_present = types_present - self.snapshot_datasets = snapshot_datasets - self.softening_of_parttype = softening_of_parttype - self.cosmology = cosmology - self.boxsize = boxsize - self.compute_basics() - self.compute_mass_profile() - - def get_dataset(self, name: str) -> unyt.unyt_array: - """ - Local wrapper for SnapshotDatasets.get_dataset(). - """ - return self.snapshot_datasets.get_dataset(name, self.data) - - def compute_basics(self): - """ - Compute some properties that are always needed, regardless of which - properties we actually want to compute. - """ - self.centre = self.input_halo["cofp"] - self.index = self.input_halo["index"] - - # Make an array of particle masses, radii and positions - mass = [] - radius = [] - position = [] - velocity = [] - types = [] - groupnr = [] - fofid = [] - softening = [] - for ptype in self.types_present: - if ptype == "PartType6": - # add neutrinos separately, since we need to treat them - # differently - continue - mass.append(self.get_dataset(f"{ptype}/{mass_dataset(ptype)}")) - pos = self.get_dataset(f"{ptype}/Coordinates") - self.centre[None, :] - position.append(pos) - r = np.sqrt(np.sum(pos**2, axis=1)) - radius.append(r) - velocity.append(self.get_dataset(f"{ptype}/Velocities")) - typearr = int(ptype[-1]) * np.ones(r.shape, dtype=np.int32) - types.append(typearr) - groupnr.append(self.get_dataset(f"{ptype}/GroupNr_bound")) - fofid.append(self.get_dataset(f"{ptype}/FOFGroupIDs")) - s = np.ones(r.shape, dtype=np.float64) * self.softening_of_parttype[ptype] - softening.append(s) - self.mass = np.concatenate(mass) - self.radius = np.concatenate(radius) - self.position = np.concatenate(position) - self.velocity = np.concatenate(velocity) - self.types = np.concatenate(types) - self.groupnr = np.concatenate(groupnr) - self.fofid = np.concatenate(fofid) - self.softening = np.concatenate(softening) - - def compute_mass_profile(self): - """ - Compute the cumulative mass profile used to determine the SO radius. - - Adds the contribution from neutrinos (if present) to the masses and - radii, sorts the particles by radius, and computes the cumulative mass - profile and the mean density within the radius of each particle. Also - determines the FOF ID of this object from its central particle, and - uses that to flag the particles which are bound to another halo. - - None of this depends on the density threshold of an individual SO - variation, so it is computed once and used by all of them. - """ - # add neutrinos - if self.has_neutrinos: - numass = self.get_dataset("PartType6/Masses") * self.get_dataset( - "PartType6/Weights" - ) - pos = self.get_dataset("PartType6/Coordinates") - self.centre[None, :] - nur = np.sqrt(np.sum(pos**2, axis=1)) - self.nu_mass = numass - self.nu_radius = nur - self.nu_softening = ( - np.ones_like(nur) * self.softening_of_parttype["PartType6"] - ) - all_mass = np.concatenate([self.mass, numass / unyt.dimensionless]) - all_r = np.concatenate([self.radius, nur]) - else: - all_mass = self.mass - all_r = self.radius - - # Sort by radius - order = np.argsort(all_r) - ordered_radius = all_r[order] - cumulative_mass = np.cumsum(all_mass[order], dtype=np.float64).astype( - self.mass.dtype - ) - # add mean neutrino mass - cumulative_mass += ( - self.cosmology["nu_density"] * 4.0 / 3.0 * np.pi * ordered_radius**3 - ) - # Determine FOF ID of object using the central non-neutrino particle - non_neutrino_order = order[order < self.radius.shape[0]] - fofid = self.fofid[non_neutrino_order[0]] - - # Compute density within radius of each particle. - # Will need to skip any at zero radius. - # Note that because of the definition of the centre of potential, the first - # particle *should* be at r=0. We need to manually exclude it, in case round - # off error places it at a very small non-zero radius. - nskip = max(1, np.argmax(ordered_radius > 0.0 * ordered_radius.units)) - self.ordered_radius = ordered_radius[nskip:] - self.cumulative_mass = cumulative_mass[nskip:] - self.nr_parts = len(self.ordered_radius) - self.density = self.cumulative_mass / ( - 4.0 / 3.0 * np.pi * self.ordered_radius**3 - ) - - # figure out which particles in the list are bound to a halo that is not the - # central halo - self.is_bound_to_satellite = ( - (self.groupnr >= 0) & (self.groupnr != self.index) & (self.fofid == fofid) - ) - self.is_bound_to_external = ( - (self.groupnr >= 0) & (self.groupnr != self.index) & (self.fofid != fofid) - ) - - class SOParticleData: """ Halo calculation class. @@ -421,20 +244,23 @@ class SOParticleData: def __init__( self, - shared: "SOSharedParticleData", + shared: SharedHaloParticleData, recently_heated_gas_filter: RecentlyHeatedGasFilter, observer_position: unyt.unyt_array, core_excision_fraction: float, virial_definition: bool, search_radius: unyt.unyt_quantity, + cosmology: dict, + boxsize: unyt.unyt_quantity, ): """ Constructor. Parameters: - - shared: SOSharedParticleData - Object holding the particle quantities that are the same for every - SO variation of this halo. + - shared: SharedHaloParticleData + Object holding the concatenated particle arrays for this halo, + shared with the other SO variations and with the inclusive + apertures. - recently_heated_gas_filter: RecentlyHeatedGasFilter Filter used to mask out gas particles that were recently heated by AGN feedback. @@ -450,8 +276,17 @@ def __init__( - search_radius: unyt.unyt_quantity Current search radius. Particles are guaranteed to be included up to this radius. + - cosmology: dict + Cosmological parameters required for SO calculation + - boxsize: unyt.unyt_quantity + Boxsize for correcting periodic boundary conditions """ self.shared = shared + # The radial profile needs the cosmology, which the aperture + # calculations sharing this object do not have, so it is requested + # here rather than built with the object. Only the first SO variation + # of a halo actually computes it. + shared.compute_mass_profile(cosmology) # Quantities that are the same for every SO variation of this halo. # Note that compute_SO_radius_and_mass() only ever rebinds these arrays @@ -462,8 +297,8 @@ def __init__( self.types_present = shared.types_present self.snapshot_datasets = shared.snapshot_datasets self.softening_of_parttype = shared.softening_of_parttype - self.cosmology = shared.cosmology - self.boxsize = shared.boxsize + self.cosmology = cosmology + self.boxsize = boxsize self.centre = shared.centre self.index = shared.index self.mass = shared.mass @@ -499,7 +334,7 @@ def compute_SO_radius_and_mass( Compute the SO radius from the density profile of the particles. Uses the cumulative mass profile computed once for this halo by - SOSharedParticleData, and calls find_SO_radius_and_mass(), unless a + SharedHaloParticleData, and calls find_SO_radius_and_mass(), unless a radius multiple is used as aperture radius. Particles outside the SO radius are then removed. @@ -3322,6 +3157,9 @@ class SOProperties(HaloProperty): sattelites. """ + # SOs always use every particle in the search radius + shared_inclusive = True + """ List of properties from the table that we want to compute. Each property should have a corresponding method/property/lazy_property in @@ -3670,7 +3508,7 @@ def calculate( search_radius: unyt.unyt_quantity, data: Dict, halo_result: Dict, - shared_particle_data: SharedParticleData = None, + shared_particle_data: ParticleDataCache = None, ): """ Compute spherical masses and overdensities for a halo @@ -3735,22 +3573,19 @@ def calculate( # types are part of the cache key because they determine the order # in which the particle arrays are concatenated. def make_shared(): - return SOSharedParticleData( + return SharedHaloParticleData( input_halo, data, types_present, + True, self.snapshot_datasets, self.softening_of_parttype, - self.cosmology, - self.boxsize, ) if shared_particle_data is None: shared = make_shared() else: - shared = shared_particle_data.get( - ("SOSharedParticleData", tuple(types_present)), make_shared - ) + shared = shared_particle_data.get(self.shared_key(data), make_shared) part_props = SOParticleData( shared, @@ -3759,6 +3594,8 @@ def make_shared(): self.core_excision_fraction, self.virial_definition, search_radius, + self.cosmology, + self.boxsize, ) # we need to make sure the physical radius uses the correct unit @@ -4009,7 +3846,7 @@ def calculate( search_radius: unyt.unyt_quantity, data: Dict, halo_result: Dict, - shared_particle_data: SharedParticleData = None, + shared_particle_data: ParticleDataCache = None, ): """ Calculate the properties of an SO of which the radius is the multiple of @@ -4027,7 +3864,7 @@ def calculate( Dictionary in which halo properties for this halo are stored. Should contain a valid result for the "parent" SO, i.e. the SO that determines the radius of this SO. - - shared_particle_data: SharedParticleData + - shared_particle_data: ParticleDataCache Cache of particle quantities shared with the other property calculations for this halo. If None, the quantities this calculation needs are computed for its own use only. diff --git a/SOAP/particle_selection/aperture_properties.py b/SOAP/particle_selection/aperture_properties.py index bc8fed8c..49a67720 100644 --- a/SOAP/particle_selection/aperture_properties.py +++ b/SOAP/particle_selection/aperture_properties.py @@ -172,7 +172,7 @@ def MetalFracStar(self): from SOAP.core.category_filter import CategoryFilter from SOAP.core.parameter_file import ParameterFile from SOAP.core.snapshot_datasets import SnapshotDatasets -from SOAP.core.shared_particle_data import SharedParticleData +from SOAP.core.shared_particle_data import ParticleDataCache from SOAP.particle_selection.shared_halo_particle_data import ( SharedHaloParticleData, ) @@ -3912,6 +3912,8 @@ def __init__( self.aperture_physical_radius_kpc = aperture_physical_radius_kpc self.aperture_property = aperture_property self.inclusive = inclusive + # which particles this aperture uses, for the shared particle arrays + self.shared_inclusive = inclusive if self.aperture_physical_radius_kpc is not None: self.physical_radius_mpc = 0.001 * self.aperture_physical_radius_kpc @@ -3978,7 +3980,7 @@ def calculate( search_radius: unyt.unyt_quantity, data: Dict, halo_result: Dict, - shared_particle_data: SharedParticleData = None, + shared_particle_data: ParticleDataCache = None, ): """ Compute centre of mass etc of bound particles @@ -4103,8 +4105,7 @@ def calculate( # Every aperture with the same value of "inclusive" sees the same # particles, so the concatenated arrays are computed once and shared # (with the bound subhalo and the projected apertures too, for the - # exclusive ones). The particle types are part of the cache key - # because they determine the order of the concatenated arrays. + # exclusive ones). def make_shared(): return SharedHaloParticleData( input_halo, @@ -4118,10 +4119,7 @@ def make_shared(): if shared_particle_data is None: shared = make_shared() else: - shared = shared_particle_data.get( - ("SharedHaloParticleData", self.inclusive, tuple(types_present)), - make_shared, - ) + shared = shared_particle_data.get(self.shared_key(data), make_shared) part_props = ApertureParticleData( shared, diff --git a/SOAP/particle_selection/halo_properties.py b/SOAP/particle_selection/halo_properties.py index 1bca96b8..0f65768b 100644 --- a/SOAP/particle_selection/halo_properties.py +++ b/SOAP/particle_selection/halo_properties.py @@ -34,6 +34,36 @@ def __init__(self, cellgrid): self.mean_density_multiple = None self.critical_density_multiple = None + # Whether the particles this calculation uses are all of them (True) or only + # those bound to the halo (False). Calculations which share a + # SharedHaloParticleData object must agree on this. None means the + # calculation does not use one. + shared_inclusive = None + + def shared_key(self, data): + """ + Return the key identifying the SharedHaloParticleData object this + calculation would use for the given particle data, or None if it does + not use one. + + The particle types are part of the key because they determine the order + in which the arrays are concatenated. Neutrinos are excluded because + they are never part of those arrays, which is what allows the SO + calculations to share an object with the inclusive apertures. + + Parameters: + - data: Dict + Dictionary containing particle data. + """ + if self.shared_inclusive is None: + return None + types_present = tuple( + ptype + for ptype in self.particle_properties + if ptype in data and ptype != "PartType6" + ) + return ("SharedHaloParticleData", self.shared_inclusive, types_present) + def expected_dataset_names(self): """ Return the set of HDF5 dataset names that this calculation will add diff --git a/SOAP/particle_selection/projected_aperture_properties.py b/SOAP/particle_selection/projected_aperture_properties.py index f59a1c18..b1b096bd 100644 --- a/SOAP/particle_selection/projected_aperture_properties.py +++ b/SOAP/particle_selection/projected_aperture_properties.py @@ -32,7 +32,7 @@ from SOAP.core.category_filter import CategoryFilter from SOAP.core.parameter_file import ParameterFile from SOAP.core.snapshot_datasets import SnapshotDatasets -from SOAP.core.shared_particle_data import SharedParticleData +from SOAP.core.shared_particle_data import ParticleDataCache from SOAP.particle_selection.shared_halo_particle_data import ( SharedHaloParticleData, ) @@ -1559,6 +1559,9 @@ class ProjectedApertureProperties(HaloProperty): the halo along the projection axis. """ + # projected apertures always use the particles bound to the halo + shared_inclusive = False + base_halo_type = "ProjectedApertureProperties" # Properties to calculate. The key is the name of the property, # the value indicates the property has a direct dependence on aperture size. @@ -1791,7 +1794,7 @@ def calculate( search_radius: unyt.unyt_quantity, data: Dict, halo_result: Dict, - shared_particle_data: SharedParticleData = None, + shared_particle_data: ParticleDataCache = None, ): """ Compute centre of mass etc of bound particles @@ -1900,8 +1903,7 @@ def calculate( # The concatenated arrays for the bound particles of this halo are # also used by the bound subhalo and the exclusive apertures, so they - # are computed once and shared. The particle types are part of the - # cache key because they determine the order of the arrays. + # are computed once and shared. def make_shared(): return SharedHaloParticleData( input_halo, @@ -1915,9 +1917,7 @@ def make_shared(): if shared_particle_data is None: shared = make_shared() else: - shared = shared_particle_data.get( - ("SharedHaloParticleData", False, tuple(types_present)), make_shared - ) + shared = shared_particle_data.get(self.shared_key(data), make_shared) part_props = ProjectedApertureParticleData( shared, diff --git a/SOAP/particle_selection/shared_halo_particle_data.py b/SOAP/particle_selection/shared_halo_particle_data.py index c8436388..5a8a9707 100644 --- a/SOAP/particle_selection/shared_halo_particle_data.py +++ b/SOAP/particle_selection/shared_halo_particle_data.py @@ -3,8 +3,8 @@ """ shared_halo_particle_data.py -Particle arrays shared between the aperture, projected aperture and subhalo -property calculations of a single halo. +Particle arrays shared between the aperture, projected aperture, subhalo and +spherical overdensity property calculations of a single halo. All of those calculations start by concatenating the same per particle quantities (masses, positions, velocities and particle types) for either the @@ -78,7 +78,13 @@ def __init__( """ self.input_halo = input_halo self.data = data - self.types_present = types_present + # Neutrinos are never part of the concatenated arrays: they only + # contribute to the spherical overdensity radius and to neutrino + # specific properties, and are handled separately below. Dropping them + # here is also what lets the SO calculations share this object with the + # inclusive apertures, whose particle types never include PartType6. + self.types_present = [t for t in types_present if t != "PartType6"] + self.has_neutrinos = "PartType6" in data self.inclusive = inclusive self.snapshot_datasets = snapshot_datasets self.softening_of_parttype = softening_of_parttype @@ -193,3 +199,105 @@ def radius_projz(self) -> unyt.unyt_array: """ pos = self.position return np.sqrt(pos[:, 0] ** 2 + pos[:, 1] ** 2) + + @lazy_property + def groupnr(self) -> unyt.unyt_array: + """ + Index of the subhalo each particle is bound to, or a negative value for + unbound particles. + """ + return np.concatenate( + [ + self.get_dataset(f"{ptype}/GroupNr_bound")[self.in_halo_mask(ptype)] + for ptype, _ in self.nr_part_of_type + ] + ) + + @lazy_property + def fofid(self) -> unyt.unyt_array: + """ + FOF group ID of each particle. + """ + return np.concatenate( + [ + self.get_dataset(f"{ptype}/FOFGroupIDs")[self.in_halo_mask(ptype)] + for ptype, _ in self.nr_part_of_type + ] + ) + + def compute_mass_profile(self, cosmology: Dict): + """ + Compute the cumulative mass profile used to determine the SO radius. + + Adds the contribution from neutrinos (if present) to the masses and + radii, sorts the particles by radius, and computes the cumulative mass + profile and the mean density within the radius of each particle. Also + determines the FOF ID of this object from its central particle, and + uses that to flag the particles which are bound to another halo. + + None of this depends on the density threshold of an individual SO + variation, so it is computed once and used by all of them. It is a + method rather than a lazy property because it needs the cosmology, + which the aperture calculations that also use this object do not have. + Repeated calls after the first are no-ops. + + Parameters: + - cosmology: dict + Cosmological parameters required for the SO calculation. + """ + if getattr(self, "have_mass_profile", False): + return + # add neutrinos + if self.has_neutrinos: + numass = self.get_dataset("PartType6/Masses") * self.get_dataset( + "PartType6/Weights" + ) + pos = self.get_dataset("PartType6/Coordinates") - self.centre[None, :] + nur = np.sqrt(np.sum(pos**2, axis=1)) + self.nu_mass = numass + self.nu_radius = nur + self.nu_softening = ( + np.ones_like(nur) * self.softening_of_parttype["PartType6"] + ) + all_mass = np.concatenate([self.mass, numass / unyt.dimensionless]) + all_r = np.concatenate([self.radius, nur]) + else: + all_mass = self.mass + all_r = self.radius + + # Sort by radius + order = np.argsort(all_r) + ordered_radius = all_r[order] + cumulative_mass = np.cumsum(all_mass[order], dtype=np.float64).astype( + self.mass.dtype + ) + # add mean neutrino mass + cumulative_mass += ( + cosmology["nu_density"] * 4.0 / 3.0 * np.pi * ordered_radius**3 + ) + # Determine FOF ID of object using the central non-neutrino particle + non_neutrino_order = order[order < self.radius.shape[0]] + fofid = self.fofid[non_neutrino_order[0]] + + # Compute density within radius of each particle. + # Will need to skip any at zero radius. + # Note that because of the definition of the centre of potential, the first + # particle *should* be at r=0. We need to manually exclude it, in case round + # off error places it at a very small non-zero radius. + nskip = max(1, np.argmax(ordered_radius > 0.0 * ordered_radius.units)) + self.ordered_radius = ordered_radius[nskip:] + self.cumulative_mass = cumulative_mass[nskip:] + self.nr_parts = len(self.ordered_radius) + self.density = self.cumulative_mass / ( + 4.0 / 3.0 * np.pi * self.ordered_radius**3 + ) + + # figure out which particles in the list are bound to a halo that is not the + # central halo + self.is_bound_to_satellite = ( + (self.groupnr >= 0) & (self.groupnr != self.index) & (self.fofid == fofid) + ) + self.is_bound_to_external = ( + (self.groupnr >= 0) & (self.groupnr != self.index) & (self.fofid != fofid) + ) + self.have_mass_profile = True diff --git a/SOAP/particle_selection/subhalo_properties.py b/SOAP/particle_selection/subhalo_properties.py index 27ca5db1..f2b49563 100644 --- a/SOAP/particle_selection/subhalo_properties.py +++ b/SOAP/particle_selection/subhalo_properties.py @@ -2324,6 +2324,9 @@ class SubhaloProperties(HaloProperty): gravitationally bound particles. """ + # the bound subhalo uses the particles bound to the halo + shared_inclusive = False + """ List of properties from the table that we want to compute. Each property should have a corresponding method/property/lazy_property in @@ -2581,8 +2584,7 @@ def calculate( # The concatenated arrays for the bound particles of this halo are also # used by the exclusive and projected aperture calculations, so they are - # computed once and shared. The particle types are part of the cache key - # because they determine the order in which the arrays are concatenated. + # computed once and shared. def make_shared(): return SharedHaloParticleData( input_halo, @@ -2596,9 +2598,7 @@ def make_shared(): if shared_particle_data is None: shared = make_shared() else: - shared = shared_particle_data.get( - ("SharedHaloParticleData", False, tuple(types_present)), make_shared - ) + shared = shared_particle_data.get(self.shared_key(data), make_shared) part_props = SubhaloParticleData( shared, From 567106ff7a88ae7c73213c2d3a6fa1c7a0833ba6 Mon Sep 17 00:00:00 2001 From: robjmcgibbon Date: Fri, 11 Sep 2026 10:28:04 +0100 Subject: [PATCH 04/11] Don't hold fof ids or density --- SOAP/compute_halo_properties.py | 7 +- SOAP/particle_selection/SO_properties.py | 2 - .../shared_halo_particle_data.py | 84 +++++++++++++------ 3 files changed, 63 insertions(+), 30 deletions(-) diff --git a/SOAP/compute_halo_properties.py b/SOAP/compute_halo_properties.py index ecab86a6..7a6729c8 100644 --- a/SOAP/compute_halo_properties.py +++ b/SOAP/compute_halo_properties.py @@ -491,12 +491,17 @@ def compute_halo_properties(): # search radius. Nothing outside the SO calculations reads an SO result, # which is what lets them move after the apertures; the assertions above # keep that true. + # - The SO calculations come last, after the inclusive apertures they share + # their particle arrays with. The SO calculations add quantities to those + # arrays which no aperture uses (the group and FOF IDs, the sorted mass + # profile), and for the largest halos those are several GB. Running them + # last means nothing else is still holding the arrays while they exist. halo_prop_list = ( subhalo_props + exclusive_apertures + projected_apertures - + so_props + inclusive_apertures + + so_props ) if len(halo_prop_list) < 1: diff --git a/SOAP/particle_selection/SO_properties.py b/SOAP/particle_selection/SO_properties.py index 3a70a948..cd634b1b 100644 --- a/SOAP/particle_selection/SO_properties.py +++ b/SOAP/particle_selection/SO_properties.py @@ -306,8 +306,6 @@ def __init__( self.position = shared.position self.velocity = shared.velocity self.types = shared.types - self.groupnr = shared.groupnr - self.fofid = shared.fofid self.softening = shared.softening if shared.has_neutrinos: self.nu_mass = shared.nu_mass diff --git a/SOAP/particle_selection/shared_halo_particle_data.py b/SOAP/particle_selection/shared_halo_particle_data.py index 5a8a9707..01b00e8c 100644 --- a/SOAP/particle_selection/shared_halo_particle_data.py +++ b/SOAP/particle_selection/shared_halo_particle_data.py @@ -200,30 +200,62 @@ def radius_projz(self) -> unyt.unyt_array: pos = self.position return np.sqrt(pos[:, 0] ** 2 + pos[:, 1] ** 2) - @lazy_property - def groupnr(self) -> unyt.unyt_array: - """ - Index of the subhalo each particle is bound to, or a negative value for - unbound particles. + def fofid_of_particle(self, index: int) -> int: """ - return np.concatenate( - [ - self.get_dataset(f"{ptype}/GroupNr_bound")[self.in_halo_mask(ptype)] - for ptype, _ in self.nr_part_of_type - ] - ) + FOF group ID of a single particle in the concatenated arrays. - @lazy_property - def fofid(self) -> unyt.unyt_array: + Only the particle type that particle belongs to is read, so the full + length array of FOF IDs never has to be built just to look up one value. + + Parameters: + - index: int + Position of the particle in the concatenated arrays. """ - FOF group ID of each particle. + # np.argsort() on a unyt_array returns the indices as a unyt_array + # carrying the units of the array that was sorted, so make sure we have + # a plain integer before doing any arithmetic with it + index = int(index) + offset = 0 + for ptype, nr_part in self.nr_part_of_type: + if index < offset + nr_part: + fofid = self.get_dataset(f"{ptype}/FOFGroupIDs") + if self.inclusive: + # every particle is included, so the position in the + # concatenated array is also the position in the raw array + return fofid[index - offset] + return fofid[self.in_halo_mask(ptype)][index - offset] + offset += nr_part + raise IndexError(f"particle {index} is not in the concatenated arrays") + + def compute_bound_masks(self, fofid_central: int): """ - return np.concatenate( - [ - self.get_dataset(f"{ptype}/FOFGroupIDs")[self.in_halo_mask(ptype)] - for ptype, _ in self.nr_part_of_type - ] - ) + Flag the particles which are bound to a halo other than this one, + separating those in the same FOF group from those in another one. + + The group numbers and FOF IDs are read one particle type at a time and + reduced to masks immediately. They are 8 bytes per particle each, so for + the largest halos holding both of them over the whole search radius + costs several GB, while the masks that are actually wanted are 1 byte + per particle. + + Parameters: + - fofid_central: int + FOF group ID of this halo. + """ + satellite = [] + external = [] + for ptype, _ in self.nr_part_of_type: + in_halo = self.in_halo_mask(ptype) + groupnr = self.get_dataset(f"{ptype}/GroupNr_bound")[in_halo] + bound_elsewhere = (groupnr >= 0) & (groupnr != self.index) + del groupnr + fofid = self.get_dataset(f"{ptype}/FOFGroupIDs")[in_halo] + same_fof = fofid == fofid_central + del fofid + satellite.append(bound_elsewhere & same_fof) + external.append(bound_elsewhere & ~same_fof) + self.is_bound_to_satellite = np.concatenate(satellite) + self.is_bound_to_external = np.concatenate(external) def compute_mass_profile(self, cosmology: Dict): """ @@ -271,13 +303,16 @@ def compute_mass_profile(self, cosmology: Dict): cumulative_mass = np.cumsum(all_mass[order], dtype=np.float64).astype( self.mass.dtype ) + del all_mass, all_r # add mean neutrino mass cumulative_mass += ( cosmology["nu_density"] * 4.0 / 3.0 * np.pi * ordered_radius**3 ) # Determine FOF ID of object using the central non-neutrino particle non_neutrino_order = order[order < self.radius.shape[0]] - fofid = self.fofid[non_neutrino_order[0]] + fofid_central = self.fofid_of_particle(non_neutrino_order[0]) + # The sort order is 8 bytes per particle and is not needed again + del order, non_neutrino_order # Compute density within radius of each particle. # Will need to skip any at zero radius. @@ -294,10 +329,5 @@ def compute_mass_profile(self, cosmology: Dict): # figure out which particles in the list are bound to a halo that is not the # central halo - self.is_bound_to_satellite = ( - (self.groupnr >= 0) & (self.groupnr != self.index) & (self.fofid == fofid) - ) - self.is_bound_to_external = ( - (self.groupnr >= 0) & (self.groupnr != self.index) & (self.fofid != fofid) - ) + self.compute_bound_masks(fofid_central) self.have_mass_profile = True From 1cc0aeb82c17688803ddab356554cd7de6117e74 Mon Sep 17 00:00:00 2001 From: robjmcgibbon Date: Mon, 14 Sep 2026 15:45:22 +0100 Subject: [PATCH 05/11] Strip logging --- SOAP/core/halo_tasks.py | 40 ------------------------------- SOAP/core/memory_use.py | 19 --------------- SOAP/core/shared_particle_data.py | 6 ----- 3 files changed, 65 deletions(-) diff --git a/SOAP/core/halo_tasks.py b/SOAP/core/halo_tasks.py index b33cbea3..5b1ae66b 100644 --- a/SOAP/core/halo_tasks.py +++ b/SOAP/core/halo_tasks.py @@ -5,8 +5,6 @@ import numpy as np import unyt -from mpi4py import MPI - 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 @@ -22,10 +20,6 @@ # Radius in Mpc at which we report halos which have a large search radius REPORT_RADIUS = 20.0 -# TEMPORARY (issue 64): number of shared particle data objects created for each -# halo processed by this rank. Strip this out after testing. -shared_data_created = [] - def process_single_halo( mesh, @@ -196,8 +190,6 @@ def process_single_halo( # If we computed all of the properties, we're done with this halo if np.all(halo_prop_done): - # TEMPORARY (issue 64) - shared_data_created.append(shared_particle_data.nr_created) break # Either the density is still too high or the property calculation failed. @@ -460,38 +452,6 @@ def process_halos( comm.barrier() nr_halos_left = comm.allreduce(np.sum(halo_arrays["done"].local.value == 0)) - # TEMPORARY (issue 64): how many shared particle data objects each halo - # needed. One per distinct set of particles is expected (two for a central, - # which also does the SO calculations, one for a satellite); anything more - # means an entry was evicted while a later calculation still wanted it. - # Strip this out after testing. - local_hist = np.zeros(5, dtype=np.int64) - for nr in shared_data_created: - local_hist[min(nr, 4)] += 1 - hist = comm.allreduce(local_hist, op=MPI.SUM) - if comm.Get_rank() == 0 and hist.sum() > 0: - print( - "SHARED_DATA_CREATED_PER_HALO " - + " ".join( - f"{n if n < 4 else '4+'}={hist[n]}" for n in range(5) if hist[n] - ), - flush=True, - ) - - # TEMPORARY (issue 64): report peak per-rank memory, to check the effect of - # sharing and evicting the particle arrays. Strip this out after testing. - peak_gb = memory_use.get_peak_rss_gb() - if peak_gb is not None: - peak_max = comm.allreduce(peak_gb, op=MPI.MAX) - peak_sum = comm.allreduce(peak_gb, op=MPI.SUM) - if comm.Get_rank() == 0: - print( - f"PEAK_RSS_PER_RANK max={peak_max:.3f}GB " - f"mean={peak_sum / comm.Get_size():.3f}GB " - f"over {comm.Get_size()} ranks", - flush=True, - ) - # Stop the clock comm.barrier() t1_all = time.time() diff --git a/SOAP/core/memory_use.py b/SOAP/core/memory_use.py index c667d219..22943f53 100644 --- a/SOAP/core/memory_use.py +++ b/SOAP/core/memory_use.py @@ -22,22 +22,3 @@ def get_memory_use(): free_mem_gb = mem.available / GB return total_mem_gb, free_mem_gb - - -def get_peak_rss_gb(): - """ - Peak resident set size of this process, in GB. - - TEMPORARY (issue 64): used to check the effect of sharing and evicting the - particle arrays on per rank memory. Strip this out after testing. - - Returns None if VmHWM cannot be read. - """ - try: - with open("/proc/self/status") as f: - for line in f: - if line.startswith("VmHWM:"): - return float(line.split()[1]) / 1024**2 - except OSError: - pass - return None diff --git a/SOAP/core/shared_particle_data.py b/SOAP/core/shared_particle_data.py index d45f3f75..4adc7f76 100644 --- a/SOAP/core/shared_particle_data.py +++ b/SOAP/core/shared_particle_data.py @@ -36,11 +36,6 @@ def __init__(self): Constructor. Creates an empty cache. """ self.cache = {} - # TEMPORARY (issue 64): how many entries this cache has had to create. - # One per distinct set of particles is expected; more means an entry was - # dropped while a later calculation still needed it, which is a waste of - # time rather than a correctness problem. Strip this out after testing. - self.nr_created = 0 def get(self, key: Hashable, factory: Callable[[], Any]) -> Any: """ @@ -59,7 +54,6 @@ def get(self, key: Hashable, factory: Callable[[], Any]) -> Any: """ if key not in self.cache: self.cache[key] = factory() - self.nr_created += 1 # TEMPORARY (issue 64) return self.cache[key] def keep_only(self, keys: Iterable[Hashable]): From c17bfa87303bbaafcf28c0d1cb7db6915ca68e5e Mon Sep 17 00:00:00 2001 From: robjmcgibbon Date: Mon, 14 Sep 2026 16:54:28 +0100 Subject: [PATCH 06/11] Review first pass --- SOAP/compute_halo_properties.py | 17 ++----- SOAP/core/halo_tasks.py | 11 +---- SOAP/core/shared_particle_data.py | 14 ------ SOAP/particle_selection/SO_properties.py | 18 ++------ .../particle_selection/aperture_properties.py | 19 ++------ SOAP/particle_selection/halo_properties.py | 45 +++++++++++++++++++ .../projected_aperture_properties.py | 19 ++------ SOAP/particle_selection/subhalo_properties.py | 19 +------- 8 files changed, 61 insertions(+), 101 deletions(-) diff --git a/SOAP/compute_halo_properties.py b/SOAP/compute_halo_properties.py index 7a6729c8..b6e8d361 100644 --- a/SOAP/compute_halo_properties.py +++ b/SOAP/compute_halo_properties.py @@ -205,10 +205,6 @@ def compute_halo_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 # 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. @@ -362,8 +358,6 @@ def compute_halo_properties(): if "radius_in_kpc" in aperture_variations[variation]: continue assert "property" in aperture_variations[variation] - # Apertures are computed before the SO calculations, so they cannot - # be defined in terms of an SO property assert not aperture_variations[variation]["property"].startswith( "SO/" ), "Apertures cannot be defined by an SO property" @@ -445,8 +439,6 @@ def compute_halo_properties(): if "radius_in_kpc" in projected_aperture_variations[variation]: continue assert "property" in projected_aperture_variations[variation] - # Projected apertures are computed before the SO calculations, so they - # cannot be defined in terms of an SO property assert not projected_aperture_variations[variation]["property"].startswith( "SO/" ), "Projected apertures cannot be defined by an SO property" @@ -488,14 +480,11 @@ def compute_halo_properties(): # 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. Nothing outside the SO calculations reads an SO result, - # which is what lets them move after the apertures; the assertions above - # keep that true. + # search radius. # - The SO calculations come last, after the inclusive apertures they share # their particle arrays with. The SO calculations add quantities to those - # arrays which no aperture uses (the group and FOF IDs, the sorted mass - # profile), and for the largest halos those are several GB. Running them - # last means nothing else is still holding the arrays while they exist. + # arrays which no aperture uses Running them last means nothing else is + # still holding the arrays once they're not needed exist. halo_prop_list = ( subhalo_props + exclusive_apertures diff --git a/SOAP/core/halo_tasks.py b/SOAP/core/halo_tasks.py index 5b1ae66b..be374119 100644 --- a/SOAP/core/halo_tasks.py +++ b/SOAP/core/halo_tasks.py @@ -119,15 +119,10 @@ def process_single_halo( pos[:, :] = ((pos - offset) % boxsize) + offset # Cache for quantities derived from these particles which more than - # one property calculation needs. It is created here, inside the - # search radius loop, so that it is discarded as soon as the set of - # particles changes. + # 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. - # The keys depend only on which particle types are present, so they - # are the same for every halo in this chunk. if shared_keys is None: shared_keys = [hp.shared_key(particle_data) for hp in halo_prop_list] @@ -345,9 +340,7 @@ def process_halos( if target_density is None or density < target_density: target_density = density - # The shared particle data key each calculation uses depends only on which - # particle types were read in, so it is the same for every halo in this - # chunk and can be worked out once here. + # 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 diff --git a/SOAP/core/shared_particle_data.py b/SOAP/core/shared_particle_data.py index 4adc7f76..5e258d2e 100644 --- a/SOAP/core/shared_particle_data.py +++ b/SOAP/core/shared_particle_data.py @@ -5,18 +5,6 @@ Cache of particle quantities that are shared between the property calculations of a single halo. - -process_single_halo() in halo_tasks.py hands the same set of particles to -every property calculation it runs for a halo. Several of those calculations -begin by deriving the same quantities from that set (concatenated masses and -radii, sorted radial profiles, ...), which is wasted work when it is repeated -once per calculation. - -A ParticleDataCache lets those calculations look up quantities that have -already been derived from the same particles. It is created inside the search -radius loop of process_single_halo(), so a new (empty) cache is used whenever -the set of particles changes, and entries are dropped as soon as no remaining -calculation needs them. """ from typing import Any, Callable, Hashable, Iterable @@ -24,8 +12,6 @@ class ParticleDataCache: """ - Cache of quantities derived from the particles of a single halo. - 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. diff --git a/SOAP/particle_selection/SO_properties.py b/SOAP/particle_selection/SO_properties.py index cd634b1b..acd03bb9 100644 --- a/SOAP/particle_selection/SO_properties.py +++ b/SOAP/particle_selection/SO_properties.py @@ -3564,26 +3564,14 @@ def calculate( # SOs only exist for central galaxies # Determine whether to skip this halo because of filter if input_halo["is_central"] and do_calculation[self.halo_filter]: - types_present = [type for type in self.particle_properties if type in data] # Quantities which are the same for every SO variation of this halo # are computed once and reused by the other variations. The particle # types are part of the cache key because they determine the order # in which the particle arrays are concatenated. - def make_shared(): - return SharedHaloParticleData( - input_halo, - data, - types_present, - True, - self.snapshot_datasets, - self.softening_of_parttype, - ) - - if shared_particle_data is None: - shared = make_shared() - else: - shared = shared_particle_data.get(self.shared_key(data), make_shared) + shared = self.get_shared_particle_data( + input_halo, data, shared_particle_data + ) part_props = SOParticleData( shared, diff --git a/SOAP/particle_selection/aperture_properties.py b/SOAP/particle_selection/aperture_properties.py index 49a67720..3cdc315a 100644 --- a/SOAP/particle_selection/aperture_properties.py +++ b/SOAP/particle_selection/aperture_properties.py @@ -4100,26 +4100,13 @@ def calculate( "Search radius is smaller than aperture" ) - types_present = [type for type in self.particle_properties if type in data] - # Every aperture with the same value of "inclusive" sees the same # particles, so the concatenated arrays are computed once and shared # (with the bound subhalo and the projected apertures too, for the # exclusive ones). - def make_shared(): - return SharedHaloParticleData( - input_halo, - data, - types_present, - self.inclusive, - self.snapshot_datasets, - self.softening_of_parttype, - ) - - if shared_particle_data is None: - shared = make_shared() - else: - shared = shared_particle_data.get(self.shared_key(data), make_shared) + shared = self.get_shared_particle_data( + input_halo, data, shared_particle_data + ) part_props = ApertureParticleData( shared, diff --git a/SOAP/particle_selection/halo_properties.py b/SOAP/particle_selection/halo_properties.py index 0f65768b..21743d5f 100644 --- a/SOAP/particle_selection/halo_properties.py +++ b/SOAP/particle_selection/halo_properties.py @@ -64,6 +64,51 @@ def shared_key(self, data): ) return ("SharedHaloParticleData", self.shared_inclusive, types_present) + def get_shared_particle_data(self, input_halo, data, cache): + """ + Return the SharedHaloParticleData object this calculation should use, + taking it from the cache if another calculation has already built the + same one. + + The object is built from the key, so which calculation happens to + create it cannot change what it contains. That matters because several + calculations share a key: the bound subhalo, the exclusive apertures + and the projected apertures all use one object, and the inclusive + apertures and the SO calculations another. + + Parameters: + - input_halo: Dict + Dictionary containing properties of the halo read from the halo + catalogue. + - data: Dict + Dictionary containing particle data. + - cache: ParticleDataCache or None + Cache shared with the other calculations for this halo. If None, the + object is built for this calculation's own use. + """ + # imported here rather than at module scope because + # SnapshotDatasets imports the property table, which imports this module + from SOAP.particle_selection.shared_halo_particle_data import ( + SharedHaloParticleData, + ) + + key = self.shared_key(data) + + def build(): + _, inclusive, types_present = key + return SharedHaloParticleData( + input_halo, + data, + list(types_present), + inclusive, + self.snapshot_datasets, + self.softening_of_parttype, + ) + + if cache is None: + return build() + return cache.get(key, build) + def expected_dataset_names(self): """ Return the set of HDF5 dataset names that this calculation will add diff --git a/SOAP/particle_selection/projected_aperture_properties.py b/SOAP/particle_selection/projected_aperture_properties.py index b1b096bd..73bf98aa 100644 --- a/SOAP/particle_selection/projected_aperture_properties.py +++ b/SOAP/particle_selection/projected_aperture_properties.py @@ -1899,25 +1899,12 @@ def calculate( * halo_result[self.aperture_property[0]][0] ) - types_present = [type for type in self.particle_properties if type in data] - # The concatenated arrays for the bound particles of this halo are # also used by the bound subhalo and the exclusive apertures, so they # are computed once and shared. - def make_shared(): - return SharedHaloParticleData( - input_halo, - data, - types_present, - False, - self.snapshot_datasets, - self.softening_of_parttype, - ) - - if shared_particle_data is None: - shared = make_shared() - else: - shared = shared_particle_data.get(self.shared_key(data), make_shared) + shared = self.get_shared_particle_data( + input_halo, data, shared_particle_data + ) part_props = ProjectedApertureParticleData( shared, diff --git a/SOAP/particle_selection/subhalo_properties.py b/SOAP/particle_selection/subhalo_properties.py index f2b49563..7b9f5e3b 100644 --- a/SOAP/particle_selection/subhalo_properties.py +++ b/SOAP/particle_selection/subhalo_properties.py @@ -53,10 +53,10 @@ from SOAP.core.category_filter import CategoryFilter from SOAP.core.parameter_file import ParameterFile from SOAP.core.snapshot_datasets import SnapshotDatasets +from SOAP.core.swift_cells import SWIFTCellGrid from SOAP.particle_selection.shared_halo_particle_data import ( SharedHaloParticleData, ) -from SOAP.core.swift_cells import SWIFTCellGrid class SubhaloParticleData: @@ -2580,25 +2580,10 @@ def calculate( Input particle data arrays are unyt_arrays. """ - types_present = [type for type in self.particle_properties if type in data] - # The concatenated arrays for the bound particles of this halo are also # used by the exclusive and projected aperture calculations, so they are # computed once and shared. - def make_shared(): - return SharedHaloParticleData( - input_halo, - data, - types_present, - False, - self.snapshot_datasets, - self.softening_of_parttype, - ) - - if shared_particle_data is None: - shared = make_shared() - else: - shared = shared_particle_data.get(self.shared_key(data), make_shared) + shared = self.get_shared_particle_data(input_halo, data, shared_particle_data) part_props = SubhaloParticleData( shared, From c7d8383a1dc6bf4441f9f1777d960b179e5d442f Mon Sep 17 00:00:00 2001 From: robjmcgibbon Date: Tue, 15 Sep 2026 10:30:15 +0100 Subject: [PATCH 07/11] Update comments --- SOAP/compute_halo_properties.py | 39 ++++++++++++++++++++++++--------- 1 file changed, 29 insertions(+), 10 deletions(-) diff --git a/SOAP/compute_halo_properties.py b/SOAP/compute_halo_properties.py index b6e8d361..02dd7940 100644 --- a/SOAP/compute_halo_properties.py +++ b/SOAP/compute_halo_properties.py @@ -202,6 +202,18 @@ def compute_halo_properties(): ) filters = parameter_file.get_filters() + # The SO calculations run after everything else (see where halo_prop_list is + # assembled below), so their results are not available to the filters which + # the other calculations are selected with. Reject this here rather than + # letting it fail with a KeyError part way through the first halo. + for filter_name, filter_info in filters.items(): + for prop in filter_info.get("properties", []): + if prop.startswith("SO/"): + raise ValueError( + f'Filter "{filter_name}" uses "{prop}", but SO properties ' + "are computed after the calculations which are selected " + "using filters. Use a BoundSubhalo property instead." + ) category_filter = CategoryFilter(filters, dmo=args.dmo) # Get the full list of property calculations we can do @@ -358,9 +370,6 @@ def compute_halo_properties(): if "radius_in_kpc" in aperture_variations[variation]: continue assert "property" in aperture_variations[variation] - assert not aperture_variations[variation]["property"].startswith( - "SO/" - ), "Apertures cannot be defined by an SO property" radius_multiple = aperture_variations[variation].get("radius_multiple", 1) # Only allow integer radius mutiples, otherwise swiftsimio will # struggle to handle the group names @@ -439,9 +448,13 @@ def compute_halo_properties(): if "radius_in_kpc" in projected_aperture_variations[variation]: continue assert "property" in projected_aperture_variations[variation] - assert not projected_aperture_variations[variation]["property"].startswith( - "SO/" - ), "Projected apertures cannot be defined by an SO property" + # Only BoundSubhalo properties are available: it is the one calculation + # guaranteed to have run first. ApertureProperties asserts the same in + # its constructor, but ProjectedApertureProperties does not. + 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 ) @@ -468,7 +481,7 @@ def compute_halo_properties(): parameter_file.write_parameters(args.output_parameters) # Assemble the calculations in the order they will be run for each halo. - # This order matters, for three separate reasons: + # 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. @@ -482,9 +495,15 @@ def compute_halo_properties(): # 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. The SO calculations add quantities to those - # arrays which no aperture uses Running them last means nothing else is - # still holding the arrays once they're not needed exist. + # 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. + # + # Note that nothing outside the SO calculations may therefore depend on an + # SO result: not as an aperture radius (asserted above), and not as a + # category filter property (asserted where the filters are read). halo_prop_list = ( subhalo_props + exclusive_apertures From 747cf2285e9f7d924eb9f88f09ad1f88cbc2da12 Mon Sep 17 00:00:00 2001 From: robjmcgibbon Date: Tue, 15 Sep 2026 10:48:54 +0100 Subject: [PATCH 08/11] Review 2 --- SOAP/particle_selection/SO_properties.py | 2 +- SOAP/particle_selection/aperture_properties.py | 2 -- SOAP/particle_selection/halo_properties.py | 6 +++--- SOAP/particle_selection/projected_aperture_properties.py | 2 +- SOAP/particle_selection/subhalo_properties.py | 2 +- 5 files changed, 6 insertions(+), 8 deletions(-) diff --git a/SOAP/particle_selection/SO_properties.py b/SOAP/particle_selection/SO_properties.py index acd03bb9..eecaf480 100644 --- a/SOAP/particle_selection/SO_properties.py +++ b/SOAP/particle_selection/SO_properties.py @@ -3156,7 +3156,7 @@ class SOProperties(HaloProperty): """ # SOs always use every particle in the search radius - shared_inclusive = True + inclusive = True """ List of properties from the table that we want to compute. diff --git a/SOAP/particle_selection/aperture_properties.py b/SOAP/particle_selection/aperture_properties.py index 3cdc315a..f3a175b8 100644 --- a/SOAP/particle_selection/aperture_properties.py +++ b/SOAP/particle_selection/aperture_properties.py @@ -3912,8 +3912,6 @@ def __init__( self.aperture_physical_radius_kpc = aperture_physical_radius_kpc self.aperture_property = aperture_property self.inclusive = inclusive - # which particles this aperture uses, for the shared particle arrays - self.shared_inclusive = inclusive if self.aperture_physical_radius_kpc is not None: self.physical_radius_mpc = 0.001 * self.aperture_physical_radius_kpc diff --git a/SOAP/particle_selection/halo_properties.py b/SOAP/particle_selection/halo_properties.py index 21743d5f..6800003a 100644 --- a/SOAP/particle_selection/halo_properties.py +++ b/SOAP/particle_selection/halo_properties.py @@ -38,7 +38,7 @@ def __init__(self, cellgrid): # those bound to the halo (False). Calculations which share a # SharedHaloParticleData object must agree on this. None means the # calculation does not use one. - shared_inclusive = None + inclusive = None def shared_key(self, data): """ @@ -55,14 +55,14 @@ def shared_key(self, data): - data: Dict Dictionary containing particle data. """ - if self.shared_inclusive is None: + if self.inclusive is None: return None types_present = tuple( ptype for ptype in self.particle_properties if ptype in data and ptype != "PartType6" ) - return ("SharedHaloParticleData", self.shared_inclusive, types_present) + return ("SharedHaloParticleData", self.inclusive, types_present) def get_shared_particle_data(self, input_halo, data, cache): """ diff --git a/SOAP/particle_selection/projected_aperture_properties.py b/SOAP/particle_selection/projected_aperture_properties.py index 73bf98aa..55e027ea 100644 --- a/SOAP/particle_selection/projected_aperture_properties.py +++ b/SOAP/particle_selection/projected_aperture_properties.py @@ -1560,7 +1560,7 @@ class ProjectedApertureProperties(HaloProperty): """ # projected apertures always use the particles bound to the halo - shared_inclusive = False + inclusive = False base_halo_type = "ProjectedApertureProperties" # Properties to calculate. The key is the name of the property, diff --git a/SOAP/particle_selection/subhalo_properties.py b/SOAP/particle_selection/subhalo_properties.py index 7b9f5e3b..2d03cf9e 100644 --- a/SOAP/particle_selection/subhalo_properties.py +++ b/SOAP/particle_selection/subhalo_properties.py @@ -2325,7 +2325,7 @@ class SubhaloProperties(HaloProperty): """ # the bound subhalo uses the particles bound to the halo - shared_inclusive = False + inclusive = False """ List of properties from the table that we want to compute. From 55923d1f3bfa16c458a0cdcb48c310acb3864063 Mon Sep 17 00:00:00 2001 From: robjmcgibbon Date: Tue, 15 Sep 2026 12:51:55 +0100 Subject: [PATCH 09/11] Review 3 --- SOAP/compute_halo_properties.py | 20 +++---------- SOAP/core/category_filter.py | 1 - SOAP/core/shared_particle_data.py | 11 +------- SOAP/particle_selection/SO_properties.py | 6 ---- .../particle_selection/aperture_properties.py | 5 +--- SOAP/particle_selection/halo_properties.py | 6 ---- .../shared_halo_particle_data.py | 28 ++++--------------- parameter_files/README.md | 2 ++ tests/dummy_halo_generator.py | 8 ------ 9 files changed, 13 insertions(+), 74 deletions(-) diff --git a/SOAP/compute_halo_properties.py b/SOAP/compute_halo_properties.py index 02dd7940..6e083ef1 100644 --- a/SOAP/compute_halo_properties.py +++ b/SOAP/compute_halo_properties.py @@ -202,18 +202,12 @@ def compute_halo_properties(): ) filters = parameter_file.get_filters() - # The SO calculations run after everything else (see where halo_prop_list is - # assembled below), so their results are not available to the filters which - # the other calculations are selected with. Reject this here rather than - # letting it fail with a KeyError part way through the first halo. for filter_name, filter_info in filters.items(): for prop in filter_info.get("properties", []): - if prop.startswith("SO/"): - raise ValueError( - f'Filter "{filter_name}" uses "{prop}", but SO properties ' - "are computed after the calculations which are selected " - "using filters. Use a BoundSubhalo property instead." - ) + 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 @@ -448,9 +442,6 @@ def compute_halo_properties(): if "radius_in_kpc" in projected_aperture_variations[variation]: continue assert "property" in projected_aperture_variations[variation] - # Only BoundSubhalo properties are available: it is the one calculation - # guaranteed to have run first. ApertureProperties asserts the same in - # its constructor, but ProjectedApertureProperties does not. assert ( projected_aperture_variations[variation]["property"].split("/")[0] == "BoundSubhalo" @@ -501,9 +492,6 @@ def compute_halo_properties(): # several GB. Running SO last means they only exist while the # calculations which need them are running. # - # Note that nothing outside the SO calculations may therefore depend on an - # SO result: not as an aperture radius (asserted above), and not as a - # category filter property (asserted where the filters are read). halo_prop_list = ( subhalo_props + exclusive_apertures diff --git a/SOAP/core/category_filter.py b/SOAP/core/category_filter.py index edfec534..e93bbe29 100644 --- a/SOAP/core/category_filter.py +++ b/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" diff --git a/SOAP/core/shared_particle_data.py b/SOAP/core/shared_particle_data.py index 5e258d2e..40e338e9 100644 --- a/SOAP/core/shared_particle_data.py +++ b/SOAP/core/shared_particle_data.py @@ -30,10 +30,7 @@ def get(self, key: Hashable, factory: Callable[[], Any]) -> Any: Parameters: - key: Hashable - Identifies the quantity being requested. Calculations that want to - share an entry have to agree on the key, so it needs to include - everything the entry depends on (e.g. the particle types that were - used to compute it). See HaloProperty.shared_key(). + 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. @@ -46,12 +43,6 @@ def keep_only(self, keys: Iterable[Hashable]): """ Drop every entry whose key is not in keys. - Called after each calculation with the keys the remaining calculations - for this halo still need, so that particle arrays are not kept alive - for longer than they are used. Dropping an entry too early is a - performance problem rather than a correctness one: the next calculation - that wants it simply rebuilds it. - Parameters: - keys: Iterable Keys to keep. Anything else is discarded. diff --git a/SOAP/particle_selection/SO_properties.py b/SOAP/particle_selection/SO_properties.py index eecaf480..cd9d5d73 100644 --- a/SOAP/particle_selection/SO_properties.py +++ b/SOAP/particle_selection/SO_properties.py @@ -282,15 +282,9 @@ def __init__( Boxsize for correcting periodic boundary conditions """ self.shared = shared - # The radial profile needs the cosmology, which the aperture - # calculations sharing this object do not have, so it is requested - # here rather than built with the object. Only the first SO variation - # of a halo actually computes it. shared.compute_mass_profile(cosmology) # Quantities that are the same for every SO variation of this halo. - # Note that compute_SO_radius_and_mass() only ever rebinds these arrays - # (it never modifies them in place), so it is safe to share them. self.input_halo = shared.input_halo self.data = shared.data self.has_neutrinos = shared.has_neutrinos diff --git a/SOAP/particle_selection/aperture_properties.py b/SOAP/particle_selection/aperture_properties.py index f3a175b8..625579c4 100644 --- a/SOAP/particle_selection/aperture_properties.py +++ b/SOAP/particle_selection/aperture_properties.py @@ -4098,10 +4098,7 @@ def calculate( "Search radius is smaller than aperture" ) - # Every aperture with the same value of "inclusive" sees the same - # particles, so the concatenated arrays are computed once and shared - # (with the bound subhalo and the projected apertures too, for the - # exclusive ones). + # Concatenated arrayss are computed once for inclusive/exclusive apertures shared = self.get_shared_particle_data( input_halo, data, shared_particle_data ) diff --git a/SOAP/particle_selection/halo_properties.py b/SOAP/particle_selection/halo_properties.py index 6800003a..191270f9 100644 --- a/SOAP/particle_selection/halo_properties.py +++ b/SOAP/particle_selection/halo_properties.py @@ -70,12 +70,6 @@ def get_shared_particle_data(self, input_halo, data, cache): taking it from the cache if another calculation has already built the same one. - The object is built from the key, so which calculation happens to - create it cannot change what it contains. That matters because several - calculations share a key: the bound subhalo, the exclusive apertures - and the projected apertures all use one object, and the inclusive - apertures and the SO calculations another. - Parameters: - input_halo: Dict Dictionary containing properties of the halo read from the halo diff --git a/SOAP/particle_selection/shared_halo_particle_data.py b/SOAP/particle_selection/shared_halo_particle_data.py index 01b00e8c..afc15bc8 100644 --- a/SOAP/particle_selection/shared_halo_particle_data.py +++ b/SOAP/particle_selection/shared_halo_particle_data.py @@ -78,11 +78,9 @@ def __init__( """ self.input_halo = input_halo self.data = data - # Neutrinos are never part of the concatenated arrays: they only + # Neutrinos are not part of the concatenated arrays: they only # contribute to the spherical overdensity radius and to neutrino - # specific properties, and are handled separately below. Dropping them - # here is also what lets the SO calculations share this object with the - # inclusive apertures, whose particle types never include PartType6. + # specific properties, and are handled separately below. self.types_present = [t for t in types_present if t != "PartType6"] self.has_neutrinos = "PartType6" in data self.inclusive = inclusive @@ -103,11 +101,7 @@ def in_halo_mask(self, ptype: str) -> NDArray[bool]: """ Mask which selects the particles of ptype that are included in the calculations: only the particles bound to this halo for exclusive - calculations, all of them for inclusive ones. This mask needs to be - applied _first_ to raw "PartTypeX" datasets. - - The mask is computed once per particle type and then reused, since - every calculation sharing this object needs the same one. + calculations, all of them for inclusive ones. Parameters: - ptype: str @@ -126,9 +120,6 @@ def compute_basics(self): """ Concatenate the quantities which every calculation sharing this object needs, over all particle types that are present. - - Also records the number of particles of each type, which the lazy - softening below uses to rebuild a per type quantity in the same order. """ mass = [] position = [] @@ -211,9 +202,6 @@ def fofid_of_particle(self, index: int) -> int: - index: int Position of the particle in the concatenated arrays. """ - # np.argsort() on a unyt_array returns the indices as a unyt_array - # carrying the units of the array that was sorted, so make sure we have - # a plain integer before doing any arithmetic with it index = int(index) offset = 0 for ptype, nr_part in self.nr_part_of_type: @@ -232,12 +220,6 @@ def compute_bound_masks(self, fofid_central: int): Flag the particles which are bound to a halo other than this one, separating those in the same FOF group from those in another one. - The group numbers and FOF IDs are read one particle type at a time and - reduced to masks immediately. They are 8 bytes per particle each, so for - the largest halos holding both of them over the whole search radius - costs several GB, while the masks that are actually wanted are 1 byte - per particle. - Parameters: - fofid_central: int FOF group ID of this halo. @@ -311,12 +293,11 @@ def compute_mass_profile(self, cosmology: Dict): # Determine FOF ID of object using the central non-neutrino particle non_neutrino_order = order[order < self.radius.shape[0]] fofid_central = self.fofid_of_particle(non_neutrino_order[0]) - # The sort order is 8 bytes per particle and is not needed again del order, non_neutrino_order # Compute density within radius of each particle. # Will need to skip any at zero radius. - # Note that because of the definition of the centre of potential, the first + # Note that because of the definition of the halo centre, the first # particle *should* be at r=0. We need to manually exclude it, in case round # off error places it at a very small non-zero radius. nskip = max(1, np.argmax(ordered_radius > 0.0 * ordered_radius.units)) @@ -331,3 +312,4 @@ def compute_mass_profile(self, cosmology: Dict): # central halo self.compute_bound_masks(fofid_central) self.have_mass_profile = True + diff --git a/parameter_files/README.md b/parameter_files/README.md index d351abf6..981b684b 100644 --- a/parameter_files/README.md +++ b/parameter_files/README.md @@ -195,6 +195,8 @@ For each alias the key is the name of the property that SOAP expects, and the va SOAP uses filters to determine whether to skip the calculation of an aperture or property based on the number of bound particles. This section of the parameter file defines the particle limits for each filter. New filters can be added if required. +Filter properties must be `BoundSubhalo` properties, since filters are evaluated +before other calculations run. There are no default filters. Every filter referenced by a property or a halo type variation must be defined here, with the sole exception of the implicit `basic` diff --git a/tests/dummy_halo_generator.py b/tests/dummy_halo_generator.py index a726da1f..e2903cd6 100644 --- a/tests/dummy_halo_generator.py +++ b/tests/dummy_halo_generator.py @@ -579,14 +579,6 @@ def get_halo_result_template(particle_numbers): ), "Dummy Nbh for filter", ), - f"SO/200_crit/{PropertyTable.full_property_list['Ngas'].name}": ( - unyt.unyt_array( - particle_numbers["PartType0"], - dtype=PropertyTable.full_property_list["Ngas"].dtype, - units="dimensionless", - ), - "Dummy SO Ngas for filter", - ), f"BoundSubhalo/EncloseRadius": ( unyt.unyt_array( 100, From f5b26940da8e17eac614a6770399e27f7584373e Mon Sep 17 00:00:00 2001 From: robjmcgibbon Date: Tue, 15 Sep 2026 13:08:19 +0100 Subject: [PATCH 10/11] Trigger workflow when job ready for review --- .github/workflows/main.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 19107013..4844a5a0 100644 --- a/.github/workflows/main.yml +++ b/.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: From ac8d96299a4d108edd272ba1f42058f337346678 Mon Sep 17 00:00:00 2001 From: robjmcgibbon Date: Tue, 15 Sep 2026 13:10:06 +0100 Subject: [PATCH 11/11] Format --- SOAP/particle_selection/shared_halo_particle_data.py | 1 - 1 file changed, 1 deletion(-) diff --git a/SOAP/particle_selection/shared_halo_particle_data.py b/SOAP/particle_selection/shared_halo_particle_data.py index afc15bc8..f275c0b6 100644 --- a/SOAP/particle_selection/shared_halo_particle_data.py +++ b/SOAP/particle_selection/shared_halo_particle_data.py @@ -312,4 +312,3 @@ def compute_mass_profile(self, cosmology: Dict): # central halo self.compute_bound_masks(fofid_central) self.have_mass_profile = True -