dcoupler 0.2.0__py3-none-any.whl

dcoupler/__init__.py

@@ -0,0 +1,52 @@
+from dcoupler.core import (
+    CouplingGraph,
+    DifferentiableComponent,
+    FluxDirection,
+    FluxSpec,
+    GradientMethod,
+    ParameterSpec,
+    FluxConnection,
+    SpatialRemapper,
+    TemporalOrchestrator,
+    ConservationChecker,
+    BMIMixin,
+)
+from dcoupler.optimization import (
+    ParameterManager,
+    MultiObservationLoss,
+    LossTerm,
+    Trainer,
+    TrainingResult,
+)
+from dcoupler import observers as _observers
+from dcoupler import diagnostics as _diagnostics
+from dcoupler import utils as _utils
+from dcoupler.observers import *  # noqa: F401,F403
+from dcoupler.diagnostics import *  # noqa: F401,F403
+from dcoupler import losses as _losses
+from dcoupler import components as _components
+from dcoupler import wrappers as _wrappers
+from dcoupler.losses import *  # noqa: F401,F403
+from dcoupler.components import *  # noqa: F401,F403
+from dcoupler.wrappers import *  # noqa: F401,F403
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "CouplingGraph",
+    "DifferentiableComponent",
+    "FluxDirection",
+    "FluxSpec",
+    "GradientMethod",
+    "ParameterSpec",
+    "FluxConnection",
+    "SpatialRemapper",
+    "TemporalOrchestrator",
+    "ConservationChecker",
+    "BMIMixin",
+    "ParameterManager",
+    "MultiObservationLoss",
+    "LossTerm",
+    "Trainer",
+    "TrainingResult",
+] + _losses.__all__ + _components.__all__ + _wrappers.__all__ + _observers.__all__ + _diagnostics.__all__ + _utils.__all__

dcoupler/components/__init__.py

@@ -0,0 +1,13 @@
+__all__ = []
+
+try:
+    from .fuse import FUSEComponent
+    __all__.append("FUSEComponent")
+except Exception:
+    FUSEComponent = None  # type: ignore
+
+try:
+    from .routing import MuskingumCungeRouting
+    __all__.append("MuskingumCungeRouting")
+except Exception:
+    MuskingumCungeRouting = None  # type: ignore

