tipbox 0.1.0__py3-none-any.whl

tipbox/__init__.py

@@ -0,0 +1,13 @@
+"""Box models and reproducible tipping-point benchmarks."""
+
+from . import generate
+from .registry import get_system_info, list_systems
+from .types import BoxModel, TrajectoryResult
+
+__all__ = [
+    "BoxModel",
+    "TrajectoryResult",
+    "generate",
+    "get_system_info",
+    "list_systems",
+]

tipbox/_amoc.py

@@ -0,0 +1,126 @@
+"""Shared physical parameters and numerical kernels for the two-state AMOC model."""
+
+from dataclasses import dataclass
+
+import jax.numpy as jnp
+
+from .utils import parameter, register_dataclass_pytree
+
+
+@register_dataclass_pytree
+@dataclass(frozen=True)
+class AMOCPhysics:
+    """Calibrated 2×CO₂ box-model parameters; salinities are mass fractions.
+
+    The states are 100 times salinity anomalies. Their drift therefore includes
+    100 times the seconds-per-year conversion. Model time is years; box volumes
+    use m³ and transports use m³/s.
+    """
+
+    S0: float = parameter(0.035, "Reference salinity mass fraction.", "mass fraction")
+    alpha: float = parameter(0.12, "Thermal-density coefficient.", "kg / m^3 / K")
+    beta: float = parameter(
+        790.0, "Haline-density coefficient per salinity mass fraction.", "kg / m^3"
+    )
+    seconds_per_year: float = parameter(
+        31500000.0,
+        "Seconds per model year; the separate factor100 converts salinity fractions to anomaly coordinates.",
+        "s / year",
+    )
+    V_N: float = parameter(3.683e16, "North Atlantic box volume.", "m^3")
+    V_T: float = parameter(5.418e16, "Tropical Atlantic box volume.", "m^3")
+    V_S: float = parameter(6.097e16, "Southern Ocean box volume.", "m^3")
+    V_IP: float = parameter(1.486e17, "Indo-Pacific box volume.", "m^3")
+    V_B: float = parameter(9.925e17, "Bottom-water box volume.", "m^3")
+    T_S: float = parameter(
+        7.919, "Southern reference temperature in the transport closure.", "degC"
+    )
+    T_0: float = parameter(
+        3.87, "North Atlantic reference temperature in the transport closure.", "degC"
+    )
+    total_salt: float = parameter(
+        4.4735e16, "Conserved salinity-volume integral.", "m^3"
+    )
+    lambda_q: float = parameter(
+        16200000.0, "Density-to-volume-transport coefficient.", "m^6 / kg / s"
+    )
+    gamma: float = parameter(
+        0.36, "Southern contribution to northward advective inflow.", "dimensionless"
+    )
+    mu: float = parameter(
+        2.2e-07,
+        "Temperature-transport coupling in the transport denominator.",
+        "K s / m^3",
+    )
+    K_N: float = parameter(1762000.0, "Northern gyre exchange transport.", "m^3 / s")
+    K_S: float = parameter(1872000.0, "Southern gyre exchange transport.", "m^3 / s")
+    S_S: float = parameter(0.034427, "Fixed Southern Ocean salinity.", "mass fraction")
+    S_B: float = parameter(0.034538, "Fixed bottom-water salinity.", "mass fraction")
+    F_N0: float = parameter(0.486, "Northern freshwater flux at zero hosing.", "Sv")
+    F_NH: float = parameter(
+        0.1311, "Northern freshwater-flux response to hosing.", "dimensionless"
+    )
+    F_T0: float = parameter(-0.997, "Tropical freshwater flux at zero hosing.", "Sv")
+    F_TH: float = parameter(
+        0.6961, "Tropical freshwater-flux response to hosing.", "dimensionless"
+    )
+
+
+def strength_sv(SN, p: AMOCPhysics):
+    """AMOC transport in Sverdrups from North Atlantic salinity anomaly."""
+    south = p.S_S - p.S0
+    return (
+        p.lambda_q
+        * (p.alpha * (p.T_S - p.T_0) + p.beta * (SN / 100 - south))
+        / (1 + p.lambda_q * p.alpha * p.mu)
+        * 1e-6
+    )
+
+
+def drift(state, H, p: AMOCPhysics):
+    """Salinity-anomaly drift, preserving the calibrated sign-dependent transport."""
+    SN, ST = state
+    SS, SB = (p.S_S - p.S0) * 100, (p.S_B - p.S0) * 100
+    SIP = (
+        100
+        * (
+            p.total_salt
+            - (p.V_N * SN + p.V_T * ST + p.V_S * SS + p.V_B * SB) / 100
+            - p.S0 * (p.V_B + p.V_N + p.V_T + p.V_IP + p.V_S)
+        )
+        / p.V_IP
+    )
+    q = strength_sv(SN, p) * 1e6
+    aq = jnp.abs(q)
+    FN, FT = (p.F_N0 + p.F_NH * H) * 1e6, (p.F_T0 + p.F_TH * H) * 1e6
+    north_exchange = p.K_N * (ST - SN) / 100
+    tropical_exchange = p.K_S * (SS - ST) / 100 + p.K_N * (SN - ST) / 100
+    north_adv = jnp.where(q >= 0, q * (ST - SN) / 100, aq * (SB - SN) / 100)
+    tropical_adv = jnp.where(
+        q >= 0,
+        q * (p.gamma * SS + (1 - p.gamma) * SIP - ST) / 100,
+        aq * (SN - ST) / 100,
+    )
+    return jnp.array(
+        [
+            100 * p.seconds_per_year / p.V_N * (north_adv + north_exchange - FN * p.S0),
+            100
+            * p.seconds_per_year
+            / p.V_T
+            * (tropical_adv + tropical_exchange - FT * p.S0),
+        ]
+    )
+
+
+def rk4_step(state, h_cur, h_mid, h_next, dt, physics):
+    """RK4 drift update; diffusion is applied separately by the caller."""
+    k1 = drift(state, h_cur, physics)
+    k2 = drift(state + 0.5 * dt * k1, h_mid, physics)
+    k3 = drift(state + 0.5 * dt * k2, h_mid, physics)
+    k4 = drift(state + dt * k3, h_next, physics)
+    return state + dt / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
+
+
+def observe(path, physics):
+    """Observed transport followed by the two physical salinity-anomaly states."""
+    return jnp.column_stack([strength_sv(path[:, 0], physics), path])

