REF: multimethod -> multipledispatch

dolo/algos/improved_time_iteration.py

@@ -230,7 +230,7 @@ def radius_jac(res,dres,jres,fut_S,dumdr,tol=1e-10,maxit=1000,verbose=False):
 from .results import AlgoResult, ImprovedTimeIterationResult
 
 def improved_time_iteration(model, method='jac', dr0=None, dprocess=None,
-            interp_method='spline', mu=2, maxbsteps=10, verbose=False,
+            interp_method='cubic', mu=2, maxbsteps=10, verbose=False,
             tol=1e-8, smaxit=500, maxit=1000,
             complementarities=True, compute_radius=False, invmethod='iti',
             details=True):
@@ -270,13 +270,15 @@ def vprint(*args, **kwargs):
 
     grid = model.get_endo_grid()
 
-    if interp_method in ('spline', 'linear'):
+    if interp_method in ('cubic', 'linear'):
         ddr = DecisionRule(dp.grid, grid, dprocess=dp, interp_method=interp_method)
         ddr_filt = DecisionRule(dp.grid, grid, dprocess=dp, interp_method=interp_method)
     elif interp_method == 'smolyak':
         ddr = SmolyakDecisionRule(n_m, grid.min, grid.max, mu)
         ddr_filt = SmolyakDecisionRule(n_m, grid.min, grid.max, mu)
         derivative_type = 'numerical'
+    else:
+        raise Exception("Unsupported interpolation method.")
 
     # s = ddr.endo_grid
     s = grid.nodes

dolo/numeric/decision_rule.py

