Model#

jeanspy.model_jax.Model

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.Model(submodels=None)[source]#

Bases: object

Functional JAX model container with explicit physical-parameter mappings.

Notes

Inputs and units. submodels maps the exact role names declared by the concrete class; no physical values are stored here. Numerical methods receive explicit parameters.

Returns and shape. A component container with __getitem__ role access and sampling_identity without JIT cache.

Validity. Physical parameter dictionaries hold scalar values; radius arrays are broadcast independently. Use vmap to batch parameter dictionaries.

Errors. Missing/extra submodel roles raise ValueError.

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

Differentiation. Composition is static; differentiation applies to array-valued numerical method arguments.

Examples. examples/docs_jax_spherical.py

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

required_param_names: tuple[str, ...] = ()#

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

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

sampling_identity()[source]#: Model configuration without the derived compilation cache.