Refactor DecisionRule to Approximator strategy and remove multipledispatch

Author: albop

pixi.lock

@@ -46,7 +46,6 @@ matplotlib = ">=3.6.0,<4"
 xarray = ">=2024.1.0,<2025"
 quantecon = ">=0.7.2,<0.8"
 interpolation = ">=2.2.7,<3"
-multipledispatch = ">=0.6.0,<0.7"
 numba = ">=0.60.0,<0.70"
 dolang = ">=0.0.21,<0.1"
 # conda-forge packages the `dynopy` project (the PyPI name pyproject.toml
@@ -128,7 +127,6 @@ matplotlib = "==3.8.4"
 xarray = "==2024.7.0"
 quantecon = "==0.7.2"
 interpolation = "==2.2.7"
-multipledispatch = "==0.6.0"
 numba = "==0.60.0"
 dolang = "==0.0.21"
 dyno = "==0.1.15"

pixi.toml

@@ -28,7 +28,6 @@ dependencies = [
     "jax>=0.4.38,<0.5",
     "matplotlib>=3.6.0,<4",
     "numba>=0.60.0,<0.70",
-    "multipledispatch>=0.6.0,<0.7",
     "numpy>=2.0.0,<3",
     "pandas>=2.2.2,<4",
     "pyyaml>=6.0,<7",

pyproject.toml

@@ -1,8 +1,5 @@
 import numpy as np
 
-# doesn't seem to support keyword arguments yet
-# from multipledispatch import dispatch
-# multimethod = dispatch()
 from dolo.misc.multimethod import multimethod
 
 # from numba import generated_jit

src/dolo/algos/ergodic.py

@@ -24,6 +24,13 @@
     ProductApproximator,
     get_approximator,
 )
+from dolo.numeric.interpolation_jax import (
+    prefilter_cubic_jax,
+    eval_cubic_jax,
+    eval_cubic_point_jax,
+    eval_linear_jax,
+    eval_linear_point_jax,
+)
 
 __all__ = [
     "GArray",
@@ -47,4 +54,9 @@
     "SmolyakApproximator",
     "ProductApproximator",
     "get_approximator",
+    "prefilter_cubic_jax",
+    "eval_cubic_jax",
+    "eval_cubic_point_jax",
+    "eval_linear_jax",
+    "eval_linear_point_jax",
 ]

src/dolo/numeric/__init__.py

@@ -61,15 +61,26 @@ def evaluate_batch(self, coeffs: Any, batch: QualifiedBatch) -> np.ndarray:
 class LinearApproximator(Approximator):
     """Piecewise multilinear interpolation on CartesianGrid."""
 
-    def __init__(self, grid: CartesianGrid, space: Optional[CartesianSpace] = None):
+    def __init__(
+        self,
+        grid: CartesianGrid,
+        space: Optional[CartesianSpace] = None,
+        backend: str = "jax",
+    ):
         self.grid = grid
         self.space = space
+        self.backend = backend
 
     def fit(self, values: Union[GArray, np.ndarray]) -> np.ndarray:
         data = values.data if isinstance(values, GArray) else np.asarray(values)
         return data.reshape(tuple(self.grid.n) + (-1,)).copy()
 
     def evaluate_coords(self, coeffs: np.ndarray, coords: np.ndarray) -> np.ndarray:
+        if self.backend == "jax":
+            from dolo.numeric.interpolation_jax import eval_linear_jax
+
+            return np.asarray(eval_linear_jax(self.grid, coeffs, coords))
+
         from interpolation.splines import eval_linear
 
         pts = np.asarray(coords, dtype=float)
@@ -81,18 +92,34 @@ def evaluate_coords(self, coeffs: np.ndarray, coords: np.ndarray) -> np.ndarray:
 class CubicSplineApproximator(Approximator):
     """Prefiltered cubic spline interpolation on CartesianGrid."""
 
-    def __init__(self, grid: CartesianGrid, space: Optional[CartesianSpace] = None):
+    def __init__(
+        self,
+        grid: CartesianGrid,
+        space: Optional[CartesianSpace] = None,
+        backend: str = "jax",
+    ):
         self.grid = grid
         self.space = space
