core

Core module for generating models and their gradients.

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

model(xyz, n_structs, n_rbins, dz, beam, *pars)

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

Source code in witch/core.py

def model(
    xyz: tuple[jax.Array, jax.Array, jax.Array, float, float],
    n_structs: tuple[int, ...],
    n_rbins: tuple[int],
    dz: float,
    beam: jax.Array,
    *pars: Unpack[tuple[float, ...]],
):
    """
    Generically create models with substructure.

    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.
    """
    params = jnp.array(pars)
    params = jnp.ravel(params)  # Fixes strange bug with params having dim (1,n)

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

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

    # Stage 2, add non-beam convolved to integrated profile
    for n_struct, struct in zip(n_structs, ORDER):
        if STRUCT_STAGE[struct] != 2:
            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):
            ip = jnp.add(ip, STRUCT_FUNCS[struct](*struct_pars[i], xyz))

    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)

    # Stage 3, add beam convolved to the integrated profile
    for n_struct, struct in zip(n_structs, ORDER):
        if STRUCT_STAGE[struct] != 3:
            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):
            ip = jnp.add(ip, STRUCT_FUNCS[struct](*struct_pars[i], xyz))

    return ip

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

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