"""Read and write potentials to text (YAML) files."""
import os
import astropy.units as u
import numpy as np
import yaml
from gala.units import DimensionlessUnitSystem
__all__ = ["load", "save"]
def _unpack_params(p):
params = p.copy()
for key, item in p.items():
if "_unit" in key:
continue
if np.iterable(item) and not isinstance(item, str):
params[key] = np.array(item).astype(float)
elif isinstance(item, str):
params[key] = item
else:
try:
params[key] = float(item)
except Exception:
params[key] = item
if key + "_unit" in params:
params[key] *= u.Unit(params[key + "_unit"])
del params[key + "_unit"]
return params
def _parse_component(component, module):
# need this here for circular import
from .. import potential as gala_potential
try:
class_name = component["class"]
except KeyError as e:
raise KeyError(
"Potential dictionary must contain a key 'class' for "
"specifying the name of the Potential class."
) from e
if "units" not in component:
unitsys = None
else:
try:
unitsys = [u.Unit(unit) for ptype, unit in component["units"].items()]
except KeyError as e:
raise KeyError(
"Potential dictionary must contain a key 'units' "
"with a list of strings specifying the unit system."
) from e
params = component.get("parameters", {})
# need to crawl the dictionary structure and unpack quantities
params = _unpack_params(params)
potential = gala_potential if module is None else module
try:
Potential = getattr(potential, class_name)
except AttributeError: # HACK: this might be bad to assume
Potential = getattr(gala_potential, class_name)
# Add any extra potential arguments to the params kwargs
params = {**params, **component.get("extra_args", {})}
return Potential(units=unitsys, **params)
def from_dict(d, module=None):
"""
Convert a dictionary potential specification into a potential object.
This function parses a dictionary representation of a potential and
creates the corresponding :class:`~gala.potential.PotentialBase`
subclass instance. Supports both simple potentials and composite
potentials with multiple components.
Parameters
----------
d : dict
Dictionary specification of a potential. Must contain at minimum
a 'class' key specifying the potential class name. For composite
potentials, should include 'type': 'composite' and a 'components'
list of component dictionaries.
module : module, optional
Python module namespace to search for potential classes. If not
provided, uses `gala.potential`.
Returns
-------
potential : `~gala.potential.PotentialBase`
The instantiated potential object.
Examples
--------
Create a simple Hernquist potential::
>>> pot_dict = {'class': 'HernquistPotential', 'm': 1e11, 'c': 2.0}
>>> pot = from_dict(pot_dict)
Create a composite potential::
>>> comp_dict = {
... 'type': 'composite',
... 'class': 'CompositePotential',
... 'components': [
... {'class': 'HernquistPotential', 'm': 1e10, 'c': 1.0},
... {'class': 'NFWPotential', 'm': 1e12, 'r_s': 20.0}
... ]
... }
>>> comp_pot = from_dict(comp_dict)
"""
# need this here for circular import issues
import gala.potential as gala_potential
potential = gala_potential if module is None else module
if "type" in d and d["type"] == "composite":
p = getattr(potential, d["class"])()
for i, component in enumerate(d["components"]):
c = _parse_component(component, module)
name = component.get("name", str(i))
p[name] = c
elif "type" in d and d["type"] == "custom":
param_groups = {}
for component in d["components"]:
c = _parse_component(component, module)
try:
name = component["name"]
except KeyError as e:
raise KeyError(
"For custom potentials, component specification "
"must include the component name (e.g., name: "
"'blah')"
) from e
params = component.get("parameters", {})
params = _unpack_params(params) # unpack quantities
param_groups[name] = params
# Append any extra arguments
param_groups = {**param_groups, **d.get("extra_args", {})}
p = getattr(potential, d["class"])(**param_groups)
else:
p = _parse_component(d, module)
return p
# ----------------------------------------------------------------------------
def _pack_params(p):
params = p.copy()
for key, item in p.items():
if hasattr(item, "unit"):
params[key] = item.value
params[key + "_unit"] = str(item.unit)
if hasattr(params[key], "tolist"): # convert array to list
params[key] = params[key].tolist()
return params
def _to_dict_help(potential):
d = {}
d["class"] = potential.__class__.__name__
if not isinstance(potential.units, DimensionlessUnitSystem):
d["units"] = {str(k): str(v) for k, v in potential.units.to_dict().items()}
if len(potential.parameters) > 0:
params = _pack_params(potential.parameters)
d["parameters"] = params
if potential._extra_serialize_args:
d["extra_args"] = {}
for arg in potential._extra_serialize_args:
d["extra_args"][arg] = getattr(potential, arg)
return d
def to_dict(potential):
"""
Convert a potential object into a dictionary representation.
This function serializes a :class:`~gala.potential.PotentialBase`
object into a dictionary that fully specifies the potential's state,
including all parameters, units, and structure. The resulting
dictionary can be used to recreate the potential using
:func:`~gala.potential.io.from_dict`.
Parameters
----------
potential : :class:`~gala.potential.PotentialBase`
The instantiated potential object to convert to dictionary form.
Returns
-------
pot_dict : dict
Dictionary representation of the potential containing the class
name, parameters, units, and (for composite potentials) component
structure.
Examples
--------
Convert a simple potential to dictionary::
>>> pot = HernquistPotential(m=1e11*u.Msun, c=2*u.kpc)
>>> pot_dict = to_dict(pot)
Convert a composite potential::
>>> comp_pot = CompositePotential(bulge=hernquist, halo=nfw)
>>> comp_dict = to_dict(comp_pot)
See Also
--------
from_dict : Create potential object from dictionary representation.
save : Save potential to file.
"""
from .. import potential as gp
if isinstance(potential, gp.CompositePotential):
d = {}
d["class"] = potential.__class__.__name__
if potential._extra_serialize_args:
d["extra_args"] = {}
for arg in potential._extra_serialize_args:
d["extra_args"][arg] = getattr(potential, arg)
d["components"] = []
for k, p in potential.items():
comp_dict = _to_dict_help(p)
comp_dict["name"] = k
d["components"].append(comp_dict)
if potential.__class__.__name__ in {
"CompositePotential",
"CCompositePotential",
}:
d["type"] = "composite"
else:
d["type"] = "custom"
else:
d = _to_dict_help(potential)
return d
# ----------------------------------------------------------------------------
[docs]
def load(f, module=None):
"""
Load a potential from a YAML specification file.
This function reads a YAML file containing a potential specification
and creates the corresponding :class:`~gala.potential.PotentialBase`
object. The file format should match the dictionary structure expected
by :func:`~gala.potential.io.from_dict`.
Parameters
----------
f : str, file-like
Path to a YAML file, a block of YAML text, or a file-like object
containing the potential specification to parse and load.
module : module, optional
Python module namespace to search for potential classes. If not
provided, uses `gala.potential`.
Returns
-------
potential : :class:`~gala.potential.PotentialBase`
The loaded potential object.
Examples
--------
Load a potential from a YAML file::
>>> pot = load('my_potential.yml')
Load from a YAML string::
>>> yaml_spec = '''
... class: HernquistPotential
... parameters:
... m: 1.0e11
... c: 2.0
... '''
>>> pot = load(yaml_spec)
See Also
--------
save : Save potential to YAML file.
from_dict : Create potential from dictionary specification.
"""
if hasattr(f, "read"):
p_dict = yaml.load(f.read(), Loader=yaml.Loader)
else:
with open(os.path.abspath(f), encoding="utf-8") as fil:
p_dict = yaml.load(fil.read(), Loader=yaml.Loader)
return from_dict(p_dict, module=module)
[docs]
def save(potential, f):
"""
Save a potential object to a YAML file.
This function serializes a :class:`~gala.potential.PotentialBase`
object to YAML format and writes it to a file. The resulting file
can be loaded using :func:`~gala.potential.io.load`.
Parameters
----------
potential : :class:`~gala.potential.PotentialBase`
The potential object to save.
f : str, file-like
Output filename or file-like object to write the potential
specification to.
Examples
--------
Save a potential to file::
>>> pot = HernquistPotential(m=1e11*u.Msun, c=2*u.kpc)
>>> save(pot, 'hernquist_potential.yml')
Save to a string buffer::
>>> from io import StringIO
>>> buffer = StringIO()
>>> save(pot, buffer)
>>> yaml_content = buffer.getvalue()
See Also
--------
load : Load potential from YAML file.
to_dict : Convert potential to dictionary representation.
"""
d = to_dict(potential)
if hasattr(f, "write"):
yaml.dump(d, f, default_flow_style=None)
else:
with open(f, "w", encoding="utf-8") as f2:
yaml.dump(d, f2, default_flow_style=None)