+        self.backend = backend
 
     def fit(self, values: Union[GArray, np.ndarray]) -> np.ndarray:
-        from interpolation.splines.prefilter_cubic import prefilter_cubic
-
         data = values.data if isinstance(values, GArray) else np.asarray(values)
         reshaped = data.reshape(tuple(self.grid.n) + (-1,))
+        if self.backend == "jax":
+            from dolo.numeric.interpolation_jax import prefilter_cubic_jax
+
+            return np.asarray(prefilter_cubic_jax(self.grid, reshaped))
+
+        from interpolation.splines.prefilter_cubic import prefilter_cubic
+
         return prefilter_cubic(self.grid.__numba_repr__(), reshaped)
 
     def evaluate_coords(self, coeffs: np.ndarray, coords: np.ndarray) -> np.ndarray:
+        if self.backend == "jax":
+            from dolo.numeric.interpolation_jax import eval_cubic_jax
+
+            return np.asarray(eval_cubic_jax(self.grid, coeffs, coords))
+
         from interpolation.splines import eval_cubic
 
         pts = np.asarray(coords, dtype=float)
@@ -172,6 +199,8 @@ def __init__(
     def fit(self, values: Union[GArray, np.ndarray]) -> Any:
         data = values.data if isinstance(values, GArray) else np.asarray(values)
         if self.is_empty_exo:
+            if data.ndim == 3 and data.shape[0] == 1:
+                data = data[0]
             return self.sub_approximator.fit(data)
         elif self.is_discrete_exo:
             n_regimes = data.shape[0]
@@ -183,6 +212,48 @@ def fit(self, values: Union[GArray, np.ndarray]) -> Any:
                 f"Unsupported exo grid type: {type(self.exo_grid)}"
             )
 
+    def evaluate_s(self, coeffs: Any, s: np.ndarray) -> np.ndarray:
+        """Evaluate continuous endogenous states `s` when exogenous state is empty or single-regime."""
+        s_arr = np.asarray(s, dtype=float)
+        if self.is_empty_exo:
+            return self.sub_approximator.evaluate_coords(coeffs, s_arr)
+        if self.is_discrete_exo and self.exo_grid.n_nodes == 1:
+            return self.sub_approximator.evaluate_coords(coeffs[0], s_arr)
+        raise ValueError(
+            "Cannot evaluate 's' alone when exogenous states are non-empty."
+        )
+
+    def evaluate_ms(self, coeffs: Any, m: np.ndarray, s: np.ndarray) -> np.ndarray:
+        """Evaluate continuous endogenous states `s` and exogenous states `m`."""
+        m_arr = np.asarray(m, dtype=float)
+        s_arr = np.asarray(s, dtype=float)
+        if self.is_empty_exo:
+            return self.sub_approximator.evaluate_coords(coeffs, s_arr)
+        elif self.is_cartesian_exo:
+            combined = np.hstack([m_arr, s_arr])
+            return self.joint_approximator.evaluate_coords(coeffs, combined)
+        elif self.is_discrete_exo:
+            if m_arr.ndim == 1:
+                m_arr = m_arr[None, :]
+            if s_arr.ndim == 1:
+                s_arr = s_arr[None, :]
+            if m_arr.shape[0] == 1 and s_arr.shape[0] > 1:
+                m_arr = np.repeat(m_arr, s_arr.shape[0], axis=0)
+            diffs = np.linalg.norm(
+                self.exo_grid.nodes[:, None, :] - m_arr[None, :, :], axis=-1
+            )
+            regimes = np.argmin(diffs, axis=0)
+            n_x = coeffs[0].shape[-1] if hasattr(coeffs[0], "shape") else 1
+            out = np.empty((s_arr.shape[0], n_x), dtype=float)
+            for r in np.unique(regimes):
+                mask = regimes == r
+                out[mask] = self.sub_approximator.evaluate_coords(
+                    coeffs[r], s_arr[mask]
+                )
+            return out
+        else:
+            raise NotImplementedError()
+
     def evaluate_coords(self, coeffs: Any, coords: np.ndarray) -> np.ndarray:
         if self.is_empty_exo:
             return self.sub_approximator.evaluate_coords(coeffs, coords)