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 |
required |
n_struct |
tuple[int, ...]
|
Number of each structure to use.
Should be in the same order as |
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 |