Skip to content

Updraft

Updraft speed specification. ParcelModel accepts a scalar (constant speed), a ConstantV instance, or an InterpolatedUpdraft for time-varying profiles. All types are JAX pytrees and are compatible with jax.grad and jax.vmap.

AbstractUpdraft

Bases: Module

Base class for updraft models. Subclasses implement __call__(t) -> V.


ConstantV

ConstantV(V: ArrayLike)

Bases: AbstractUpdraft

A constant updraft speed V (m/s).

Source code in pyrcel/updraft.py
def __init__(self, V: ArrayLike) -> None:
    # Store as a float64 scalar array so it is a differentiable pytree leaf.
    self.V = jnp.asarray(V, dtype=jnp.float64)

InterpolatedUpdraft

InterpolatedUpdraft(ts: ArrayLike, vs: ArrayLike)

Bases: AbstractUpdraft

A tabulated, piecewise-linear V(t) profile.

Parameters:

Name Type Description Default
ts array

Strictly increasing knot times (s).

required
vs array

Updraft speeds (m/s) at ts. Outside [ts[0], ts[-1]] the endpoint values are held constant (the jax.numpy.interp default).

required
Source code in pyrcel/updraft.py
def __init__(self, ts: ArrayLike, vs: ArrayLike) -> None:
    self.ts = jnp.asarray(ts, dtype=jnp.float64)
    self.vs = jnp.asarray(vs, dtype=jnp.float64)

as_updraft

as_updraft(V) -> AbstractUpdraft

Coerce V to an AbstractUpdraft.

Accepts an existing updraft (returned as-is) or a scalar (wrapped in ConstantV). Plain Python callables are not accepted here: model an arbitrary profile as an AbstractUpdraft so the vector field stays a pure, jit-able pytree.

Parameters:

Name Type Description Default
V float or AbstractUpdraft

Updraft speed (m/s) or an existing AbstractUpdraft instance.

required

Returns:

Type Description
AbstractUpdraft

The original updraft, or V wrapped in ConstantV.

Raises:

Type Description
TypeError

If V is a plain Python callable; use InterpolatedUpdraft instead.

Source code in pyrcel/updraft.py
def as_updraft(V) -> AbstractUpdraft:
    """Coerce ``V`` to an [AbstractUpdraft][pyrcel.updraft.AbstractUpdraft].

    Accepts an existing updraft (returned as-is) or a scalar (wrapped in
    [ConstantV][pyrcel.updraft.ConstantV]). Plain Python callables are *not* accepted here: model an
    arbitrary profile as an [AbstractUpdraft][pyrcel.updraft.AbstractUpdraft] so the vector field
    stays a pure,
    ``jit``-able pytree.

    Parameters
    ----------
    V : float or AbstractUpdraft
        Updraft speed (m/s) or an existing [AbstractUpdraft][pyrcel.updraft.AbstractUpdraft]
        instance.

    Returns
    -------
    AbstractUpdraft
        The original updraft, or ``V`` wrapped in [ConstantV][pyrcel.updraft.ConstantV].

    Raises
    ------
    TypeError
        If ``V`` is a plain Python callable; use
        [InterpolatedUpdraft][pyrcel.updraft.InterpolatedUpdraft] instead.
    """
    if isinstance(V, AbstractUpdraft):
        return V
    if callable(V):
        raise TypeError(
            "pass V as an AbstractUpdraft (e.g. ConstantV / InterpolatedUpdraft), "
            "not a bare Python callable"
        )
    return ConstantV(V)