3. Define and extend a model

Download this notebook (.ipynb) · Read the saved outputs here, or run the cells in Jupyter.

A spherical Jeans model combines three components: a tracer number density, a gravitating halo, and velocity anisotropy. A sampler and its priors are separate objects. Start by constructing a forward model and checking its predictions at fixed physical parameters.

This notebook is self-contained. Install JeansPy with the numpyro_cpu and plotting extras as described in the installation guide, then select that environment as your Jupyter kernel and run cells from top to bottom. Saved outputs are an example run; timings and short-chain results can vary.

Compose the JAX model

The JAX interface used by the Quickstart separates the components from their physical values. The keys in submodels are role names, so preserve their spelling even when replacing the concrete classes.

import os
os.environ.setdefault("JEANSPY_JAX_PLATFORM", "cpu")
os.environ.setdefault("JEANSPY_JAX_ENABLE_X64", "true")
from jeanspy.model_jax import (
    DSphModel, PlummerModel, NFWModel, ConstantAnisotropyModel, get_runtime_config,
)
import jax
import jax.numpy as jnp
import numpy as np
import matplotlib.pyplot as plt

tracer, halo, anisotropy = PlummerModel(), NFWModel(), ConstantAnisotropyModel()
model = DSphModel(submodels={"StellarModel": tracer, "DMModel": halo,
                             "AnisotropyModel": anisotropy})
params = dict(re_pc=300., rs_pc=500., rhos_Msunpc3=.1,
              r_t_pc=5000., beta_ani=-.2)
R_pc = jnp.array([30., 100., 300.])
variance = jax.block_until_ready(model.sigmalos2(
    R_pc, params=params, solver="kernel", n_u=128, n_kernel=64))
assert variance.shape == (3,) and np.all(np.asarray(variance) > 0)
sampled_R = tracer.sample_R(jax.random.PRNGKey(20260913), 8, re_pc=300.)
assert sampled_R.shape == (8,)
print(get_runtime_config(), variance)
#

{'jax_platform_requested': 'cpu', 'jax_platform_effective_env': 'cpu', 'jax_backend_active': 'cpu', 'jax_enable_x64': True, 'kernel_backend_default': 'jax', 'constant_kernel_n_quad_default': 32, 'sigmalos2_solver_default': 'auto', 'sigmalos2_jit_default': True, 'sigmalos2_kernel_outer_transform_default': 'sqrtlog', 'baes_kernel_n_quad_default': 32, 'baes_eta_recommended_max': 10.0, 'sigmalos2_n_u_default': 128, 'sigmalos2_n_r_default': 768, 'sigmalos2_u_max_default': 10000.0} [91.44464827 89.19126387 74.77863921]

This notebook includes its own imports and runtime setup above. params holds scalar physical values. A new dictionary changes the prediction without mutating the model:

denser = {**params, "rhos_Msunpc3": 2 * params["rhos_Msunpc3"]}
variance_denser = model.sigmalos2(R_pc, params=denser, solver="kernel")

Logarithms and other sampling transforms belong in SamplingParameter (NumPy/SciPy) or ParameterSpec (NumPyro), not in the physical dictionary. For instance, rs_pc is a radius in pc even when the sampler uses a coordinate called log10_rs_pc.

Generalize the NFW inner slope

NFWModel fixes the inner and outer density slopes to 1 and 3. The Quickstart instead uses ZhaoModel with alpha=1, beta=3 and a variable gamma:

\[ \rho(r)=\frac{\rho_s}{(r/r_s)^gamma(1+r/r_s)^{3-gamma}}. \]

Here gamma=0 has a finite central density, gamma=1 recovers NFW and gamma=2 has a steeper cusp. rhos_Msunpc3 is the normalization \(\rho_s\); the density at \(r_s\) is \(\rho_s/2^{3-gamma}\). Keep a and b in the fixed physical dictionary, and sample gamma in SamplingParameter (NumPy/SciPy) or ParameterSpec (NumPyro). Use the differentiable numerical Zhao mass when fitting gamma; switching to NFWModel fixes this slope instead. The inference notebook demonstrates this fit.

Use the NumPy/SciPy interface

NumPy/SciPy spherical components store their parameters. All required physical values must be supplied; omitted values can remain NaN. This example supplies every parameter and evaluates the assembled model:

import numpy as np
from jeanspy.model import ConstantAnisotropyModel, DSphModel, NFWModel, PlummerModel