dcoupler/components/fuse.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+from typing import Dict, List, Optional
+import torch
+import torch.nn as nn
+
+from dcoupler.core.component import (
+    DifferentiableComponent,
+    FluxDirection,
+    FluxSpec,
+    GradientMethod,
+    ParameterSpec,
+)
+
+try:
+    import cfuse
+    import cfuse_core
+    from cfuse.torch import DifferentiableFUSEBatch
+except ImportError:
+    cfuse = None
+    cfuse_core = None
+    DifferentiableFUSEBatch = None
+
+
+class FUSEComponent(DifferentiableComponent, nn.Module):
+    """FUSE component wrapper using DifferentiableFUSEBatch."""
+
+    def __init__(
+        self,
+        name: str,
+        fuse_config,
+        n_hrus: int,
+        dt: float = 86400.0,
+        spatial_params: bool = True,
+        n_states: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+        if DifferentiableFUSEBatch is None or cfuse is None:
+            raise RuntimeError("cfuse and DifferentiableFUSEBatch are required for FUSEComponent")
+        if fuse_config is None:
+            raise ValueError("fuse_config is required")
+
+        self._name = name
+        self.fuse_config = fuse_config
+        self.config_dict = fuse_config.to_dict() if hasattr(fuse_config, "to_dict") else fuse_config
+        self.n_hrus = n_hrus
+        self.dt_seconds = float(dt)
+        self.dt_days = self.dt_seconds / 86400.0
+        self.spatial_params = spatial_params
+
+        if n_states is None:
+            if cfuse_core is None:
+                raise RuntimeError("cfuse_core required to infer state size")
+            self.n_states = int(cfuse_core.get_num_active_states(self.config_dict))
+        else:
+            self.n_states = int(n_states)
+
+        self.param_names = list(cfuse.PARAM_NAMES)
+        self.n_params = len(self.param_names)
+        lowers = torch.tensor([cfuse.PARAM_BOUNDS[n][0] for n in self.param_names], dtype=torch.float32)
+        uppers = torch.tensor([cfuse.PARAM_BOUNDS[n][1] for n in self.param_names], dtype=torch.float32)
+        self.register_buffer("param_lower", lowers)
+        self.register_buffer("param_upper", uppers)
+
+        self._raw_params = nn.ParameterList()
+        for _ in range(self.n_params):
+            if self.spatial_params:
+                init = torch.zeros(self.n_hrus) + torch.randn(self.n_hrus) * 0.2
+            else:
+                init = torch.zeros(())
+            self._raw_params.append(nn.Parameter(init))
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def input_fluxes(self) -> List[FluxSpec]:
+        return [
+            FluxSpec(
+                name="forcing",
+                units="mixed",
+                direction=FluxDirection.INPUT,
+                spatial_type="hru",
+                temporal_resolution=self.dt_seconds,
+                dims=("time", "hru", "var"),
+                optional=False,
+            )
+        ]
+
+    @property
+    def output_fluxes(self) -> List[FluxSpec]:
+        return [
+            FluxSpec(
+                name="runoff",
+                units="mm/day",
+                direction=FluxDirection.OUTPUT,
+                spatial_type="hru",
+                temporal_resolution=self.dt_seconds,
+                dims=("time", "hru"),
+                conserved_quantity="water_mass",
+            )
+        ]
+
+    @property
+    def parameters(self) -> List[ParameterSpec]:
+        specs: List[ParameterSpec] = []
+        for i, name in enumerate(self.param_names):
+            specs.append(
+                ParameterSpec(
+                    name=name,
+                    lower_bound=float(self.param_lower[i].item()),
+                    upper_bound=float(self.param_upper[i].item()),
+                    spatial=self.spatial_params,
+                    n_spatial=self.n_hrus if self.spatial_params else None,
+                    log_transform=False,
+                )
+            )
+        return specs
+
+    @property
+    def gradient_method(self) -> GradientMethod:
+        return GradientMethod.ENZYME
+
+    @property
+    def state_size(self) -> int:
+        return self.n_states
+
+    @property
+    def requires_batch(self) -> bool:
+        return True
+
+    def get_initial_state(self) -> torch.Tensor:
+        state = torch.zeros(self.n_hrus, self.n_states)
+        if self.n_states > 0:
+            state[:, 0] = 50.0
+        if self.n_states > 1:
+            state[:, 1] = 20.0
+        if self.n_states > 2:
+            state[:, 2] = 200.0
+        return state
+
+    def _raw_param_matrix(self) -> torch.Tensor:
+        if self.spatial_params:
+            return torch.stack(list(self._raw_params), dim=1)
+        return torch.stack(list(self._raw_params), dim=0)
+
+    def get_physical_parameters(self) -> Dict[str, torch.Tensor]:
+        raw = self._raw_param_matrix()
+        if self.spatial_params:
+            phys = self.param_lower + (self.param_upper - self.param_lower) * torch.sigmoid(raw)
+        else:
+            phys = self.param_lower + (self.param_upper - self.param_lower) * torch.sigmoid(raw)
+        return {name: phys[:, i] if self.spatial_params else phys[i] for i, name in enumerate(self.param_names)}
+
+    def _physical_param_tensor(self) -> torch.Tensor:
+        raw = self._raw_param_matrix()
+        if self.spatial_params:
+            return self.param_lower + (self.param_upper - self.param_lower) * torch.sigmoid(raw)
+        return self.param_lower + (self.param_upper - self.param_lower) * torch.sigmoid(raw)
+
+    def step(
+        self,
+        inputs: Dict[str, torch.Tensor],
+        state: torch.Tensor,
+        dt: float,
+    ):
+        raise RuntimeError("FUSEComponent requires batch execution")
+
+    def run(
+        self,
+        inputs: Dict[str, torch.Tensor],
+        state: torch.Tensor,
+        dt: float,
+        n_timesteps: int,
+    ):
+        forcing = inputs.get("forcing")
+        if forcing is None:
+            raise ValueError("Missing required input 'forcing'")
+        phys_params = self._physical_param_tensor()
+        if state is None:
+            state = self.get_initial_state().to(forcing.device)
+        runoff = DifferentiableFUSEBatch.apply(
+            phys_params,
+            state,
+            forcing,
+            self.config_dict,
+            self.dt_days,
+        )
+        return {"runoff": runoff}, state
+
+    def get_torch_parameters(self) -> List[torch.nn.Parameter]:
+        return list(self._raw_params)

dcoupler/components/routing.py

@@ -0,0 +1,219 @@
+from __future__ import annotations
+
+from typing import Dict, List, Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+import xarray as xr
+
+import droute as dmc
+
+from dcoupler.core.component import (
+    DifferentiableComponent,
+    FluxDirection,
+    FluxSpec,
+    GradientMethod,
+    ParameterSpec,
+)
+from dcoupler.wrappers.enzyme import DifferentiableRouting
+
+
+class MuskingumCungeRouting(DifferentiableComponent, nn.Module):
+    """Muskingum-Cunge routing component with Enzyme AD."""
+
+    def __init__(
+        self,
+        name: str,
+        topology_file: str,
+        hru_areas: np.ndarray,
+        dt: float = 86400.0,
+        outlet_reach_id: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+        self._name = name
+        self.dt_seconds = float(dt)
+
+        self.network = self._load_network(topology_file)
+        self.n_reaches = self.network.num_reaches()
+        self.n_hrus = len(hru_areas)
+
+        topo_order = list(self.network.topological_order())
+        self.reach_ids = topo_order
+        self.id_to_idx = {rid: i for i, rid in enumerate(topo_order)}
+
+        if outlet_reach_id is None:
+            outlet_reach_id = int(topo_order[-1])
+        self.outlet_reach_id = int(outlet_reach_id)
+
+        config = dmc.RouterConfig()
+        config.dt = self.dt_seconds
+        config.num_substeps = 4
+        config.enable_gradients = False
+
+        self.router = dmc.MuskingumCungeRouter(self.network, config)
+
+        initial_log_n = torch.full((self.n_reaches,), np.log(0.035))
+        initial_log_n = initial_log_n + torch.randn(self.n_reaches) * 0.1
+        self.log_manning_n = nn.Parameter(initial_log_n)
+
+        self.register_buffer(
+            "mapping_matrix",
+            self._build_mapping_matrix(topology_file, hru_areas),
+        )
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def input_fluxes(self) -> List[FluxSpec]:
+        return [
+            FluxSpec(
+                name="lateral_inflow",
+                units="m3/s",
+                direction=FluxDirection.INPUT,
+                spatial_type="reach",
+                temporal_resolution=self.dt_seconds,
+                dims=("time", "reach"),
+                optional=False,
+            )
+        ]
+
+    @property
+    def output_fluxes(self) -> List[FluxSpec]:
+        return [
+            FluxSpec(
+                name="discharge",
+                units="m3/s",
+                direction=FluxDirection.OUTPUT,
+                spatial_type="point",
+                temporal_resolution=self.dt_seconds,
+                dims=("time",),
+                conserved_quantity="water_mass",
+            )
+        ]
+
+    @property
+    def parameters(self) -> List[ParameterSpec]:
+        return [
+            ParameterSpec(
+                name="manning_n",
+                lower_bound=1e-4,
+                upper_bound=1.0,
+                spatial=True,
+                n_spatial=self.n_reaches,
+                log_transform=True,
+            )
+        ]
+
+    @property
+    def gradient_method(self) -> GradientMethod:
+        return GradientMethod.ENZYME
+
+    @property
+    def state_size(self) -> int:
+        return 0
+
+    @property
+    def requires_batch(self) -> bool:
+        return True
+
+    def get_initial_state(self) -> torch.Tensor:
+        return torch.empty(0)
+
+    def get_physical_parameters(self) -> Dict[str, torch.Tensor]:
+        return {"manning_n": torch.exp(self.log_manning_n)}
+
+    def step(self, inputs: Dict[str, torch.Tensor], state: torch.Tensor, dt: float):
+        raise RuntimeError("MuskingumCungeRouting requires batch execution")
+
+    def run(
+        self,
+        inputs: Dict[str, torch.Tensor],
+        state: torch.Tensor,
+        dt: float,
+        n_timesteps: int,
+    ):
+        lateral = inputs.get("lateral_inflow")
+        if lateral is None:
+            raise ValueError("Missing required input 'lateral_inflow'")
+        manning_n = torch.exp(self.log_manning_n)
+        discharge = DifferentiableRouting.apply(
+            lateral,
+            manning_n,
+            self.router,
+            self.network,
+            self.outlet_reach_id,
+            self.dt_seconds,
+        )
+        return {"discharge": discharge}, state
+
+    def get_torch_parameters(self) -> List[torch.nn.Parameter]:
+        return [self.log_manning_n]
+
+    def _load_network(self, topology_file: str) -> dmc.Network:
+        ds = xr.open_dataset(topology_file)
+
+        seg_ids = ds["segId"].values.astype(int)
+        down_seg_ids = ds["downSegId"].values.astype(int)
+        lengths = ds["length"].values.astype(float)
+        slopes = ds["slope"].values.astype(float)
+        mann_n = ds["mann_n"].values if "mann_n" in ds else np.full(len(seg_ids), 0.035)
+
+        network = dmc.Network()
+        seg_id_set = set(seg_ids)
+
+        upstream_map = {int(sid): [] for sid in seg_ids}
+        for i, down_id in enumerate(down_seg_ids):
+            if int(down_id) in seg_id_set:
+                upstream_map[int(down_id)].append(int(seg_ids[i]))
+
+        for i, sid in enumerate(seg_ids):
+            reach = dmc.Reach()
+            reach.id = int(sid)
+            reach.length = float(lengths[i])
+            reach.slope = max(float(slopes[i]), 0.0001)
+            reach.manning_n = float(mann_n[i])
+            reach.geometry.width_coef = 7.2
+            reach.geometry.width_exp = 0.5
+            reach.geometry.depth_coef = 0.27
+            reach.geometry.depth_exp = 0.3
+            reach.upstream_junction_id = int(sid)
+            down_id = int(down_seg_ids[i])
+            reach.downstream_junction_id = down_id if down_id in seg_id_set else -1
+            network.add_reach(reach)
+
+        for i, sid in enumerate(seg_ids):
+            junc = dmc.Junction()
+            junc.id = int(sid)
+            junc.upstream_reach_ids = upstream_map[int(sid)]
+            junc.downstream_reach_ids = [int(sid)]
+            network.add_junction(junc)
+
+        network.build_topology()
+        ds.close()
+        return network
+
+    def _build_mapping_matrix(self, topology_file: str, hru_areas: np.ndarray) -> torch.Tensor:
+        ds = xr.open_dataset(topology_file)
+        hru_to_seg = ds["hruToSegId"].values.astype(int)
+        ds.close()
+
+        topo_order = self.network.topological_order()
+        id_to_idx = {rid: i for i, rid in enumerate(topo_order)}
+
+        n_hrus = len(hru_to_seg)
+        n_reaches = len(topo_order)
+        mapping = torch.zeros((n_reaches, n_hrus), dtype=torch.float32)
+
+        for h_idx, seg_id in enumerate(hru_to_seg):
+            if seg_id in id_to_idx:
+                r_idx = id_to_idx[seg_id]
+                conversion = hru_areas[h_idx] / 1000.0 / 86400.0
+                mapping[r_idx, h_idx] = conversion
+
+        indices = mapping.nonzero(as_tuple=False).T
+        values = mapping[indices[0], indices[1]]
+        sparse = torch.sparse_coo_tensor(indices, values, mapping.shape)
+        return sparse

dcoupler/core/__init__.py

@@ -0,0 +1,26 @@
+from .component import (
+    DifferentiableComponent,
+    FluxDirection,
+    FluxSpec,
+    GradientMethod,
+    ParameterSpec,
+)
+from .connection import FluxConnection, SpatialRemapper
+from .temporal import TemporalOrchestrator
+from .conservation import ConservationChecker
+from .graph import CouplingGraph
+from .bmi import BMIMixin
+
+__all__ = [
+    "DifferentiableComponent",
+    "FluxDirection",
+    "FluxSpec",
+    "GradientMethod",
+    "ParameterSpec",
+    "FluxConnection",
+    "SpatialRemapper",
+    "TemporalOrchestrator",
+    "ConservationChecker",
+    "CouplingGraph",
+    "BMIMixin",
+]

dcoupler/core/bmi.py

@@ -0,0 +1,76 @@
+"""BMI-aligned lifecycle mixin for dCoupler components.
+
+The Basic Model Interface (BMI) is a standardized interface for model
+coupling.  This mixin maps BMI lifecycle methods to the DifferentiableComponent
+protocol so that existing components can be used in BMI-based workflows
+without modifying their core implementation.
+"""
+
+from __future__ import annotations
+
+import abc
+from typing import Any, Dict, List
+
+
+class BMIMixin(abc.ABC):
+    """BMI-aligned lifecycle for dCoupler components.
+
+    Subclasses implement these methods to expose a BMI-style interface
+    while still participating in the dCoupler CouplingGraph.
+
+    BMI ↔ DifferentiableComponent mapping:
+        bmi_initialize(config)    → initialize(config) + get_initial_state()
+        bmi_update(inputs, dt)    → step(inputs, state, dt)
+        bmi_update_batch(…, n)    → run(inputs, state, dt, n)
+        bmi_get_state()           → return internal state tensor
+        bmi_set_state(state)      → set internal state tensor
+        bmi_get_value(name)       → index into last output dict
+        bmi_finalize()            → cleanup hook
+    """
+
+    @abc.abstractmethod
+    def bmi_initialize(self, config: dict) -> None:
+        """One-time setup: parse config, allocate state, load data."""
+
+    @abc.abstractmethod
+    def bmi_update(self, inputs: Dict[str, Any], dt: float) -> Dict[str, Any]:
+        """Advance one timestep, return outputs dict."""
+
+    def bmi_update_batch(
+        self, inputs: Dict[str, Any], dt: float, n_timesteps: int
+    ) -> Dict[str, Any]:
+        """Advance multiple timesteps. Default: loop bmi_update."""
+        all_outputs: Dict[str, list] = {}
+        for _ in range(n_timesteps):
+            outputs = self.bmi_update(inputs, dt)
+            for k, v in outputs.items():
+                all_outputs.setdefault(k, []).append(v)
+        return all_outputs
+
+    @abc.abstractmethod
+    def bmi_finalize(self) -> None:
+        """Release resources."""
+
+    @abc.abstractmethod
+    def bmi_get_state(self) -> Any:
+        """Return the current internal state."""
+
+    @abc.abstractmethod
+    def bmi_set_state(self, state: Any) -> None:
+        """Overwrite the current internal state."""
+
+    @abc.abstractmethod
+    def bmi_get_value(self, name: str) -> Any:
+        """Get a named output/state variable by name."""
+
+    @abc.abstractmethod
+    def bmi_set_value(self, name: str, value: Any) -> None:
+        """Set a named input variable by name."""
+
+    @abc.abstractmethod
+    def bmi_get_output_var_names(self) -> List[str]:
+        """Return list of output variable names."""
+
+    @abc.abstractmethod
+    def bmi_get_input_var_names(self) -> List[str]:
+        """Return list of input variable names."""