Skip to content

Ensemble

Convenience functions for running batched ensembles over Gaussian updraft-speed distributions using jax.vmap.

run_updraft_ensemble 

run_updraft_ensemble(
    aerosols: list[Any],
    T0: float,
    S0: float,
    P0: float,
    *,
    mean: float,
    std: float,
    n: int,
    seed: int = 0,
    v_min: float = 0.01,
    accom: float = c.ac,
    z_cap: float = 400.0,
    t_end: float | None = None,
    max_steps: int = 100000,
) -> dict[str, Any]

Convenience: sample a Gaussian updraft distribution and run the ensemble.

Equilibrates y0 once for the given aerosols/conditions, samples n updraft speeds from Normal(mean, std) (clipped at v_min), and returns the per-member S_max/N_act (see smax_nact_ensemble). t_end defaults to z_cap / min(V) so even the slowest member reaches its supersaturation maximum.

Source code in pyrcel/ensemble.py 
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
156
157
158
159
160
def run_updraft_ensemble(
    aerosols: list[Any],
    T0: float,
    S0: float,
    P0: float,
    *,
    mean: float,
    std: float,
    n: int,
    seed: int = 0,
    v_min: float = 1e-2,
    accom: float = c.ac,
    z_cap: float = 400.0,
    t_end: float | None = None,
    max_steps: int = 100_000,
) -> dict[str, Any]:
    """Convenience: sample a Gaussian updraft distribution and run the ensemble.

    Equilibrates ``y0`` once for the given aerosols/conditions, samples ``n`` updraft
    speeds from ``Normal(mean, std)`` (clipped at ``v_min``), and returns the per-member
    ``S_max``/``N_act`` (see [smax_nact_ensemble][pyrcel.ensemble.smax_nact_ensemble]).
    ``t_end`` defaults to
    ``z_cap / min(V)`` so even the slowest member reaches its supersaturation maximum.
    """
    species, r_drys, kappas, Nis = [], [], [], []
    for aer in aerosols:
        r_drys.extend(aer.r_drys)
        kappas.extend([aer.kappa] * aer.nr)
        Nis.extend(aer.Nis)
        species.extend([aer.species] * aer.nr)
    r_drys = np.asarray(r_drys)
    kappas = np.asarray(kappas)
    Nis = np.asarray(Nis)

    y0 = np.asarray(equilibrate_initial_state(T0, S0, P0, r_drys, kappas, Nis))
    V_samples = sample_gaussian_updrafts(mean, std, n, seed=seed, v_min=v_min)
    if t_end is None:
        t_end = float(z_cap / np.min(V_samples))

    result = smax_nact_ensemble(
        y0, r_drys, Nis, kappas, float(accom), V_samples, t_end, max_steps=max_steps
    )
    result["t_end"] = t_end
    return result

smax_nact_ensemble 

smax_nact_ensemble(
    y0: ArrayLike,
    r_drys: ArrayLike,
    Nis: ArrayLike,
    kappas: ArrayLike,
    accom: float,
    V_samples: ArrayLike,
    t_end: float,
    *,
    rtol: float = STATE_RTOL,
    atol: ArrayLike | None = None,
    max_steps: int = 100000,
) -> dict[str, NDArray[np.floating[Any]] | NDArray[np.bool_] | float]

Peak supersaturation and activated number for a batch of updraft speeds.

Parameters:

Name Type Description Default
y0 array, shape ``(7 + nr,)``

Equilibrated initial state (shared across the ensemble).

 required
r_drys array, shape ``(nr,)``

Aerosol dry radii (m), number concentrations (m^-3), hygroscopicities.

 required
Nis array, shape ``(nr,)``

Aerosol dry radii (m), number concentrations (m^-3), hygroscopicities.

 required
kappas array, shape ``(nr,)``

Aerosol dry radii (m), number concentrations (m^-3), hygroscopicities.

 required
accom float

Accommodation coefficient.

 required
V_samples array, shape ``(n,)``

Updraft speeds (m/s).

 required
t_end float

Upper bound on integration time; must exceed the slowest member's time-to-peak (use z_cap / min(V) for a target height z_cap).

 required

Returns:

