"""Prior and inference utilities for the NumPy/SciPy backend."""
from __future__ import annotations
from abc import ABCMeta, abstractmethod
from functools import cached_property
from multiprocessing.shared_memory import SharedMemory
import os
import numpy as np
import pandas as pd
from scipy.stats import norm, truncnorm
from ..parameters import SamplingParameter, _validate_parameter_specs, _photometry_coordinate
from .core import Model, logger
from .profiles import (ConstantAnisotropyModel, NFWModel, PlummerModel,
ProjectedExponentialModel, _EXP_HALF_LIGHT_FACTOR)
from .solver import DSphModel
class FittableModel(Model, metaclass=ABCMeta):
"""Subclassing interface for stateful likelihoods and prior terms.
Parameters
----------
args_load_data : list
Positional arguments forwarded to the concrete load_data method.
kwargs_load_data : dict or None, optional
Keyword arguments forwarded to load_data; None means an empty mapping.
*args, **kwargs
Model component and physical parameter initialization arguments.
Raises
------
TypeError
The data-loading arguments have the wrong container types, or an
abstract subclass has not implemented the required interface.
AttributeError
The initialized concrete model does not declare prior_names.
Notes
-----
Concrete subclasses define observation shapes and units, sampling-vector
order, conversion to physical parameters, and likelihood/prior terms.
Calling a target method updates the stateful components. lnposterior
returns the total log posterior followed by log likelihood and individual
log priors for emcee blobs. WBIC requires more than one observation.
This NumPy/SciPy host interface does not support physical-parameter JAX
tracing. See SphericalDSphEstimationModel for the spherical kinematic target
and ``examples/docs_inference.py`` for a complete short storage example.
"""
def __init__(self, args_load_data=None, kwargs_load_data=None, *args, **kwargs):
super().__init__(*args, **kwargs)
self.logger.info("Fittable Model: args_load_data: %r", args_load_data)
if not isinstance(args_load_data, list):
raise TypeError("args_load_data must be a list.")
if kwargs_load_data is None:
kwargs_load_data = {}
self.logger.info("Fittable Model: kwargs_load_data: %r", kwargs_load_data)
if not isinstance(kwargs_load_data, dict):
raise TypeError("kwargs_load_data must be a dict.")
self.load_data(*args_load_data, **kwargs_load_data)
if not hasattr(self, "prior_names"):
raise AttributeError("FittableModel must have the prior_names attribute.")
[docs]
@abstractmethod
def convert_params(self, p):
"""Map a sampling-coordinate vector p to named physical parameters.
Subclasses define vector order, transforms and units. This abstract
interface raises NotImplementedError.
"""
raise NotImplementedError
[docs]
@abstractmethod
def load_data(self, *args, **kwargs):
"""Load observations in a concrete estimation model.
Subclasses define the accepted arguments and validation. This abstract
interface raises NotImplementedError.
"""
raise NotImplementedError
@cached_property
def inverse_temperature(self):
"""Return the WBIC inverse temperature ``1/log(N_data)``."""
n_data = self.n_data if hasattr(self, "n_data") else len(self.data)
if n_data <= 1:
raise ValueError("WBIC requires at least two observations")
return 1 / np.log(n_data)
@abstractmethod
def _lnlikelihoods(self, *args, **kwargs):
raise NotImplementedError
[docs]
def lnlikelihoods(self, p, *args, **kwargs):
r"""Evaluate the explicit kinematic inference target.
Notes
-----
**Inputs and units.** p is one parameter vector of shape (ndim,) in
``p_names_lnprob`` order. Uses already loaded observations.
**Returns and shape.** Per-star log densities, shape (N,).
"""
params = self.convert_params(p)
self.update(params)
return self._lnlikelihoods(*args, **kwargs)
def _lnlikelihood(self, *args, **kwargs):
value = np.sum(self._lnlikelihoods(*args, **kwargs))
return -np.inf if np.isnan(value) else value
[docs]
def lnlikelihood(self, p, *args, **kwargs):
r"""Evaluate the explicit kinematic inference target.
Notes
-----
**Inputs and units.** p is one parameter vector of shape (ndim,) in
``p_names_lnprob`` order. Uses already loaded observations.
**Returns and shape.** Scalar sum of log likelihoods.
"""
params = self.convert_params(p)
self.update(params)
return self._lnlikelihood(*args, **kwargs)
@abstractmethod
def _lnpriors(self, p, *args, **kwargs):
raise NotImplementedError
[docs]
def lnpriors(self, p, *args, **kwargs):
r"""Evaluate the explicit kinematic inference target.
Notes
-----
**Inputs and units.** p is one parameter vector of shape (ndim,) in
``p_names_lnprob`` order. Uses already loaded observations.
**Returns and shape.** Sequence of log prior contributions in
``prior_names`` order.
"""
params = self.convert_params(p)
self.update(params)
return self._lnpriors(p, *args, **kwargs)
@property
def blobs_dtype(self):
"""Return emcee blob fields for log likelihood and individual log priors.
The list contains (name, float) pairs in the same order as the values
returned by lnposterior after its first element.
"""
return [("lnl", float), *((name, float) for name in self.prior_names)]
[docs]
def lnposterior(self, p, *args, **kwargs):
r"""Evaluate the explicit kinematic inference target.
Notes
-----
**Inputs and units.** p is one parameter vector of shape (ndim,) in
``p_names_lnprob`` order. Uses already loaded observations.
**Returns and shape.** Tuple (logposterior, loglikelihood,
individual prior terms) for emcee blobs.
"""
params = self.convert_params(p)
self.update(params)
lnl = -np.inf
lnp_list = self._lnpriors(p, *args, **kwargs)
if np.all([lnp > -np.inf for lnp in lnp_list]):
lnl = self._lnlikelihood(*args, **kwargs)
result = (lnl + np.sum(lnp_list), lnl, *lnp_list)
if np.isnan(result[0]):
self.logger.error("lnposterior is nan. lnl:%s, lnp_list:%s", lnl, lnp_list)
self.logger.error("p:%s", p)
self.logger.error("args:%s", args)
self.logger.error("kwargs:%s", kwargs)
self.logger.error("params:%s", params)
raise ValueError(
[
f"lnposterior is nan. lnl:{lnl}, lnp_list:{lnp_list}",
f"p:{p}",
f"args:{args}",
f"kwargs:{kwargs}",
f"params:{params}",
]
)
return result
[docs]
def lnposterior_wbic(self, p, *args, **kwargs):
r"""Evaluate the explicit kinematic inference target.
Notes
-----
**Inputs and units.** p is one parameter vector of shape (ndim,) in
``p_names_lnprob`` order. Uses already loaded observations.
**Returns and shape.** Tuple using loglikelihood/log(N) plus the original
prior; requires N>1.
"""
params = self.convert_params(p)
self.update(params)
lnl = -np.inf
lnp_list = self._lnpriors(p, *args, **kwargs)
if np.all([lnp > -np.inf for lnp in lnp_list]):
lnl = self._lnlikelihood(*args, **kwargs) * self.inverse_temperature
result = (lnl + np.sum(lnp_list), lnl, *lnp_list)
if np.isnan(result[0]):
raise ValueError(
"lnposterior_wbic is nan. "
f"lnl:{lnl}, lnp_list:{lnp_list}\np:{p}\n"
f"args:{args}\nkwargs:{kwargs}\nparams:{params}"
)
return result
@cached_property
def ndim(self):
"""Number of flattened physical parameters, cached on first access."""
return len(self.params_all)
class FlatPriorModel(Model):
r"""Finite uniform bounds in explicitly named sampling coordinates.
The DataFrame is the single source of truth for evaluation and sampling.
A generated template must be filled in before constructing this model.
Notes
-----
**Inputs and units.** config is a pandas DataFrame indexed by ordered
parameter names, with finite lower/upper columns, or a CSV path.
sample(size) uses NumPy's random state. ``write_config_template``
writes a CSV template.
**Returns and shape.** A validated prior object; sample returns coordinates
with trailing parameter axis. lower/upper are array copies.
``extract_value_by_name`` expects exactly one parameter vector.
**Validity.** Unique nonempty names and lower<upper. Bounds apply before
log/power transforms. Unfilled default NaN bounds are intentionally unusable
for inference.
**Errors.** Invalid schema/bounds/vector shape raise ValueError or
TypeError; missing CSV raises FileNotFoundError.
**Backend.** NumPy/SciPy CPU; stateful components, with no JAX tracing.
**Differentiation.** No physical-parameter automatic differentiation on this
API.
**Examples.** ``examples/docs_inference.py``
"""
required_param_names = []
required_models = {}
def __init__(self, config, show_init=False, submodels=None, **params):
super().__init__(show_init, submodels or {}, **params)
self.load_config(config)
[docs]
def load_config(self, config):
"""Load and copy uniform-prior bounds from a DataFrame or CSV path.
CSV input uses its first column as the parameter-name index. The bounds
are checked by validate_config before replacing stored data. Returns
None; file, parse and validation errors propagate.
"""
self.fname_config = (
os.fspath(config) if isinstance(config, (str, os.PathLike)) else None
)
if self.fname_config is not None:
try:
data = pd.read_csv(self.fname_config, index_col=0)
except FileNotFoundError:
logger.error("config file '%s' is not found.", config)
raise
else:
data = config
self.validate_config(data)
self.data = data.copy(deep=True)
[docs]
@staticmethod
def validate_config(data):
"""Validate a DataFrame of explicit finite uniform-prior bounds.
The nonempty index contains unique parameter names; columns must include
unique lower and upper bounds with lower < upper in each row. Returns
None. A non-DataFrame raises TypeError; invalid schema or bounds raise
ValueError. Use load_config to read a CSV path first.
"""
if not isinstance(data, pd.DataFrame):
raise TypeError("Prior config must be a DataFrame or CSV path.")
if data.empty or not data.index.is_unique or not data.columns.is_unique:
raise ValueError("Prior config must have nonempty, unique parameter names and columns.")
if any(not isinstance(name, str) or not name for name in data.index):
raise ValueError("Prior parameter names must be nonempty strings.")
if not {"lower", "upper"}.issubset(data.columns):
raise ValueError("Prior config needs lower and upper columns.")
bounds = data[["lower", "upper"]].to_numpy(dtype=float)
valid = np.isfinite(bounds).all(axis=1) & (bounds[:, 0] < bounds[:, 1])
if not valid.all():
raise ValueError(
"Supply explicit finite prior bounds with lower < upper for: "
+ ", ".join(data.index[~valid])
+ ". Fill in the prior template before inference."
)
@property
def lower(self):
"""Return a float array copy of lower bounds, shape (ndim,), in prior order."""
return self.data["lower"].to_numpy(dtype=float, copy=True)
@property
def upper(self):
"""Return a float array copy of upper bounds, shape (ndim,), in prior order."""
return self.data["upper"].to_numpy(dtype=float, copy=True)
[docs]
def get_index(self, param_name):
"""Return the index of the named parameter in the validated prior table.
An unknown param_name raises KeyError.
"""
return self.data.index.get_loc(param_name)
[docs]
def sample(self, size=None):
r"""Draw from the finite sampling-coordinate bounds.
Notes
-----
**Inputs and units.** size is a sample count, tuple of sample axes or None;
uses NumPy's global random state.
**Returns and shape.** Uniform coordinates with trailing parameter axis;
size=None returns one vector.
"""
self.validate_config(self.data)
size = (size,) if isinstance(size, int) else size
size = size + (len(self.lower),) if isinstance(size, tuple) else size
try:
return np.random.uniform(self.lower, self.upper, size=size)
except OverflowError as exc:
message = f"OverflowError: lower:{self.lower}, upper:{self.upper}, size:{size}"
exc.args = (message,) + exc.args
raise
def _lnprior(self, p):
self.validate_config(self.data)
if np.shape(p) != (len(self.data),):
raise ValueError(f"Parameters must have shape ({len(self.data)},).")
lower = self.lower
upper = self.upper
return 0.0 if np.all((lower <= p) & (p <= upper)) else -np.inf
[docs]
@staticmethod
def write_config_template(fname, param_names, lower=np.nan, upper=np.nan):
"""Write a template; unspecified bounds deliberately cannot be sampled."""
df = pd.DataFrame({"lower": lower, "upper": upper}, index=param_names)
df.to_csv(fname)
logger.info("generated %s.", fname)
return df
class PhotometryPriorModel(Model):
r"""Gaussian prior for ``log10(re_pc)``.
Notes
-----
**Inputs and units.** loc and scale are location and standard deviation in
log10(pc); sample(size) uses SciPy's random state.
``reset_prior(loc,scale)`` replaces that distribution.
**Returns and shape.** A prior object; sample returns log10 radii, not
physical pc.
**Validity.** Finite loc and positive finite scale are required by the
estimation model.
**Errors.** Invalid prior values are rejected when composing
SphericalDSphEstimationModel.
**Backend.** NumPy/SciPy CPU; stateful components, with no JAX tracing.
**Differentiation.** No physical-parameter automatic differentiation on this
API.
**Examples.** ``examples/docs_inference.py``
"""
required_param_names = []
required_models = {}
def __init__(self, loc, scale, show_init=False, submodels=None, **params):
super().__init__(show_init, submodels or {}, **params)
self.logger.info(
"%s:%r",
self.__class__.__name__,
{"log10_re_pc": loc, "e_log10_re_pc": scale},
)
self.reset_prior(loc, scale)
[docs]
def reset_prior(self, loc, scale):
"""Replace the Gaussian prior on log10 half-light radius in pc.
``loc`` and positive ``scale`` are the mean and standard deviation in
log10(pc). Returns None and replaces the stored log-PDF and sampler.
This helper delegates domain behavior to scipy.stats.norm.
"""
self.loc, self.scale = loc, scale
self._lnprior_func = norm(loc=loc, scale=scale).logpdf
self._sample = norm(loc=loc, scale=scale).rvs
def _lnprior(self, log10_re_pc):
return self._lnprior_func(log10_re_pc)
[docs]
def sample(self, size):
r"""Draw a log-radius prior value.
Notes
-----
**Inputs and units.** size is a sample count/shape or None, following SciPy
normal-distribution sampling.
**Returns and shape.** Samples in log10(pc), not physical radii; shape
follows size.
"""
return self._sample(size=size)
[docs]
def sampling_identity(self, sampled_names=()):
"""Return the Gaussian photometric-prior location and scale in log10(pc).
The host metadata dictionary excludes random/runtime state.
sampled_names is accepted for the common interface and is unused.
"""
return {"loc": self.loc, "scale": self.scale}
class DotDict(dict):
r"""Dictionary with attribute access to existing keys.
Notes
-----
**Inputs and units.** An optional mapping plus keyword values; keys are
strings and units belong to the stored values.
**Returns and shape.** A dict subclass: mapping operations and keys()/values()
retain dict behavior. Reading an attribute falls back to an existing key;
assigning or deleting an attribute changes an existing key if present.
**Validity.** Container operations have no physical validation. Assignment
to a new attribute creates an ordinary instance attribute, not a mapping key.
Use item assignment to insert keys. Dict method names take precedence over
attribute access to identically named keys.
**Errors.** Missing mapping keys raise KeyError; missing attributes raise
AttributeError.
**Backend.** Python host-side containers.
**Differentiation.** No physical-parameter automatic differentiation on this
API.
**Examples.** ``DotDict({"R_pc": [10., 20.]}).R_pc`` returns the stored list.
"""
def __getattr__(self, key):
if key in self:
return self[key]
raise AttributeError(key)
def __setattr__(self, key, value):
if key in self:
self[key] = value
else:
super().__setattr__(key, value)
def __delattr__(self, key):
if key in self:
del self[key]
else:
super().__delattr__(key)
class SphericalDSphEstimationModel(FittableModel, Model):
r"""Spherical Jeans inference with a Gaussian LOS-velocity likelihood.
Notes
-----
**Inputs and units.** SphericalDSphEstimationModel composes DSphModel,
FlatPriorModel and PhotometryPriorModel. ``args_load_data=[data]`` supplies
a DataFrame with ``R_pc``, ``vlos_kms`` and ``e_vlos_kms``;
``kwargs_load_data`` may contain shared=True. Parameter vectors follow
``p_names_lnprob`` exactly. ``parameter_specs`` is an ordered sequence of
:class:`jeanspy.parameters.SamplingParameter` objects specifying physical
names and transforms. Without specifications, names map by identity only.
**Returns and shape.** lnlikelihoods gives (N,) log densities; lnlikelihood
sums them. lnpriors returns prior terms. lnposterior returns (logposterior,
loglikelihood, individual prior terms) for emcee blobs. sample draws starting
coordinates; ``sample_data`` simulates velocities at supplied positions.
**Validity.** Nonempty finite 1-D data, R>0, error>=0; mean/error in km/s.
``dtype=None`` preserves the common floating dtype of the three input
columns (integer-only data use float64). An explicit floating ``dtype``
selects storage precision. Shared buffers use that same dtype and cannot
change dtype on reset. Numerical solvers may promote arithmetic precision.
``vmem_prior_from_data`` defaults to False. WBIC uses
``inverse_temperature = 1/log(N)`` and requires N>1. Shared data cannot be
resized.
**Errors.** Invalid prior order/schema/data raise ValueError. FittableModel
requires a list ``args_load_data``. Shared buffers must be released with
``release_shared_memory`` after workers stop.
**Backend.** NumPy/SciPy CPU; stateful components, with no JAX tracing.
**Differentiation.** No physical-parameter automatic differentiation on this
API.
**Examples.** ``examples/docs_inference.py``
"""
required_param_names = []
required_models = {
"DSphModel": DSphModel,
"FlatPriorModel": FlatPriorModel,
"PhotometryPriorModel": PhotometryPriorModel,
}
prior_names = ["flat_prior", "photometry_prior"]
def __init__(self, *args, parameter_specs=None, dtype=None,
vmem_prior_from_data=False, **kwargs):
self.parameter_specs = None if parameter_specs is None else tuple(parameter_specs)
self._requested_dtype = None if dtype is None else np.dtype(dtype)
if self._requested_dtype is not None and self._requested_dtype.kind != "f":
raise ValueError("Observation dtype must be a real floating dtype")
self.dtype = self._requested_dtype
self.vmem_prior_from_data = vmem_prior_from_data
super().__init__(*args, **kwargs)
self._validate_prior_schema()
[docs]
def sampling_identity(self):
"""Describe the persisted target using observations, priors and parameter order.
Returns host metadata for the sampler's identity checks. Shared-memory
handles, loggers and cached runtime state are excluded; observation
values are included for both shared and ordinary storage.
"""
# All physical coordinates are sampled by this model's validated schema.
# Shared-memory handles and cached WBIC temperature are runtime state;
# the observations themselves are hashed for shared and ordinary models.
ignored = {"logger", "_parammap", "params", "submodels", "_data",
"shared", "shared_shape", "buffer_size", "inverse_temperature",
"shm_R_pc", "shm_vlos_kms", "shm_e_vlos_kms"}
state = {k: v for k, v in vars(self).items() if k not in ignored}
state["data"] = self.data
state["parameter_order"] = self.p_names_lnprob
state["submodels"] = {
k: (type(v), v.sampling_identity(self.required_param_names_combined))
for k, v in self.submodels.items()
}
return state
def _validate_prior_schema(self):
prior = self["FlatPriorModel"]
prior.validate_config(prior.data)
names = self.p_names_lnprob
self.parameter_specs = _validate_parameter_specs(self.parameter_specs, names)
physical = [spec.param_name for spec in self.parameter_specs]
if set(physical) != set(self.required_param_names_combined):
raise ValueError(
"Parameter specifications must match model parameters exactly: "
f"expected {self.required_param_names_combined}, got {physical}."
)
exponential = isinstance(self["DSphModel"]["StellarModel"], ProjectedExponentialModel)
radius_name = "r_exp_pc" if exponential else "re_pc"
self._photometry_index = _photometry_coordinate(self.parameter_specs, radius_name)
self._photometry_offset = np.log10(_EXP_HALF_LIGHT_FACTOR) if exponential else 0.0
photometry = self["PhotometryPriorModel"]
if not np.isfinite(photometry.loc) or not np.isfinite(photometry.scale) or photometry.scale <= 0:
raise ValueError("Photometry prior needs a finite location and positive finite scale.")
@property
def p_names_lnprob(self):
"""Return sampling-coordinate names in the exact order expected by lnposterior."""
return self["FlatPriorModel"].data.index.tolist()
[docs]
def convert_params(self, p):
r"""Map sampling coordinates to physical parameters.
Notes
-----
**Inputs and units.** One parameter vector in exact prior order;
``parameter_specs`` explicitly supplies each transformation.
**Returns and shape.** A Series indexed by physical parameter names,
with units defined by the composed spherical model. Priors remain in
the sampled coordinates; conversion does not add a Jacobian.
"""
self._validate_prior_schema()
if np.shape(p) != (len(self.parameter_specs),):
raise ValueError(f"Parameters must have shape ({len(self.parameter_specs)},) in prior config order.")
return pd.Series({spec.param_name: spec.to_physical(value).item()
for spec, value in zip(self.parameter_specs, p)})
[docs]
def load_data(self, data, shared=False):
"""Load explicitly supplied observed kinematic data."""
# Reject schema mistakes before allocating shared observation buffers.
self._validate_prior_schema()
previous_shared = getattr(self, "shared", False)
self.shared = shared
try:
self.reset_data(data)
except Exception:
self.shared = previous_shared
raise
[docs]
def reset_data(self, data):
"""Replace observations while sampler workers are idle.
Shared models keep their buffer shape so existing readers stay attached.
Construct a new model to use a different number of observations.
Explicit velocity-prior bounds are preserved unless the model was
constructed with vmem_prior_from_data=True (an empirical-prior choice).
"""
# Validate optional empirical bounds before committing a data update.
prior = self["FlatPriorModel"]
updated_prior = prior.data.copy(deep=True)
if self.vmem_prior_from_data:
velocity_spec = next((spec for spec in self.parameter_specs
if spec.param_name == "vmem_kms" and spec.transform == "identity"), None)
if velocity_spec is None:
raise ValueError("Data-derived velocity bounds require an identity vmem_kms specification.")
velocities = np.asarray(data["vlos_kms"], dtype=self._observation_dtype(data))
updated_prior.loc[velocity_spec.sample_name, ["lower", "upper"]] = [
velocities.min(), velocities.max()
]
prior.validate_config(updated_prior)
self.data = data
if self.vmem_prior_from_data:
prior.data = updated_prior
self.__dict__.pop("inverse_temperature", None)
@property
def shared_memory_basename(self):
"""Return the observation-buffer name for this instance, or None if unshared."""
if not self.shared:
return None
return f"SphericalDSphEstimationModel_{id(self)}"
@property
def data(self):
"""Access stored radii, velocities and velocity errors as named arrays.
Columns R_pc, vlos_kms and e_vlos_kms have shape (N,) and units pc, km/s
and km/s. Shared mode returns views of the shared buffers and raises
FileNotFoundError or AttributeError if they have not been initialized.
"""
if not self.shared:
return self._data
try:
return DotDict(
{
"R_pc": np.ndarray(
self.shared_shape,
dtype=self.dtype,
buffer=self.shm_R_pc.buf,
),
"vlos_kms": np.ndarray(
self.shared_shape,
dtype=self.dtype,
buffer=self.shm_vlos_kms.buf,
),
"e_vlos_kms": np.ndarray(
self.shared_shape,
dtype=self.dtype,
buffer=self.shm_e_vlos_kms.buf,
),
}
)
except (FileNotFoundError, AttributeError):
self.logger.error(
"SharedMemory '%s' is not initialized yet.",
self.shared_memory_basename,
)
raise
@property
def n_data(self):
"""Return the number of observed stars as an integer."""
return self._n_data
def _observation_dtype(self, data):
if self._requested_dtype is not None:
return self._requested_dtype
dtype = np.result_type(*(data[field].dtype for field in ("R_pc", "vlos_kms", "e_vlos_kms")))
if dtype.kind not in "fiu":
raise ValueError("Kinematic columns must have real numeric dtypes")
return dtype if dtype.kind == "f" else np.dtype(np.float64)
@data.setter
def data(self, data: pd.DataFrame):
fields = ("R_pc", "vlos_kms", "e_vlos_kms")
shape = data["R_pc"].shape
if self.shared and hasattr(self, "shared_shape"):
if shape != self.shared_shape:
raise ValueError(
"Cannot resize shared kinematic data; construct a new model "
"for a different number of observations."
)
dtype = self._observation_dtype(data)
if self.shared and hasattr(self, "shared_shape") and dtype != self.dtype:
raise ValueError("Cannot change shared observation dtype; construct a new model "
"or explicitly select the existing dtype when creating the model.")
values = {field: data[field].to_numpy(dtype=dtype, copy=True) for field in fields}
if any(array.shape != shape for array in values.values()):
raise ValueError("Kinematic columns must have matching shapes.")
if len(shape) != 1 or not len(data) or not all(np.isfinite(v).all() for v in values.values()):
raise ValueError("Kinematic data must contain nonempty finite one-dimensional columns.")
if np.any(values["R_pc"] <= 0) or np.any(values["e_vlos_kms"] < 0):
raise ValueError("Kinematic data require R_pc > 0 and e_vlos_kms >= 0.")
if not self.shared:
self.dtype = dtype
self._data = DotDict(values)
self._n_data = len(data)
return
buffer_size = values["R_pc"].nbytes
handles = {}
arrays = {}
opened = []
created = []
try:
# Validate every segment before changing any observations or metadata.
for field in fields:
shm_name = self.shared_memory_basename + "_" + field
shm = getattr(self, f"shm_{field}", None)
if shm is None:
try:
shm = SharedMemory(name=shm_name, create=True, size=buffer_size)
created.append(shm)
except FileExistsError:
shm = SharedMemory(name=shm_name, create=False)
opened.append(shm)
if shm.size != buffer_size:
raise ValueError(
f"Shared memory {shm.name!r} has size {shm.size} bytes; "
f"expected {buffer_size} bytes for {field}."
)
handles[field] = shm
arrays[field] = np.ndarray(shape, dtype=dtype, buffer=shm.buf)
except Exception:
arrays.clear()
for shm in opened:
shm.close()
for shm in created:
shm.unlink()
raise
for field in fields:
arrays[field][:] = values[field]
setattr(self, f"shm_{field}", handles[field])
self.dtype = dtype
self._n_data = len(data)
self.shared_shape = shape
self.buffer_size = buffer_size
def _release_shared_memory(self, suffix):
if not self.shared:
return
name = self.shared_memory_basename + suffix
if not hasattr(self, "shared_shape"):
raise ValueError(
f"{self.__class__.__name__}: try to release shared memory "
f"{name} before initialization."
)
try:
shm = getattr(self, f"shm{suffix}")
shm.close()
shm.unlink()
self.logger.info("shared memory '%s' is released.", name)
except FileNotFoundError:
self.logger.info("shared memory '%s' is already released.", name)
[docs]
def release_shared_memory(self):
"""Close and release this model's three shared observation buffers.
Returns None. Existing array views must no longer be used after their
buffers are released.
"""
self._release_shared_memory("_R_pc")
self._release_shared_memory("_vlos_kms")
self._release_shared_memory("_e_vlos_kms")
def _lnlikelihoods(self):
s2 = self["DSphModel"].sigmalos2(self.data.R_pc)
err2 = self.data.e_vlos_kms**2
vmem_kms = self["DSphModel"].params.vmem_kms
return norm.logpdf(
self.data.vlos_kms,
loc=vmem_kms,
scale=np.sqrt(s2 + err2),
)
def _lnpriors(self, p_before_conversion):
idx_log10_re_pc = self._photometry_index
log10_re_pc = p_before_conversion[idx_log10_re_pc] + self._photometry_offset
return [
self["FlatPriorModel"]._lnprior(p_before_conversion),
self["PhotometryPriorModel"]._lnprior(log10_re_pc),
]
[docs]
def sample(self, size=None):
"""Draw sampling-coordinate vectors from the specified joint prior.
Uniform bounds apply to every coordinate; the log-radius coordinate
is drawn from the product of those bounds and the Gaussian photometric
prior. ``size=None`` returns (ndim,); a sample count/shape precedes that
parameter axis. Uses the NumPy/SciPy global random state. Invalid prior
schemas raise ValueError before sampling.
"""
self._validate_prior_schema()
p = self["FlatPriorModel"].sample(size)
idx_log10_re_pc = self._photometry_index
prior = self["FlatPriorModel"]
photometry = self["PhotometryPriorModel"]
loc, scale = photometry.loc - self._photometry_offset, photometry.scale
# Draw from the product of the Gaussian photometry prior and finite
# uniform support, so generated walkers always satisfy both priors.
a = (prior.lower[idx_log10_re_pc] - loc) / scale
b = (prior.upper[idx_log10_re_pc] - loc) / scale
p[..., idx_log10_re_pc] = truncnorm.rvs(a, b, loc=loc, scale=scale, size=size)
return p
[docs]
def sample_data(self, size=None):
"""Draw conditional Gaussian LOS velocities at the stored positions.
The mean is vmem_kms and the variance is the predicted LOS variance plus
the squared stored measurement error, in (km/s)^2. ``size=None`` returns
a broadcast vector of shape (N,); explicit sizes must be compatible with
that per-star shape. Uses SciPy's global random state. Positions and
measurement errors remain fixed; this is not a phase-space DF sampler.
"""
s2 = self["DSphModel"].sigmalos2(self.data.R_pc)
err2 = self.data.e_vlos_kms**2
vmem_kms = self["DSphModel"].params.vmem_kms
return norm.rvs(
loc=vmem_kms,
scale=np.sqrt(s2 + err2),
size=size,
)
def plummer_nfw_constant_anisotropy_model(
data,
photometry_prior_loc,
photometry_prior_scale,
config,
*,
vmem_prior_from_data=False,
dtype=None,
):
r"""Compose Plummer + NFW + constant anisotropy with explicit finite priors.
``config`` is a DataFrame or CSV in this order: vmem_kms, log10_re_pc,
log10_rs_pc, log10_rhos_Msunpc3, log10_r_t_pc, log10_one_minus_beta_ani.
The caller must supply config. A missing CSV raises FileNotFoundError;
this constructor never writes a prior template.
Caller velocity bounds are preserved unless vmem_prior_from_data is True.
The preset supplies explicit SamplingParameter objects for these names;
``dtype`` follows SphericalDSphEstimationModel's observation-storage contract.
Notes
-----
**Inputs and units.** data is an observed DataFrame; config is a finite
ordered prior DataFrame or CSV path; ``photometry_prior_loc``/scale specify
the Gaussian in log10(``re_pc``).
**Returns and shape.** SphericalDSphEstimationModel with the standard
Plummer/NFW/constant-anisotropy components.
**Validity.** A convenience constructor does not choose scientifically
justified prior bounds for the caller.
**Errors.** Missing CSV files raise FileNotFoundError. Invalid data or
prior bounds/order raise ValueError; omitted config raises TypeError.
**Backend.** NumPy/SciPy CPU; stateful components, with no JAX tracing.
**Differentiation.** No physical-parameter automatic differentiation on this
API.
**Examples.** ``examples/docs_inference.py``
"""
dsph_model = DSphModel(
submodels={
"StellarModel": PlummerModel(),
"DMModel": NFWModel(),
"AnisotropyModel": ConstantAnisotropyModel(),
}
)
names = ["vmem_kms", "log10_re_pc", "log10_rs_pc", "log10_rhos_Msunpc3",
"log10_r_t_pc", "log10_one_minus_beta_ani"]
prior = FlatPriorModel(config=config)
if prior.data.index.tolist() != names:
raise ValueError(f"Plummer/NFW/constant-anisotropy prior names/order must be {names}.")
return SphericalDSphEstimationModel(
args_load_data=[data],
dtype=dtype,
parameter_specs=[
SamplingParameter("vmem_kms", "vmem_kms"),
SamplingParameter("log10_re_pc", "re_pc", "pow10"),
SamplingParameter("log10_rs_pc", "rs_pc", "pow10"),
SamplingParameter("log10_rhos_Msunpc3", "rhos_Msunpc3", "pow10"),
SamplingParameter("log10_r_t_pc", "r_t_pc", "pow10"),
SamplingParameter("log10_one_minus_beta_ani", "beta_ani", "one_minus_pow10"),
],
vmem_prior_from_data=vmem_prior_from_data,
submodels={
"DSphModel": dsph_model,
"FlatPriorModel": prior,
"PhotometryPriorModel": PhotometryPriorModel(
loc=photometry_prior_loc,
scale=photometry_prior_scale,
),
},
)
__all__ = [
"DotDict",
"FittableModel",
"FlatPriorModel",
"PhotometryPriorModel",
"SphericalDSphEstimationModel",
"plummer_nfw_constant_anisotropy_model",
]