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_structs tuple[int, ...]

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

required
n_rbins tuple[int]

Number of rbins for each non-parametric model

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).

model3D(xyz, n_structs, n_rbins, params)

Generate a 3D profile from params on xyz.

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_structs tuple[int, ...]

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

required
n_rbins tuple[int]

Number of rbins for each non-parametric model

required
params tuple[float, ...]

1D container of model parameters.

required

Returns:

Name Type Description
pressure Array

The 3D model with the specified substructure evaluated on the grid.

start int

Current Total npar.

Source code in witch/core.py
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
def model3D(
    xyz: tuple[jax.Array, jax.Array, jax.Array, float, float],
    n_structs: tuple[int, ...],
    n_rbins: tuple[int],
    params: tuple[float, ...],  # TODO: not sure this is a tuple
) -> jax.Array:
    """
    Generate a 3D profile from params on xyz.

    Parameters
    ----------
    xyz : tuple[jax.Array, jax.Array, jax.Array, float, float]
        Grid to compute model on.
        See `containers.Model.xyz` for details.
    n_structs : tuple[int, ...]
        Number of each structure to use.
        Should be in the same order as `order`.
    n_rbins : tuple[int]
        Number of rbins for each non-parametric model
    params : tuple[float,...]
        1D container of model parameters.

    Returns
    -------
    pressure : jax.Array
        The 3D model with the specified substructure evaluated on the grid.
    start : int
        Current Total npar.
    """
    pressure = jnp.zeros((xyz[0].shape[0], xyz[1].shape[1], xyz[2].shape[2]))
    start = 0

    for i, (n_struct, struct) in enumerate(zip(n_structs, ORDER)):
        if STRUCT_STAGE[struct] != -1:
            continue
        if not n_struct:
            continue
        delta = n_struct * (
            n_rbins[i] * STRUCT_N_NONPARA[struct]
            + STRUCT_N_PAR[struct]
            - STRUCT_N_NONPARA[struct]
        )
        struct_pars = params[start : start + delta].reshape(
            (n_struct, int(delta / n_struct))
        )
        start += delta
        for j in range(n_struct):
            cur_struct_pars = struct_pars[j]
            nonpara_struct_pars = cur_struct_pars[
                : n_rbins[i] * STRUCT_N_NONPARA[struct]
            ].reshape((STRUCT_N_NONPARA[struct], n_rbins[i]))
            cur_struct_pars = cur_struct_pars[n_rbins[i] * STRUCT_N_NONPARA[struct] :]
            # pressure = jnp.add(pressure, STRUCT_FUNCS[struct](*nonpara_struct_pars, *struct_pars, xyz))
            cur_pars = [
                nonpara_struct_pars[k] for k in range(STRUCT_N_NONPARA[struct])
            ] + [
                cur_struct_pars[k]
                for k in range(STRUCT_N_PAR[struct] - STRUCT_N_NONPARA[struct])
            ]
            # pressure = jnp.add(pressure, STRUCT_FUNCS[struct](*nonpara_struct_pars, *struct_pars, xyz))
            pressure = jnp.add(pressure, STRUCT_FUNCS[struct](*cur_pars, xyz))

    # Stage 0, add to the 3d grid
    for n_struct, struct in zip(n_structs, ORDER):
        if STRUCT_STAGE[struct] != 0:
            continue
        if not n_struct:
            continue
        delta = n_struct * STRUCT_N_PAR[struct]
        struct_pars = params[start : start + delta].reshape(
            (n_struct, STRUCT_N_PAR[struct])
        )
        start += delta
        for i in range(n_struct):
            pressure = jnp.add(pressure, STRUCT_FUNCS[struct](*struct_pars[i], xyz))

    # Stage 1, modify the 3d grid
    for n_struct, struct in zip(n_structs, ORDER):
        if STRUCT_STAGE[struct] != 1:
            continue
        if not n_struct:
            continue
        delta = n_struct * STRUCT_N_PAR[struct]
        struct_pars = params[start : start + delta].reshape(
            (n_struct, STRUCT_N_PAR[struct])
        )
        start += delta
        for i in range(n_struct):
            pressure = STRUCT_FUNCS[struct](pressure, xyz, *struct_pars[i])

    return pressure, start

stage2_model(xyz, n_structs, dz, beam, *pars)

Only returns the second stage of the model. Used for visualizing shocks, etc. that can otherwise be hard to see in a model plot

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_structs 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.

()

Returns:

Name Type Description
model Array

The model with the specified substructure evaluated on the grid. No stage 3 structures are included.

Source code in witch/core.py
 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
147
148
149
150
151
152
153
154
155
def stage2_model(
    xyz: tuple[jax.Array, jax.Array, jax.Array, float, float],
    n_structs: tuple[int, ...],
    dz: float,
    beam: jax.Array,
    *pars: Unpack[tuple[float, ...]],
):
    """
    Only returns the second stage of the model. Used for visualizing shocks, etc.
    that can otherwise be hard to see in a model plot

    Parameters
    ----------
    xyz : tuple[jax.Array, jax.Array, jax.Array, float, float]
        Grid to compute model on.
        See `containers.Model.xyz` for details.
    n_structs : tuple[int, ...]
        Number of each structure to use.
        Should be in the same order as `order`.
    dz : float
        Factor to scale by while integrating.
        Should at least include the pixel size along the LOS.
    beam : jax.Array
        Beam to convolve by, should be a 2d array.
    *pars : Unpack[tuple[float,...]]
        1D container of model parameters.

    Returns
    -------
    model : jax.Array
        The model with the specified substructure evaluated on the grid.
        No stage 3 structures are included.
    """
    params = jnp.array(pars)
    params = jnp.ravel(params)  # Fixes strange bug with params having dim (1,n)

    pressure = jnp.ones((xyz[0].shape[0], xyz[1].shape[1], xyz[2].shape[2]))
    start = 0

    # Stage 0, track delta but don't add anything
    for n_struct, struct in zip(n_structs, ORDER):
        if STRUCT_STAGE[struct] != 0:
            continue
        if not n_struct:
            continue
        delta = n_struct * STRUCT_N_PAR[struct]
        struct_pars = params[start : start + delta].reshape(
            (n_struct, STRUCT_N_PAR[struct])
        )
        start += delta

    # Stage 1, modify the 3d grid
    for n_struct, struct in zip(n_structs, ORDER):
        if STRUCT_STAGE[struct] != 1:
            continue
        if not n_struct:
            continue

        delta = n_struct * STRUCT_N_PAR[struct]
        struct_pars = params[start : start + delta].reshape(
            (n_struct, STRUCT_N_PAR[struct])
        )

        start += delta
        for i in range(n_struct):
            pressure = STRUCT_FUNCS[struct](pressure, xyz, *struct_pars[i])

    # Integrate along line of site
    ip = trapz(pressure, dx=dz, axis=-1)

    bound0, bound1 = int((ip.shape[0] - beam.shape[0]) / 2), int(
        (ip.shape[1] - beam.shape[1]) / 2
    )
    beam = jnp.pad(
        beam,
        (
            (bound0, ip.shape[0] - beam.shape[0] - bound0),
            (bound1, ip.shape[1] - beam.shape[1] - bound1),
        ),
    )

    ip = fft_conv(ip, beam)

    return ip