Added implementation of non cartesian grids.

dolark/dolo_improvements.py

@@ -2,6 +2,11 @@
 from dolo.numeric.decision_rule import CallableDecisionRule, cat_grids
 import numpy as np
 import tqdm
+from dolo.numeric.misc import cartesian
+
+
+class Vector:
+    pass  # not really used
 
 class Linear:
     pass
@@ -19,6 +24,89 @@ class Chebychev:
     'chebychev': Chebychev()
 }
 
+
+from dolark.dolo_improvements import *
+from typing import List, Union
+
+functions = {
+    '': [lambda u: u, lambda u: u],
+    'exp': [np.exp, np.log],
+    'log': [np.log, np.exp],
+    'exp(exp)': [lambda x: np.exp(np.exp(x)), lambda x: np.log(np.log(x))],
+    'log(log)': [lambda x: np.log(np.log(x)), lambda x: np.exp(np.exp(x))],
+    'exp(exp(exp))': [lambda x: np.exp(np.exp(np.exp(x))), lambda x: np.log(np.log(np.log(x)))],
+    'log(log(log))': [lambda x: np.log(np.log(np.log(x))), lambda x: np.exp(np.exp(np.exp(x)))]
+}
+
+class WarpGrid(Grid):
+
+    base: Grid
+    warp: List[str]
+
+    def __init__(self, base:Grid, warp: Union[str, List[str]]):
+        if isinstance(warp, str):
+            warp = [str]
+        self.warp = warp
+        self.base = base
+
+        assert(base.d == len(warp))
+
+        self.warp_functions = []
+        self.warp_ifunctions = []
+        for w in self.warp:
+            if w not in functions:
+                raise Exception("Unsupported warp function")
+            else:
+                f, i_f = functions[w]
+                self.warp_functions.append(f)
+                self.warp_ifunctions.append(i_f)
+
+    @property
+    def nodes(self):
+        nn = self.base.nodes.copy()
+        for i in range(nn.shape[1]):
+            nn[:,i] = self.warp_functions[i](nn[:,i])
+        return nn
+
+    @property
+    def n_nodes(self):
+        return self.base.n_nodes
+
+    def node(self, i):
+        nn = self.base.node(i)
+        return np.array([self.warp_functions[i](nn[i]) for i in range(len(self.warp))])
+
+
+class ICartesianGrid(Grid):
+
+    axes: List[Vector]
+
+    def __init__(self, axes: List[Vector]):
+        self.axes = tuple([np.array(e) for e in axes])
+        self.d = len(axes)
+        self.__nodes__ = None
+
+    @property
+    def n(self):
+        return tuple([len(e) for e in self.axes])
+
+    @property
+    def nodes(self):
+        if self.__nodes__ is None:
+            self.__nodes__ = cartesian(self.axes)
+
+        return self.__nodes__
+
+    @property
+    def n_nodes(self):
+        return np.prod([len(e) for e in self.axes])
+
+    def node(self, i):
+        # TODO: improve this!
+        return self.nodes[i,:]
+
+
+
 # we keep the same user-facing-API
 
 class CallableDecisionRule:
@@ -180,7 +268,7 @@ def eval_is(itp, exo_grid: CartesianGrid, endo_grid: CartesianGrid, interp_type:
 
 @multimethod
 def get_coefficients(itp, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_type: Linear, x):
-    return [x[i].copy() for i in range(x.shape[0])]
+    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):
@@ -196,6 +284,26 @@ def eval_is(itp, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid, interp_ty
     return eval_linear(gg, coeffs, s)
 
 
+
+# UnstructuredGrid x ICartesian x Linear
+
+@multimethod
+def get_coefficients(itp, exo_grid: UnstructuredGrid, endo_grid: ICartesianGrid, interp_type: Linear, x):
+    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: ICartesianGrid, interp_type: Linear, i, s):
+
+    from interpolation.splines import eval_linear
+    assert(s.ndim==2)
+
+    grid = endo_grid # one single CartesianGrid
+    coeffs = itp.coefficients[i]
+    gg = endo_grid.axes
+
+    return eval_linear(gg, coeffs, s)
+
+
 # UnstructuredGrid x Cartesian x Cubic
 
 @multimethod
