AxisymmetricDSphModel#

jeanspy.model_jax.AxisymmetricDSphModel

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.AxisymmetricDSphModel(n_force=96, n_vertical=96, n_los=96, submodels=None)[source]#

Bases: object

Compose functional JAX axisymmetric stellar, halo and anisotropy components.

submodels uses StellarModel=AxisymmetricPlummerModel(), DMModel=AxisymmetricZhaoModel() and AnisotropyModel=AxisymmetricConstantAnisotropyModel(). These are also the defaults. Physical parameters remain explicit in params on every call. Components are immutable and participate in JIT and sampling identity.

Fixed quadrature settings are static under JIT. lax.map with rematerialized per-star evaluation bounds intermediate memory when differentiating catalogs.

Notes

Inputs and units. params requires re_pc, rs_pc (pc), rhos_Msunpc3 (Msun/pc^3), and exactly one of q or q_projected. Optional Q, alpha, beta, gamma, beta_z and inclination (radians) have the defaults shown in the axisymmetric guide. r_t_pc is a positive ellipsoidal cutoff (pc). Use alpha/beta/gamma, matching the spherical Zhao profile. Physical parameter dictionaries hold scalar values; radius arrays are broadcast independently. Use vmap to batch parameter dictionaries. Constructor node counts n_force/n_vertical/n_los are static integers >=16.

Returns and shape. sigmalos2 and intrinsic_moments return (km/s)^2; density_3d is normalized pc^-3, surface_density pc^-2, mass_density_3d Msun/pc^3, enclosed_mass Msun inside an ellipsoid; potential_gradient is (km/s)^2/pc. Coordinates broadcast; intrinsic moments and forces are tuples of matching arrays.

Validity. Cylindrical alignment with constant beta_z; same physical restrictions as the NumPy axisymmetric solver. Scalar sky inputs yield scalars; centers are allowed for projected moments.

Errors. Invalid dynamic parameters yield NaN, including under jit. Schema/shape/configuration errors raise before evaluation; likelihood rejects invalid variances.

Backend. JAX arrays on the configured CPU/GPU, with dtype set before import.

Differentiation. Physical scalar parameters and supported coordinate values are differentiable in admissible smooth regions. Node counts, schema decisions, rejection masks and hard-cutoff boundaries are not continuous model parameters. J/D methods are not part of this JAX class.

Examples. examples/docs_axisymmetric.py; examples/docs_jax.py

Parameters:

n_force (int)
n_vertical (int)
n_los (int)
submodels (object)

n_force: int = 96#

n_vertical: int = 96#

n_los: int = 96#

submodels: object = None#

sampling_identity()[source]#

Return the three fixed quadrature orders used to identify a sampling target.

The host dictionary contains n_force, n_vertical and n_los. This metadata helper has no physical-parameter derivative.

intrinsic_moments(R_pc, z_pc, *, params)[source]#

Evaluate the intrinsic Jeans second moments.

Notes

Inputs and units. R_pc >= 0 and signed z_pc in pc, broadcastable; params supplies the physical dictionary.

Returns and shape. Tuple (vR2,vz2,vphi2), each in (km/s)^2 with the broadcast coordinate shape. vphi2 is the total azimuthal second moment; no rotation/dispersion split is assigned.

potential_gradient(R_pc, z_pc, *, params)[source]#

Evaluate derivatives of the gravitational potential.

Notes

Inputs and units. R_pc >= 0, signed z_pc in pc and explicit params.

Returns and shape. Tuple (dPhi/dR,dPhi/dz) in (km/s)^2/pc; gravitational acceleration has the opposite sign.

surface_density(x_pc, y_pc, *, params)[source]#

Evaluate the projected spheroidal Plummer tracer.

Notes

Inputs and units. Signed x_pc/y_pc in pc and explicit params; coordinates broadcast.

Returns and shape. Normalized surface density in pc^-2 with the broadcast coordinate shape.

density_3d(R_pc, z_pc, *, params)[source]#: Unit-normalized stellar density in pc^-3.

mass_density_3d(R_pc, z_pc, *, params)[source]#: Halo density, including the optional ellipsoidal cutoff.

enclosed_mass(m_pc, *, params, n_steps=128)[source]#

Mass in Msun inside the spheroid m <= m_pc (with truncation).

Notes

Inputs and units. m_pc >= 0 is the ellipsoidal radius in pc; params supplies halo scales, slopes, Q and cutoff.

Returns and shape. Msun inside R^2+z^2/Q^2<=``min(m_pc, r_t_pc)``^2, matching m_pc shape.

sigmalos2(x_pc, y_pc, *, params)[source]#

Project a cylindrically aligned second moment.

Notes

Inputs and units. Signed x_pc/y_pc in pc, broadcastable scalar/arrays; params is the explicit physical dictionary.

Returns and shape. LOS second moment in (km/s)^2 with the broadcast coordinate shape, including scalar output.