NFWModel#

jeanspy.model_jax.NFWModel

See methods and properties for individual lookup pages, or the alphabetical API dictionary to search all classes. The full class contract and existing member anchors are retained below.

class jeanspy.model_jax.NFWModel(submodels=None)[source]#

Bases: jeanspy.model_jax.DMModel

Functional JAX NFW density and mass with a hard cutoff at r_t_pc.

params contains positive rs_pc and r_t_pc in pc and rhos_Msunpc3 in Msun/pc^3. rs and rhos must be finite; r_t_pc may be infinite for finite-radius predictions. mass_density_3d is rho_s/(x*(1+x)**2) inside the cutoff (including its boundary), zero outside, and divergent at the center. resolve_params also returns the mass normalization nfw_mass_coeff in Msun.

enclosed_mass(r_pc, params=..., method=...) returns Msun inside min(r_pc,r_t_pc). Both auto and analytic use the stable analytic NFW mass; numeric uses a radial quadrature controlled by n_steps. Density and mass preserve radius shape. Invalid dynamic domains yield NaN; unsupported method names raise ValueError.

JAX tracing and physical-parameter differentiation work in smooth valid regions. The hard cutoff is not differentiable at its boundary. The runtime JAX precision/platform configuration applies to this class. See examples/docs_jax_spherical.py.

Parameters:

submodels (Dict[str, jeanspy.model_jax.Model])

required_param_names: tuple[str, ...] = ('rs_pc', 'rhos_Msunpc3', 'r_t_pc')#
analytic_enclosed_mass_autodiff_safe = True#
resolve_params(params)[source]#

Resolve NFW scales and cache their mass coefficient as JAX arrays.

Requires rs_pc and r_t_pc in pc, and rhos_Msunpc3 in Msun/pc^3; a missing key raises KeyError. Returns those entries and nfw_mass_coeff in Msun. Continuous expressions are differentiable; physical-domain validation is performed by the mass/likelihood paths that consume them.

Parameters:

params (Mapping[str, Any])

Return type:

dict[str, jax.Array]

mass_density_3d(r_pc, *, params)[source]#

Evaluate functional spherical halo density.

Notes

Inputs and units. r_pc is scalar/array in pc; params is a physical scalar dictionary.

Returns and shape. Density in Msun/pc^3 with radius shape; zero for r > r_t_pc, including the boundary r = r_t_pc in the halo. Invalid dynamic domains yield NaN; cusps can diverge at r=0.

Parameters:
Return type:

jax.Array

enclosed_mass_analytic(r_pc, *, params)[source]#

Evaluate the analytic NFW mass, with a stable small-radius limit.

Notes

Inputs and units. r_pc is scalar/array in pc; params contains rs_pc, rhos_Msunpc3 and r_t_pc. This direct analytic method takes no quadrature or method-selection arguments.

Returns and shape. Mass in Msun within min(r_pc, r_t_pc), matching radius shape. Invalid dynamic proposals yield NaN.

Parameters:
Return type:

jax.Array

enclosed_mass(r_pc, method='auto', *, params, n_steps=None)#

Return enclosed mass with a selectable analytic/numeric method.

The default auto method uses each model’s autodiff-safe default: analytic for NFW and numeric for Zhao. Pass method="analytic" to request a closed form explicitly, or method="numeric" for the autodiff-friendly numerical integral.

n_steps sets numerical mass resolution (Zhao: Gauss nodes per segment). None uses the model default; analytic methods ignore it.

Notes

Inputs and units. r_pc is scalar/array in pc; params is the physical dictionary. enclosed_mass accepts method=auto/analytic/numeric; numerical methods accept n_steps.

Returns and shape. Mass in Msun within min(r_pc, r_t_pc), matching radius shape. Invalid dynamic proposals yield NaN.

Parameters:
Return type:

jax.Array

enclosed_mass_numeric(r_pc, *, params, n_steps=256, t_min=1e-06)#

Numerically compute enclosed mass via 4π∫ρ(r)r²dr.

This default path is AD-friendly and avoids special-function gradient issues. If r_t_pc exists in params, radius is truncated at that value.

Notes

Inputs and units. r_pc is scalar/array in pc; params is the physical dictionary. n_steps is an integer >= 2 giving the trapezoidal grid size. t_min is the lower radial fraction, strictly between 0 and 1; density is integrated from t_min * min(r_pc, r_t_pc) to that outer radius. Out-of-range grid settings raise ValueError. To select a different mass method, call enclosed_mass instead.

Returns and shape. Mass in Msun within min(r_pc, r_t_pc), matching radius shape. Invalid dynamic proposals yield NaN.

Parameters:
Return type:

jax.Array

property has_analytic_enclosed_mass: bool#

Whether this halo exposes a callable analytic enclosed-mass implementation.

required_models: Mapping[str, type[jeanspy.model_jax.Model]] = {}#
sampling_identity()#

Model configuration without the derived compilation cache.

valid_mass_domain(r_pc, *, params)#

Dynamic validity mask; custom profiles may impose stricter domains.

submodels: Dict[str, jeanspy.model_jax.Model]#