@@ -205,22 +313,45 @@ def get_coefficients(itp, exo_grid: UnstructuredGrid, endo_grid: CartesianGrid,
     d = len(grid.n)
     # this gg could be stored as a member of itp
     gg = tuple( [(grid.min[i], grid.max[i], grid.n[i]) for i in range(d)] )
-    return [prefilter_cubic(gg, x[i]) 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):
 
     from interpolation.splines import eval_cubic
     assert(s.ndim==2)
-
     grid = endo_grid # one single CartesianGrid
     coeffs = itp.coefficients[i]
     d = len(grid.n)
     gg = tuple( [(grid.min[i], grid.max[i], grid.n[i]) for i in range(d)] )
 
     return eval_cubic(gg, coeffs, s)
 
+
+
+# deal with warped Grids
+@multimethod
+def get_coefficients(itp, exo_grid, endo_grid: WarpGrid, interp_type, x):
+    return get_coefficients(itp, exo_grid, endo_grid.base, interp_type, x)
+
+@multimethod
+def eval_is(itp, exo_grid, endo_grid: WarpGrid, interp_type, i, s):
+    base = endo_grid.base
+    ss = s.copy()
+    for k,f in enumerate(endo_grid.warp_ifunctions):
+        ss[:,k] = f(ss[:,k])
+    return eval_is(itp, exo_grid, base, interp_type, i, ss)
+
+@multimethod
+def eval_ms(itp, exo_grid, endo_grid: WarpGrid, interp_type, m, s):
+    base = endo_grid.base
+    ss = s.copy()
+    for i,f in enumerate(endo_grid.warp_ifunctions):
+        ss[:,i] = f(ss[:,i])
+    return eval_ms(itp, exo_grid, base, interp_type, m, ss)
+
+
 ###### Test
 
 if __name__ == "__main__":

dolark/equilibrium.py

@@ -13,7 +13,7 @@ def __init__(self, aggmodel, m, μ, dr, y):
         self.m = m
         self.μ = μ
         self.dr = dr
-        self.x = np.concatenate([e[None,:,:] for e in [dr(i,dr.endo_grid.nodes()) for i in range(max(dr.exo_grid.n_nodes(),1))] ], axis=0)
+        self.x = np.concatenate([e[None,:,:] for e in [dr(i,dr.endo_grid.nodes) for i in range(max(dr.exo_grid.n_nodes,1))] ], axis=0)
         self.y = y
         self.c = dr.coefficients
 
@@ -24,12 +24,12 @@ def __init__(self, aggmodel, m, μ, dr, y):
     def as_df(self):
         model = self.aggmodel.model
         eq = self
-        exg = np.column_stack([range(eq.dr.exo_grid.n_nodes()), eq.dr.exo_grid.nodes()])
-        edg = np.column_stack([eq.dr.endo_grid.nodes()])
+        exg = np.column_stack([range(eq.dr.exo_grid.n_nodes, eq.dr.exo_grid.n_nodes)])
+        edg = np.column_stack([eq.dr.endo_grid.nodes])
         N_m = exg.shape[0]
         N_s = edg.shape[0]
         ssg = np.concatenate([exg[:,None,:].repeat(N_s, axis=1), edg[None,:,:].repeat(N_m, axis=0)], axis=2).reshape((N_m*N_s,-1))
-        x = np.concatenate([eq.dr(i, edg) for i in range(max(eq.dr.exo_grid.n_nodes(),1))], axis=0)
+        x = np.concatenate([eq.dr(i, edg) for i in range(max(eq.dr.exo_grid.n_nodes,1))], axis=0)
         import pandas as pd
         cols = ['i_m'] + model.symbols['exogenous'] + model.symbols['states'] + ['μ'] + model.symbols['controls']
         df = pd.DataFrame(np.column_stack([ssg, eq.μ.ravel(), x]), columns=cols)
@@ -55,12 +55,12 @@ def equilibrium(hmodel, m0: 'vector', y0: 'vector', p=None, dr0=None, grids=None
 
     Π0, μ0 = ergodic_distribution(hmodel.model, dr, exg, edg, dp)
 
-    s = edg.nodes()
-    if exg.n_nodes()==0:
+    s = edg.nodes
+    if exg.n_nodes==0:
         nn = 1
         μμ0 = μ0.data[None,:]
     else:
-        nn = exg.n_nodes()
+        nn = exg.n_nodes
         μμ0 = μ0.data
 
     xx0 = np.concatenate([e[None,:,:] for e in [dr(i,s) for i in range(nn)] ], axis=0)
@@ -111,8 +111,8 @@ def fun(u):
         if verbose: print(colored("done", "green"))
 
 
-    # grid_m = model.exogenous.discretize(to='mc', options=[{},{'N':N_mc}]).nodes()
-    # grid_s = model.get_grid().nodes()
+    # grid_m = model.exogenous.discretize(to='mc', options=[{},{'N':N_mc}]).nodes
+    # grid_s = model.get_grid().nodes
     #
     y_ss = solution.x # vector of aggregate endogenous variables
     m_ss = m0 # vector fo aggregate exogenous

dolark/model.py

@@ -223,8 +223,8 @@ def 𝒜(self, grids, m0: 'n_e', μ0: "n_m.N" , xx0: "n_m.N.n_x", y0: "n_y", p:
             # this is so sad
             mi = self.model.calibration['exogenous'][None,:] # not used anyway...
         else:
-            mi = exg.nodes()
-        s = eng.nodes()
+            mi = exg.nodes
+        s = eng.nodes
         res = sum( [μ0[i,:] @ ℰ(mi[i,:],s,xx0[i,:,:],m0,y0,p) for i in range(xx0.shape[0]) ])
         return res
 

dolark/perturbation.py

@@ -81,8 +81,8 @@ def F(hmodel, equilibrium, states, controls, states_f, controls_f, p):
     dr1.set_values(x1)
 
     dr0 = time_iteration(hmodel.model, dr0=dr1, verbose=False, maxit=1, dprocess=tmc)
-    s = dr0.endo_grid.nodes()
-    n_m = _mc.n_nodes()
+    s = dr0.endo_grid.nodes
+    n_m = _mc.n_nodes
     xx0 = np.concatenate([e[None,:,:] for e in [dr0(i,s) for i in range(n_m)] ], axis=0)
 
     res_0 = xx0-x0

examples/bfs_2017.yaml

@@ -12,11 +12,11 @@ equations:
         - m = ((D+r)/Đ)*(m(-1)-c(-1))*(1/ψ) + ξ
 
     arbitrage:
-        - ( β*(c(1)/c)^(-ρ) )*(ψ(1)^(-ρ))*(D+r(1))/Đ - 1 | 0.0<=c<=m
+        - ( β*(c(1)/c)^(-ρ) )*(ψ(1)^(-ρ))*(D+r(1)) - 1 | 0.0<=c<=m
 
 calibration:
     ρ: 1
-    β: 0.99
+    β: 0.98
     δ: 0.025
     D: 1-δ
     Đ: 1-0.00625

examples/equilibrium.py

@@ -55,7 +55,7 @@
 eqs = find_steady_state(hmodel2)
 eqs # results is (for now) a list of equilibrium objects
 
-s = eqs[0][1].dr.endo_grid.nodes().ravel()
+s = eqs[0][1].dr.endo_grid.nodes.ravel()
 i=0
 for (w,eq) in eqs:
     dens = eq.μ.sum(axis=0) # \mu is a 2d array: exogenous x endogenous states

examples/reiter_example.py

@@ -23,6 +23,7 @@
 aggmodel.features
 
 # %%
+aggmodel.agent
 
 # %% [markdown]
 # First we can check whether the one-agent sub-part of it works, or whether we will need to find better initial guess.
@@ -39,10 +40,13 @@
 eq = find_steady_state(aggmodel)
 eq
 
+# %%
+from matplotlib import pyplot as plt
+
 # %%
 # lot's look at the aggregate equilibrium
 for i in range(eq.μ.shape[0]):
-    s = eq.dr.endo_grid.nodes() # grid for states (temporary)
+    s = eq.dr.endo_grid.nodes # grid for states (temporary)
     plt.plot(s, eq.μ[i,:]*(eq.μ[i,:].sum()), label=f"y={eq.dr.exo_grid.node(i)[2]: .2f}")
 plt.plot(s, eq.μ.sum(axis=0), label='total', color='black')
 plt.grid()

experiments/jagman/heterogeneous_agents_example.ipynb

@@ -429,7 +429,7 @@
    "source": [
     "%matplotlib inline\n",
     "\n",
-    "egrid = np.unique(dr.endo_grid.nodes()[:,0])\n",
+    "egrid = np.unique(dr.endo_grid.nodes[:,0])\n",
     "kss = model.calibration['states'][1]\n",
     "\n",
     "plt.figure(figsize=(15, 7.5))\n",

experiments/krusell_smith/krusell_smith.ipynb

@@ -269,17 +269,17 @@
    "source": [
     "import matplotlib.pyplot as plt\n",
     "plt.figure()\n",
-    "plt.plot(sol.dr.endo_grid.nodes(),sol.dr(0,sol.dr.endo_grid.nodes()))\n",
-    "plt.plot(sol.dr.endo_grid.nodes(),sol.dr(1,sol.dr.endo_grid.nodes()))\n",
-    "plt.plot(sol.dr.endo_grid.nodes(),sol.dr.endo_grid.nodes(),'k--')\n",
+    "plt.plot(sol.dr.endo_grid.nodes,sol.dr(0,sol.dr.endo_grid.nodes))\n",
+    "plt.plot(sol.dr.endo_grid.nodes,sol.dr(1,sol.dr.endo_grid.nodes))\n",
+    "plt.plot(sol.dr.endo_grid.nodes,sol.dr.endo_grid.nodes,'k--')\n",
     "plt.xlabel(\"$a_t$\")\n",
     "plt.ylabel(\"$a_{t+1}$\")\n",
     "plt.legend([\"$e = 0$\",\"$e = 1$\",\"$a_{t+1} = a_{t}$\"])\n",
     "\n",
     "plt.figure()\n",
-    "plt.plot(sol.dr.endo_grid.nodes(),sol.dr(n_K*n_z*n_e-2,sol.dr.endo_grid.nodes()))\n",
-    "plt.plot(sol.dr.endo_grid.nodes(),sol.dr(n_K*n_z*n_e-1,sol.dr.endo_grid.nodes()))\n",
-    "plt.plot(sol.dr.endo_grid.nodes(),sol.dr.endo_grid.nodes(),'k--')\n",
+    "plt.plot(sol.dr.endo_grid.nodes,sol.dr(n_K*n_z*n_e-2,sol.dr.endo_grid.nodes))\n",
+    "plt.plot(sol.dr.endo_grid.nodes,sol.dr(n_K*n_z*n_e-1,sol.dr.endo_grid.nodes))\n",
+    "plt.plot(sol.dr.endo_grid.nodes,sol.dr.endo_grid.nodes,'k--')\n",
     "plt.xlabel(\"$a_t$\")\n",
     "plt.ylabel(\"$a_{t+1}$\")\n",
     "plt.legend([\"$e = 0$\",\"$e = 1$\",\"$a_{t+1} = a_{t}$\"])"

experiments/krusell_smith/krusell_smith_wout_agg_shocks.ipynb

@@ -248,9 +248,9 @@
    ],
    "source": [
     "import matplotlib.pyplot as plt\n",
-    "plt.plot(sol.endo_grid.nodes(),sol(0,sol.endo_grid.nodes()))\n",
-    "plt.plot(sol.endo_grid.nodes(),sol(1,sol.endo_grid.nodes()))\n",
-    "plt.plot(sol.endo_grid.nodes(),sol.endo_grid.nodes(),'k--')\n",
+    "plt.plot(sol.endo_grid.nodes,sol(0,sol.endo_grid.nodes))\n",
+    "plt.plot(sol.endo_grid.nodes,sol(1,sol.endo_grid.nodes))\n",
+    "plt.plot(sol.endo_grid.nodes,sol.endo_grid.nodes,'k--')\n",
     "plt.xlabel(\"$a_t$\")\n",
     "plt.ylabel(\"$a_{t+1}$\")\n",
     "plt.legend([\"$z = 0$\",\"$z = 1$\",\"$a_{t+1} = a_{t}$\"])"
@@ -454,9 +454,9 @@
    ],
    "source": [
     "#Policy function\n",
-    "plt.plot(s.endo_grid.nodes(),sol(0,s.endo_grid.nodes()))\n",
-    "plt.plot(s.endo_grid.nodes(),sol(1,s.endo_grid.nodes()))\n",
-    "plt.plot(s.endo_grid.nodes(),sol.endo_grid.nodes(),'k--')\n",
+    "plt.plot(s.endo_grid.nodes,sol(0,s.endo_grid.nodes))\n",
+    "plt.plot(s.endo_grid.nodes,sol(1,s.endo_grid.nodes))\n",
+    "plt.plot(s.endo_grid.nodes,sol.endo_grid.nodes,'k--')\n",
     "plt.xlabel(\"$a_t$\")\n",
     "plt.ylabel(\"$a_{t+1}$\")\n",
     "plt.legend([\"$z = 0$\",\"$z = 1$\",\"$a_{t+1} = a_{t}$\"])"
@@ -523,7 +523,7 @@
     }
    ],
    "source": [
-    "plt.plot(s.endo_grid.nodes(),hist)\n",
+    "plt.plot(s.endo_grid.nodes,hist)\n",
     "plt.xlabel(\"$a$\")\n",
     "plt.title(\"Asymptotic density of assets\")"
    ]