"""Support for moving a halo catalogue between processes, or storing it independently of its files.
A halo catalogue normally depends on the files written by a halo finder, and on the machinery in the
corresponding :class:`~pynbody.halo.HaloCatalogue` subclass for interpreting them. That is inconvenient when a
catalogue needs to be handed to another process (for example by a parallel task manager such as `tangos
<https://pynbody.github.io/tangos/>`_, which loads a catalogue once and shares it between workers).
To make this possible, :meth:`~pynbody.halo.HaloCatalogue.get_portable_state` expresses everything that a
catalogue knows -- its halo numbering, the particles belonging to each halo, what it knows about particles
that are missing from the snapshot, and any halo finder properties -- as a dictionary of numpy arrays and
python primitives. That includes the limits of its knowledge: a catalogue which cannot detect missing
particles says so, rather than appearing to have become able to. The recipient does not need to understand
the role of any of these arrays; it needs only to get them to the other end (e.g. by putting them into shared
memory), where :meth:`~pynbody.halo.HaloCatalogue.from_portable_state` turns them back into a live catalogue,
attached to a specified simulation, in the form of a :class:`PortableHaloCatalogue`.
.. versionadded:: 2.7.0
"""
from __future__ import annotations
import warnings
import numpy as np
from .. import array, units
from . import HaloCatalogue
from .details.number_mapping import HaloNumberMapper
from .details.particle_indices import HaloParticleIndices
#: The version of the portable state format generated by this version of pynbody. It is incremented if a
#: change is made that older versions of pynbody cannot interpret.
PORTABLE_STATE_VERSION = 1
_PRIMITIVE_TYPES = (str, bytes, bool, int, float, complex, np.number, np.bool_, type(None))
[docs]
class PortableHaloCatalogue(HaloCatalogue):
"""A halo catalogue reconstructed from the portable state of another halo catalogue.
Instances are normally created by :meth:`~pynbody.halo.HaloCatalogue.from_portable_state` rather than
directly. They behave like any other halo catalogue, except that they never refer back to the halo finder's
files; all the information they have is what was in the state dictionary they were created from.
.. versionadded:: 2.7.0
"""
_uses_file_position_addressing = False
"""Never true, whatever the originating catalogue did.
Halo membership arrives here as offsets into the snapshot the state was generated from, already resolved
by that catalogue; nothing is expressed relative to a file. The requirement is instead that *sim* presents
the same particles in the same order, which the caller asserts by passing it."""
[docs]
def __init__(self, sim, state: dict):
"""Create a halo catalogue for *sim* from the *state* dictionary of another catalogue.
Parameters
----------
sim : :class:`~pynbody.snapshot.simsnap.SimSnap`
The simulation to attach the catalogue to. It must have the same particle ordering as the
simulation the state was generated from, since halo membership is expressed as offsets into it.
state : dict
A dictionary previously returned by :meth:`~pynbody.halo.HaloCatalogue.get_portable_state`.
"""
version = state.get('version', PORTABLE_STATE_VERSION)
if version > PORTABLE_STATE_VERSION:
raise ValueError(f"This halo catalogue was exported in format version {version}, but this version "
f"of pynbody understands only up to version {PORTABLE_STATE_VERSION}")
super().__init__(sim, number_mapper=HaloNumberMapper.from_portable_state(state['number_mapper']))
self._index_lists = HaloParticleIndices.from_portable_state(state['particle_indices'])
self._original_class_name = state.get('catalogue_class')
# A catalogue which could not tell whether particles were missing does not become able to by being
# transferred, so the capability travels with the state. Without this, such a catalogue would report
# every halo as complete at the far end, and without the warning that says why.
self._can_determine_completeness = state.get('can_determine_completeness', True)
properties = state.get('properties')
if properties:
self._properties = properties_from_portable_state(properties)
def _get_all_particle_indices(self) -> HaloParticleIndices:
return self._index_lists
[docs]
def get_properties_all_halos(self, with_units=True) -> dict:
if self._properties is None:
return {}
elif with_units:
return self._properties
else:
return {k: _strip_units(v) for k, v in self._properties.items()}
def __repr__(self):
if self._original_class_name is None:
return super().__repr__()
else:
return f"<{type(self).__name__} from {self._original_class_name}, length {len(self)}>"
[docs]
def map_arrays(state, mapper, types=np.ndarray):
"""Return a copy of a portable state, with every numpy array in it replaced by ``mapper(array)``.
This is a convenience for code that transfers a state elsewhere, e.g. via shared memory. It saves the
caller from having to know the structure of the state, let alone the meaning of the individual arrays.
.. versionadded:: 2.7.0
Parameters
----------
state : dict | list | object
A portable state, as returned by :meth:`~pynbody.halo.HaloCatalogue.get_portable_state`, or any part
of one.
mapper : callable
A function which is called with each array in turn, and returns its replacement.
types : type | tuple[type]
The types to be passed to *mapper*; anything else is left alone. This defaults to numpy arrays, but
may be changed to reverse a mapping that has replaced them with something else.
"""
if isinstance(state, types):
return mapper(state)
elif isinstance(state, dict):
return {key: map_arrays(value, mapper, types) for key, value in state.items()}
elif isinstance(state, list):
return [map_arrays(value, mapper, types) for value in state]
else:
return state
[docs]
def properties_to_portable_state(properties: dict) -> dict:
"""Express a dictionary of halo properties using only numpy arrays and python primitives.
Properties which cannot be expressed in this way are omitted, with a warning. The result can be turned
back into a properties dictionary by :func:`properties_from_portable_state`."""
state = {}
for key, value in properties.items():
encoded = _property_to_portable_state(value)
if encoded is None:
warnings.warn(f"Halo property '{key}' is of a type that cannot be transferred between processes, "
f"and has been omitted", RuntimeWarning)
else:
state[key] = encoded
return state
[docs]
def properties_from_portable_state(state: dict) -> dict:
"""Recreate a dictionary of halo properties from one returned by :func:`properties_to_portable_state`."""
return {key: _property_from_portable_state(value) for key, value in state.items()}
def _property_to_portable_state(value):
"""Encode a single halo property, or return None if it cannot be encoded"""
unit_name = str(value.units) if units.has_unit(value) else None
if isinstance(value, np.ndarray) and value.dtype != object:
return {'kind': 'array', 'data': np.asarray(value), 'units': unit_name}
if isinstance(value, _PRIMITIVE_TYPES):
return {'kind': 'primitive', 'value': value, 'units': unit_name}
encoded = _sequence_to_portable_state(value)
if encoded is not None:
encoded['units'] = unit_name
return encoded
def _sequence_to_portable_state(value):
"""Encode a list, tuple or object array, or return None if it cannot be encoded.
Properties like 'children' are a ragged collection of arrays, one per halo. These are stored the same way
as halo membership itself, i.e. concatenated together with a list of boundaries. A simple sequence of
scalars is stored as a single array."""
if not isinstance(value, (list, tuple)) and not (isinstance(value, np.ndarray) and value.dtype == object):
return None
try:
rows = [np.asarray(row) for row in value]
except (TypeError, ValueError):
return None
if all(row.ndim == 0 and row.dtype != object for row in rows):
return {'kind': 'sequence', 'data': np.asarray(value)}
if any(row.ndim != 1 or row.dtype == object for row in rows):
return None
lengths = np.array([len(row) for row in rows], dtype=np.intp)
stops = np.cumsum(lengths)
boundaries = np.vstack((stops - lengths, stops)).T
non_empty = [row for row in rows if len(row) > 0]
if len(non_empty) > 0:
data = np.concatenate(non_empty)
elif len(rows) > 0:
data = rows[0]
else:
data = np.zeros(0, dtype=np.intp)
return {'kind': 'ragged', 'data': data, 'boundaries': boundaries}
def _property_from_portable_state(state: dict):
kind = state['kind']
unit_name = state.get('units')
if kind == 'primitive':
return _with_units(state['value'], unit_name)
elif kind == 'array':
return _with_units(np.asarray(state['data']), unit_name)
elif kind == 'sequence':
return [_with_units(value, unit_name) for value in np.asarray(state['data'])]
elif kind == 'ragged':
data = np.asarray(state['data'])
return [_with_units(data[start:stop], unit_name) for start, stop in state['boundaries']]
else:
raise ValueError(f"Unknown halo property encoding '{kind}'")
def _with_units(value, unit_name):
if unit_name is None:
return value
elif isinstance(value, np.ndarray):
return array.SimArray(value, units.Unit(unit_name))
else:
return value * units.Unit(unit_name)
def _strip_units(value):
if isinstance(value, list):
return [_strip_units(v) for v in value]
elif isinstance(value, np.ndarray):
return np.asarray(value)
else:
return value