tipbox/bifurcation/amoc.py

@@ -0,0 +1,307 @@
+"""
+AMOC 2-box model, bifurcation-induced tipping (Ritchie et al. 2023; Alkhayuon et al. 2019).
+
+Port from dews/experiments/amoc_boxmodel.py. State (SN, ST); control H(t) (freshwater hosing).
+Saddle-node (fold) at H_fold; ramp H(t) with cap for tipping vs null.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import jax
+import jax.numpy as jnp
+from jax import lax
+
+from .._amoc import AMOCPhysics, drift, observe, rk4_step, strength_sv
+from ..types import BoxModel
+from ..utils import (
+    parameter,
+    register_dataclass_pytree,
+    sampling_stride,
+)
+
+# 2x CO2 parameters from BoxModel_2DH_IVP.m (rtippingreview)
+
+
+def _forcing_bifurcation_ramp_cap(
+    t: jnp.ndarray, H0: float, rate: float, H_cap: float
+) -> jnp.ndarray:
+    """H(t) = min(H0 + rate*t, H_cap) — ramp then constant at H_cap (same as amoc_boxmodel.forcing_bifurcation_ramp_cap)."""
+    return jnp.minimum(H0 + rate * t, H_cap)
+
+
+def _bifurcation_ramp_cap_params(
+    t_end: float,
+    H0: float,
+    H_fold: float,
+    *,
+    delta_cap: float = 0.05,
+) -> tuple:
+    """Rate and caps so H(t) reaches H_fold at t_end (like other box models: crossing at t_total/param_end_multiplier)."""
+    rate = (H_fold - H0) / t_end
+    H_cap_nontip = H_fold - delta_cap
+    H_cap_tip = H_fold + delta_cap
+    return rate, H_cap_nontip, H_cap_tip
+
+
+@register_dataclass_pytree
+@dataclass(frozen=True)
+class AMOCFoldConfig:
+    """Configuration for AMOC 2-box model (bifurcation-induced)."""
+
+    physics: AMOCPhysics = parameter(
+        AMOCPhysics(), "Shared calibrated physical AMOC parameters, expanded below.", ""
+    )
+    initial_SN: float = parameter(
+        0.0324, "Mean initial North Atlantic salinity anomaly.", "100 × mass fraction"
+    )
+    initial_ST: float = parameter(
+        0.1434,
+        "Mean initial tropical Atlantic salinity anomaly.",
+        "100 × mass fraction",
+    )
+    H0: float = parameter(0.0, "Initial freshwater-hosing level.", "Sv")
+    H_fold: float = parameter(
+        0.4236,
+        "Explicit fold marker for the default physical calibration; recalibrate after changing physics.",
+        "Sv",
+    )
+    delta_cap: float = parameter(
+        0.05, "Distance from H_fold to each forced/null cap.", "Sv"
+    )
+    sigma_SN: float = parameter(
+        0.0001,
+        "North Atlantic diffusion coefficient before the noise multiplier.",
+        "state / sqrt(year)",
+    )
+    sigma_ST: float = parameter(
+        0.0001,
+        "Tropical Atlantic diffusion coefficient before the noise multiplier.",
+        "state / sqrt(year)",
+    )
+    sigma_ic: float = parameter(
+        0.01,
+        "Initial-state standard deviation in make_trajectory; simulate uses the configured mean.",
+        "state",
+    )
+    dt: float = parameter(
+        5.0,
+        "Default output spacing; dt_inner controls RK4 integration.",
+        "year",
+    )
+    t_start: float = parameter(
+        0.0,
+        "Simulation start time.",
+        "year",
+    )
+    t_end: float = parameter(
+        5200.0,
+        "End time when make_trajectory is called without num_samples; explicit num_samples sets the end time.",
+        "year",
+    )
+    dt_inner: float = parameter(0.2, "AMOC RK4 integration step.", "year")
+
+    @classmethod
+    def default(cls) -> "AMOCFoldConfig":
+        return cls()
+
+
+class AMOCFoldModel(BoxModel):
+    """AMOC 2-box model, bifurcation-induced (H(t) ramp with fold at H_fold)."""
+
+    config_class = AMOCFoldConfig
+
+    @property
+    def critical_parameter(self):
+        """Configured fold marker; calibration depends on the physical parameters."""
+        return self._config.H_fold
+
+    state_names = ("north_salinity_anomaly", "tropical_salinity_anomaly")
+    observable_names = ("amoc_strength",) + state_names
+    noise_dimension = 2
+    integration_step_field = "dt_inner"
+    integration_method = "RK4 drift with additive Euler noise"
+    uses_midpoint_forcing = True
+    forcing_name = "freshwater_hosing"
+    time_unit = "year"
+
+    def _initial_state(self, key, null_simulation):
+        return self._sample_initial_state(key, null_simulation)
+
+    def _forcing(self, time, null_simulation, param_end_multiplier):
+        return self._H_schedule(time, null_simulation, param_end_multiplier)
+
+    def _integrate(
+        self, time, forcing, dt, noise, state, innovations, midpoint_forcing
+    ):
+        scale = (
+            jnp.array([self._config.sigma_SN, self._config.sigma_ST])
+            * noise
+            * jnp.sqrt(dt)
+        )
+        path = self._simulate_path(
+            forcing[:-1], midpoint_forcing, forcing[1:], dt, state, scale, innovations
+        )
+        return path, scale
+
+    def _observe(self, path):
+        return observe(path, self._config.physics)
+
+    def _H_schedule(
+        self,
+        t: jnp.ndarray,
+        null_simulation: bool = False,
+        param_end_multiplier: float = 1.1,
+    ) -> jnp.ndarray:
+        """H(t) = forcing_bifurcation_ramp_cap(t, H0, rate, H_cap) — same as dews amoc_boxmodel."""
+        H0 = self._config.H0
+        H_fold = self._config.H_fold
+        delta_cap = self._config.delta_cap
+        t_total = t[-1] - t[0]
+        t_end_eff = t_total / param_end_multiplier
+        rate, H_cap_nontip, H_cap_tip = _bifurcation_ramp_cap_params(
+            t_end_eff, H0, H_fold, delta_cap=delta_cap
+        )
+        H_cap = H_cap_nontip if null_simulation else H_cap_tip
+        t_rel = t - t[0]  # time from 0 for forcing_bifurcation_ramp_cap
+        return _forcing_bifurcation_ramp_cap(t_rel, H0, rate, H_cap)
+
+    def _sample_initial_state(
+        self,
+        key: jax.random.PRNGKey,
+        null_simulation: bool = False,
+    ) -> jnp.ndarray:
+        """Gaussian anomalies about the configured mean; negative anomalies are valid."""
+        key, _ = jax.random.split(key)
+        noise = jax.random.normal(key, (2,))
+        x0 = (
+            jnp.array([self._config.initial_SN, self._config.initial_ST])
+            + self._config.sigma_ic * noise
+        )
+        return x0
+
+    def _drift(self, state: jnp.ndarray, H: jnp.ndarray) -> jnp.ndarray:
+        return drift(state, H, self._config.physics)
+
+    def _simulate_path(
+        self,
+        H_cur: jnp.ndarray,
+        H_mid: jnp.ndarray,
+        H_next: jnp.ndarray,
+        dt: float,
+        initial_state: jnp.ndarray,
+        noise_scale: jnp.ndarray,
+        noise: jnp.ndarray,
+    ) -> jnp.ndarray:
+        """Single RK4 integrator: state_next = rk4_step(...) + noise_scale * noise_t. Use noise_scale=0 for deterministic."""
+
+        def step(state, inputs):
+            h_cur, h_mid, h_next, noise_t = inputs
+            state_next = (
+                rk4_step(state, h_cur, h_mid, h_next, dt, self._config.physics)
+                + noise_scale * noise_t
+            )
+            return state_next, state_next
+
+        _, states = lax.scan(step, initial_state, (H_cur, H_mid, H_next, noise))
+        return jnp.concatenate([initial_state[None, :], states], axis=0)
+
+    def simulate(
+        self,
+        key: jax.random.PRNGKey,
+        total_len: int,
+        sigma_tilde: float = 1.0,
+        null_simulation: bool = False,
+        dt_sample: Optional[float] = None,
+        param_end_multiplier: float = 1.1,
+        initial_state: Optional[jnp.ndarray] = None,
+        noise: Optional[jnp.ndarray] = None,
+    ):
+        key_noise, _ = jax.random.split(key)
+        total_time = self._config.t_end
+        dt_inner = self._config.dt_inner
+        dt_out = self._config.dt if dt_sample is None else dt_sample
+        sample_every = sampling_stride(dt_out, dt_inner)
+        n_steps = sampling_stride(total_time - self._config.t_start, dt_inner)
+        # Single inner grid: RK4 for both branches (stochastic adds noise per step)
+        time_inner = jnp.linspace(self._config.t_start, total_time, n_steps + 1)
+        H_inner = self._H_schedule(
+            time_inner,
+            null_simulation=null_simulation,
+            param_end_multiplier=param_end_multiplier,
+        )
+        H_cur = H_inner[:-1]
+        H_next = H_inner[1:]
+        H_mid = (H_cur + H_next) / 2.0
+
+        if initial_state is None:
+            initial_state = jnp.array(
+                [self._config.initial_SN, self._config.initial_ST]
+            )
+        else:
+            initial_state = jnp.asarray(initial_state)
+
+        deterministic = sigma_tilde <= 1e-9
+        n_steps = H_cur.shape[0]
+        if deterministic:
+            noise_scale = jnp.zeros(2)
+            noise_arr = jnp.zeros((n_steps, 2))
+        else:
+            noise_scale = jnp.array(
+                [
+                    self._config.sigma_SN * jnp.sqrt(dt_inner) * sigma_tilde,
+                    self._config.sigma_ST * jnp.sqrt(dt_inner) * sigma_tilde,
+                ]
+            )
+            noise_arr = (
+                jax.random.normal(key_noise, (n_steps, 2))
+                if noise is None
+                else jnp.asarray(noise)
+            )
+        trajectory_fine = self._simulate_path(
+            H_cur, H_mid, H_next, dt_inner, initial_state, noise_scale, noise_arr
+        )
+
+        n_inner = trajectory_fine.shape[0]
+        indices = jnp.arange(0, n_inner, sample_every)
+        time = time_inner[indices]
+        path_sampled = trajectory_fine[indices]
+
+        param_values = H_inner[indices]
+
+        # Output trajectory with both physical state (S_N, S_T) and AMOC strength Q(t).
+        # Column convention:
+        #   trajectory[:, 0] = Q(t)  (AMOC strength in Sv, for 1D EWS)
+        #   trajectory[:, 1] = S_N(t)  (North Atlantic salinity)
+        #   trajectory[:, 2] = S_T(t)  (Tropical Atlantic salinity)
+        SN_sampled = path_sampled[:, 0]
+        ST_sampled = path_sampled[:, 1]
+        Q_t = strength_sv(SN_sampled, self._config.physics)
+        trajectory = jnp.stack([Q_t, SN_sampled, ST_sampled], axis=1)
+
+        sigma_float = float(jnp.max(noise_scale))
+        return time, trajectory, param_values, sigma_float
+
+    def is_tipped(
+        self,
+        trajectory: jnp.ndarray,
+        param_values: jnp.ndarray,
+    ):
+        """True if H crossed the bifurcation point (max H >= H_fold)."""
+        return float(jnp.max(param_values)) >= self._config.H_fold
+
+    def locate_transition(
+        self,
+        trajectory: jnp.ndarray,
+        null_simulation: bool,
+        param_values: jnp.ndarray,
+        param_end_multiplier: float = 1.1,
+        time=None,
+    ):
+        """Midpoint for null; first hosing fold crossing for forced, or None."""
+        if null_simulation:
+            return len(param_values) // 2
+        crossed = param_values >= self._config.H_fold
+        if bool(jnp.any(crossed)):
+            return int(jnp.argmax(crossed))
+        return None

tipbox/bifurcation/bifurcation_simulator.py

@@ -0,0 +1,140 @@
+from dataclasses import asdict
+from typing import Any, Dict
+
+import jax
+
+from ..types import TrajectoryResult
+from ..registry import get_bifurcation_real_system_spec
+from ..utils import calculate_total_len
+
+
+def available_bifurcation_real_systems():
+    """Get the available real systems with bifurcation tipping (internal names)."""
+    from ..registry import SYSTEM_REGISTRY
+
+    return tuple(
+        spec.bifurcation_system
+        for spec in SYSTEM_REGISTRY.values()
+        if spec.bifurcation_system is not None
+    )
+
+
+def _select_system(key: jax.random.PRNGKey, system_name: str = None):
+    """Resolve a named bifurcation model and its typed configuration."""
+    from ..registry import SYSTEM_REGISTRY, SystemType
+
+    if system_name is not None:
+        return get_bifurcation_real_system_spec(system_name)
+
+    # Random selection
+    real_systems = [
+        spec
+        for spec in SYSTEM_REGISTRY.values()
+        if spec.system_type == SystemType.BIFURCATION
+        and spec.bifurcation_system is not None
+    ]
+    if not real_systems:
+        raise ValueError("No bifurcation real systems available")
+    idx = int(jax.random.randint(key, (), 0, len(real_systems)))
+    return real_systems[idx]
+
+
+def simulate_real_bifurcation_trajectory(
+    key: jax.random.PRNGKey,
+    context_samples: int = 500,
+    excluded_samples: int = 1,
+    buffer_samples: int = 1000,
+    sigma_tilde: float = 0.01,
+    null_simulation: bool = False,
+    system_name: str = None,
+    param_end_multiplier: float = 1.1,
+    config=None,
+    **generate_kwargs,
+):
+    """Simulate and extract bifurcation context. Return None on outcome/context rejection."""
+
+    key_select, key_sim = jax.random.split(key)
+
+    if excluded_samples > buffer_samples:
+        raise ValueError(
+            "excluded_samples exceeds buffer_samples. "
+            "Increase buffer_samples to maintain constant parameter rate."
+        )
+
+    # Calculate total simulation length (extended if param_end_multiplier > 1.0)
+    total_len = calculate_total_len(
+        context_samples, buffer_samples, param_end_multiplier
+    )
+
+    # Select, setup and initialize model using unified registry
+    system_spec = _select_system(key_select, system_name=system_name)
+    config = system_spec.default_config() if config is None else config
+    model = system_spec.create_model(config)
+
+    # Simulate
+    time, trajectory, param_values, sigma = model.simulate(
+        key_sim,
+        total_len=total_len,
+        sigma_tilde=sigma_tilde,
+        null_simulation=null_simulation,
+        param_end_multiplier=param_end_multiplier,
+        **generate_kwargs,
+    )
+
+    # Check if the trajectory crossed the bifurcation point
+    is_tipped = model.is_tipped(trajectory, param_values=param_values)
+    if is_tipped == null_simulation:
+        return None
+
+    # Locate transition index (pass param_values for parameter-based detection)
+    trans_idx = model.locate_transition(
+        trajectory,
+        null_simulation=null_simulation,
+        param_values=param_values,
+    )
+
+    if trans_idx is None:
+        return None
+
+    # Extract segment
+    time_seg, traj_seg, param_seg, trans_idx_out, success = model.extract_segment(
+        time,
+        trajectory,
+        param_values,
+        trans_idx,
+        excluded_samples=excluded_samples,
+        context_samples=context_samples,
+    )
+
+    if not success:
+        return None
+
+    # Return results
+    metadata: Dict[str, Any] = {
+        "label": "null" if null_simulation else system_spec.bifurcation_system,
+        "bifurcation_real_system": system_spec.bifurcation_system,
+        "excluded_samples": excluded_samples,
+        "system_config": asdict(config),
+        "transition_index": int(trans_idx_out),
+        "param_value_transition": float(
+            param_values[min(trans_idx_out, param_values.shape[0] - 1)]
+        ),
+    }
+
+    full_trajectory = TrajectoryResult(
+        time=time,
+        trajectory=trajectory,
+        param_values=param_values,
+        sigma=sigma,
+        metadata=metadata,
+    )
+
+    part_trajectory = TrajectoryResult(
+        time=time_seg,
+        trajectory=traj_seg,
+        param_values=param_seg,
+        sigma=sigma,
+        metadata=metadata,
+    )
+
+    return (full_trajectory, part_trajectory)