Source code for py21cmfast.input_serialization
"""Functions for handling InputParameters for various purposes, e.g. I/O and unstructuring."""
import warnings
from typing import Any, Literal
import attrs
from .utils import recursive_difference
from .wrapper._utils import camel_to_snake, snake_to_camel
from .wrapper.inputs import CosmoTables, InputParameters, InputStruct, Table1D
def _deserialize_cosmo_table(value: dict[str, Any] | Table1D | Any) -> Table1D | Any:
"""Deserialize a CosmoTables subfield value to Table1D when needed."""
return Table1D(**value) if isinstance(value, dict) else value
[docs]
def convert_inputs_to_dict(
inputs: InputParameters,
mode: Literal["full", "minimal"] = "full",
only_structs: bool = True,
camel: bool = True,
only_cstruct_params: bool = False,
use_aliases: bool = True,
include_cosmo_tables: Literal["always", "if_cached", "never"] = "always",
) -> dict[str, dict[str, Any]]:
"""Convert an InputParameters object to a dictionary, with various options.
While the :class:`~py21cmfast.wrapper.inputs.InputParameters.asdict` method allows
several options for "unstructuring" the inputs to a dictionary, this function is
much more flexible, having options that cover a number of use-cases.
Parameters
----------
inputs
The input parameters to convert to a dict.
mode
Either 'minimal' to get the minimal dictionary required to specify the
parameters (on top of defaults) or 'full' to keep all parameters.
only_structs
Whether to only return InputStruct objects (unstructured to dicts),
otherwise also return other attributes (e.g. random_seed).
camel
Whether the keys of the returned dict should be camel-case, e.g.
SimulationOptions. Otherwise, return as snake_case. Only applies to
InputStruct attributes of the inputs (i.e. not node_redshifts or random_seed).
only_cstruct_params
Only return parameters that are part of the Cstruct, rather than all
fields of the class. This is useful for pretty-printing.
use_aliases
If True, use correct aliases for parameters, which allows the dictionary
to be passed back into the InputStruct constructors (e.g. use DIM instead
of _DIM).
include_cosmo_tables
Controls whether derived ``cosmo_tables`` are emitted.
- ``"always"`` includes them and may materialize CLASS-derived tables.
- ``"if_cached"`` includes them only if already materialized on ``inputs``.
- ``"never"`` omits them.
"""
kw = {
"only_structs": only_structs,
"camel": camel,
"only_cstruct_params": only_cstruct_params,
"use_aliases": use_aliases,
"include_cosmo_tables": include_cosmo_tables,
}
all_inputs = inputs.asdict(**kw)
if mode == "minimal":
defaults = InputParameters(random_seed=0)
default_dct = defaults.asdict(**kw)
cosmo_tables_key = "CosmoTables" if camel else "cosmo_tables"
# we still want to keep cosmo_tables, even in minimal mode (since we want
# to keep ps_norm and USE_SIGMA_8), if they are present.
cosmo_tables_dct = None
if cosmo_tables_key in all_inputs:
cosmo_tables_dct = all_inputs[cosmo_tables_key].copy()
all_inputs = recursive_difference(all_inputs, default_dct)
if cosmo_tables_dct is not None:
all_inputs[cosmo_tables_key] = cosmo_tables_dct
return all_inputs
[docs]
def prepare_inputs_for_serialization(
inputs: InputParameters,
mode: Literal["full", "minimal"] = "full",
only_structs: bool = True,
camel: bool = True,
include_cosmo_tables: Literal["always", "if_cached", "never"] = "always",
) -> dict[str, dict[str, Any]]:
"""Prepare an inputs class for serialization (to e.g. TOML, YAML or HDF5).
This is a thin wrapper around :func:`~convert_inputs_to_dict` that also
ensures that 'None' values are removed (so long as their default is also None)
and that the parameter names map back to aliases of InputStruct attributes.
Parameters
----------
inputs
The input parameters to convert.
mode
Either ``"full"`` or ``"minimal"``. ``"minimal"`` keeps only differences
from defaults (plus any included ``cosmo_tables``).
only_structs
If True, include only input-struct fields.
camel
If True, use CamelCase top-level struct keys.
include_cosmo_tables
Forwarded to :func:`convert_inputs_to_dict`.
Use ``"if_cached"`` for cache writes to avoid triggering expensive CLASS
computations, and ``"never"`` for template files that should not carry
derived tables.
"""
dct = convert_inputs_to_dict(
inputs,
mode=mode,
only_structs=only_structs,
camel=camel,
use_aliases=False, # convert to aliases below instead.
include_cosmo_tables=include_cosmo_tables,
)
# dct is a dict of dicts, where each subdict represents a
# single input parameter struct (SimulationOptions, AstroParams etc).
# The rules are that _all_ of the input structs are in this dict, even if they are
# empty. The keys allowed in each struct are such that:
#
# 1. They are valid keys to pass to the class constructor, not necessarily
# all the fields required for the Cstruct (since sometimes these are
# auto-calcualted).
# 2. They are not necessarily exhaustive, especially if mode=="minimal".
# 3. The values are those that are required to *build the isntance*, not their
# final transformed values for use in C.
#
# However, some attributes of the InputStruct classes are not necessarily
# required for building the Cstruct. This can occur if there are two pathways of
# setting a particular necessary parameter (e.g. DIM can be set either by setting
# DIM directly, or by setting HII_DIM and HIRES_TO_LOWRES_FACTOR). In this case,
# at least one of the attributes will have a default value of None. If something
# has an actual value of None, we require that its default is None. This way,
# we can simply leave it out of the written TOML (or HDF5) when its value is None,
# and it will anyway be set to its own default if read back in.
out = {}
if "random_seed" in dct:
out["random_seed"] = dct.pop("random_seed")
if "node_redshifts" in dct:
out["node_redshifts"] = dct.pop("node_redshifts")
for structname, structvals in dct.items():
this = {}
clsname = snake_to_camel(structname)
if clsname in InputStruct._subclasses:
fields = attrs.fields_dict(InputStruct._subclasses[clsname])
elif clsname == "CosmoTables":
fields = attrs.fields_dict(CosmoTables)
for key, val in structvals.items():
if val is None:
if fields[key].default is not None: # pragma: nocover
# This should not be reachable because setting a required parameter
# to None should error on validation, rather than reaching here.
raise RuntimeError(
f"Detected that {structname} has {key}=None but it is not an optional parameter!"
)
else:
# Convert to use the alias, rather than the name (i.e. DIM instead of _DIM)
this[fields[key].alias] = val
if this:
out[structname] = this
# Furthermore, some
return out
[docs]
def deserialize_inputs(
dict_of_structdicts: dict[str, Any],
safe: bool = True,
include_cosmo_tables: bool = False,
**loose_params,
) -> dict[str, InputStruct]:
"""Construct a dictionary of InputStructs ready to be converted to InputParameters.
Parameters
----------
dict_of_structdicts
A dictionary whose keys are names of InputStruct attributes of
InputParameters (e.g. CosmoParams, SimulationOptions), and whose values
are dictionaries of parameters specific to each struct. Not every parameter
of every struct is required.
safe
Whether to raise on unrecognized/excess parameters.
include_cosmo_tables
Whether to deserialize the optional ``CosmoTables`` block if present.
This is intentionally a boolean because deserialization does not need the
tri-state behavior used when serializing.
Other Parameters
----------------
All other parameters are considered "loose" parameters of one of the InputStructs,
and will *override* the parameter matching their name if it is found in one of the
structs.
Returns
-------
dict_of_structs
A dictionary whose keys are attribute names of the InputParameters class,
and whose values are InputStruct instances, ready to instantiate an
InputParameters class, e.g. ``InputParameters(**dict_of_structs, random_seed=1)``.
"""
# It's possible that names in dict_of_structdicts are in camel case or snake
# case. We normalise to camel case here. This also de-references the input
# so when we modify it in-place (via .pop()) later, we don't mess with the user's
# input.
dict_of_structdicts = {
snake_to_camel(name): dct for name, dct in dict_of_structdicts.items()
}
input_dict = {}
extra_params = {}
for structname, kls in InputStruct._subclasses.items():
# Use field.alias instead of field.name because the alias is what needs
# to be passed to the class constructor (e.g. "DIM" instead of "_DIM")
fieldnames = [field.alias for field in attrs.fields(kls)]
kw_dict = {kk: loose_params.pop(kk) for kk in fieldnames if kk in loose_params}
these_all = dict_of_structdicts.pop(structname, {})
if structname == "CosmoTables":
these = {
kk: Table1D(**these_all[kk])
if isinstance(these_all[kk], dict)
else these_all[kk]
for kk in these_all
if kk in fieldnames
}
else:
these = {kk: these_all[kk] for kk in these_all if kk in fieldnames}
extra = {kk: these_all[kk] for kk in these_all if kk not in fieldnames}
# Here, if structname is not in the input dict_of_structdicts, it is OK,
# because we just assume an empty set of parameters, potentially added to
# by the loose params.
arg_dict = {**these, **kw_dict}
input_struct = kls.new(arg_dict)
input_dict[camel_to_snake(structname)] = input_struct
if extra:
extra_params[structname] = extra
if include_cosmo_tables:
structname = "CosmoTables"
fieldnames = [field.alias for field in attrs.fields(CosmoTables)]
kw_dict = {kk: loose_params.pop(kk) for kk in fieldnames if kk in loose_params}
these_all = dict_of_structdicts.pop(structname, {})
# Table1D instances may arrive either as serialized dicts (e.g. from HDF5/TOML)
# or as already-instantiated Table1D objects (e.g. in-memory call sites).
these = {
kk: _deserialize_cosmo_table(these_all[kk])
for kk in these_all
if kk in fieldnames
}
extra = {kk: these_all[kk] for kk in these_all if kk not in fieldnames}
arg_dict = {**these, **kw_dict}
if arg_dict:
input_dict[camel_to_snake(structname)] = CosmoTables.new(arg_dict)
if extra:
extra_params[structname] = extra
else:
dict_of_structdicts.pop("CosmoTables", None)
if dict_of_structdicts:
warnings.warn(
f"The following keys were not recognized for deserializing to InputParameters: {tuple(dict_of_structdicts.keys())}",
stacklevel=2,
)
if extra_params or loose_params:
all_extra = {**extra_params, **loose_params}
msg = f"Excess arguments exist: {all_extra}"
if safe:
raise ValueError(msg)
else:
warnings.warn(msg, stacklevel=2)
return input_dict