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
to_run	tuple[bool, bool, bool, bool]	Which stages to run. Should be a 5 element tuple of bools, see make_to_run for details.	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).

make_to_run(stage_m1=True, stage_0=True, stage_1=True, stage_2=True)

Constructs the to_run tuple needed for calls to model and model_grad in the correct order. This is genrally prefered over manual construction, if you need to contsruct manualy the order of parameters in this function is the order of the to_run tuple. Any stages that are set to False here will not be run when computing the model.

Parameters:

Name	Type	Description	Default
stage_m1	bool	If True run stage -1. This is non-parametric models.	True
stage_0	bool	If True run stage 0. This is 3D parametric models.	True
stage_1	bool	If True run stage 1. This is 3D sub-structure.	True
stage_2	bool	If True run stage 2. This is 2D structure.	True

Returns:

Name	Type	Description
to_run	tuple[bool, bool, bool, bool, bool, bool]	Which stages to run in the order needed for model and model_grad.

Source code in witch/core.py

def make_to_run(
    stage_m1: bool = True,
    stage_0: bool = True,
    stage_1: bool = True,
    stage_2: bool = True,
) -> tuple[bool, bool, bool, bool]:
    """
    Constructs the `to_run` tuple needed for calls to `model` and `model_grad` in the correct order.
    This is genrally prefered over manual construction, if you need to contsruct manualy the order of parameters in
    this function is the order of the `to_run` tuple.
    Any stages that are set to `False` here will not be run when computing the model.

    Parameters
    ----------
    stage_m1 : bool, default: True
        If True run stage -1.
        This is non-parametric models.
    stage_0 : bool, default: True
        If True run stage 0.
        This is 3D parametric models.
    stage_1 : bool, default: True
        If True run stage 1.
        This is 3D sub-structure.
    stage_2 : bool, default: True
        If True run stage 2.
        This is 2D structure.

    Returns
    -------
    to_run : tuple[bool, bool, bool, bool, bool, bool]
        Which stages to run in the order needed for `model` and `model_grad`.
    """
    return (stage_m1, stage_0, stage_1, stage_2)

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	Array	1D container of model parameters.	required

Returns:

Name	Type	Description
pressure	Array	The 3D model with the specified substructure evaluated on the grid.

Source code in witch/core.py

def model3D(
    xyz: tuple[jax.Array, jax.Array, jax.Array, float, float],
    n_structs: tuple[int, ...],
    n_rbins: tuple[int, ...],
    params: jax.Array,
) -> 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 : jax.Array
        1D container of model parameters.

    Returns
    -------
    pressure : jax.Array
        The 3D 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.zeros((xyz[0].shape[0], xyz[1].shape[1], xyz[2].shape[2]))
    start = 0

    # Stage -1, non para
    pressure, start = _stage_m1(xyz, n_structs, n_rbins, params, start, pressure, True)

    # Stage 0, add to the 3d grid
    pressure, start = _stage_0(xyz, n_structs, params, start, pressure, True)

    # Stage 1, modify the 3d grid
    pressure, start = _stage_1(xyz, n_structs, params, start, pressure, True)

    return pressure

stage2_model(xyz, n_structs, dz, *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
*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,
    *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.
    *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, but just track param start
    pressure, start = _stage_0(xyz, n_structs, params, start, pressure, False)

    # Stage 1, modify the 3d grid
    pressure, start = _stage_1(xyz, n_structs, params, start, pressure)

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

    return ip