Unify DecisionRule and Approximator evaluation and deprecate eval_* methods

Author: albop

src/dolo/algos/newton.py

@@ -428,7 +428,7 @@ def newton(
 
     if dr0 is not None:
         for i_m in range(n_m):
-            x0[i_m, :, :] = dr0.eval_is(i_m, s)
+            x0[i_m, :, :] = dr0(i_m, s)
     ddr.set_values(x0)
 
     if x_lb is not None and complementarities:

src/dolo/algos/simulations.py

@@ -131,15 +131,15 @@ def simulate(
             sim_type = "continuous"
             if m0 is None:
                 m0 = model.calibration["exogenous"]
-            x_simul[0, :, :] = dr.eval_ms(m0[None, :], s0[None, :])[0, :]
+            x_simul[0, :, :] = dr(m0[None, :], s0[None, :])[0, :]
         elif len(driving_process.shape) == 2:
             i_simul = driving_process
             nodes = dr.exo_grid.nodes
             m_simul = nodes[i_simul]
             # inds = i_simul.ravel()
             # m_simul = np.reshape( np.concatenate( [nodes[i,:][None,:] for i in inds.ravel()], axis=0 ), inds.shape + (-1,) )
             sim_type = "discrete"
-            x_simul[0, :, :] = dr.eval_is(i0, s0[None, :])[0, :]
+            x_simul[0, :, :] = dr(i0, s0[None, :])[0, :]
         else:
             raise Exception("Incorrect specification of driving values.")
         m0 = m_simul[0, :, :]
@@ -165,15 +165,15 @@ def simulate(
             m_simul = dp.simulate(N, T, i0=i0, stochastic=stochastic)
             i_simul = find_index(m_simul, dp.values)
             m0 = dp.node(i0)
-            x0 = dr.eval_is(i0, s0[None, :])[0, :]
+            x0 = dr(i0, s0[None, :])[0, :]
         else:
             m_simul = process.simulate(N, T, m0=m0, stochastic=stochastic)
             if isinstance(m_simul, xr.DataArray):
                 m_simul = m_simul.data
             sim_type = "continuous"
             if m0 is None:
                 m0 = model.calibration["exogenous"]
-            x0 = dr.eval_ms(m0[None, :], s0[None, :])[0, :]
+            x0 = dr(m0[None, :], s0[None, :])[0, :]
             x_simul[0, :, :] = x0[None, :]
 
     f = model.functions["arbitrage"]
@@ -187,12 +187,10 @@ def simulate(
         s = s_simul[i, :, :]
         if sim_type == "discrete":
             i_m = i_simul[i, :]
-            xx = [
-                dr.eval_is(i_m[ii], s[ii, :][None, :])[0, :] for ii in range(s.shape[0])
-            ]
+            xx = [dr(i_m[ii], s[ii, :][None, :])[0, :] for ii in range(s.shape[0])]
             x = np.vstack(xx)
         else:
-            x = dr.eval_ms(m, s)
+            x = dr(m, s)
 
         x_simul[i, :, :] = x
 
@@ -284,10 +282,10 @@ def tabulate(
 
     if i0 is not None:
         m = dp.node(i0)
-        xvec = dr.eval_is(i0, svec)
+        xvec = dr(i0, svec)
     elif m0 is not None:
         m = m0
-        xvec = dr.eval_ms(m0, svec)
+        xvec = dr(m0, svec)
 
     mm = numpy.vstack([m] * n_steps)
     l = [mm, svec, xvec]
@@ -336,7 +334,7 @@ def tabulate_2d(model, dr, states=None, i0=0, s0=None, n=[12, 13]):
         s[i_x] = xx
         for yy in lps[1]:
             s[i_y] = yy
-            x = dr.eval_is(i0, s)
+            x = dr(i0, s)
             vv.append(numpy.concatenate([s, x]))
         vals.append(vv)
     vv = numpy.array(vals)

src/dolo/algos/value_iteration.py

@@ -103,7 +103,7 @@ def value_iteration(
         m = np.repeat(dprocess.node(i_ms)[None, :], N, axis=0)
         lb_i = controls_lb(m, s, p_batched)
         ub_i = controls_ub(m, s, p_batched)
-        x_i = mdr.eval_is(i_ms, s)
+        x_i = mdr(i_ms, s)
         # the calibrated controls may violate the bounds: project onto
         # [lb, ub] (feasible constant policy, keeps transitioned states sane)
         controls_0[i_ms, :, :] = np.clip(x_i, lb_i, ub_i)
@@ -150,7 +150,7 @@ def value_iteration(
 
         mdrv = ev.solution
         for i_ms in range(n_ms):
-            values_0[i_ms, :, :] = mdrv.eval_is(i_ms, s)
+            values_0[i_ms, :, :] = mdrv(i_ms, s)
 
         values = values_0.copy()
         controls = controls_0.copy()
@@ -315,15 +315,15 @@ def evaluate_policy(
     else:
         controls = np.zeros((n_ms, N, n_x))
         for i_m in range(n_ms):
-            controls[i_m, :, :] = mdr.eval_is(i_m, s)
+            controls[i_m, :, :] = mdr(i_m, s)
 
     values_0 = np.zeros((n_ms, N, n_v))
     if dr0 is None:
         for i_m in range(n_ms):
             values_0[i_m, :, :] = v0[None, :]
     else:
         for i_m in range(n_ms):
-            values_0[i_m, :, :] = dr0.eval_is(i_m, s)
+            values_0[i_m, :, :] = dr0(i_m, s)
 
     val = model.functions["value"]
     g = model.functions["transition"]

src/dolo/numeric/approximators.py

@@ -45,13 +45,29 @@ def evaluate_coords(self, coeffs: Any, coords: np.ndarray) -> np.ndarray:
         """Evaluate approximation at continuous coordinate matrix of shape (N, d)."""
         pass
 
-    def evaluate(self, coeffs: Any, qp: QualifiedPoint) -> np.ndarray:
-        """Evaluate approximation at a single QualifiedPoint."""
-        coords = np.asarray(qp.val, dtype=float)
-        if coords.ndim == 1:
-            res = self.evaluate_coords(coeffs, coords[None, :])
-            return res[0]
-        return self.evaluate_coords(coeffs, coords)
+    def evaluate(self, coeffs: Any, *args) -> np.ndarray:
+        """Universal evaluation method for approximation schemes.
+
+        Supports:
+        - `evaluate(coeffs, qp)`: single QualifiedPoint / QP
+        - `evaluate(coeffs, batch)`: QualifiedBatch / SliceQP
+        - `evaluate(coeffs, coords)`: coordinate array (N, d) or (d,)
+        - `evaluate(coeffs, *args)`: multi-argument forms implemented by subclasses
+        """
+        if len(args) == 1:
+            arg = args[0]
+            if isinstance(arg, QualifiedBatch):
+                return self.evaluate_batch(coeffs, arg)
+            if isinstance(arg, QualifiedPoint):
+                coords = np.asarray(arg.val, dtype=float)
+                if coords.ndim == 1:
+                    return self.evaluate_coords(coeffs, coords[None, :])[0]
+                return self.evaluate_coords(coeffs, coords)
+            coords = np.asarray(arg, dtype=float)
+            if coords.ndim == 1:
+                return self.evaluate_coords(coeffs, coords[None, :])[0]
+            return self.evaluate_coords(coeffs, coords)
+        raise ValueError(f"Approximator.evaluate cannot handle {len(args)} arguments.")
 
     def evaluate_batch(self, coeffs: Any, batch: QualifiedBatch) -> np.ndarray:
         """Evaluate approximation over a QualifiedBatch / SliceQP."""
@@ -264,9 +280,55 @@ def evaluate_coords(self, coeffs: Any, coords: np.ndarray) -> np.ndarray:
                 "Discrete exogenous states require regime index or QualifiedPoint for evaluation."
             )
 
-    def evaluate(self, coeffs: Any, qp: QualifiedPoint) -> np.ndarray:
+    def evaluate(self, coeffs: Any, *args) -> np.ndarray:
+        """Universal evaluation on ProductApproximator.
+
+        Supports:
+        - `evaluate(coeffs, qp)`: single QualifiedPoint / QP
+        - `evaluate(coeffs, batch)`: QualifiedBatch / SliceQP
+        - `evaluate(coeffs, s)`: continuous endogenous states alone (empty exo or single-regime)
+        - `evaluate(coeffs, i, s)`: discrete regime index `i` + continuous states `s`
+        - `evaluate(coeffs, m, s)`: continuous exogenous `m` + endogenous `s`
+        """
+        if len(args) == 1:
+            arg = args[0]
+            if isinstance(arg, QualifiedBatch):
+                return self.evaluate_batch(coeffs, arg)
+            if isinstance(arg, QualifiedPoint):
+                return self._evaluate_qp(coeffs, arg)
+            # Coordinates array: continuous endogenous states alone
+            return self.evaluate_s(coeffs, np.asarray(arg, dtype=float))
+
+        elif len(args) == 2:
+            a0, a1 = args
+            # Check for discrete regime index `i`
+            if isinstance(a0, (int, np.integer)):
+                s_arr = np.asarray(a1, dtype=float)
+                if s_arr.ndim == 1:
+                    return self.evaluate_regime(coeffs, int(a0), s_arr[None, :])[0]
+                return self.evaluate_regime(coeffs, int(a0), s_arr)
+
+            a0_arr = np.asarray(a0)
+            a1_arr = np.asarray(a1, dtype=float)
+            if a0_arr.dtype.kind in ("i", "u") and a0_arr.ndim == 0:
+                if a1_arr.ndim == 1:
+                    return self.evaluate_regime(coeffs, int(a0_arr), a1_arr[None, :])[0]
+                return self.evaluate_regime(coeffs, int(a0_arr), a1_arr)
+
+            # Continuous (m, s)
+            if a0_arr.ndim == 1 and a1_arr.ndim == 1:
+                return self.evaluate_ms(coeffs, a0_arr[None, :], a1_arr[None, :])[0]
+            return self.evaluate_ms(coeffs, a0_arr, a1_arr)
+
+        raise ValueError(
+            f"ProductApproximator.evaluate cannot handle {len(args)} arguments."
+        )
+
+    def _evaluate_qp(self, coeffs: Any, qp: QualifiedPoint) -> np.ndarray:
         if self.is_empty_exo:
             coords = np.asarray(qp.val, dtype=float)
+            if coords.ndim == 1:
+                return self.sub_approximator.evaluate_coords(coeffs, coords[None, :])[0]
             return self.sub_approximator.evaluate_coords(coeffs, coords)
         elif self.is_discrete_exo:
             i = qp.index
@@ -277,9 +339,12 @@ def evaluate(self, coeffs: Any, qp: QualifiedPoint) -> np.ndarray:
             )
             if isinstance(qp.loc, tuple) and len(qp.loc) >= 2:
                 endo_coords = qp.loc[1]
-            return self.sub_approximator.evaluate_coords(
-                coeffs[i], np.asarray(endo_coords, dtype=float)
-            )
+            c_arr = np.asarray(endo_coords, dtype=float)
+            if c_arr.ndim == 1:
+                return self.sub_approximator.evaluate_coords(coeffs[i], c_arr[None, :])[
+                    0
+                ]
+            return self.sub_approximator.evaluate_coords(coeffs[i], c_arr)
         elif self.is_cartesian_exo:
             return self.joint_approximator.evaluate(coeffs, qp)
         else:

src/dolo/numeric/decision_rule.py

@@ -54,15 +54,17 @@ class Chebychev:
 
 class CallableDecisionRule:
     def __call__(self, *args):
+        # 1. Fast path: delegate directly to approximator when present
+        if hasattr(self, "approximator") and self.approximator is not None:
+            return self.approximator.evaluate(self.coefficients, *args)
+
+        # 2. Fallback for custom / legacy decision rules without Approximator
         from dolo.numeric.garray import QP, SliceQP
         from dolo.numeric.grids import EmptyGrid
 
-        # 1. Single argument: SliceQP, QP, or coordinates array
         if len(args) == 1:
             arg = args[0]
             if isinstance(arg, SliceQP):
-                if hasattr(self, "approximator") and self.approximator is not None:
-                    return self.approximator.evaluate_batch(self.coefficients, arg)
                 is_empty_exo = hasattr(self, "exo_grid") and isinstance(
                     self.exo_grid, EmptyGrid
                 )
@@ -71,8 +73,6 @@ def __call__(self, *args):
                 return self.eval_s(arg.endo)
 
             if isinstance(arg, QP):
-                if hasattr(self, "approximator") and self.approximator is not None:
-                    return self.approximator.evaluate(self.coefficients, arg)
                 is_empty_exo = hasattr(self, "exo_grid") and isinstance(
                     self.exo_grid, EmptyGrid
                 )
@@ -91,13 +91,11 @@ def __call__(self, *args):
                     return self.eval_s(val_arr[None, :])[0, :]
                 return self.eval_s(val_arr)
 
-            # Continuous coords array
             s_arr = np.asarray(arg)
             if s_arr.ndim == 1:
                 return self.eval_s(s_arr[None, :])[0, :]
             return self.eval_s(s_arr)
 
-        # 2. Two arguments: (i, s) or (m, s)
         elif len(args) == 2:
             a0, a1 = args
             if isinstance(a0, (int, np.integer)):
@@ -108,13 +106,11 @@ def __call__(self, *args):
 
             a0_arr = np.asarray(a0)
             a1_arr = np.asarray(a1)
-            if a0_arr.dtype.kind in ("i", "u"):
-                if a0_arr.ndim == 0:
-                    if a1_arr.ndim == 1:
-                        return self.eval_is(int(a0_arr), a1_arr[None, :])[0, :]
-                    return self.eval_is(int(a0_arr), a1_arr)
+            if a0_arr.dtype.kind in ("i", "u") and a0_arr.ndim == 0:
+                if a1_arr.ndim == 1:
+                    return self.eval_is(int(a0_arr), a1_arr[None, :])[0, :]
+                return self.eval_is(int(a0_arr), a1_arr)
 
-            # (m, s) coordinates
             if a0_arr.ndim == 1 and a1_arr.ndim == 1:
                 return self.eval_ms(a0_arr[None, :], a1_arr[None, :])[0, :]
             return self.eval_ms(a0_arr, a1_arr)
@@ -205,34 +201,47 @@ def fit(self, x):
     set_values = fit
 
     def eval_ms(self, m, s):
-        m_arr = np.asarray(m, dtype=float)
-        s_arr = np.asarray(s, dtype=float)
-        if s_arr.ndim == 1:
-            s_arr = s_arr[None, :]
-        if m_arr.ndim == 1:
-            m_arr = m_arr[None, :]
-        return self.approximator.evaluate_ms(self.coefficients, m_arr, s_arr)
+        """Deprecated: use `dr(m, s)` instead."""
+        import warnings
+
+        warnings.warn(
+            "dr.eval_ms(m, s) is deprecated, use dr(m, s) instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self(m, s)
 
     def eval_is(self, i, s):
-        s_arr = np.asarray(s, dtype=float)
-        if s_arr.ndim == 1:
-            s_arr = s_arr[None, :]
-        return self.approximator.evaluate_regime(self.coefficients, int(i), s_arr)
+        """Deprecated: use `dr(i, s)` instead."""
+        import warnings
+
+        warnings.warn(
+            "dr.eval_is(i, s) is deprecated, use dr(i, s) instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self(i, s)
 
     def eval_s(self, s):
-        s_arr = np.asarray(s, dtype=float)
-        if s_arr.ndim == 1:
-            s_arr = s_arr[None, :]
-        return self.approximator.evaluate_s(self.coefficients, s_arr)
+        """Deprecated: use `dr(s)` instead."""
+        import warnings
+
+        warnings.warn(
+            "dr.eval_s(s) is deprecated, use dr(s) instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self(s)
 
     def eval_ijs(self, i, j, s):
+        """Evaluate policy at transitioned state for regime transition (i -> j)."""
         if isinstance(self.exo_grid, UnstructuredGrid):
-            return self.eval_is(j, s)
+            return self(j, s)
         elif isinstance(self.exo_grid, EmptyGrid):
-            return self.eval_s(s)
+            return self(s)
         elif isinstance(self.exo_grid, CartesianGrid):
             m = self.dprocess.inode(i, j)[None, :].repeat(s.shape[0], axis=0)
-            return self.eval_ms(m, s)
+            return self(m, s)
         else:
             raise NotImplementedError()
 

src/dolo/numeric/grids.py

@@ -129,8 +129,8 @@ def __add__(self, g):
             raise Exception("Not implemented.")
 
         n = np.array(tuple(self.n) + tuple(g.n))
-        min = np.array(tuple(self.min) + tuple(self.min))
-        max = np.array(tuple(self.max) + tuple(self.max))
+        min = np.array(tuple(self.min) + tuple(g.min))
+        max = np.array(tuple(self.max) + tuple(g.max))
 
         return UniformCartesianGrid(min, max, n)