model = DSphModel(
    vmem_kms=0.,
    submodels={
        "StellarModel": PlummerModel(re_pc=200.),
        "DMModel": NFWModel(rs_pc=1000., rhos_Msunpc3=.01, r_t_pc=10000.),
        "AnisotropyModel": ConstantAnisotropyModel(beta_ani=0.),
    },
)
R_pc = np.array([50., 100., 300.])
variance = model.sigmalos2(R_pc)
sigma_kms = np.sqrt(variance)
assert variance.shape == R_pc.shape
assert np.all(np.isfinite(sigma_kms) & (sigma_kms > 0))
print(sigma_kms)
#

[4.42046919 4.21841733 4.22716713]

Use model["DMModel"] to access a component, and model.update(rhos_Msunpc3=.02) to change its density scale through the composite model. This mutation affects subsequent predictions. Use separate model instances when comparing parameter choices concurrently.

The profile guide describes tracer scales, deprojection domains, anisotropy families and halo cutoffs. The API catalogue separates the NumPy and JAX implementations.

Implement your own tracer

The following minimal NumPy/SciPy tracer implements the Plummer formula again so that it can be checked against a built-in reference. Replace these two consistent projected and spatial densities with your own profile once their normalization and projection relation are established.

import numpy as np
from jeanspy.model import StellarModel, PlummerModel, DSphModel, NFWModel, ConstantAnisotropyModel

class CustomPlummer(StellarModel):
    """Unit-normalized Plummer tracer with projected half-light radius re_pc."""
    required_param_names = ["re_pc"]
    required_models = {}

    def density_2d(self, R_pc):
        """Projected number density in pc^-2; accepts scalar or array radii."""
        re = self.params.re_pc
        return (1. + (np.asarray(R_pc) / re)**2)**-2 / (np.pi * re**2)

    def density_3d(self, r_pc):
        """Spatial number density in pc^-3; accepts scalar or array radii."""
        re = self.params.re_pc
        return 3. * (1. + (np.asarray(r_pc) / re)**2)**-2.5 / (4. * np.pi * re**3)

custom = CustomPlummer(re_pc=200.)
model = DSphModel(vmem_kms=0., submodels={
    "StellarModel": custom,
    "DMModel": NFWModel(rs_pc=500., rhos_Msunpc3=.1, r_t_pc=5000.),
    "AnisotropyModel": ConstantAnisotropyModel(beta_ani=0.),
})
R_pc = np.geomspace(10., 2000., 20)
variance = model.sigmalos2(R_pc, n=256, n_kernel=64)
#

required_param_names declares the scalar parameters maintained by the NumPy/SciPy Model container. required_models={} means that the component has no nested components. The LOS solver needs both density_2d and density_3d; implement them for scalar and broadcast array inputs. Extra helpers such as a radial CDF or sampling method are needed only when your own workflow uses them.

Check a custom implementation before inference

from scipy.integrate import quad
builtin = PlummerModel(re_pc=200.)
np.testing.assert_allclose(custom.density_2d(R_pc), builtin.density_2d(R_pc))
np.testing.assert_allclose(custom.density_3d(R_pc), builtin.density_3d(R_pc))
projected_total = quad(lambda R: 2*np.pi*R*custom.density_2d(R), 0., np.inf)[0]
spatial_total = quad(lambda r: 4*np.pi*r**2*custom.density_3d(r), 0., np.inf)[0]
np.testing.assert_allclose([projected_total, spatial_total], 1., rtol=1e-7)
assert np.all(np.isfinite(variance) & (variance > 0))
#

The checks compare the two densities against the built-in Plummer expression and integrate each to unit tracer number. For a new profile, also check that projecting the 3-D density reproduces the 2-D density, inspect central and outer limits, and refine the Jeans quadrature. These checks do not establish the existence of a physical distribution function.

For a custom NumPy/SciPy halo, implement a consistent mass density and enclosed mass with an explicit cutoff convention; consult jeanspy.model.DMModel. A custom anisotropy must implement the beta, f and kernel interfaces consistently; see jeanspy.model.AnisotropyModel.

JAX extensions use JAX operations and explicit parameters. The current spherical JAX solver calls tracer densities as density_2d(R, re_pc=...) and density_3d(r, re_pc=...): it does not pass an arbitrary tracer parameter dictionary. A tracer requiring additional sampled shape parameters therefore needs a compatible solver extension. A custom JAX halo must satisfy the DMModel mass interface; an all-JAX anisotropy kernel is needed for gradients. Defining an arbitrary Python class alone does not make it NUTS-compatible.

Extend the geometry

Axisymmetric components store their physical values in immutable objects; use dataclasses.replace or the supported per-call parameter mapping to change them. Their flattening and cylindrical anisotropy differ from the spherical roles above. Follow the axisymmetric guide and worked analysis when changing geometry.

Next: predict observables and generate mocks.