@@ -90,32 +90,11 @@ def __init__(self, exo_grid: Grid, endo_grid: Grid, interp_method='cubic', dproc
 
         self.__interp_method__ = interp_methods[interp_method]
 
-        args = (self, exo_grid, endo_grid, interp_methods[interp_method])
-
-        try:
-            aa = args + (None, None)
-            fun = eval_ms[tuple(map(type, aa))]
-            self.__eval_ms__ = fun
-        except Exception as exc:
-            pass
-
-        try:
-            aa = args + (None, None)
-            fun = eval_is[tuple(map(type, aa))]
-            self.__eval_is__ = fun
-        except Exception as exc:
-            pass
-
-        try:
-            aa = args + (None, None)
-            fun = eval_s[tuple(map(type, aa))]
-            self.__eval_s__ = fun
-            # self.__eval_is__ = lambda i, s: fun(s)
-        except Exception as exc:
-            pass
-
-        fun = get_coefficients[tuple(map(type, args))]
-        self.__get_coefficients__ = fun
+        # here we could replace with a caching mechanism resolving dispatch in advance
+        self.__eval_ms__ = eval_ms
+        self.__eval_is__ = eval_is
+        self.__eval_s__ = eval_s
+        self.__get_coefficients__ = get_coefficients
 
         if values is not None:
             self.set_values(values)
@@ -149,19 +128,21 @@ def eval_ijs(self, i, j, s):
 
 # this is *not* meant to be used by users
 
-from multimethod import multimethod
+from multipledispatch import dispatch
+namespace = dict()
+multimethod = dispatch(namespace=namespace)
 
 # Cartesian x Cartesian x Linear
 
 @multimethod
-def get_coefficients(itp, exo_grid: CartesianGrid, endo_grid: CartesianGrid, interp_type: Linear, x):
+def get_coefficients(itp: object, exo_grid: CartesianGrid, endo_grid: CartesianGrid, interp_type: Linear, x: object):
     grid = exo_grid + endo_grid
     xx = x.reshape(tuple(grid.n)+(-1,))
     return xx
 
 
 @multimethod
-def eval_ms(itp, exo_grid: CartesianGrid, endo_grid: CartesianGrid, interp_type: Linear, m, s):
+def eval_ms(itp: object, exo_grid: CartesianGrid, endo_grid: CartesianGrid, interp_type: Linear, m: object, s: object):
 
     assert(m.ndim==s.ndim==2)
 
@@ -178,15 +159,15 @@ def eval_ms(itp, exo_grid: CartesianGrid, endo_grid: CartesianGrid, interp_type:
 
 
 @multimethod
-def eval_is(itp, exo_grid: CartesianGrid, endo_grid: CartesianGrid, interp_type: Linear, i, s):
+def eval_is(itp: object, exo_grid: CartesianGrid, endo_grid: CartesianGrid, interp_type: Linear, i: object, s: object):
     m = exo_grid.node(i)[None,:]
     return eval_ms(itp, exo_grid, endo_grid, interp_type, m, s)
 
 
 # Cartesian x Cartesian x Cubic
 
 @multimethod
-def get_coefficients(itp, exo_grid: CartesianGrid, endo_grid: CartesianGrid, interp_type: Cubic, x):
+def get_coefficients(itp: object, exo_grid: CartesianGrid, endo_grid: CartesianGrid, interp_type: Cubic, x: object):
 
     from interpolation.splines.prefilter_cubic import prefilter_cubic
     grid = exo_grid + endo_grid # one single CartesianGrid
@@ -195,7 +176,7 @@ def get_coefficients(itp, exo_grid: CartesianGrid, endo_grid: CartesianGrid, int
     return prefilter_cubic(gg, x)
 
 @multimethod
-def eval_ms(itp, exo_grid: CartesianGrid, endo_grid: CartesianGrid, interp_type: Cubic, m, s):
+def eval_ms(itp: object, exo_grid: CartesianGrid, endo_grid: CartesianGrid, interp_type: Cubic, m: object, s: object):
 
     from interpolation.splines import eval_cubic
 
@@ -212,7 +193,7 @@ def eval_ms(itp, exo_grid: CartesianGrid, endo_grid: CartesianGrid, interp_type:
 
 
 @multimethod
-def eval_is(itp, exo_grid: CartesianGrid, endo_grid: CartesianGrid, interp_type: Cubic, i, s):
+def eval_is(itp: object, exo_grid: CartesianGrid, endo_grid: CartesianGrid, interp_type: Cubic, i: object, s: object):
     m = exo_grid.node(i)[None,:]
     return eval_ms(itp, exo_grid, endo_grid, interp_type, m, s)
 
@@ -221,11 +202,11 @@ def eval_is(itp, exo_grid: CartesianGrid, endo_grid: CartesianGrid, interp_type:
 # UnstructuredGrid x Cartesian x Linear
 
 @multimethod
-def get_coefficients(itp, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_type: Linear, x):
+def get_coefficients(itp: object, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_type: Linear, x: object):
     return [x[i].reshape( tuple(endo_grid.n) + (-1,)).copy() for i in range(x.shape[0])]
 
 @multimethod
-def eval_is(itp, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_type: Linear, i, s):
+def eval_is(itp: object, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_type: Linear, i: object, s: object):
 
     from interpolation.splines import eval_linear
     assert(s.ndim==2)
@@ -238,14 +219,14 @@ def eval_is(itp, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_ty
 # UnstructuredGrid x Cartesian x Cubic
 
 @multimethod
-def get_coefficients(itp, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_type: Cubic, x):
+def get_coefficients(itp: object, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_type: Cubic, x: object):
     from interpolation.splines.prefilter_cubic import prefilter_cubic
     gg = endo_grid.__numba_repr__()
-    return [prefilter_cubic(gg, x[i].reshape( tuple(grid.n) + (-1,))) for i in range(x.shape[0])]
+    return [prefilter_cubic(gg, x[i].reshape( tuple(endo_grid.n) + (-1,))) for i in range(x.shape[0])]
 
 
 @multimethod
-def eval_is(itp, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_type: Cubic, i, s):
+def eval_is(itp: object, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_type: Cubic, i: object, s: object):
 
     from interpolation.splines import eval_cubic
     assert(s.ndim==2)
@@ -257,11 +238,11 @@ def eval_is(itp, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_ty
 # UnstructuredGrid x Cartesian x Linear
 
 @multimethod
-def get_coefficients(itp, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_type: Linear, x):
+def get_coefficients(itp: object, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_type: Linear, x: object):
     return [x[i].copy() for i in range(x.shape[0])]
 
 @multimethod
-def eval_is(itp, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_type: Linear, i, s):
+def eval_is(itp: object, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_type: Linear, i: object, s: object):
 
     from interpolation.splines import eval_linear
     assert(s.ndim==2)
@@ -274,13 +255,13 @@ def eval_is(itp, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_ty
 # Empty x Cartesian x Linear
 
 @multimethod
-def get_coefficients(itp, exo_grid: EmptyGrid, endo_grid: CartesianGrid, interp_type: Linear, x):
+def get_coefficients(itp: object, exo_grid: EmptyGrid, endo_grid: CartesianGrid, interp_type: Linear, x: object):
     grid = exo_grid + endo_grid
     xx = x.reshape(tuple(grid.n)+(-1,))
     return xx
 
 @multimethod
-def eval_s(itp, exo_grid: EmptyGrid, endo_grid: CartesianGrid, interp_type: Linear, s):
+def eval_s(itp: object, exo_grid: EmptyGrid, endo_grid: CartesianGrid, interp_type: Linear, s: object):
     from interpolation.splines import eval_linear
     assert(s.ndim==2)
     coeffs = itp.coefficients
@@ -290,7 +271,7 @@ def eval_s(itp, exo_grid: EmptyGrid, endo_grid: CartesianGrid, interp_type: Line
 # Empty x Cartesian x Cubic
 
 @multimethod
-def get_coefficients(itp, exo_grid: EmptyGrid, endo_grid: CartesianGrid, interp_type: Cubic, x):
+def get_coefficients(itp: object, exo_grid: EmptyGrid, endo_grid: CartesianGrid, interp_type: Cubic, x: object):
     from interpolation.splines.prefilter_cubic import prefilter_cubic
     grid = endo_grid # one single CartesianGrid
     gg = endo_grid.__numba_repr__()
@@ -299,19 +280,19 @@ def get_coefficients(itp, exo_grid: EmptyGrid, endo_grid: CartesianGrid, interp_
 
 
 @multimethod
-def eval_s(itp, exo_grid: EmptyGrid, endo_grid: CartesianGrid, interp_type: Cubic, s):
+def eval_s(itp: object, exo_grid: EmptyGrid, endo_grid: CartesianGrid, interp_type: Cubic, s: object):
     from interpolation.splines import eval_cubic
     assert(s.ndim==2)
     coeffs = itp.coefficients
     gg = endo_grid.__numba_repr__()
     return eval_cubic(gg, coeffs, s)
 
 @multimethod
-def eval_is(itp, exo_grid: EmptyGrid, endo_grid: CartesianGrid, interp_type: Cubic, i, s):
+def eval_is(itp: object, exo_grid: EmptyGrid, endo_grid: CartesianGrid, interp_type: Cubic, i: object, s: object):
     return eval_s(itp, exo_grid, endo_grid, interp_type, s)
 
 @multimethod
-def eval_ms(itp, exo_grid: EmptyGrid, endo_grid: CartesianGrid, interp_type: Cubic, m, s):
+def eval_ms(itp: object, exo_grid: EmptyGrid, endo_grid: CartesianGrid, interp_type: Cubic, m: object, s: object):
     return eval_s(itp, exo_grid, endo_grid, interp_type, s)
 
 

dolo/tests/test_grid.py

@@ -1,10 +1,10 @@
 def test_grids():
 
-    from dolo.numeric.grids import CartesianGrid, UnstructuredGrid, NonUniformCartesianGrid, SmolyakGrid
+    from dolo.numeric.grids import UniformCartesianGrid, UnstructuredGrid, NonUniformCartesianGrid, SmolyakGrid
     from dolo.numeric.grids import nodes, n_nodes, node
 
     print("Cartsian Grid")
-    grid = CartesianGrid([0.1, 0.3], [9, 0.4], [50, 10])
+    grid = UniformCartesianGrid([0.1, 0.3], [9, 0.4], [50, 10])
     print(grid.nodes)
     print(nodes(grid))
 

dolo/tests/test_grids.py

@@ -1,9 +1,9 @@
 def test_grids():
 
-    from dolo.numeric.grids import CartesianGrid, UnstructuredGrid, NonUniformCartesianGrid, SmolyakGrid
+    from dolo.numeric.grids import UniformCartesianGrid, UnstructuredGrid, NonUniformCartesianGrid, SmolyakGrid
     from dolo.numeric.grids import n_nodes, nodes, node
     print("Cartsian Grid")
-    grid = CartesianGrid([0.1, 0.3], [9, 0.4], [50, 10])
+    grid = UniformCartesianGrid([0.1, 0.3], [9, 0.4], [50, 10])
     print(grid.nodes)
     print(nodes(grid))
 

environment.yml

@@ -17,5 +17,6 @@ dependencies:
 - statsmodels
 - matplotlib
 - xarray
+- multipledispatch
 # - pip:
 #   - sphinx-better-theme==0.13

requirements.txt

@@ -1,2 +1,2 @@
 dolang
-multimethod
+multipledispatch

setup.py

@@ -15,7 +15,7 @@
     ],
     install_requires=[
         "dolang", "pyyaml", "numba", "numpy", "sympy", "scipy", "quantecon", "pandas",
-        "interpolation", "ruamel.yaml", "xarray", "altair", "multimethod"
+        "interpolation", "ruamel.yaml", "xarray", "altair", "multipledispatch"
     ],
     extras_require={
         'interactive': ['ipython'],