hcx 0.1.0__tar.gz

hcx-0.1.0/PKG-INFO

@@ -0,0 +1,87 @@
+Metadata-Version: 2.3
+Name: hcx
+Version: 0.1.0
+Summary: hcx
+Requires-Dist: numpy
+Requires-Dist: torch
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+
+# hcx
+
+hcx defines typed, runtime-usable contracts for hydrological model batches,
+forecasts, output specifications, models, and model factories.
+
+The [contract specification](docs/spec.md) is normative. Public docstrings and
+this README are explanatory summaries of that specification.
+
+The distribution combines three parts: the normative specification, an
+inline-typed Python contract marked with `py.typed`, and a standalone
+conformance harness. Independent model packages can import the same contract
+and prove their behavior without an application, training engine, or external
+dataset.
+
+## Authoring a model package
+
+1. Add `hcx` as a runtime dependency of the model package.
+2. Implement a `torch.nn.Module` whose `forward(batch: Batch) -> Forecast`
+   behavior satisfies `ForecastModel`, including the normative shape,
+   dtype/device, and metadata identity rules.
+3. Implement a callable matching `ModelFactory`. It accepts a model-specific
+   `dict[str, object]` and the keyword-only ordered `dynamic_inputs` and
+   `static_inputs`, resolved `input_size`, `static_size`, and `output_size`, and
+   a resolved `OutputSpecification[object]`. It returns the module. Preserve
+   name order; the sizes are resolved from the first batch.
+4. Declare the factory in the model package's `pyproject.toml`:
+
+   ```toml
+   [project.entry-points.'hcx.models']
+   my_model = "my_model_package.factory:create_model"
+   ```
+
+5. Exercise the complete contract with synthetic data:
+
+   ```python
+   from importlib.metadata import entry_points
+
+   from hcx import Point, assert_conforms, make_synthetic_batch
+
+   batch = make_synthetic_batch()
+   assert batch.scalar_dynamic is not None
+   assert batch.scalar_static is not None
+
+   factory = entry_points(group="hcx.models", name="my_model")[0].load()
+   dynamic_inputs = [
+       f"dynamic_{index}" for index in range(batch.scalar_dynamic.shape[-1])
+   ]
+   static_inputs = [
+       f"static_{index}" for index in range(batch.scalar_static.shape[-1])
+   ]
+   model = factory(
+       {},
+       dynamic_inputs=dynamic_inputs,
+       static_inputs=static_inputs,
+       input_size=batch.scalar_dynamic.shape[-1],
+       static_size=batch.scalar_static.shape[-1],
+       output_size=batch.target.shape[-1],
+       output_specification=Point(),
+   )
+   assert_conforms(model, batch)
+   ```
+
+The factory can also be imported and called directly in this smoke test. hcx
+packages `scalar_lstm` as a reference entry point, not as a required base class.
+Entry-point names are consumer-facing identifiers and should remain stable and
+unique within an environment.
+
+See the normative [factory and entry-point clauses](docs/spec.md#6-model-factory-and-hcxmodels-entry-points)
+and the [changelog](CHANGELOG.md) for version history.
+
+## Releases
+
+Maintainers bump the version with bump-my-version and update the changelog in
+the same commit. A human then creates and publishes the GitHub Release, which
+creates the `vX.Y.Z` tag. The GitHub Actions workflow builds distributions with
+`uv build` and publishes them through OIDC trusted publishing. Prereleases and
+manual workflow dispatches target TestPyPI; published non-prereleases target
+PyPI. Local publishing and hand-made tags are forbidden.

hcx-0.1.0/README.md

@@ -0,0 +1,78 @@
+# hcx
+
+hcx defines typed, runtime-usable contracts for hydrological model batches,
+forecasts, output specifications, models, and model factories.
+
+The [contract specification](docs/spec.md) is normative. Public docstrings and
+this README are explanatory summaries of that specification.
+
+The distribution combines three parts: the normative specification, an
+inline-typed Python contract marked with `py.typed`, and a standalone
+conformance harness. Independent model packages can import the same contract
+and prove their behavior without an application, training engine, or external
+dataset.
+
+## Authoring a model package
+
+1. Add `hcx` as a runtime dependency of the model package.
+2. Implement a `torch.nn.Module` whose `forward(batch: Batch) -> Forecast`
+   behavior satisfies `ForecastModel`, including the normative shape,
+   dtype/device, and metadata identity rules.
+3. Implement a callable matching `ModelFactory`. It accepts a model-specific
+   `dict[str, object]` and the keyword-only ordered `dynamic_inputs` and
+   `static_inputs`, resolved `input_size`, `static_size`, and `output_size`, and
+   a resolved `OutputSpecification[object]`. It returns the module. Preserve
+   name order; the sizes are resolved from the first batch.
+4. Declare the factory in the model package's `pyproject.toml`:
+
+   ```toml
+   [project.entry-points.'hcx.models']
+   my_model = "my_model_package.factory:create_model"
+   ```
+
+5. Exercise the complete contract with synthetic data:
+
+   ```python
+   from importlib.metadata import entry_points
+
+   from hcx import Point, assert_conforms, make_synthetic_batch
+
+   batch = make_synthetic_batch()
+   assert batch.scalar_dynamic is not None
+   assert batch.scalar_static is not None
+
+   factory = entry_points(group="hcx.models", name="my_model")[0].load()
+   dynamic_inputs = [
+       f"dynamic_{index}" for index in range(batch.scalar_dynamic.shape[-1])
+   ]
+   static_inputs = [
+       f"static_{index}" for index in range(batch.scalar_static.shape[-1])
+   ]
+   model = factory(
+       {},
+       dynamic_inputs=dynamic_inputs,
+       static_inputs=static_inputs,
+       input_size=batch.scalar_dynamic.shape[-1],
+       static_size=batch.scalar_static.shape[-1],
+       output_size=batch.target.shape[-1],
+       output_specification=Point(),
+   )
+   assert_conforms(model, batch)
+   ```
+
+The factory can also be imported and called directly in this smoke test. hcx
+packages `scalar_lstm` as a reference entry point, not as a required base class.
+Entry-point names are consumer-facing identifiers and should remain stable and
+unique within an environment.
+
+See the normative [factory and entry-point clauses](docs/spec.md#6-model-factory-and-hcxmodels-entry-points)
+and the [changelog](CHANGELOG.md) for version history.
+
+## Releases
+
+Maintainers bump the version with bump-my-version and update the changelog in
+the same commit. A human then creates and publishes the GitHub Release, which
+creates the `vX.Y.Z` tag. The GitHub Actions workflow builds distributions with
+`uv build` and publishes them through OIDC trusted publishing. Prereleases and
+manual workflow dispatches target TestPyPI; published non-prereleases target
+PyPI. Local publishing and hand-made tags are forbidden.

hcx-0.1.0/pyproject.toml

@@ -0,0 +1,89 @@
+[project]
+name = "hcx"
+version = "0.1.0"
+description = "hcx"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "numpy",
+    "torch",
+]
+
+[project.entry-points.'hcx.models']
+scalar_lstm = "hcx.models.lstm:factory"
+
+
+[build-system]
+requires = ["uv_build>=0.9.13,<0.10.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "pytest>=8.3", # tests
+    "pytest-cov>=6.2.1",
+    "ruff>=0.12", # linter + formatter
+    "bump-my-version>=1.2.7",
+    "ty>=0.0.58",
+    "pyyaml>=6.0.3",
+]
+
+# -----------------------
+# Ruff (lint + format)
+# -----------------------
+[tool.ruff]
+line-length = 120
+target-version = "py311"
+# extend-exclude = ["notebooks/", "experiments/"]
+
+[tool.ruff.lint]
+# Errors, warnings, naming, imports, modern py, bugbear, comprehensions, simplify
+select = ["E", "F", "W", "N", "I", "UP", "B", "C4", "SIM"]
+ignore = [
+    "E501", # let formatter handle wrapping
+    "N803", # allow X, Y for ML-style params
+    "N806", # allow X, Y for local vars
+]
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+
+# # -----------------------
+# # ty (fast static type checker)
+# # -----------------------
+# [tool.ty]
+# [tool.ty.rules]
+# possibly-unresolved-reference = "warn"
+
+# -----------------------
+# Pytest
+# -----------------------
+[tool.pytest.ini_options]
+pythonpath = ["src"]
+testpaths = ["tests"]
+python_files = ["test_*.py", "*_test.py"]
+python_classes = ["Test*"]
+python_functions = ["test_*"]
+addopts = ["--strict-markers", "--strict-config", "-ra"]
+
+# -----------------------
+# Bump My Version
+# -----------------------
+[tool.bumpversion]
+current_version = "0.1.0"
+commit = false
+tag = false
+allow_dirty = true
+
+[[tool.bumpversion.files]]
+filename = "pyproject.toml"
+search = 'version = "{current_version}"'
+replace = 'version = "{new_version}"'
+
+[[tool.bumpversion.files]]
+filename = "src/hcx/__init__.py"
+search = '__version__ = "{current_version}"'
+replace = '__version__ = "{new_version}"'
+
+[tool.uv.sources]
+# ctrl-freak = { git = "https://github.com/hydrosolutions/ctrl-freak.git", rev = "main" } # example of how to add a git dependency

hcx-0.1.0/src/hcx/__init__.py

@@ -0,0 +1,30 @@
+__version__ = "0.1.0"
+
+from hcx.batch import Batch, BatchMetadata, GriddedDynamic, GriddedStatic
+from hcx.conformance import ConformanceCheck, ConformanceError, assert_conforms, check_conformance
+from hcx.output import Forecast
+from hcx.protocol import MODEL_ENTRY_POINT_GROUP, ForecastModel, ModelFactory
+from hcx.specifications import Gaussian, GaussianParameters, OutputSpecification, Point, PointParameters
+from hcx.synthetic import make_synthetic_batch
+
+__all__ = [
+    "MODEL_ENTRY_POINT_GROUP",
+    "Batch",
+    "BatchMetadata",
+    "ConformanceCheck",
+    "ConformanceError",
+    "Forecast",
+    "ForecastModel",
+    "Gaussian",
+    "GaussianParameters",
+    "GriddedDynamic",
+    "GriddedStatic",
+    "ModelFactory",
+    "OutputSpecification",
+    "Point",
+    "PointParameters",
+    "assert_conforms",
+    "check_conformance",
+    "make_synthetic_batch",
+    "__version__",
+]

hcx-0.1.0/src/hcx/batch.py

@@ -0,0 +1,35 @@
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+
+@dataclass(frozen=True)
+class GriddedDynamic:
+    values: torch.Tensor
+    coordinates: torch.Tensor
+    padding_mask: torch.Tensor
+
+
+@dataclass(frozen=True)
+class GriddedStatic:
+    values: torch.Tensor
+    coordinates: torch.Tensor
+    padding_mask: torch.Tensor
+
+
+@dataclass(frozen=True)
+class BatchMetadata:
+    sample_ids: tuple[str, ...]
+    input_end_indices: np.ndarray
+    target_fill_mask: np.ndarray
+
+
+@dataclass(frozen=True)
+class Batch:
+    scalar_dynamic: torch.Tensor | None
+    scalar_static: torch.Tensor | None
+    gridded_dynamic: dict[str, GriddedDynamic]
+    gridded_static: dict[str, GriddedStatic]
+    target: torch.Tensor
+    metadata: BatchMetadata

hcx-0.1.0/src/hcx/conformance.py

@@ -0,0 +1,220 @@
+from contextlib import suppress
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+from hcx.batch import Batch
+from hcx.output import Forecast
+from hcx.protocol import ForecastModel
+
+_SKIPPED = "skipped: forward did not return a Forecast"
+_QUADRANTS = ("scalar_dynamic", "scalar_static", "gridded_dynamic", "gridded_static")
+
+
+@dataclass(frozen=True)
+class ConformanceCheck:
+    name: str
+    passed: bool
+    detail: str
+
+
+class ConformanceError(AssertionError):
+    pass
+
+
+def _consumed_quadrants(model: torch.nn.Module, batch: Batch) -> ConformanceCheck:
+    declared = getattr(model, "consumed_quadrants", None)
+    if declared is None:
+        return ConformanceCheck("consumed_quadrants_present", True, "no consumed quadrants declared")
+    populated = {
+        "scalar_dynamic": batch.scalar_dynamic is not None,
+        "scalar_static": batch.scalar_static is not None,
+        "gridded_dynamic": bool(batch.gridded_dynamic),
+        "gridded_static": bool(batch.gridded_static),
+    }
+    try:
+        names = tuple(declared)
+        invalid = [name for name in names if name not in _QUADRANTS]
+        absent = [name for name in names if name in populated and not populated[name]]
+    except Exception as exc:
+        return ConformanceCheck("consumed_quadrants_present", False, f"declaration iteration failed: {exc!r}")
+    if invalid:
+        return ConformanceCheck("consumed_quadrants_present", False, f"invalid quadrants: {invalid!r}")
+    if absent:
+        return ConformanceCheck("consumed_quadrants_present", False, f"absent quadrants: {absent!r}")
+    return ConformanceCheck("consumed_quadrants_present", True, "all declared quadrants are populated")
+
+
+def _identity(name: str, got: object, original: object, snapshot: object) -> ConformanceCheck:
+    same_object = got is original
+    if isinstance(snapshot, tuple):
+        same_value = got == snapshot
+    elif isinstance(got, np.ndarray) and isinstance(snapshot, np.ndarray):
+        same_value = bool(np.array_equal(got, snapshot))
+    else:
+        same_value = False
+    passed = bool(same_object and same_value)
+    detail = "carried verbatim (original object, unchanged values)" if passed else "rebuilt or changed"
+    return ConformanceCheck(name, passed, detail)
+
+
+def _finiteness(out: Forecast) -> ConformanceCheck:
+    prediction = out.prediction
+    valid = isinstance(prediction, torch.Tensor) and bool(torch.isfinite(prediction).all())
+    detail = "prediction is finite"
+    if out.variance is not None:
+        variance = out.variance
+        valid_variance = (
+            isinstance(variance, torch.Tensor)
+            and isinstance(prediction, torch.Tensor)
+            and variance.shape == prediction.shape
+            and variance.dtype == prediction.dtype
+            and variance.device == prediction.device
+            and bool(torch.isfinite(variance).all())
+            and bool((variance > 0).all())
+        )
+        valid = valid and valid_variance
+        detail = "prediction finite; variance must match prediction and be finite and strictly positive"
+    return ConformanceCheck("finiteness", bool(valid), detail)
+
+
+def _trainability(model: torch.nn.Module, batch: Batch) -> ConformanceCheck:
+    try:
+        model.train()
+        model.zero_grad(set_to_none=True)
+        out = model.forward(batch)
+        if not isinstance(out, Forecast):
+            return ConformanceCheck("trainability_grad_flow", False, f"train forward returned {type(out).__name__}")
+        requires_grad = isinstance(out.prediction, torch.Tensor) and bool(out.prediction.requires_grad)
+        if requires_grad:
+            out.prediction.sum().backward()
+        gradient = any(
+            parameter.grad is not None
+            and bool(torch.isfinite(parameter.grad).all())
+            and float(parameter.grad.abs().sum()) > 0
+            for parameter in model.parameters()
+        )
+        return ConformanceCheck(
+            "trainability_grad_flow",
+            requires_grad and gradient,
+            f"requires_grad={requires_grad}, finite nonzero gradient={gradient}",
+        )
+    except Exception as exc:
+        return ConformanceCheck("trainability_grad_flow", False, f"train forward/backward raised {exc!r}")
+    finally:
+        with suppress(Exception):
+            model.zero_grad(set_to_none=True)
+
+
+def check_conformance(model: torch.nn.Module, batch: Batch) -> tuple[list[ConformanceCheck], Forecast | None]:
+    was_training = model.training
+    metadata = batch.metadata
+    sample_ids = metadata.sample_ids
+    input_end_indices = metadata.input_end_indices
+    target_fill_mask = metadata.target_fill_mask
+    input_end_snapshot = input_end_indices.copy()
+    target_fill_snapshot = target_fill_mask.copy()
+    checks = [
+        ConformanceCheck("structural_protocol", isinstance(model, ForecastModel), "isinstance(model, ForecastModel)"),
+        _consumed_quadrants(model, batch),
+    ]
+    output: Forecast | None = None
+    try:
+        model.eval()
+        try:
+            with torch.no_grad():
+                produced = model.forward(batch)
+            if isinstance(produced, Forecast):
+                output = produced
+                checks.append(ConformanceCheck("returns_forecast", True, "forward returned Forecast"))
+            else:
+                checks.append(
+                    ConformanceCheck("returns_forecast", False, f"forward returned {type(produced).__name__}")
+                )
+        except Exception as exc:
+            checks.append(ConformanceCheck("returns_forecast", False, f"forward raised {exc!r}"))
+
+        if output is None:
+            for name in (
+                "prediction_tensor_dtype_device",
+                "prediction_shape_exact",
+                "identity_sample_ids_verbatim",
+                "identity_input_end_indices_verbatim",
+                "identity_target_fill_mask_verbatim",
+                "finiteness",
+            ):
+                checks.append(ConformanceCheck(name, False, _SKIPPED))
+            checks.append(ConformanceCheck("stable_across_two_eval_forwards", False, _SKIPPED))
+        else:
+            prediction = output.prediction
+            tensor_match = (
+                isinstance(prediction, torch.Tensor)
+                and prediction.dtype == batch.target.dtype
+                and prediction.device == batch.target.device
+            )
+            checks.append(
+                ConformanceCheck("prediction_tensor_dtype_device", tensor_match, "must match target dtype/device")
+            )
+            shape = tuple(prediction.shape) if isinstance(prediction, torch.Tensor) else None
+            expected = (batch.target.shape[0], batch.target.shape[-1])
+            checks.append(
+                ConformanceCheck("prediction_shape_exact", shape == expected, f"expected {expected}, got {shape}")
+            )
+            checks.append(_identity("identity_sample_ids_verbatim", output.sample_ids, sample_ids, sample_ids))
+            checks.append(
+                _identity(
+                    "identity_input_end_indices_verbatim",
+                    output.input_end_indices,
+                    input_end_indices,
+                    input_end_snapshot,
+                )
+            )
+            checks.append(
+                _identity(
+                    "identity_target_fill_mask_verbatim",
+                    output.target_fill_mask,
+                    target_fill_mask,
+                    target_fill_snapshot,
+                )
+            )
+            checks.append(_finiteness(output))
+            try:
+                model.eval()
+                with torch.no_grad():
+                    again = model.forward(batch)
+                stable = (
+                    isinstance(again, Forecast)
+                    and torch.equal(output.prediction, again.prediction)
+                    and (
+                        (output.variance is None and again.variance is None)
+                        or (
+                            isinstance(output.variance, torch.Tensor)
+                            and isinstance(again.variance, torch.Tensor)
+                            and torch.equal(output.variance, again.variance)
+                        )
+                    )
+                )
+                detail = "adjacent eval forwards equal; cross-run/device determinism is documented, not enforced"
+                checks.append(ConformanceCheck("stable_across_two_eval_forwards", bool(stable), detail))
+            except Exception as exc:
+                checks.append(
+                    ConformanceCheck("stable_across_two_eval_forwards", False, f"second eval forward raised {exc!r}")
+                )
+        checks.append(_trainability(model, batch))
+        return checks, output
+    finally:
+        model.train(was_training)
+
+
+def assert_conforms(model: torch.nn.Module, batch: Batch) -> Forecast:
+    checks, forecast = check_conformance(model, batch)
+    for check in checks:
+        if not check.passed:
+            raise ConformanceError(f"{check.name}: {check.detail}")
+    if forecast is None:
+        raise ConformanceError("harness error: all checks passed without a forecast")
+    return forecast
+
+
+__all__ = ["ConformanceCheck", "ConformanceError", "assert_conforms", "check_conformance"]

hcx-0.1.0/src/hcx/models/__init__.py

@@ -0,0 +1,3 @@
+from hcx.models.lstm import LSTMConfig, ScalarLSTM, factory
+
+__all__ = ["LSTMConfig", "ScalarLSTM", "factory"]

hcx-0.1.0/src/hcx/models/lstm.py

@@ -0,0 +1,165 @@
+"""A scalar-input LSTM reference model."""
+
+import math
+from dataclasses import dataclass, fields
+from numbers import Real
+
+import torch
+from torch import nn
+
+from hcx.batch import Batch
+from hcx.output import Forecast
+from hcx.specifications import OutputSpecification
+
+
+@dataclass(frozen=True)
+class LSTMConfig:
+    hidden_size: int = 64
+    num_layers: int = 1
+    dropout: float = 0.0
+    decoder_hidden_size: int = 32
+
+
+def _parse_config(values: dict[str, object]) -> LSTMConfig:
+    allowed = {field.name for field in fields(LSTMConfig)}
+    for key, value in values.items():
+        if key not in allowed:
+            raise ValueError(f"unknown model config key {key!r} with value {value!r}")
+
+    defaults = LSTMConfig()
+    for key in ("hidden_size", "num_layers", "decoder_hidden_size"):
+        value = values.get(key, getattr(defaults, key))
+        if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+            raise ValueError(f"{key} must be a positive integer; got {value!r}")
+    dropout_value = values.get("dropout", defaults.dropout)
+    if isinstance(dropout_value, bool) or not isinstance(dropout_value, Real):
+        raise ValueError(f"dropout must be a finite real number in [0, 1); got {dropout_value!r}")
+    dropout = float(dropout_value)
+    if not math.isfinite(dropout) or not 0 <= dropout < 1:
+        raise ValueError(f"dropout must be a finite real number in [0, 1); got {dropout_value!r}")
+    hidden_size = values.get("hidden_size", defaults.hidden_size)
+    num_layers = values.get("num_layers", defaults.num_layers)
+    decoder_hidden_size = values.get("decoder_hidden_size", defaults.decoder_hidden_size)
+    assert isinstance(hidden_size, int) and not isinstance(hidden_size, bool)
+    assert isinstance(num_layers, int) and not isinstance(num_layers, bool)
+    assert isinstance(decoder_hidden_size, int) and not isinstance(decoder_hidden_size, bool)
+    return LSTMConfig(
+        hidden_size=hidden_size,
+        num_layers=num_layers,
+        dropout=dropout,
+        decoder_hidden_size=decoder_hidden_size,
+    )
+
+
+def _validate_size(key: str, value: object, *, allow_zero: bool = False) -> None:
+    valid_range = value >= 0 if isinstance(value, int) else False
+    if isinstance(value, bool) or not isinstance(value, int) or not valid_range or (not allow_zero and value == 0):
+        requirement = "a nonnegative integer" if allow_zero else "a positive integer"
+        raise ValueError(f"{key} must be {requirement}; got {value!r}")
+
+
+def _validate_names(key: str, names: list[str], size: int) -> tuple[str, ...]:
+    if len(names) != size:
+        raise ValueError(f"{key} count must equal its supplied size {size}; got {len(names)}")
+    for name in names:
+        if not isinstance(name, str):
+            raise ValueError(f"{key} names must be strings; got {name!r}")
+    if len(set(names)) != len(names):
+        raise ValueError(f"{key} names must be unique; got {names!r}")
+    return tuple(names)
+
+
+class ScalarLSTM(nn.Module):
+    def __init__(
+        self,
+        config: LSTMConfig,
+        input_size: int,
+        static_size: int,
+        output_size: int,
+        dynamic_inputs: tuple[str, ...],
+        static_inputs: tuple[str, ...],
+        output_specification: OutputSpecification[object],
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.input_size = input_size
+        self.static_size = static_size
+        self.output_size = output_size
+        self.dynamic_inputs = dynamic_inputs
+        self.static_inputs = static_inputs
+        self.output_specification = output_specification
+        self.consumed_quadrants = ("scalar_dynamic",) if static_size == 0 else ("scalar_dynamic", "scalar_static")
+
+        self.embedding = nn.Linear(input_size + static_size, config.hidden_size)
+        self.lstm = nn.LSTM(
+            config.hidden_size,
+            config.hidden_size,
+            config.num_layers,
+            dropout=config.dropout if config.num_layers > 1 else 0,
+            batch_first=True,
+        )
+        self.decoder = nn.Sequential(
+            nn.Linear(config.hidden_size, config.decoder_hidden_size),
+            nn.ReLU(),
+            nn.Dropout(config.dropout),
+            nn.Linear(config.decoder_hidden_size, output_size * output_specification.raw_head_width),
+        )
+
+    def forward(self, batch: Batch) -> Forecast:
+        x = batch.scalar_dynamic
+        if x is None:
+            raise ValueError("scalar dynamic input is required")
+        if x.shape[-1] != self.input_size:
+            raise ValueError(f"scalar dynamic width must be {self.input_size}; got {x.shape[-1]}")
+        batch_size, input_length, _ = x.shape
+
+        if self.static_size > 0:
+            static = batch.scalar_static
+            if static is None:
+                raise ValueError(f"scalar static input of width {self.static_size} is required")
+            if static.shape[-1] != self.static_size:
+                raise ValueError(f"scalar static width must be {self.static_size}; got {static.shape[-1]}")
+            x = torch.cat([x, static.unsqueeze(1).expand(-1, input_length, -1)], dim=-1)
+        elif batch.scalar_static is not None:
+            raise ValueError("scalar static input was supplied when static_size is zero")
+
+        embedded = self.embedding(x)
+        recurrent, _ = self.lstm(embedded)
+        decoded = self.decoder(recurrent[:, -1, :])
+        raw = decoded.view(batch_size, self.output_size, self.output_specification.raw_head_width)
+        parameters = self.output_specification.parameterize(raw)
+        return self.output_specification.populate_forecast(
+            parameters,
+            sample_ids=batch.metadata.sample_ids,
+            input_end_indices=batch.metadata.input_end_indices,
+            target_fill_mask=batch.metadata.target_fill_mask,
+        )
+
+
+def factory(
+    model_config: dict[str, object],
+    *,
+    dynamic_inputs: list[str],
+    static_inputs: list[str],
+    input_size: int,
+    static_size: int,
+    output_size: int,
+    output_specification: OutputSpecification[object],
+) -> torch.nn.Module:
+    _validate_size("input_size", input_size)
+    _validate_size("static_size", static_size, allow_zero=True)
+    _validate_size("output_size", output_size)
+    dynamic_names = _validate_names("dynamic_inputs", dynamic_inputs, input_size)
+    static_names = _validate_names("static_inputs", static_inputs, static_size)
+    return ScalarLSTM(
+        _parse_config(model_config),
+        input_size,
+        static_size,
+        output_size,
+        dynamic_names,
+        static_names,
+        output_specification,
+    )
+
+
+__all__ = ["LSTMConfig", "ScalarLSTM", "factory"]

hcx-0.1.0/src/hcx/output.py

@@ -0,0 +1,15 @@
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+
+@dataclass(frozen=True)
+class Forecast:
+    """A forecast whose values are in target space."""
+
+    prediction: torch.Tensor
+    sample_ids: tuple[str, ...]
+    input_end_indices: np.ndarray
+    target_fill_mask: np.ndarray
+    variance: torch.Tensor | None = None

hcx-0.1.0/src/hcx/protocol.py

@@ -0,0 +1,29 @@
+from typing import Protocol, runtime_checkable
+
+import torch
+
+from hcx.batch import Batch
+from hcx.output import Forecast
+from hcx.specifications import OutputSpecification
+
+MODEL_ENTRY_POINT_GROUP = "hcx.models"
+
+
+@runtime_checkable
+class ForecastModel(Protocol):
+    def forward(self, batch: Batch) -> Forecast: ...
+
+
+@runtime_checkable
+class ModelFactory(Protocol):
+    def __call__(
+        self,
+        model_config: dict[str, object],
+        *,
+        dynamic_inputs: list[str],
+        static_inputs: list[str],
+        input_size: int,
+        static_size: int,
+        output_size: int,
+        output_specification: OutputSpecification[object],
+    ) -> torch.nn.Module: ...

hcx-0.1.0/src/hcx/specifications.py

@@ -0,0 +1,92 @@
+from dataclasses import dataclass
+from typing import ClassVar, Protocol, TypeVar, runtime_checkable
+
+import numpy as np
+import torch
+import torch.nn.functional as functional
+
+from hcx.output import Forecast
+
+ParametersT = TypeVar("ParametersT")
+
+
+@runtime_checkable
+class OutputSpecification(Protocol[ParametersT]):
+    raw_head_width: int
+
+    def parameterize(self, raw: torch.Tensor) -> ParametersT: ...
+
+    def populate_forecast(
+        self,
+        parameters: ParametersT,
+        *,
+        sample_ids: tuple[str, ...],
+        input_end_indices: np.ndarray,
+        target_fill_mask: np.ndarray,
+    ) -> Forecast: ...
+
+
+def _validate_raw_head(raw: torch.Tensor, expected_width: int) -> None:
+    expected = f"[B, T_out, {expected_width}]"
+    actual = tuple(raw.shape)
+    if raw.ndim != 3 or raw.shape[-1] != expected_width:
+        raise ValueError(f"expected shape {expected}; actual shape {actual}")
+
+
+@dataclass(frozen=True)
+class PointParameters:
+    prediction: torch.Tensor
+
+
+@dataclass(frozen=True)
+class GaussianParameters:
+    prediction: torch.Tensor
+    variance: torch.Tensor
+
+
+@dataclass(frozen=True)
+class Point:
+    raw_head_width: ClassVar[int] = 1
+
+    def parameterize(self, raw: torch.Tensor) -> PointParameters:
+        _validate_raw_head(raw, self.raw_head_width)
+        return PointParameters(prediction=raw[..., 0])
+
+    def populate_forecast(
+        self,
+        parameters: PointParameters,
+        *,
+        sample_ids: tuple[str, ...],
+        input_end_indices: np.ndarray,
+        target_fill_mask: np.ndarray,
+    ) -> Forecast:
+        return Forecast(parameters.prediction, sample_ids, input_end_indices, target_fill_mask)
+
+
+@dataclass(frozen=True)
+class Gaussian:
+    raw_head_width: ClassVar[int] = 2
+
+    def parameterize(self, raw: torch.Tensor) -> GaussianParameters:
+        _validate_raw_head(raw, self.raw_head_width)
+        raw_variance = raw[..., 1]
+        return GaussianParameters(
+            prediction=raw[..., 0],
+            variance=functional.softplus(raw_variance) + torch.finfo(raw_variance.dtype).tiny,
+        )
+
+    def populate_forecast(
+        self,
+        parameters: GaussianParameters,
+        *,
+        sample_ids: tuple[str, ...],
+        input_end_indices: np.ndarray,
+        target_fill_mask: np.ndarray,
+    ) -> Forecast:
+        return Forecast(
+            parameters.prediction,
+            sample_ids,
+            input_end_indices,
+            target_fill_mask,
+            parameters.variance,
+        )

hcx-0.1.0/src/hcx/synthetic.py

@@ -0,0 +1,90 @@
+import numpy as np
+import torch
+
+from hcx.batch import Batch, BatchMetadata, GriddedDynamic, GriddedStatic
+
+
+def make_synthetic_batch(
+    *,
+    batch_size: int = 4,
+    input_length: int = 6,
+    output_length: int = 2,
+    scalar_dynamic_features: int = 3,
+    scalar_static_features: int = 2,
+    grid_cells: int = 5,
+    gridded_dynamic_features: int = 2,
+    gridded_static_features: int = 3,
+    dtype: torch.dtype = torch.float32,
+    device: torch.device | str = "cpu",
+    seed: int = 1729,
+    include_scalar_dynamic: bool = True,
+    include_scalar_static: bool = True,
+    include_gridded_dynamic: bool = True,
+    include_gridded_static: bool = True,
+) -> Batch:
+    dimensions = {
+        "batch_size": batch_size,
+        "input_length": input_length,
+        "output_length": output_length,
+        "grid_cells": grid_cells,
+    }
+    if invalid := [name for name, value in dimensions.items() if value <= 0]:
+        raise ValueError(f"dimensions must be positive: {invalid!r}")
+    counts = {
+        "scalar_dynamic": scalar_dynamic_features,
+        "scalar_static": scalar_static_features,
+        "gridded_dynamic": gridded_dynamic_features,
+        "gridded_static": gridded_static_features,
+    }
+    if invalid := [name for name, value in counts.items() if value < 0]:
+        raise ValueError(f"feature counts must be nonnegative: {invalid!r}")
+    included = {
+        "scalar_dynamic": include_scalar_dynamic,
+        "scalar_static": include_scalar_static,
+        "gridded_dynamic": include_gridded_dynamic,
+        "gridded_static": include_gridded_static,
+    }
+    if invalid := [name for name, enabled in included.items() if enabled and counts[name] == 0]:
+        raise ValueError(f"included quadrants must have features: {invalid!r}")
+
+    device = torch.device(device)
+    generator = torch.Generator(device=device)
+    generator.manual_seed(seed)
+
+    def randn(shape: tuple[int, ...]) -> torch.Tensor:
+        return torch.randn(shape, dtype=dtype, device=device, generator=generator)
+
+    def rand(shape: tuple[int, ...]) -> torch.Tensor:
+        return torch.rand(shape, dtype=dtype, device=device, generator=generator)
+
+    scalar_dynamic = randn((batch_size, input_length, scalar_dynamic_features)) if include_scalar_dynamic else None
+    scalar_static = randn((batch_size, scalar_static_features)) if include_scalar_static else None
+
+    gridded_dynamic: dict[str, GriddedDynamic] = {}
+    coordinates: torch.Tensor | None = None
+    padding_mask: torch.Tensor | None = None
+    if include_gridded_dynamic:
+        values = randn((batch_size, input_length, grid_cells, gridded_dynamic_features))
+        coordinates = rand((batch_size, grid_cells, 2))
+        padding_mask = torch.zeros((batch_size, grid_cells), dtype=torch.bool, device=device)
+        gridded_dynamic["meteorology"] = GriddedDynamic(values, coordinates, padding_mask)
+
+    gridded_static: dict[str, GriddedStatic] = {}
+    if include_gridded_static:
+        values = randn((batch_size, grid_cells, gridded_static_features))
+        if coordinates is None:
+            coordinates = rand((batch_size, grid_cells, 2))
+            padding_mask = torch.zeros((batch_size, grid_cells), dtype=torch.bool, device=device)
+        assert padding_mask is not None
+        gridded_static["physiography"] = GriddedStatic(values, coordinates, padding_mask)
+
+    target = randn((batch_size, output_length))
+    metadata = BatchMetadata(
+        sample_ids=tuple(f"sample-{index:03d}" for index in range(batch_size)),
+        input_end_indices=np.full((batch_size,), input_length - 1, dtype=np.int64),
+        target_fill_mask=np.zeros((batch_size, output_length), dtype=np.int8),
+    )
+    return Batch(scalar_dynamic, scalar_static, gridded_dynamic, gridded_static, target, metadata)
+
+
+__all__ = ["make_synthetic_batch"]