from functools import reduce
from operator import mul
from quantecon import cartesian
import numpy as np
from typing import TypeVar, Generic
T = TypeVar("T")
S = TypeVar("S")
[docs]
def prod(l):
return reduce(mul, l, 1.0)
[docs]
class Grid:
def __mul__(self, rgrid):
return cat_grids(self, rgrid)
@property
def nodes(self):
return self.__nodes__
@property
def n_nodes(self):
return self.__nodes__.shape[0]
[docs]
def node(self, i):
return self.__nodes__[i, :]
[docs]
class ProductGrid(Grid, Generic[T, S]):
def __init__(self, g1: T, g2: S, names=None):
self.grids = [g1, g2]
self.names = names
def __getitem__(self, v):
return self.grids[self.names.index(v)]
def __repr__(self):
return str.join(" × ", [e.__repr__() for e in self.grids])
[docs]
class EmptyGrid(Grid):
type = "empty"
@property
def nodes(self):
return None
@property
def n_nodes(self):
return 0
[docs]
def node(self, i):
return None
def __add__(self, g):
return g
[docs]
class PointGrid(Grid):
type = "point"
def __init__(self, point):
self.point = np.array(point)
@property
def nodes(self):
return None
@property
def n_nodes(self):
return 1
[docs]
def node(self, i):
return None
[docs]
class UnstructuredGrid(Grid):
type = "unstructured"
def __init__(self, nodes):
nodes = np.array(nodes, dtype=float)
self.min = nodes.min(axis=0)
self.max = nodes.max(axis=0)
self.__nodes__ = nodes
self.d = len(self.min)
[docs]
class CartesianGrid(Grid):
pass
[docs]
class SmolyakGrid(Grid):
type = "Smolyak"
def __init__(self, min, max, mu=2):
from interpolation.smolyak import SmolyakGrid as ISmolyakGrid
min = np.array(min)
max = np.array(max)
self.min = min
self.max = max
self.mu = mu
d = len(min)
sg = ISmolyakGrid(d, mu, lb=min, ub=max)
self.sg = sg
self.d = d
self.__nodes__ = sg.grid
[docs]
def cat_grids(grid_1, grid_2):
if isinstance(grid_1, EmptyGrid):
return grid_2
if isinstance(grid_1, CartesianGrid) and isinstance(grid_2, CartesianGrid):
min = np.concatenate([grid_1.min, grid_2.min])
max = np.concatenate([grid_1.max, grid_2.max])
n = np.concatenate([grid_1.n, grid_2.n])
return CartesianGrid(min, max, n)
else:
raise Exception("Not Implemented.")
# compat
[docs]
def node(grid, i):
return grid.node(i)
[docs]
def nodes(grid):
return grid.nodes
[docs]
def n_nodes(grid):
return grid.n_nodes
if __name__ == "__main__":
print("Cartsian Grid")
grid = CartesianGrid([0.1, 0.3], [9, 0.4], [50, 10])
print(grid.nodes)
print(nodes(grid))
print("UnstructuredGrid")
ugrid = UnstructuredGrid([[0.1, 0.3], [9, 0.4], [50, 10]])
print(nodes(ugrid))
print(node(ugrid, 0))
print(n_nodes(ugrid))
print("Non Uniform CartesianGrid")
ugrid = NonUniformCartesianGrid([[0.1, 0.3], [9, 0.4], [50, 10]])
print(nodes(ugrid))
print(node(ugrid, 0))
print(n_nodes(ugrid))
print("Smolyak Grid")
sg = SmolyakGrid([0.1, 0.2], [1.0, 2.0], 2)
print(nodes(sg))
print(node(sg, 1))
print(n_nodes(sg))