continuousintelligence 0.1.0__py3-none-any.whl

continuousintelligence/ci_proto/_calendar.py

@@ -0,0 +1,153 @@
+"""The calendar covariates used by ci-proto training and inference."""
+
+from __future__ import annotations
+
+import re
+
+import numpy as np
+import pandas as pd
+from numpy.typing import NDArray
+
+_DAY_NS = 86_400 * 10**9
+_HOUR_NS = 3_600 * 10**9
+_HOLIDAY_DATES = ((1, 1), (12, 24), (12, 25), (12, 26), (12, 31))
+_LEGACY_ALIASES = {
+    "T": "min",
+    "H": "h",
+    "S": "s",
+    "L": "ms",
+    "U": "us",
+    "N": "ns",
+    "M": "ME",
+    "Q": "QE",
+    "A": "YE",
+    "Y": "YE",
+}
+
+
+def frequency_step_ns(freq: str | None) -> int:
+    """Return the step of a pandas frequency in nanoseconds, or ``0`` for none or an irregular one.
+
+    A month counts 30 days and a quarter 91, so they select the monthly and quarterly calendar
+    rows. A step of several weeks, months or quarters returns ``0``.
+    """
+    if not freq:
+        return 0
+    try:
+        offset = pd.tseries.frequencies.to_offset(freq)
+    except ValueError:
+        match = re.fullmatch(r"(\d*)([A-Za-z]+)(-.*)?", freq)
+        if match is None or match[2] not in _LEGACY_ALIASES:
+            return 0
+        offset = pd.tseries.frequencies.to_offset(
+            match[1] + _LEGACY_ALIASES[match[2]] + (match[3] or "")
+        )
+    if isinstance(offset, pd.offsets.Day):
+        return int(offset.n) * _DAY_NS
+    for kind, days in (
+        (pd.offsets.Week, 7),
+        (pd.offsets.MonthBegin, 30),
+        (pd.offsets.MonthEnd, 30),
+        (pd.offsets.QuarterBegin, 91),
+        (pd.offsets.QuarterEnd, 91),
+    ):
+        if isinstance(offset, kind):
+            return days * _DAY_NS if offset.n == 1 else 0
+    if isinstance(offset, pd.offsets.Tick):
+        return int(pd.Timedelta(offset).value)
+    return 0
+
+
+def _easter_day(year: int) -> int:
+    """Days since 1970-01-01 of Easter Sunday, by the anonymous Gregorian algorithm."""
+    a, b, c = year % 19, year // 100, year % 100
+    d, e = b // 4, b % 4
+    g = (b - (b + 8) // 25 + 1) // 3
+    h = (19 * a + b - d - g + 15) % 30
+    i, k = c // 4, c % 4
+    w = (32 + 2 * e + 2 * i - h - k) % 7
+    m = (a + 11 * h + 22 * w) // 451
+    month, day = divmod(h + w - 7 * m + 114, 31)
+    return int(np.datetime64(f"{year:04d}-{month:02d}-{day + 1:02d}").astype(np.int64))
+
+
+def _holiday_flags(days: NDArray[np.int64]) -> NDArray[np.bool_]:
+    """Whether each day (days since 1970-01-01) is a holiday.
+
+    The holidays are New Year's Day, Christmas Eve to Boxing Day, New Year's Eve, Good Friday,
+    Easter Monday and US Thanksgiving; the movable ones only for years 1800 to 2300.
+    """
+    dates = days.astype("datetime64[D]")
+    months = dates.astype("datetime64[M]").astype(np.int64) % 12 + 1
+    day_of_month = (dates - dates.astype("datetime64[M]")).astype(np.int64) + 1
+    flags = np.zeros(days.shape, dtype=bool)
+    for month, day in _HOLIDAY_DATES:
+        flags |= (months == month) & (day_of_month == day)
+    years = dates.astype("datetime64[Y]").astype(np.int64) + 1970
+    for year in np.unique(years).tolist():
+        if not 1800 <= year <= 2300:
+            continue
+        easter = _easter_day(year)
+        november = int(np.datetime64(f"{year:04d}-11-01").astype(np.int64))
+        # The fourth Thursday; 1970-01-01, day 0, was a Thursday.
+        thanksgiving = november + (3 - (november + 3) % 7) % 7 + 21
+        flags |= np.isin(days, (easter - 2, easter + 1, thanksgiving))
+    return flags
+
+
+def _business_days(timestamps: NDArray[np.int64], step_ns: int) -> NDArray[np.float32] | None:
+    """Return the weekdays that are not holidays in each step's month or quarter.
+
+    Monthly steps (28 to 31 days) count their calendar month and quarterly steps (89 to 92
+    days) their calendar quarter; any other step returns ``None``.
+    """
+    months = (timestamps // _DAY_NS).astype("datetime64[D]").astype("datetime64[M]")
+    if 28 * _DAY_NS <= step_ns <= 31 * _DAY_NS:
+        starts = months
+        ends = months + np.timedelta64(1, "M")
+    elif 89 * _DAY_NS <= step_ns <= 92 * _DAY_NS:
+        starts = months - (months.astype(np.int64) % 3).astype("timedelta64[M]")
+        ends = starts + np.timedelta64(3, "M")
+    else:
+        return None
+    first, last = starts.astype("datetime64[D]"), ends.astype("datetime64[D]")
+    span = np.arange(first.min(), last.max(), dtype="datetime64[D]")
+    holidays = span[_holiday_flags(span.astype(np.int64))]
+    return np.busday_count(first, last, holidays=holidays).astype(np.float32)
+
+
+def calendar_channels(timestamps: NDArray[np.int64], step_ns: int) -> NDArray[np.float32]:
+    """``[channels, time]`` calendar features of a series sampled every ``step_ns``.
+
+    Steps that divide the day give hour of day (sub-daily steps only), day of week (Monday
+    ``0``), day of year (``0``-``365``) and a holiday flag; weekly steps give week of year
+    (``0``-``52``) and whether a holiday falls within three days; monthly steps (28 to 31 days)
+    give month of year (``0``-``11``); quarterly steps (89 to 92 days) give the quarter
+    (``0``-``3``). Both also include the period's business-day count. Any other step gives
+    no channels.
+    """
+    days = timestamps // _DAY_NS
+    dates = days.astype("datetime64[D]")
+    day_of_year = (dates - dates.astype("datetime64[Y]")).astype(np.int64)
+    rows: list[NDArray[np.int64] | NDArray[np.bool_]]
+    if 0 < step_ns <= _DAY_NS and _DAY_NS % step_ns == 0:
+        # 1970-01-01, day 0, was a Thursday.
+        rows = [(days + 3) % 7, day_of_year, _holiday_flags(days)]
+        if step_ns < _DAY_NS:
+            rows.insert(0, (timestamps % _DAY_NS) // _HOUR_NS)
+    elif step_ns == 7 * _DAY_NS:
+        window = (days[:, None] + np.arange(-3, 4)).ravel()
+        rows = [day_of_year // 7, _holiday_flags(window).reshape(days.shape[0], 7).any(axis=1)]
+    elif 28 * _DAY_NS <= step_ns <= 31 * _DAY_NS:
+        rows = [
+            dates.astype("datetime64[M]").astype(np.int64) % 12,
+            _business_days(timestamps, step_ns),
+        ]
+    elif 89 * _DAY_NS <= step_ns <= 92 * _DAY_NS:
+        rows = [
+            (dates.astype("datetime64[M]").astype(np.int64) % 12) // 3,
+            _business_days(timestamps, step_ns),
+        ]
+    else:
+        return np.empty((0, timestamps.shape[0]), dtype=np.float32)
+    return np.stack([row.astype(np.float32) for row in rows])

continuousintelligence/ci_proto/_checkpoint.py

@@ -0,0 +1,106 @@
+"""Reconstruct ci-proto weights, calibration and inference backend."""
+
+import json
+import math
+from pathlib import Path
+
+import torch
+
+from ._forecast import Forecaster
+from ._model import Model
+
+
+def _load_tta(path: str | Path) -> dict[str, float]:
+    booleans = {
+        "postprocess",
+        "tta_diff",
+        "tta_sign_flip",
+        "tta_running_totals",
+        "target_encode_categoricals",
+        "clip_nonnegative",
+        "integer_round",
+        "bfloat16",
+        "calendar_covariates",
+    }
+    constants = {
+        "trend_threshold",
+        "trend_min_r2",
+        "trend_window",
+        "diff_band_exponent",
+        "diff_band_scale",
+        "raw_band_exponent",
+        "raw_band_scale",
+        "differencing_max_horizon",
+        "static_groups",
+        "joint_neighbours",
+        "chunk_joint",
+    }
+    required = booleans | constants
+    raw = json.loads(Path(path).read_text())
+    if not isinstance(raw, dict):
+        raise ValueError("tta config must contain a JSON object")
+    missing = required - raw.keys()
+    unknown = raw.keys() - required
+    if missing or unknown:
+        raise ValueError(
+            f"invalid tta fields: missing={sorted(missing)}, unknown={sorted(unknown)}"
+        )
+    result = {}
+    for key in booleans:
+        if type(raw[key]) is not bool:
+            raise ValueError(f"{key} must be a boolean")
+        result[key] = raw[key]
+    for key in constants:
+        value = raw[key]
+        if key in ("differencing_max_horizon", "static_groups", "joint_neighbours", "chunk_joint"):
+            valid = (
+                type(value) is int
+                and value >= (1 if key == "differencing_max_horizon" else 0)
+                and not (key in ("joint_neighbours", "chunk_joint") and value == 1)
+            )
+        else:
+            if type(value) not in (int, float) or not math.isfinite(value):
+                raise ValueError(f"{key} must be a finite number")
+            valid = value >= 0 if key.endswith("exponent") else value > 0
+            if key in ("trend_min_r2", "trend_window"):
+                valid = 0 < value <= 1
+            value = float(value)
+        if not valid:
+            raise ValueError(f"invalid {key}: {value}")
+        result[key] = value
+    return result
+
+
+def from_checkpoint(payload: object, tta_config: str | Path, device: torch.device) -> Forecaster:
+    """Build ci-proto from a loaded checkpoint and its calibration."""
+    tta = _load_tta(tta_config)
+    metadata = payload.get("hyper_parameters") if isinstance(payload, dict) else None
+    if not isinstance(metadata, dict) or not isinstance(metadata.get("model"), dict):
+        raise ValueError("checkpoint must contain a hyper_parameters.model mapping")
+    try:
+        model = Model(metadata["model"])
+    except (KeyError, AttributeError, TypeError, ZeroDivisionError, OverflowError) as error:
+        raise ValueError(f"invalid checkpoint model configuration: {error}") from error
+    state = payload.get("state_dict")
+    if not isinstance(state, dict) or not state or any(not isinstance(key, str) for key in state):
+        raise ValueError("checkpoint must contain a nonempty state_dict with string keys")
+    if all(key.startswith("model.") for key in state):
+        state = {key.removeprefix("model."): value for key, value in state.items()}
+    model.load_state_dict(state, strict=True)
+    levels = model.quantiles
+    if (
+        len(levels) % 2 == 0
+        or not torch.isclose(levels[len(levels) // 2], torch.tensor(0.5), atol=1e-6, rtol=0)
+        or not torch.allclose(
+            levels,
+            1 - levels.flip(0),
+            atol=1e-6,
+            rtol=0,
+        )
+    ):
+        raise ValueError("ci-proto TTA requires symmetric quantiles centred on the median")
+    model = model.to(device).eval()
+    from ._cuda import configure_backend
+
+    configure_backend(model)
+    return Forecaster(model, tta)

continuousintelligence/ci_proto/_cuda.py

@@ -0,0 +1,135 @@
+"""Automatic selection and inference wrapper for the bundled Mamba-3 MIMO kernels.
+
+Includes Mamba-3 forward integration adapted by Continuous Intelligence:
+https://github.com/state-spaces/mamba/tree/a14b1dff0454a3bc27d9eb31355dc01e4b2490ec
+
+Copyright (c) 2026, Dao AI Lab, Goombalab.
+Upstream portions are licensed under Apache-2.0; see NOTICE.
+"""
+
+from __future__ import annotations
+
+import warnings
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import Tensor, nn
+
+from ._mamba import Mamba3
+
+
+def configure_backend(model: nn.Module) -> None:
+    """Enable fused inference only for supported hardware and cell geometry."""
+    model.backend = "torch"
+    device = next(model.parameters()).device
+    if device.type != "cuda" or torch.version.hip is not None:
+        return
+    cells = [cell for cell in model.modules() if isinstance(cell, Mamba3)]
+    if torch.cuda.get_device_capability(device)[0] < 8 or not cells:
+        return
+    # Float32 tiles need more shared memory than bf16. Use smaller chunks on
+    # Ampere and consumer GPUs.
+    chunk_budget = (
+        64
+        if torch.cuda.get_device_properties(device).shared_memory_per_block_optin >= 196608
+        else 32
+    )
+    if any(
+        not cell.is_mimo
+        or cell.mimo_rank not in (1, 2, 4, 8)
+        or cell.d_state not in (32, 64, 128)
+        or cell.headdim not in (32, 64, 128)
+        or max(8, chunk_budget // cell.mimo_rank) * cell.mimo_rank > chunk_budget
+        for cell in cells
+    ):
+        return
+    try:
+        from .._vendor.mamba3.mamba3_mimo_fwd import mamba_mimo_forward  # noqa: F401
+    except (ImportError, OSError) as error:
+        warnings.warn(
+            f"Fused Mamba-3 is unavailable; using portable PyTorch on CUDA: {error}",
+            RuntimeWarning,
+            stacklevel=2,
+        )
+        return
+    for cell in cells:
+        cell.chunk_size = max(8, chunk_budget // cell.mimo_rank)
+        cell.fused = True
+    model.backend = "mamba3"
+
+
+def mamba3_forward(cell: Mamba3, u: Tensor) -> Tensor:
+    """Run the forward path, including rank-scaled chunk padding."""
+    from .._vendor.mamba3.angle_dt import angle_dt_fwd
+    from .._vendor.mamba3.layernorm_gated import _layer_norm_fwd
+    from .._vendor.mamba3.mamba3_mimo_fwd import mamba_mimo_forward
+    from .._vendor.mamba3.mamba3_mimo_utils import compute_dacs_segsum_triton
+
+    seq = u.shape[1]
+    chunk = cell.chunk_size
+    u = F.pad(u, (0, 0, 0, (-seq) % chunk))
+    bc = cell.d_state * cell.mimo_rank
+    z, x, B, C, dd_dt, dd_A, trap, angles = torch.split(
+        cell.in_proj(u),
+        [
+            cell.d_inner,
+            cell.d_inner,
+            bc,
+            bc,
+            cell.nheads,
+            cell.nheads,
+            cell.nheads,
+            cell.num_rope_angles,
+        ],
+        dim=-1,
+    )
+    x = rearrange(x, "b l (h p) -> b l h p", p=cell.headdim)
+    z = rearrange(z, "b l (h p) -> b l h p", p=cell.headdim)
+    B = rearrange(B, "b l (r g n) -> b l r g n", r=cell.mimo_rank, g=1)
+    C = rearrange(C, "b l (r g n) -> b l r g n", r=cell.mimo_rank, g=1)
+    # Keep the kernel's reduction and rounding order, including 1 / sqrt.
+    B = _layer_norm_fwd(
+        B.reshape(-1, cell.d_state),
+        cell.B_norm.weight,
+        None,
+        cell.B_norm.eps,
+        is_rms_norm=True,
+    )[0].view_as(B)
+    C = _layer_norm_fwd(
+        C.reshape(-1, cell.d_state),
+        cell.C_norm.weight,
+        None,
+        cell.C_norm.eps,
+        is_rms_norm=True,
+    )[0].view_as(C)
+    A = (-F.softplus(dd_A.float())).clamp(max=-1e-4)
+    dt = F.softplus(dd_dt + cell.dt_bias)
+    adt = rearrange(A * dt, "b l h -> b h l").contiguous()
+    dt = rearrange(dt, "b l h -> b h l").contiguous()
+    angles = angles[:, :, None].expand(-1, -1, cell.nheads, -1).float().contiguous()
+    with torch.cuda.device(u.device):
+        cumulative, _ = angle_dt_fwd(angles, dt, chunk_size=chunk, return_output_state=True)
+        da_cs, da_cs_rev, segsum = compute_dacs_segsum_triton(adt, chunk)
+        y, _, _ = mamba_mimo_forward(
+            C.contiguous(),
+            B.contiguous(),
+            x.contiguous(),
+            cell.C_bias,
+            cell.B_bias,
+            cell.mimo_x,
+            cell.mimo_o,
+            z.contiguous(),
+            cell.D,
+            cell.mimo_z,
+            cumulative,
+            da_cs,
+            da_cs_rev,
+            dt,
+            rearrange(trap, "b l h -> b h l").contiguous(),
+            segsum,
+            chunk_size=chunk,
+            rotary_dim_divisor=4,
+            dtype=x.dtype,
+        )
+    return cell.out_proj(rearrange(y, "b l h p -> b l (h p)").to(u.dtype))[:, :seq]