diff --git a/Documentation/reference/tools/econforgeinterp.rst b/Documentation/reference/tools/econforgeinterp.rst index 9f0b7aa25..d4ede1132 100644 --- a/Documentation/reference/tools/econforgeinterp.rst +++ b/Documentation/reference/tools/econforgeinterp.rst @@ -1,7 +1,7 @@ Econforgeinterp --------------- -.. automodule:: HARK.econforgeinterp +.. automodule:: HARK.interpolation.econforgeinterp :members: :undoc-members: :show-inheritance: diff --git a/HARK/interpolation/__init__.py b/HARK/interpolation/__init__.py new file mode 100644 index 000000000..b01acabba --- /dev/null +++ b/HARK/interpolation/__init__.py @@ -0,0 +1,5 @@ +from HARK.interpolation._econforge import * +from HARK.interpolation._hark import * +from HARK.interpolation._multi import * +from HARK.interpolation._scipy import * +from HARK.interpolation._sklearn import * diff --git a/HARK/econforgeinterp.py b/HARK/interpolation/_econforge.py similarity index 100% rename from HARK/econforgeinterp.py rename to HARK/interpolation/_econforge.py diff --git a/HARK/interpolation.py b/HARK/interpolation/_hark.py similarity index 100% rename from HARK/interpolation.py rename to HARK/interpolation/_hark.py diff --git a/HARK/interpolation/_multi.py b/HARK/interpolation/_multi.py new file mode 100644 index 000000000..eba9582d0 --- /dev/null +++ b/HARK/interpolation/_multi.py @@ -0,0 +1,528 @@ +import numpy as np +from numba import njit, prange, typed +from scipy.ndimage import map_coordinates + +from HARK.metric import MetricObject + +AVAILABLE_TARGETS = ["cpu", "parallel"] + +try: + import cupy as cp + from cupyx.scipy.ndimage import map_coordinates as cupy_map_coordinates + + CUPY_AVAILABLE = True + AVAILABLE_TARGETS.append("gpu") +except ImportError: + CUPY_AVAILABLE = False + + +MC_KWARGS = { + "order": 1, # order of interpolation + "mode": "nearest", # how to handle extrapolation + "cval": 0.0, # value to use for extrapolation + "prefilter": False, # whether to prefilter input +} + + +class _RegularGridInterp(MetricObject): + """ + Abstract class for interpolating on a regular grid. Sets up + structure for using different targets (cpu, parallel, gpu). + Takes in arguments to be used by `map_coordinates`. + """ + + distance_criteria = ["values"] + + def __init__(self, values, target="cpu", **kwargs): + """ + Initialize a regular grid interpolator. + + Parameters + ---------- + values : np.ndarray + Functional values on a regular grid. + target : str, optional + Determines which target to use for interpolation. + Options are "cpu", "parallel", and "gpu". + If "cpu", uses numpy and scipy. + If "parallel", uses numba and scipy. + If "gpu", uses cupy. + + Raises + ------ + ValueError + Target is invalid. + """ + if target not in AVAILABLE_TARGETS: + raise ValueError("Invalid target.") + self.target = target + + self.mc_kwargs = MC_KWARGS.copy() + # update mc_kwargs with any kwargs that are in MC_KWARGS + self.mc_kwargs.update((k, v) for k, v in kwargs.items() if k in MC_KWARGS) + + if target in ["cpu", "parallel"]: + self.values = np.asarray(values) + elif target == "gpu": + self.values = cp.asarray(values) + + self.ndim = self.values.ndim # should match number of grids + self.shape = self.values.shape # should match points in each grid + + def __call__(self, *args): + """ + Interpolates arguments on the regular grid. + + Returns + ------- + np.ndarray + Interpolated functional values for each argument. + + Raises + ------ + ValueError + Number of argumets does not match number of dimensions. + """ + if self.target in ["cpu", "parallel"]: + args = np.asarray(args) + elif self.target == "gpu": + args = cp.asarray(args) + + if args.shape[0] != self.ndim: + raise ValueError("Number of arguments must match number of dimensions.") + + coordinates = self._get_coordinates(args) + return self._map_coordinates(coordinates) + + def _get_coordinates(self, args): + """ + Abstract method for getting coordinates for interpolation. + + Parameters + ---------- + args : np.ndarray + Arguments to be interpolated. + + Raises + ------ + NotImplementedError + Must be implemented by subclass. + """ + raise NotImplementedError("Must be implemented by subclass.") + + def _map_coordinates(self, coordinates): + """ + Uses coordinates to interpolate on the regular grid with + `map_coordinates` from scipy or cupy, depending on target. + + Parameters + ---------- + coordinates : np.ndarray + Index coordinates for interpolation. + + Returns + ------- + np.ndarray + Interpolated functional values for each coordinate. + """ + if self.target in ["cpu", "parallel"]: + # there is no parallelization for scipy map_coordinates + output = map_coordinates( + self.values, coordinates.reshape(self.ndim, -1), **self.mc_kwargs + ) + + elif self.target == "gpu": + output = cupy_map_coordinates( + self.values, coordinates.reshape(self.ndim, -1), **self.mc_kwargs + ) + + return output.reshape(coordinates[0].shape) + + +class MultivariateInterp(_RegularGridInterp): + """ + Multivariate Interpolator on a regular grid. Maps functional coordinates + to index coordinates and uses `map_coordinates` from scipy or cupy. + """ + + distance_criteria = ["values", "grids"] + + def __init__(self, values, grids, target="cpu", **kwargs): + """ + Initialize a multivariate interpolator. + + Parameters + ---------- + values : np.ndarray + Functional values on a regular grid. + grids : _type_ + 1D grids for each dimension. + target : str, optional + One of "cpu", "parallel", or "gpu". Determines + hardware to use for interpolation. + """ + + super().__init__(values, target=target, **kwargs) + + if target == "cpu": + self.grids = [np.asarray(grid) for grid in grids] + elif target == "parallel": + self.grids = typed.List(grids) + elif target == "gpu": + self.grids = [cp.asarray(grid) for grid in grids] + + if not (self.ndim == len(self.grids)): + raise ValueError("Number of grids must match number of dimensions.") + + if not all(self.shape[i] == grid.size for i, grid in enumerate(self.grids)): + raise ValueError("Values shape must match points in each grid.") + + def _get_coordinates(self, args): + """ + For each argument, finds the index coordinates for interpolation. + + Parameters + ---------- + args : np.ndarray + Arguments to be interpolated. + + Returns + ------- + np.ndarray + Index coordinates for interpolation. + """ + + if self.target == "cpu": + coordinates = np.empty_like(args) + for dim, grid in enumerate(self.grids): # for each dimension + coordinates[dim] = np.interp( # x, xp, fp (new x, x points, y values) + args[dim], grid, np.arange(self.shape[dim]) + ) + elif self.target == "parallel": + coordinates = _nb_interp(self.grids, args) + elif self.target == "gpu": + coordinates = cp.empty_like(args) + for dim, grid in enumerate(self.grids): # for each dimension + coordinates[dim] = cp.interp( # x, xp, fp (new x, x points, y values) + args[dim], grid, cp.arange(self.shape[dim]) + ) + + return coordinates + + +@njit(parallel=True, cache=True, fastmath=True) +def _nb_interp(grids, args): + """ + Just-in-time compiled function for interpolating on a regular grid. + + Parameters + ---------- + grids : np.ndarray + 1D grids for each dimension. + args : np.ndarray + Arguments to be interpolated. + + Returns + ------- + np.ndarray + Index coordinates for each argument. + """ + + coordinates = np.empty_like(args) + for dim in prange(args.shape[0]): + coordinates[dim] = np.interp(args[dim], grids[dim], np.arange(grids[dim].size)) + + return coordinates + + +class _CurvilinearGridInterp(_RegularGridInterp): + """ + Abstract class for interpolating on a curvilinear grid. + """ + + distance_criteria = ["values", "grids"] + + def __init__(self, values, grids, target="cpu"): + """ + Initialize a curvilinear grid interpolator. + + Parameters + ---------- + values : np.ndarray + Functional values on a curvilinear grid. + grids : np.ndarray + ND curvilinear grids for each dimension + target : str, optional + One of "cpu", "parallel", or "gpu". + """ + + super().__init__(values, target=target) + + if target in ["cpu", "parallel"]: + self.grids = np.asarray(grids) + elif target == "gpu": + self.grids = cp.asarray(grids) + + assert ( + self.ndim == self.grids[0].ndim + ), "Number of grids must match number of dimensions." + assert ( + self.shape == self.grids[0].shape + ), "Values shape must match points in each grid." + + +class WarpedInterpOnInterp2D(_CurvilinearGridInterp): + """ + Warped Grid Interpolation on a 2D grid. + """ + + def __call__(self, *args, axis=1): + """ + Interpolate on a warped grid using the Warped Grid Interpolation + method described in `EGM$^n$`. + + Parameters + ---------- + axis : int, 0 or 1 + Determines which axis to use for linear interpolators. + Setting to 0 may fix some issues where interpolation fails. + + Returns + ------- + np.ndarray + Interpolated values on a warped grid. + + Raises + ------ + ValueError + Number of arguments doesn't match number of dimensions. + """ + + if self.target in ["cpu", "parallel"]: + args = np.asarray(args) + elif self.target == "gpu": + args = cp.asarray(args) + + if args.shape[0] != self.ndim: + raise ValueError("Number of arguments must match number of dimensions.") + + if self.target == "cpu": + output = self._target_cpu(args, axis) + elif self.target == "parallel": + output = self._target_parallel(args, axis) + elif self.target == "gpu": + output = self._target_gpu(args, axis) + + return output + + def _target_cpu(self, args, axis): + """ + Uses numpy to interpolate on a warped grid. + + Parameters + ---------- + args : np.ndarray + Coordinates to be interpolated. + axis : int, 0 or 1 + See `WarpedInterpOnInterp2D.__call__`. + + Returns + ------- + np.ndarray + Interpolated values on arguments. + """ + + shape = args[0].shape # original shape of arguments + size = args[0].size # number of points in arguments + shape_axis = self.shape[axis] # number of points in axis + + # flatten arguments by dimension + args = args.reshape((self.ndim, -1)) + + y_intermed = np.empty((shape_axis, size)) + z_intermed = np.empty((shape_axis, size)) + + for i in range(shape_axis): + # for each dimension, interpolate the first argument + grids0 = np.take(self.grids[0], i, axis=axis) + grids1 = np.take(self.grids[1], i, axis=axis) + values = np.take(self.values, i, axis=axis) + y_intermed[i] = np.interp(args[0], grids0, grids1) + z_intermed[i] = np.interp(args[0], grids0, values) + + output = np.empty_like(args[0]) + + for j in range(size): + y_temp = y_intermed[:, j] + z_temp = z_intermed[:, j] + + if y_temp[0] > y_temp[-1]: + # reverse + y_temp = y_temp[::-1] + z_temp = z_temp[::-1] + + output[j] = np.interp(args[1][j], y_temp, z_temp) + + return output.reshape(shape) + + def _target_parallel(self, args, axis): + """ + Uses numba to interpolate on a warped grid. + + Parameters + ---------- + args : np.ndarray + Coordinates to be interpolated. + axis : int, 0 or 1 + See `WarpedInterpOnInterp2D.__call__`. + + Returns + ------- + np.ndarray + Interpolated values on arguments. + """ + + return nb_interp_piecewise(args, self.grids, self.values, axis) + + def _target_gpu(self, args, axis): + """ + Uses cupy to interpolate on a warped grid. + + Parameters + ---------- + args : np.ndarray + Coordinates to be interpolated. + axis : int, 0 or 1 + See `WarpedInterpOnInterp2D.__call__`. + + Returns + ------- + np.ndarray + Interpolated values on arguments. + """ + + shape = args[0].shape # original shape of arguments + size = args[0].size # number of points in arguments + shape_axis = self.shape[axis] # number of points in axis + + args = args.reshape((self.ndim, -1)) + + y_intermed = cp.empty((shape_axis, size)) + z_intermed = cp.empty((shape_axis, size)) + + for i in range(shape_axis): + # for each dimension, interpolate the first argument + grids0 = cp.take(self.grids[0], i, axis=axis) + grids1 = cp.take(self.grids[1], i, axis=axis) + values = cp.take(self.values, i, axis=axis) + y_intermed[i] = cp.interp(args[0], grids0, grids1) + z_intermed[i] = cp.interp(args[0], grids0, values) + + output = cp.empty_like(args[0]) + + for j in range(size): + y_temp = y_intermed[:, j] + z_temp = z_intermed[:, j] + + if y_temp[0] > y_temp[-1]: + # reverse + y_temp = y_temp[::-1] + z_temp = z_temp[::-1] + + output[j] = cp.interp(args[1][j], y_temp, z_temp) + + return output.reshape(shape) + + def warmup(self): + """ + Warms up the JIT compiler. + """ + self(*self.grids) + + return None + + +@njit(parallel=True, cache=True, fastmath=True) +def nb_interp_piecewise(args, grids, values, axis): + """ + Just-in-time compiled function to interpolate on a warped grid. + + Parameters + ---------- + args : np.ndarray + Arguments to be interpolated. + grids : np.ndarray + Curvilinear grids for each dimension. + values : np.ndarray + Functional values on a curvilinear grid. + axis : int, 0 or 1 + See `WarpedInterpOnInterp2D.__call__`. + + + Returns + ------- + np.ndarray + Interpolated values on arguments. + """ + + shape = args[0].shape # original shape of arguments + size = args[0].size # number of points in arguments + shape_axis = values.shape[axis] # number of points in axis + + # flatten arguments by dimension + args = args.reshape((values.ndim, -1)) + + y_intermed = np.empty((shape_axis, size)) + z_intermed = np.empty((shape_axis, size)) + + for i in prange(shape_axis): + # for each dimension, interpolate the first argument + grids0 = grids[0][i] if axis == 0 else grids[0][:, i] + grids1 = grids[1][i] if axis == 0 else grids[1][:, i] + vals = values[i] if axis == 0 else values[:, i] + y_intermed[i] = np.interp(args[0], grids0, grids1) + z_intermed[i] = np.interp(args[0], grids0, vals) + + output = np.empty_like(args[0]) + + for j in prange(size): + y_temp = y_intermed[:, j] + z_temp = z_intermed[:, j] + + if y_temp[0] > y_temp[-1]: + # reverse + y_temp = y_temp[::-1] + z_temp = z_temp[::-1] + + output[j] = np.interp(args[1][j], y_temp, z_temp) + + return output.reshape(shape) + + +class _UnstructuredGridInterp(_CurvilinearGridInterp): + """ + Abstract class for interpolation on unstructured grids. + """ + + def __init__(self, values, grids, target="cpu"): + """ + Initialize interpolation on unstructured grids. + + Parameters + ---------- + values : np.ndarray + Functional values on an unstructured grid. + grids : np.ndarray + ND unstructured grids for each dimension. + target : str, optional + One of "cpu", "parallel", or "gpu". + """ + + super().__init__(values, grids, target=target) + # remove non finite values that might result from + # sequential endogenous grid method + condition = np.logical_and.reduce([np.isfinite(grid) for grid in self.grids]) + condition = np.logical_and(condition, np.isfinite(self.values)) + self.values = self.values[condition] + self.grids = self.grids[:, condition] + self.ndim = self.grids.shape[0] diff --git a/HARK/interpolation/_scipy.py b/HARK/interpolation/_scipy.py new file mode 100644 index 000000000..6a6c9837a --- /dev/null +++ b/HARK/interpolation/_scipy.py @@ -0,0 +1,120 @@ +import numpy as np +from scipy.interpolate import ( + CloughTocher2DInterpolator, + LinearNDInterpolator, + NearestNDInterpolator, + RBFInterpolator, +) + +from HARK.interpolation._multi import _UnstructuredGridInterp + +LNDI_KWARGS = {"fill_value": np.nan, "rescale": False} # linear +NNDI_KWARGS = {"rescale": False, "tree_options": None} # nearest +CT2DI_KWARGS = { # cubic + "fill_value": np.nan, + "tol": 1e-06, + "maxiter": 400, + "rescale": False, +} +RBFI_KWARGS = { # rbf (radial basis function) + "neighbors": None, + "smoothing": 0.0, + "kernel": "linear", + "epsilon": None, + "degree": None, +} + +AVAILABLE_METHODS = ["nearest", "linear", "cubic", "rbf"] + + +class UnstructuredInterp(_UnstructuredGridInterp): + """ + Multivariate interpolation on an unstructured grid. + This class wraps various scipy unstructured interpolation + methods to provide a common interface. Additionally, it + can be used with meshgrids and returns meshgrids, which + are not supported by scipy but are the default used in + HARK. + """ + + distance_criteria = ["values", "grids"] + + def __init__( + self, + values, + grids, + method="linear", + **kwargs, + ): + """ + Initialize an Unstructured Grid Interpolator. + + Parameters + ---------- + values : np.ndarray + Functional values on an unstructured grid. + grids : np.ndarray + Points on an unstructured grid. + method : str, optional + One of "nearest", "linear", "cubic", "rbf". Determines + which scipy interpolation method to use. The interpolators + are "nearest" for NearestNDInterpolator, "linear" for + LinearNDInterpolator, "cubic" for CloughTocher2DInterpolator + and "rbf" for RBFInterpolator. The default is "linear". + + Raises + ------ + ValueError + The interpolation method is not valid. + """ + + # scipy can only do target = cpu + super().__init__(values, grids, target="cpu") + + # Check for valid interpolation method + if method not in AVAILABLE_METHODS: + raise ValueError("Invalid interpolation method.") + + self.method = method + + interpolator_mapping = { + "nearest": (NNDI_KWARGS, NearestNDInterpolator), + "linear": (LNDI_KWARGS, LinearNDInterpolator), + "cubic": (CT2DI_KWARGS, CloughTocher2DInterpolator) + if self.ndim == 2 + else (None, None), + "rbf": (RBFI_KWARGS, RBFInterpolator), + } + + interp_kwargs, interpolator_class = interpolator_mapping.get( + method, (None, None) + ) + + if not interp_kwargs: + raise ValueError( + f"Unknown interpolation method {method} for {self.ndim} dimensional data." + ) + + self.interp_kwargs = interp_kwargs.copy() + self.interp_kwargs.update( + (k, v) for k, v in kwargs.items() if k in interp_kwargs + ) + self.interpolator = interpolator_class( + np.moveaxis(self.grids, -1, 0), self.values, **self.interp_kwargs + ) + + def __call__(self, *args): + """ + Interpolates function on arguments. + + Returns + ------- + np.ndarray + Interpolated values. + """ + + if self.method == "rbf": + coords = np.asarray(args).reshape(self.ndim, -1).T + return self.interpolator(coords).reshape(args[0].shape) + + return self.interpolator(*args) diff --git a/HARK/interpolation/_sklearn.py b/HARK/interpolation/_sklearn.py new file mode 100644 index 000000000..e400500db --- /dev/null +++ b/HARK/interpolation/_sklearn.py @@ -0,0 +1,345 @@ +import numpy as np +from sklearn.gaussian_process import GaussianProcessRegressor +from sklearn.kernel_ridge import KernelRidge +from sklearn.linear_model import ElasticNet, ElasticNetCV, SGDRegressor +from sklearn.pipeline import make_pipeline +from sklearn.preprocessing import PolynomialFeatures, SplineTransformer, StandardScaler +from sklearn.svm import SVR + +from HARK.interpolation._multi import _CurvilinearGridInterp, _UnstructuredGridInterp + + +class PipelineCurvilinearInterp(_CurvilinearGridInterp): + """ + Curvilinear Interpolator using a pipeline of sklearn models. + """ + + def __init__(self, values, grids, pipeline, **kwargs): + """ + Initialize a PipelineCurvilinearInterp object. + + Parameters + ---------- + values : np.ndarray + Functional values on a curvilinear grid. + grids : np.ndarray + Functional coordinates on a curvilinear grid. + pipeline : sklearn.pipeline.Pipeline + Pipeline of sklearn models. + """ + # for now, only support cpu + super().__init__(values, grids, target="cpu", **kwargs) + + self.pipeline = pipeline + + X_train = np.reshape(self.grids, (self.ndim, -1)) + y_train = np.mgrid[[slice(0, dim) for dim in self.shape]] + y_train = np.reshape(y_train, (self.ndim, -1)) + + self.models = [make_pipeline(*pipeline) for _ in range(self.ndim)] + for dim in range(self.ndim): + self.models[dim].fit(X_train, y_train[dim]) + + def _get_coordinates(self, args): + """ + Apply the sklearn pipeline to each dimension of arguments. + + Parameters + ---------- + args : np.ndarray + Values to interpolate for each dimension. + + Returns + ------- + np.ndarray + Interpolated values. + """ + X_test = np.reshape(args, (self.ndim, -1)) + return np.array([m.predict(X_test).reshape(args[0].shape) for m in self.models]) + + +class _PreprocessingCurvilinearInterp(PipelineCurvilinearInterp): + """ + Abstract class for PipelineCurvilinearInterp with preprocessing. + """ + + def __init__( + self, + values, + grids, + pipeline, + std=False, + preprocessing_options=None, + ): + """ + Initialize a _PreprocessingCurvilinearInterp object. Preprocessing options + includes standardization, polynomial features, and spline features. + + Parameters + ---------- + values : np.ndarray + Functional values on a curvilinear grid. + grids : np.ndarray + Functional coordinates on a curvilinear grid. + pipeline : sklearn.pipeline.Pipeline + Pipeline of sklearn models. + std : bool, optional + Standardize data by removing the mean and scaling to unit variance, + by default False + preprocessing_options : dict, optional + Can be None, or a dictionary with key "feature". + If "feature" is "pol", then "degree" must be specified. + If "feature" is "spl", then "degree" and "n_knots" must be specified. + + Raises + ------ + AttributeError + Feature not recognized. + """ + self.std = std + + if preprocessing_options is None: + preprocessing_options = {} + + self.preprocessing_options = preprocessing_options + + feature = preprocessing_options.get("feature", None) + + if feature and isinstance(feature, str): + degree = preprocessing_options.get("degree", 3) + assert isinstance(degree, int), "Degree must be an integer." + if feature.startswith("pol"): + pipeline.insert(0, PolynomialFeatures(degree)) + elif feature.startswith("spl"): + n_knots = preprocessing_options.get("n_knots", 5) + assert isinstance(n_knots, int), "n_knots must be an integer." + pipeline.insert(0, SplineTransformer(n_knots=n_knots, degree=degree)) + else: + raise AttributeError(f"Feature {feature} not recognized.") + else: + raise AttributeError(f"Feature {feature} not recognized.") + + if std: + pipeline.insert(0, StandardScaler()) + + super().__init__(values, grids, pipeline) + + +class GeneralizedRegressionCurvilinearInterp(_PreprocessingCurvilinearInterp): + """ + Generalized Regression for each dimension of the curvilinear grid. + Use regression to map from the curvilinear grid to an index grid. + Then use map_coordinates to interpolate on the index grid. + """ + + def __init__( + self, values, grids, model="elastic-net-cv", model_kwargs=None, **kwargs + ): + """ + Initialize a GeneralizedRegressionCurvilinearInterp object. + The model determines the regression used for each dimension. + + Parameters + ---------- + values : np.ndarray + Functional values on a curvilinear grid. + grids : np.ndarray + Functional coordinates on a curvilinear grid. + model : str, optional + One of "elastic-net", "elastic-net-cv", "kernel-ridge", "svr", "sgd", + "gaussian-process", by default "elastic-net". + model_kwargs : dict, optional + Options for the model, by default None. + + Raises + ------ + AttributeError + Model is not implemented. + """ + if model_kwargs is None: + model_kwargs = {} + + self.model = model + self.model_kwargs = model_kwargs + + if model == "elastic-net": + pipeline = [ElasticNet(**model_kwargs)] + elif model == "elastic-net-cv": + pipeline = [ElasticNetCV(**model_kwargs)] + elif model == "kernel-ridge": + pipeline = [KernelRidge(**model_kwargs)] + elif model == "svr": + pipeline = [SVR(**model_kwargs)] + elif model == "sgd": + pipeline = [SGDRegressor(**model_kwargs)] + elif model == "gaussian-process": + pipeline = [GaussianProcessRegressor(**model_kwargs)] + else: + raise AttributeError( + f"Model {model} not implemented. Consider using `PipelineCurvilinearInterp`." + ) + + super().__init__(values, grids, pipeline, **kwargs) + + +class PipelineUnstructuredInterp(_UnstructuredGridInterp): + """ + Unstructured Interpolator using a pipeline of sklearn models. + """ + + def __init__(self, values, grids, pipeline): + """ + Initialize a PipelineUnstructuredInterp object. + + Parameters + ---------- + values : np.ndarray + Functional values on an unstructured grid. + grids : np.ndarray + Functional coordinates on an unstructured grid. + pipeline : sklearn.pipeline.Pipeline + Pipeline of sklearn models. + """ + # for now, only support cpu + super().__init__(values, grids, target="cpu") + X_train = np.moveaxis(self.grids, -1, 0) + y_train = self.values + self.pipeline = pipeline + self.model = make_pipeline(*self.pipeline) + self.model.fit(X_train, y_train) + + def __call__(self, *args: np.ndarray): + """ + Interpolate on the unstructured grid. + + Returns + ------- + np.ndarray + Interpolated values. + """ + + X_test = np.c_[tuple(arg.ravel() for arg in args)] + return self.model.predict(X_test).reshape(args[0].shape) + + +class _PreprocessingUnstructuredInterp(PipelineUnstructuredInterp): + """ + Abstract class for PipelineUnstructuredInterp with preprocessing. + """ + + def __init__( + self, + values, + grids, + pipeline, + std=False, + feature=None, + degree=3, + n_knots=5, + ): + """ + Initialize a _PreprocessingUnstructuredInterp object. Preprocessing options + includes standardization, polynomial features, and spline features. + + Parameters + ---------- + values : np.ndarray + Functional values on an unstructured grid. + grids : np.ndarray + Functional coordinates on an unstructured grid. + pipeline : sklearn.pipeline.Pipeline + Pipeline of sklearn models. + std : bool, optional + Standardize data by removing the mean and scaling to unit variance, + by default False + preprocessing_options : dict, optional + Can be None, or a dictionary with key "feature". + If "feature" is "pol", then "degree" must be specified. + If "feature" is "spl", then "degree" and "n_knots" must be specified. + + Raises + ------ + AttributeError + Feature not recognized. + """ + + self.std = std + self.feature = feature + self.degree = degree + self.n_knots = n_knots + + if feature is None: + pass + elif isinstance(feature, str): + assert isinstance(degree, int), "Degree must be an integer." + if feature.startswith("pol"): + pipeline = [PolynomialFeatures(degree=degree)] + pipeline + elif feature.startswith("spl"): + assert isinstance(n_knots, int), "n_knots must be an integer." + pipeline = [ + SplineTransformer(n_knots=n_knots, degree=degree) + ] + pipeline + else: + raise AttributeError(f"Feature {feature} not recognized.") + else: + raise AttributeError(f"Feature {feature} not recognized.") + + if std: + pipeline = [StandardScaler()] + pipeline + + super().__init__(values, grids, pipeline) + + +class GeneralizedRegressionUnstructuredInterp(_PreprocessingUnstructuredInterp): + """ + Generalized Regression for an unstructured grid. + """ + + def __init__( + self, values, grids, model="elastic-net-cv", model_kwargs=None, **kwargs + ): + """ + Initialize a GeneralizedRegressionUnstructuredInterp object. + The model determines the regression used. + + Parameters + ---------- + values : np.ndarray + Functional values on an unstructured grid. + grids : np.ndarray + Functional coordinates on an unstructured grid. + model : str, optional + One of "elastic-net", "elastic-net-cv", "kernel-ridge", "svr", "sgd", + "gaussian-process", by default "elastic-net". + model_kwargs : dict, optional + Options for the model, by default None. + + Raises + ------ + AttributeError + Model is not implemented. + """ + if model_kwargs is None: + model_kwargs = {} + + self.model = model + self.model_kwargs = model_kwargs + + if model == "elastic-net": + pipeline = [ElasticNet(**model_kwargs)] + elif model == "elastic-net-cv": + pipeline = [ElasticNetCV(**model_kwargs)] + elif model == "kernel-ridge": + pipeline = [KernelRidge(**model_kwargs)] + elif model == "svr": + pipeline = [SVR(**model_kwargs)] + elif model == "sgd": + pipeline = [SGDRegressor(**model_kwargs)] + elif model == "gaussian-process": + pipeline = [GaussianProcessRegressor(**model_kwargs)] + else: + raise AttributeError( + f"Model {model} not implemented. Consider using `PipelineUnstructuredInterp`." + ) + + super().__init__(values, grids, pipeline, **kwargs) diff --git a/HARK/tests/test_econforgeinterp.py b/HARK/interpolation/tests/test_econforge.py similarity index 99% rename from HARK/tests/test_econforgeinterp.py rename to HARK/interpolation/tests/test_econforge.py index 2f7fed096..418d6106a 100644 --- a/HARK/tests/test_econforgeinterp.py +++ b/HARK/interpolation/tests/test_econforge.py @@ -2,8 +2,7 @@ import numpy as np -from HARK.econforgeinterp import DecayInterp, LinearFast -from HARK.interpolation import BilinearInterp, LinearInterp +from HARK.interpolation import BilinearInterp, DecayInterp, LinearFast, LinearInterp from HARK.metric import distance_metric diff --git a/HARK/tests/test_interpolation.py b/HARK/interpolation/tests/test_hark.py similarity index 95% rename from HARK/tests/test_interpolation.py rename to HARK/interpolation/tests/test_hark.py index ad0008330..7373e076b 100644 --- a/HARK/tests/test_interpolation.py +++ b/HARK/interpolation/tests/test_hark.py @@ -11,7 +11,8 @@ class testsLinearInterp(unittest.TestCase): - """tests for LinearInterp, currently tests for uneven length of + """ + tests for LinearInterp, currently tests for uneven length of x and y with user input as lists, arrays, arrays with column orientation """ @@ -41,7 +42,8 @@ def test_same_length(self): class testsCubicInterp(unittest.TestCase): - """tests for CubicInterp, currently tests for uneven length of + """ + tests for CubicInterp, currently tests for uneven length of x, y and derivative with user input as lists, arrays, arrays with column orientation """ @@ -80,7 +82,8 @@ def test_same_length(self): class testsBilinearInterp(unittest.TestCase): - """tests for BilinearInterp, currently tests for uneven length of + """ + tests for BilinearInterp, currently tests for uneven length of x, y, f(x,y) with user input as arrays, arrays with column orientation """ @@ -108,7 +111,8 @@ def test_same_length(self): class testsTrilinearInterp(unittest.TestCase): - """tests for TrilinearInterp, currently tests for uneven length of + """ + tests for TrilinearInterp, currently tests for uneven length of x, y, z, f(x, y, z) with user input as arrays, arrays with column orientation """ @@ -151,7 +155,8 @@ def test_same_length(self): class testsQuadlinearInterp(unittest.TestCase): - """tests for TrilinearInterp, currently tests for uneven length of + """ + tests for TrilinearInterp, currently tests for uneven length of w, x, y, z, f(w, x, y, z) with user input as arrays, arrays with column orientation """ diff --git a/HARK/interpolation/tests/test_multi.py b/HARK/interpolation/tests/test_multi.py new file mode 100644 index 000000000..4fac4e2d9 --- /dev/null +++ b/HARK/interpolation/tests/test_multi.py @@ -0,0 +1,58 @@ +import unittest + +import numpy as np + +from HARK.interpolation import MultivariateInterp + + +def function(*args): + mats = np.meshgrid(*args, indexing="ij") + + return np.sum(mats, axis=0) + + +class TestMultivariateInterp(unittest.TestCase): + def setUp(self): + # create test data + + self.grids = [ + np.linspace(0, 1, 10), + np.linspace(0, 1, 11), + np.linspace(0, 1, 12), + ] + + self.args = [ + np.linspace(0, 1, 11), + np.linspace(0, 1, 12), + np.linspace(0, 1, 13), + ] + + def test_interpolation_values(self): + # check that interpolation values match expected values + + interpolator2D_cpu = MultivariateInterp( + function(*self.grids[0:2]), self.grids[0:2], target="cpu" + ) + interpolator2D_parallel = MultivariateInterp( + function(*self.grids[0:2]), self.grids[0:2], target="parallel" + ) + interpolator3D_cpu = MultivariateInterp( + function(*self.grids), self.grids, target="cpu" + ) + interpolator3D_parallel = MultivariateInterp( + function(*self.grids), self.grids, target="parallel" + ) + + val2D_cpu = interpolator2D_cpu(*np.meshgrid(*self.args[0:2], indexing="ij")) + val2D_parallel = interpolator2D_parallel( + *np.meshgrid(*self.args[0:2], indexing="ij") + ) + val3D_cpu = interpolator3D_cpu(*np.meshgrid(*self.args, indexing="ij")) + val3D_parallel = interpolator3D_parallel( + *np.meshgrid(*self.args, indexing="ij") + ) + + self.assertTrue(np.allclose(val2D_cpu, function(*self.args[0:2]))) + self.assertTrue(np.allclose(val2D_parallel, function(*self.args[0:2]))) + self.assertTrue(np.allclose(val3D_cpu, function(*self.args))) + self.assertTrue(np.allclose(val3D_parallel, function(*self.args))) diff --git a/HARK/interpolation/tests/test_scipy.py b/HARK/interpolation/tests/test_scipy.py new file mode 100644 index 000000000..220487be7 --- /dev/null +++ b/HARK/interpolation/tests/test_scipy.py @@ -0,0 +1,46 @@ +import unittest + +import numpy as np + +from HARK.interpolation import UnstructuredInterp + + +def function(*args): + mats = np.meshgrid(*args, indexing="ij") + + return np.sum(mats, axis=0) + + +class TestMultivariateInterp(unittest.TestCase): + def setUp(self): + # create test data + + self.grids = [ + np.linspace(0, 1, 10), + np.linspace(0, 1, 11), + np.linspace(0, 1, 12), + ] + + self.args = [ + np.linspace(0, 1, 11), + np.linspace(0, 1, 12), + np.linspace(0, 1, 13), + ] + + def test_interpolation_values(self): + # check that interpolation values match expected values + + interpolator2D = UnstructuredInterp( + function(*self.grids[0:2]), [*np.meshgrid(*self.grids[0:2], indexing="ij")] + ) + + interpolator3D = UnstructuredInterp( + function(*self.grids), [*np.meshgrid(*self.grids, indexing="ij")] + ) + + val2D = interpolator2D(*np.meshgrid(*self.args[0:2], indexing="ij")) + + val3D = interpolator3D(*np.meshgrid(*self.args, indexing="ij")) + + self.assertTrue(np.allclose(val2D, function(*self.args[0:2]))) + self.assertTrue(np.allclose(val3D, function(*self.args))) diff --git a/HARK/interpolation/tests/test_sklearn.py b/HARK/interpolation/tests/test_sklearn.py new file mode 100644 index 000000000..2bafdce68 --- /dev/null +++ b/HARK/interpolation/tests/test_sklearn.py @@ -0,0 +1,50 @@ +import unittest + +import numpy as np + +from HARK.interpolation import GeneralizedRegressionUnstructuredInterp + + +def function(*args): + mats = np.meshgrid(*args, indexing="ij") + + return np.sum(mats, axis=0) + + +class TestMultivariateInterp(unittest.TestCase): + def setUp(self): + # create test data + + self.grids = [ + np.linspace(0, 1, 10), + np.linspace(0, 1, 11), + np.linspace(0, 1, 12), + ] + + self.args = [ + np.linspace(0, 1, 11), + np.linspace(0, 1, 12), + np.linspace(0, 1, 13), + ] + + def test_interpolation_values(self): + # check that interpolation values match expected values + + interpolator2D = GeneralizedRegressionUnstructuredInterp( + function(*self.grids[0:2]), + [*np.meshgrid(*self.grids[0:2], indexing="ij")], + model_kwargs={"fit_intercept": False}, + ) + + interpolator3D = GeneralizedRegressionUnstructuredInterp( + function(*self.grids), + [*np.meshgrid(*self.grids, indexing="ij")], + model_kwargs={"fit_intercept": False}, + ) + + val2D = interpolator2D(*np.meshgrid(*self.args[0:2], indexing="ij")) + + val3D = interpolator3D(*np.meshgrid(*self.args, indexing="ij")) + + self.assertTrue(np.allclose(val2D, function(*self.args[0:2]), rtol=1e-2)) + self.assertTrue(np.allclose(val3D, function(*self.args), rtol=1e-2)) diff --git a/environment.yml b/environment.yml new file mode 100644 index 000000000..d989bc130 --- /dev/null +++ b/environment.yml @@ -0,0 +1,20 @@ +name: hark-dev +channels: + - conda-forge + - defaults +dependencies: + - estimagic + - interpolation>=2.2.3 + - joblib>=1.2 + - jupyterlab>=3.6 + - jupytext>=1.14 + - matplotlib>=3.6 + - networkx>=3 + - numba>=0.56 + - numpy>=1.23 + - pandas>=1.5 + - quantecon + - scipy>=1.10 + - scikit-learn + - seaborn>=0.12 + - xarray>=2023 diff --git a/examples/Interpolation/DecayInterp.ipynb b/examples/Interpolation/DecayInterp.ipynb index e3812bd61..ea3ae542d 100644 --- a/examples/Interpolation/DecayInterp.ipynb +++ b/examples/Interpolation/DecayInterp.ipynb @@ -35,7 +35,7 @@ "# Setup\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", - "from HARK.econforgeinterp import LinearFast, DecayInterp\n", + "from HARK.interpolation import LinearFast, DecayInterp\n", "from HARK.interpolation import LinearInterp" ] }, @@ -47,7 +47,7 @@ "## Basic use of the `DecayInterp` class\n", "\n", "`DecayInterp` requires two basic inputs:\n", - "- `interp`, which is the interpolator $f(\\cdot)$. It must be an instance of the `HARK.econforgeinterp.LinearFast` class.\n", + "- `interp`, which is the interpolator $f(\\cdot)$. It must be an instance of the `HARK.interpolation.LinearFast` class.\n", "- `limit_fun`, which is the limiting function $g(\\cdot)$. It must receive the same number of inputs as `interp` and be able to take `numpy` arrays as inputs.\n", "\n", "And that's it!\n", diff --git a/examples/Interpolation/DecayInterp.py b/examples/Interpolation/DecayInterp.py index dfbc4a0dc..abe564e79 100644 --- a/examples/Interpolation/DecayInterp.py +++ b/examples/Interpolation/DecayInterp.py @@ -30,14 +30,14 @@ # Setup import numpy as np import matplotlib.pyplot as plt -from HARK.econforgeinterp import LinearFast, DecayInterp +from HARK.interpolation import LinearFast, DecayInterp from HARK.interpolation import LinearInterp # %% [markdown] # ## Basic use of the `DecayInterp` class # # `DecayInterp` requires two basic inputs: -# - `interp`, which is the interpolator $f(\cdot)$. It must be an instance of the `HARK.econforgeinterp.LinearFast` class. +# - `interp`, which is the interpolator $f(\cdot)$. It must be an instance of the `HARK.interpolation.LinearFast` class. # - `limit_fun`, which is the limiting function $g(\cdot)$. It must receive the same number of inputs as `interp` and be able to take `numpy` arrays as inputs. # # And that's it! diff --git a/examples/Interpolation/README.md b/examples/Interpolation/README.md new file mode 100644 index 000000000..38891a098 --- /dev/null +++ b/examples/Interpolation/README.md @@ -0,0 +1,3 @@ +# A practical guide to Multivariate Interpolation on Irregular Grids + +More examples can be found at [alanlujan91/multinterp](http://github.com/alanlujan91/multinterp) diff --git a/requirements/base.txt b/requirements/base.txt index 5821f865e..9ddfba33f 100644 --- a/requirements/base.txt +++ b/requirements/base.txt @@ -5,7 +5,9 @@ networkx>=3 numba>=0.56 numpy>=1.23 pandas>=1.5 -quantecon +quantecon>=0.6 +recommonmark>=0.7 +scikit-learn scipy>=1.10 seaborn>=0.12 xarray>=2023