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: diff --git a/SOAP/compute_halo_properties.py b/SOAP/compute_halo_properties.py index 1cf0279f..6e083ef1 100644 --- a/SOAP/compute_halo_properties.py +++ b/SOAP/compute_halo_properties.py @@ -202,21 +202,30 @@ def compute_halo_properties(): ) filters = parameter_file.get_filters() + for filter_name, filter_info in filters.items(): + for prop in filter_info.get("properties", []): + 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 - # 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 - 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 +245,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 +260,7 @@ def compute_halo_properties(): ) ) else: - halo_prop_list.append( + so_props.append( SO_properties.SOProperties( cellgrid, parameter_file, @@ -268,7 +277,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, @@ -302,7 +311,11 @@ 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 + # Add the apertures defined with fixed physical radii, followed by those + # whose radius is defined by a SOAP property. Exclusive and inclusive + # 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 @@ -316,7 +329,7 @@ def compute_halo_properties(): if aperture_variations[variation].get("skip_gt_enclose_radius", False): radii_kpc = inclusive_radii_kpc - halo_prop_list.append( + inclusive_apertures.append( aperture_properties.InclusiveSphereProperties( cellgrid, parameter_file, @@ -331,7 +344,7 @@ def compute_halo_properties(): ) ) else: - halo_prop_list.append( + exclusive_apertures.append( aperture_properties.ExclusiveSphereProperties( cellgrid, parameter_file, @@ -346,7 +359,7 @@ def compute_halo_properties(): ) ) - # Add the apertures based on SOAP properties + # Apertures based on SOAP properties for variation in aperture_variations: if "radius_in_kpc" in aperture_variations[variation]: continue @@ -356,7 +369,7 @@ def compute_halo_properties(): # struggle to handle the group names assert int(radius_multiple) == radius_multiple if aperture_variations[variation]["inclusive"]: - halo_prop_list.append( + inclusive_apertures.append( aperture_properties.InclusiveSphereProperties( cellgrid, parameter_file, @@ -371,7 +384,7 @@ def compute_halo_properties(): ) ) else: - halo_prop_list.append( + exclusive_apertures.append( aperture_properties.ExclusiveSphereProperties( cellgrid, parameter_file, @@ -413,7 +426,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, @@ -429,11 +442,15 @@ def compute_halo_properties(): if "radius_in_kpc" in projected_aperture_variations[variation]: continue assert "property" in projected_aperture_variations[variation] + 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 ) assert int(radius_multiple) == radius_multiple - halo_prop_list.append( + projected_apertures.append( projected_aperture_properties.ProjectedApertureProperties( cellgrid, parameter_file, @@ -454,6 +471,35 @@ 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 four 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. + # - The SO calculations come last, after the inclusive apertures they share + # 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. + # + halo_prop_list = ( + subhalo_props + + exclusive_apertures + + projected_apertures + + inclusive_apertures + + so_props + ) + if len(halo_prop_list) < 1: raise Exception("Must select at least one halo property calculation!") 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/halo_tasks.py b/SOAP/core/halo_tasks.py index 4da85e23..be374119 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 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 @@ -30,6 +31,7 @@ def process_single_halo( boxsize, input_halo, target_density, + shared_keys=None, ): """ This computes properties for one halo and runs on a single @@ -116,6 +118,14 @@ 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. + 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. + 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): if halo_prop_done[prop_nr]: @@ -124,7 +134,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. @@ -161,6 +175,14 @@ 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): break @@ -318,6 +340,9 @@ def process_halos( if target_density is None or density < target_density: target_density = density + # 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 if comm.Get_rank() == 0: local_shape = (1,) @@ -381,6 +406,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 diff --git a/SOAP/core/shared_particle_data.py b/SOAP/core/shared_particle_data.py new file mode 100644 index 00000000..40e338e9 --- /dev/null +++ b/SOAP/core/shared_particle_data.py @@ -0,0 +1,52 @@ +#!/bin/env python + +""" +shared_particle_data.py + +Cache of particle quantities that are shared between the property +calculations of a single halo. +""" + +from typing import Any, Callable, Hashable, Iterable + + +class ParticleDataCache: + """ + 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. + """ + + 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. 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() + return self.cache[key] + + def keep_only(self, keys: Iterable[Hashable]): + """ + Drop every entry whose key is not in keys. + + 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 b5478558..cd9d5d73 100644 --- a/SOAP/particle_selection/SO_properties.py +++ b/SOAP/particle_selection/SO_properties.py @@ -45,6 +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 ParticleDataCache +from SOAP.particle_selection.shared_halo_particle_data import ( + SharedHaloParticleData, +) def cumulative_mass_intersection(r: float, rho_dim: float, slope_dim: float) -> float: @@ -240,14 +244,10 @@ class SOParticleData: def __init__( self, - input_halo: Dict, - data: Dict, - types_present: List[str], + shared: SharedHaloParticleData, 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, @@ -257,27 +257,19 @@ 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 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. - 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 @@ -289,20 +281,37 @@ def __init__( - 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.shared = shared + shared.compute_mass_profile(cosmology) + + # Quantities that are the same for every SO variation of this halo. + 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 = cosmology + self.boxsize = 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.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.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() def get_dataset(self, name: str) -> unyt.unyt_array: """ @@ -310,59 +319,16 @@ def get_dataset(self, name: str) -> unyt.unyt_array: """ 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_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. - 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. + Uses the cumulative mass profile computed once for this halo by + 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. Parameters: - reference_density: unyt.unyt_quantity @@ -376,48 +342,11 @@ def compute_SO_radius_and_mass( Rethrows any SearchRadiusTooSmallError thrown by find_SO_radius_and_mass(). """ - # 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)) - 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) + # 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 +386,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 @@ -3224,6 +3149,9 @@ class SOProperties(HaloProperty): sattelites. """ + # SOs always use every particle in the search radius + inclusive = True + """ List of properties from the table that we want to compute. Each property should have a corresponding method/property/lazy_property in @@ -3572,6 +3500,7 @@ def calculate( search_radius: unyt.unyt_quantity, data: Dict, halo_result: Dict, + shared_particle_data: ParticleDataCache = None, ): """ Compute spherical masses and overdensities for a halo @@ -3583,6 +3512,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. """ @@ -3625,17 +3558,20 @@ 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. + shared = self.get_shared_particle_data( + input_halo, data, shared_particle_data + ) 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, @@ -3890,6 +3826,7 @@ def calculate( search_radius: unyt.unyt_quantity, data: Dict, halo_result: Dict, + shared_particle_data: ParticleDataCache = None, ): """ Calculate the properties of an SO of which the radius is the multiple of @@ -3907,6 +3844,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: 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. Throws a RuntimeError if the "parent" SO radius cannot be obtained from halo_result. @@ -3929,5 +3870,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..625579c4 100644 --- a/SOAP/particle_selection/aperture_properties.py +++ b/SOAP/particle_selection/aperture_properties.py @@ -172,6 +172,10 @@ 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 ParticleDataCache +from SOAP.particle_selection.shared_halo_particle_data import ( + SharedHaloParticleData, +) class ApertureParticleData: @@ -200,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, ): @@ -217,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 @@ -237,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() @@ -269,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]: @@ -547,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: @@ -787,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: @@ -1714,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: @@ -4027,6 +3978,7 @@ def calculate( search_radius: unyt.unyt_quantity, data: Dict, halo_result: Dict, + shared_particle_data: ParticleDataCache = None, ): """ Compute centre of mass etc of bound particles @@ -4039,6 +3991,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. The halo_result dictionary is updated with the properties computed by this function. @@ -4142,18 +4098,17 @@ def calculate( "Search radius is smaller than aperture" ) - types_present = [type for type in self.particle_properties if type in data] + # Concatenated arrayss are computed once for inclusive/exclusive apertures + shared = self.get_shared_particle_data( + input_halo, data, shared_particle_data + ) + 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/halo_properties.py b/SOAP/particle_selection/halo_properties.py index 1bca96b8..191270f9 100644 --- a/SOAP/particle_selection/halo_properties.py +++ b/SOAP/particle_selection/halo_properties.py @@ -34,6 +34,75 @@ 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. + 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.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.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. + + 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 58e3d324..55e027ea 100644 --- a/SOAP/particle_selection/projected_aperture_properties.py +++ b/SOAP/particle_selection/projected_aperture_properties.py @@ -32,6 +32,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.shared_particle_data import ParticleDataCache +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, @@ -55,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() @@ -97,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 @@ -385,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: @@ -429,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: @@ -1137,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: @@ -1587,6 +1559,9 @@ class ProjectedApertureProperties(HaloProperty): the halo along the projection axis. """ + # projected apertures always use the particles bound to the halo + 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. @@ -1819,6 +1794,7 @@ def calculate( search_radius: unyt.unyt_quantity, data: Dict, halo_result: Dict, + shared_particle_data: ParticleDataCache = None, ): """ Compute centre of mass etc of bound particles @@ -1831,6 +1807,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. The halo_result dictionary is updated with the properties computed by this function. @@ -1919,13 +1899,16 @@ 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. + shared = self.get_shared_particle_data( + input_halo, data, shared_particle_data + ) + 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..f275c0b6 --- /dev/null +++ b/SOAP/particle_selection/shared_halo_particle_data.py @@ -0,0 +1,314 @@ +#!/bin/env python + +""" +shared_halo_particle_data.py + +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 +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 + # 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. + 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 + 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. + + 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. + """ + 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) + + def fofid_of_particle(self, index: int) -> int: + """ + FOF group ID of a single particle in the concatenated arrays. + + 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. + """ + 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): + """ + 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. + + 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): + """ + 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 + ) + 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_central = self.fofid_of_particle(non_neutrino_order[0]) + 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 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)) + 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.compute_bound_masks(fofid_central) + self.have_mass_profile = True diff --git a/SOAP/particle_selection/subhalo_properties.py b/SOAP/particle_selection/subhalo_properties.py index f5bb7e1f..2d03cf9e 100644 --- a/SOAP/particle_selection/subhalo_properties.py +++ b/SOAP/particle_selection/subhalo_properties.py @@ -54,6 +54,9 @@ 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, +) class SubhaloParticleData: @@ -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: @@ -2352,6 +2324,9 @@ class SubhaloProperties(HaloProperty): gravitationally bound particles. """ + # the bound subhalo uses the particles bound to the halo + inclusive = False + """ List of properties from the table that we want to compute. Each property should have a corresponding method/property/lazy_property in @@ -2584,7 +2559,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,20 +2572,23 @@ 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. 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. + shared = self.get_shared_particle_data(input_halo, data, shared_particle_data) 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, ) 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,