Type Description
dict

S_max (n,), N_act (n, m^-3), T_smax (n,), and activated (n, bool) -- the last flags members that reached a genuine interior maximum before t_end. Activated number uses the equilibrium criterion (modal critical supersaturation below S_max) evaluated at the peak temperature.
Source code in pyrcel/ensemble.py 
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 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
def smax_nact_ensemble(
    y0: ArrayLike,
    r_drys: ArrayLike,
    Nis: ArrayLike,
    kappas: ArrayLike,
    accom: float,
    V_samples: ArrayLike,
    t_end: float,
    *,
    rtol: float = STATE_RTOL,
    atol: ArrayLike | None = None,
    max_steps: int = 100_000,
) -> dict[str, NDArray[np.floating[Any]] | NDArray[np.bool_] | float]:
    """Peak supersaturation and activated number for a batch of updraft speeds.

    Parameters
    ----------
    y0 : array, shape ``(7 + nr,)``
        Equilibrated initial state (shared across the ensemble).
    r_drys, Nis, kappas : array, shape ``(nr,)``
        Aerosol dry radii (m), number concentrations (m^-3), hygroscopicities.
    accom : float
        Accommodation coefficient.
    V_samples : array, shape ``(n,)``
        Updraft speeds (m/s).
    t_end : float
        Upper bound on integration time; must exceed the slowest member's time-to-peak
        (use ``z_cap / min(V)`` for a target height ``z_cap``).

    Returns
    -------
    dict
        ``S_max`` (n,), ``N_act`` (n, m^-3), ``T_smax`` (n,), and ``activated`` (n, bool)
        -- the last flags members that reached a genuine interior maximum before
        ``t_end``. Activated number uses the equilibrium criterion (modal critical
        supersaturation below ``S_max``) evaluated at the peak temperature.
    """
    y0 = jnp.asarray(y0)
    r_drys = jnp.asarray(r_drys)
    Nis = jnp.asarray(Nis)
    kappas = jnp.asarray(kappas)
    nr = int(y0.shape[0] - c.N_STATE_VARS)
    if atol is None:
        atol = atol_vector(nr)
    V_samples = jnp.asarray(V_samples)

    def _member(V):
        args = (r_drys, Nis, kappas, accom, ConstantV(V))
        sol = _solve_to_smax(y0, args, t_end, rtol, atol, max_steps)
        y_peak = sol.ys[-1]
        s_max = y_peak[6]
        T_peak = y_peak[2]
        t_peak = sol.ts[-1]
        _, s_crit = kohler_crit_approx(T_peak, r_drys, kappas)
        n_act = jnp.sum(jnp.where(s_crit < s_max, Nis, 0.0))
        return s_max, n_act, T_peak, t_peak < t_end

    s_max, n_act, T_smax, activated = jax.jit(jax.vmap(_member))(V_samples)
    return {
        "S_max": np.asarray(s_max),
        "N_act": np.asarray(n_act),
        "T_smax": np.asarray(T_smax),
        "activated": np.asarray(activated),
        "V": np.asarray(V_samples),
    }

sample_gaussian_updrafts 

sample_gaussian_updrafts(
    mean: float, std: float, n: int, *, seed: int = 0, v_min: float = 0.01
) -> NDArray[np.floating[Any]]

Draw n updraft speeds from Normal(mean, std) (m/s).

Samples are clipped at v_min so every member is a physical (positive) updraft; with a low-mean/high-std request this truncation matters, so the returned array is what was actually simulated.

Source code in pyrcel/ensemble.py 
37
38
39
40
41
42
43
44
45
46
47
def sample_gaussian_updrafts(
    mean: float, std: float, n: int, *, seed: int = 0, v_min: float = 1e-2
) -> NDArray[np.floating[Any]]:
    """Draw ``n`` updraft speeds from ``Normal(mean, std)`` (m/s).

    Samples are clipped at ``v_min`` so every member is a physical (positive) updraft;
    with a low-mean/high-std request this truncation matters, so the returned array is
    what was actually simulated.
    """
    rng = np.random.default_rng(seed)
    return np.clip(rng.normal(mean, std, size=int(n)), v_min, None)