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

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 
 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
147
def 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.
    """
    params = jnp.array(params)
    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