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 |
stage2_model(xyz, n_structs, dz, beam, *params)
Only returns the second stage of the model. Used for visualizing shocks, etc. that can otherwise be hard to see in a model plot
Arguments:
xyz: Coordinate grid to compute profile on.
n_struct: Number of each structure to use.
Should be in the same order as `order`.
dz: Factor to scale by while integrating.
Since it is a global factor it can contain unit conversions.
Historically equal to y2K_RJ * dr * da * XMpc / me.
beam: Beam to convolve by, should be a 2d array.
params: 1D array of model parameters.
Returns:
model: The model with the specified substructure evaluated on the grid.
Source code in witch/core.py
66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
|