Skip to content

core

Core module for generating models and their gradients.

model = jax.jit(model, static_argnums=model_static) module-attribute

Generically create models with substructure.

Parameters:

Name Type Description Default
xyz tuple[Array, Array, Array, float, float]

Grid to compute model on. See containers.Model.xyz for details.

required
n_struct tuple[int, ...]

Number of each structure to use. Should be in the same order as order.

required
dz float

Factor to scale by while integrating. Should at least include the pixel size along the LOS.

required
beam Array

Beam to convolve by, should be a 2d array.

required
*pars Unpack[tuple[float, ...]]

1D container of model parameters.

required

Returns:

Name Type Description
model Array

The model with the specified substructure evaluated on the grid.

model_grad = jax.jit(model_grad, static_argnums=model_grad_static) module-attribute

A wrapper around model that also returns the gradients of the model. Only the additional arguments are described here, see model for the others. Note that the additional arguments are passed before the *params argument.

Parameters:

Name Type Description Default
argnums tuple[int, ...]

The indices of the arguments to evaluate the gradient at.

required

Returns:

Name Type Description
model Array

The model with the specified substructure evaluated on the grid.

grad Array

The gradient of the model with respect to the model parameters. Has shape (len(pars),) + model.shape).