diff --git a/CHANGELOG.md b/CHANGELOG.md index e3482dfb..9b112168 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,13 @@ +## v0.1.2.dev10 (NOT YET RELEASED) + +### CoFI Utils + +- [#65](https://github.com/inlab-geo/cofi/issues/54) Utility functions using findiff to generate the difference matrices + + ## v0.1.2.dev9 (13/09/2022) ### CoFI Core diff --git a/docs/_templates/autosummary/class.rst b/docs/_templates/autosummary/class.rst index ff870fa5..ce99e595 100644 --- a/docs/_templates/autosummary/class.rst +++ b/docs/_templates/autosummary/class.rst @@ -10,35 +10,33 @@ .. rubric:: Reference Details -.. ---- {% for item in members %} {% if item in ['__call__'] %} -.. automethod:: {{ objname }}.{{ item }} + .. automethod:: {{ objname }}.{{ item }} +{% endif %} +{% if "BaseRegularisation" in objname and item == '__add__' %} + .. automethod:: {{ objname }}.{{ item }} {% endif %} {% endfor %} -.. :ref:`back to top ` -.. ---- {% for item in methods %} {% if "BaseSolver" in objname and item == '__init__' %} -.. automethod:: {{ objname }}.{{ item }} + .. automethod:: {{ objname }}.{{ item }} {% endif %} {% if item != '__init__' %} -.. automethod:: {{ objname }}.{{ item }} + .. automethod:: {{ objname }}.{{ item }} {% endif %} {% endfor %} -.. :ref:`back to top ` ----- {% for item in attributes %} -.. autoattribute:: {{ objname }}.{{ item }} + .. autoattribute:: {{ objname }}.{{ item }} {% endfor %} :ref:`back to top ` @@ -51,7 +49,7 @@ .. rubric:: Reference Details -.. automethod:: cofi.SamplingResult.to_arviz + .. automethod:: cofi.SamplingResult.to_arviz :ref:`back to top ` {% endif %} diff --git a/docs/api/index.rst b/docs/api/index.rst index 0c5d9254..5dd4c00a 100644 --- a/docs/api/index.rst +++ b/docs/api/index.rst @@ -61,3 +61,15 @@ tools, simply create a subclass of :code:`BaseSolver` and implements :code:`__in :toctree: generated/ cofi.solvers.BaseSolver + + +CoFI Utils +---------- + +Module :mod:`cofi.utils` contains a set of helper classes and routines to enable +easier problem definitions. + +.. autosummary:: + :toctree: generated/ + + cofi.utils diff --git a/docs/cofi-examples b/docs/cofi-examples index 24af2027..4cb411ef 160000 --- a/docs/cofi-examples +++ b/docs/cofi-examples @@ -1 +1 @@ -Subproject commit 24af2027590568579c81b65e7a67ce55fe6e032a +Subproject commit 4cb411ef35d46291f2b8897bd1e22249c216de00 diff --git a/docs/conf.py b/docs/conf.py index 6c11b123..b96de402 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -91,6 +91,7 @@ def setup(app): "scipy": ("https://docs.scipy.org/doc/scipy/", None), "numpy": ("https://numpy.org/doc/stable/", None), "arviz": ("https://python.arviz.org/en/latest/", None), + "findiff": ("https://findiff.readthedocs.io/en/latest/", None), } # Disable including boostrap CSS for sphinx_panels since it's already included @@ -156,6 +157,7 @@ def setup(app): "dollarmath", "html_image", ] +nb_execution_mode = "cache" # -- Cutomised variables ------------------------------------------------------ diff --git a/docs/installation.rst b/docs/installation.rst index c347d731..76c657ab 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -11,6 +11,7 @@ CoFI requires Python 3.6+, and the following dependencies: - scipy>=1.0.0 - emcee>=3.1.0 - arviz>=0.9.0 +- findiff>=0.7.0 Virtual environment setup diff --git a/envs/environment_dev.yml b/envs/environment_dev.yml index f4c36320..b9a5069c 100644 --- a/envs/environment_dev.yml +++ b/envs/environment_dev.yml @@ -48,3 +48,4 @@ dependencies: - watermark - sphinxcontrib-svg2pdfconverter==1.2.0 - papermill + - findiff>=0.7.0 diff --git a/setup.py b/setup.py index 4e1a4926..3c07a70a 100644 --- a/setup.py +++ b/setup.py @@ -2,6 +2,7 @@ from io import StringIO import sys from pathlib import Path +from setuptools import find_namespace_packages try: from skbuild import setup @@ -54,8 +55,7 @@ "License :: OSI Approved :: BSD License", ] PACKAGE_DIR = {"": "src"} -# PACKAGES = find_packages("src") -PACKAGES = ["cofi"] +PACKAGES = find_namespace_packages(where='src') CMAKE_INSTALL_DIR = "src/cofi" CMAKE_ARGS=['-DSKBUILD=ON'] PYTHON_REQUIRES = ">=3.6" @@ -64,6 +64,7 @@ "scipy>=1.0.0", "emcee>=3.1.0", "arviz>=0.9.0", + "findiff>=0.7.0", ] EXTRAS_REQUIRE = { "petsc": ["petsc4py>=3.16.0"], diff --git a/src/cofi/base_problem.py b/src/cofi/base_problem.py index 24402255..d1afe470 100644 --- a/src/cofi/base_problem.py +++ b/src/cofi/base_problem.py @@ -1195,7 +1195,7 @@ def set_model_shape(self, model_shape: Tuple): np.reshape(self.initial_model, model_shape) except ValueError as err: raise DimensionMismatchError( - entered_dimenion=model_shape, + entered_dimension=model_shape, entered_name="model shape", expected_dimension=self.initial_model.shape, expected_source="initial model" diff --git a/src/cofi/exceptions.py b/src/cofi/exceptions.py index 0142104e..852f4c73 100644 --- a/src/cofi/exceptions.py +++ b/src/cofi/exceptions.py @@ -62,22 +62,23 @@ class DimensionMismatchError(CofiError, ValueError): def __init__( self, *args, - entered_dimenion: Tuple, + entered_dimension: Tuple, entered_name: str, expected_dimension: Tuple, expected_source: str, ) -> None: super().__init__(*args) - self._entered_dimension = entered_dimenion + self._entered_dimension = entered_dimension self._entered_name = entered_name self._expected_dimension = expected_dimension self._expected_source = expected_source def __str__(self) -> str: super_msg = super().__str__() - msg = f"the {self._entered_name} you've provided {self._entered_dimension}" \ - f" doesn't match and cannot be reshaped into the dimension you've set" \ - f" for {self._expected_source} which is {self._expected_dimension}" + msg = f"the {self._entered_name} you've provided (shape: " \ + f"{self._entered_dimension}) doesn't match and cannot be reshaped " \ + f"into the dimension you've set for {self._expected_source} which is " \ + f"{self._expected_dimension}" return self._form_str(super_msg, msg) diff --git a/src/cofi/utils/__init__.py b/src/cofi/utils/__init__.py new file mode 100644 index 00000000..c7e32686 --- /dev/null +++ b/src/cofi/utils/__init__.py @@ -0,0 +1,13 @@ +"""Utility classes and functions (e.g. to generate regularisation terms and more) +""" + +from .regularisation import ( + BaseRegularisation, + QuadraticReg, +) + + +__all__ = [ + "BaseRegularisation", + "QuadraticReg", +] diff --git a/src/cofi/utils/regularisation.py b/src/cofi/utils/regularisation.py new file mode 100644 index 00000000..6dba17fc --- /dev/null +++ b/src/cofi/utils/regularisation.py @@ -0,0 +1,446 @@ +from numbers import Number +from abc import abstractmethod, ABCMeta +from typing import Union, Any +import numpy as np +import findiff + +from ..exceptions import DimensionMismatchError + + +REG_TYPES = { + "damping": 0, + "flattening": 1, + "roughening": 1, + "smoothing": 2, +} + + +class BaseRegularisation(metaclass=ABCMeta): + r"""Base class for a regularisation term + + Check :class:`QuadraticReg` for a concrete example. + + .. rubric:: Basic interface + + The basic properties / methods for a regularisation term in ``cofi.utils`` + include the following: + + .. autosummary:: + BaseRegularisation.model_size + BaseRegularisation.reg + BaseRegularisation.gradient + BaseRegularisation.hessian + BaseRegularisation.__call__ + + .. rubric:: Adding two terms + + Two instances of ``BaseRegularisation`` can also be added together using the ``+`` + operator: + + .. autosummary:: + BaseRegularisation.__add__ + + """ + def __init__(self,): + pass + + @property + @abstractmethod + def model_size(self) -> Number: + """the number of unknowns that current regularisation function accepts + """ + raise NotImplementedError + + def __call__(self, model: np.ndarray) -> Number: + r"""a class instance itself can also be called as a function + + It works exactly the same as :meth:`reg`. + + In other words, the following two usages are exactly the same:: + + >>> my_reg = QuadraticReg(factor=1, model_size=3) + >>> my_reg_value = my_reg(np.array([1,2,3])) # usage 1 + >>> my_reg_value = my_reg.reg(np.array([1,2,3])) # usage 2 + """ + return self.reg(model) + + @abstractmethod + def reg(self, model: np.ndarray) -> Number: + """the regularisation function value given a model to evaluate + """ + raise NotImplementedError + + @abstractmethod + def gradient(self, model: np.ndarray) -> np.ndarray: + """the gradient of regularisation function with respect to model given a model + + The usual size for the gradient is :math:`(M,)` where :math:`M` is the number + of model parameters + """ + raise NotImplementedError + + @abstractmethod + def hessian(self, model: np.ndarray) -> np.ndarray: + """the hessian of regularisation function with respect to model given a model + + The usual size for the Hessian is :math:`(M,M)` where :math:`M` is the number + of model parameters + """ + raise NotImplementedError + + def __add__(self, other_reg): + r"""Adds two regularisation terms + + Parameters + ---------- + other_reg : BaseRegularisation + the second argument of "+" operator; must also be a + :class:`BaseRegularisation` instance + + Returns + ------- + BaseRegularisation + a regularisation term ``resRegularisation`` such that: + + - :math:`\text{resRegularisation.reg}(m)=\text{self.reg}(m)+\text{other_reg.reg}(m)` + - :math:`\text{resRegularisation.gradient}(m)=\text{self.gradient}(m)+\text{other_reg.gradient}(m)` + - :math:`\text{resRegularisation.hessian}(m)=\text{self.hessian}(m)+\text{other_reg.hessian}(m)` + + Raises + ------ + TypeError + when the ``other_reg`` is not a regularisation term generated by CoFI Utils + DimensionMismatchError + when the ``other_reg`` doesn't accept model_size that matches the one of + ``self`` + + Examples + -------- + + >>> from cofi import BaseProblem + >>> from cofi.utils import QuadraticReg + >>> reg1 = QuadraticReg(factor=1, model_size=3, reg_type="damping") + >>> reg2 = QuadraticReg(factor=2, model_size=3, reg_type="smoothing") + >>> my_problem = BaseProblem() + >>> my_problem.set_regularisation(reg1 + reg2) + + """ + if not isinstance(other_reg, BaseRegularisation): + raise TypeError( + f"unsupported operand type(s) for +: '{self.__class__.__name__}' " + f"and '{other_reg.__class__.__name__}" + ) + if self.model_size != other_reg.model_size: + raise DimensionMismatchError( + entered_name="the second regularisation term", + entered_dimension=other_reg.model_size, + expected_source="the first regularisation term", + expected_dimension=self.model_size, + ) + tmp_model_size = self.model_size + tmp_reg = self.reg + tmp_grad = self.gradient + tmp_hess = self.hessian + class NewRegularisation(BaseRegularisation): + @property + def model_size(self): + return tmp_model_size + def reg(self, model): + return tmp_reg(model) + other_reg(model) + def gradient(self, model): + return tmp_grad(model) + other_reg.gradient(model) + def hessian(self, model): + return tmp_hess(model) + other_reg.hessian(model) + return NewRegularisation() + + +class QuadraticReg(BaseRegularisation): + r"""CoFI's utility class to calculate damping, flattening (roughening), and + smoothing regularisation + + They correspond to the zeroth order, first order and second order Tikhonov + regularisation approaches respectively. + + - If ``reg_type == "damping"``, then + + .. toggle:: + + - :meth:`reg` produces :math:`\text{reg}=\text{factor}\times||m-m_0||_2^2` + - :meth:`gradient` produces + :math:`\frac{\partial\text{reg}}{\partial m}=2\times\text{factor}\times(m-m_0)` + - :meth:`hessian` produces + :math:`\frac{\partial^2\text{reg}}{\partial m}=2\times\text{factor}\times I` + - :attr:`matrix` is the identity matrix of size :math:`(M,M)` + - where + + - :math:`m_0` is a reference model that you can specify in ``ref_model`` argument + (default to zero) + - :math:`M` is the number of model parameters + + - If ``reg_type == "roughening"`` (or equivalently ``"flattening"``), + then + + .. toggle:: + + - :meth:`reg` produces :math:`\text{reg}=\text{factor}\times||Dm||_2^2` + - :meth:`gradient` produces + :math:`\frac{\partial\text{reg}}{\partial m}=2\times\text{factor}\times D^TDm` + - :meth:`hessian` produces + :math:`\frac{\partial^2\text{reg}}{\partial m}=2\times\text{factor}\times D^TD` + - :attr:`matrix` is :math:`D` + - where + + - :math:`D` matrix helps calculate the first order derivative of :math:`m`. + For 1D problems, it looks like + + :math:`\begin{pmatrix}-1.5&2&-0.5&&&\\-0.5&&0.5&&&&&\\&-0.5&&0.5&&&&\\&&...&&...&&&\\&&&-0.5&&0.5&&\\&&&&-0.5&&0.5&\\&&&&&0.5&-2&1.5\end{pmatrix}` + + .. :math:`\begin{pmatrix}-1&1&&&&\\&-1&1&&&\\&&...&...&&\\&&&-1&1&\\&&&&&-1&1\end{pmatrix}` + + While for higher dimension problems, by default it's a flattened version of + an N-D array. The actual ordering of model parameters in higher dimensions + is controlled by :class:`findiff.operators.FinDiff`. + + - If ``reg_type == "smoothing"``, then + + .. toggle:: + + - :meth:`reg` produces :math:`\text{reg}=\text{factor}\times||Dm||_2^2` + - :meth:`gradient` produces + :math:`\frac{\partial\text{reg}}{\partial m}=2\times\text{factor}\times D^TDm` + - :meth:`hessian` produces + :math:`\frac{\partial^2\text{reg}}{\partial m}=2\times\text{factor}\times D^TD` + - :attr:`matrix` is :math:`D` + - where + + - :math:`D` matrix helps calculate the second order derivatives of :math:`m`. + For 1D problems, it looks like + + :math:`\begin{pmatrix}2&-5&4&-1&&&\\1&-2&1&&&&\\&1&-2&1&&&\\&&...&...&...&&\\&&&1&-2&1&\\&&&&1&-2&1\\&&&-1&4&-5&2\end{pmatrix}` + + .. :math:`\begin{pmatrix}1&-2&1&&&&\\&1&-2&1&&&\\&&...&...&...&&\\&&&1&-2&1&\\&&&&1&-2&1\end{pmatrix}` + + While for higher dimension problems, by default it's a flattened version of + an N-D array. The actual ordering of model parameters in higher dimensions + is controlled by :class:`findiff.operators.FinDiff`. + + - If ``reg_type == None``, then we assume you want to use the argument + ``byo_matrix``, + + .. toggle:: + + - :meth:`reg` produces :math:`\text{reg}=\text{factor}\times||Dm||_2^2` + - :meth:`gradient` produces + :math:`\frac{\partial\text{reg}}{\partial m}=2\times\text{factor}\times D^TDm` + - :meth:`hessian` produces + :math:`\frac{\partial^2\text{reg}}{\partial m}=2\times\text{factor}\times D^TD` + - :attr:`matrix` is :math:`D` + - where + + - :math:`D` matrix is ``byo_matrix`` from the arguments (or identity matrix + if ``byo_matrix is None``) + + Parameters + ---------- + factor : Number + the scale for the regularisation term + model_size : Number + the number of elements in a inference model + reg_type : str + specify what kind of regularisation is to be calculated, by default + ``"damping"`` + ref_model : np.ndarray + reference model used only when ``reg_type == "damping"``, + by default None (if this is None, then reference model is assumed to be zero) + byo_matrix : np.ndarray + bring-your-own matrix, activated only when ``reg_type == None`` + + Raises + ------ + ValueError + when input arguments don't conform to the standards described above. Check + error message for details. + + Examples + -------- + + Generate a quadratic damping regularisation matrix for model of size 3: + + >>> from cofi.utils import QuadraticReg + >>> reg = QuadraticReg(factor=1, model_size=3) + >>> reg(np.array([1,2,3])) + 3.0 + + To use together with :class:`cofi.BaseProblem`: + + >>> from cofi import BaseProblem + >>> from cofi.utils import QuadraticReg + >>> reg = QuadraticReg(factor=1, model_size=3) + >>> my_problem = BaseProblem() + >>> my_problem.set_regularisation(reg) + + You may also combine two regularisation terms: + + >>> from cofi import BaseProblem + >>> from cofi.utils import QuadraticReg + >>> reg1 = QuadraticReg(factor=1, model_size=3, reg_type="damping") + >>> reg2 = QuadraticReg(factor=2, model_size=3, reg_type="smoothing") + >>> my_problem = BaseProblem() + >>> my_problem.set_regularisation(reg1 + reg2) + """ + def __init__( + self, + factor: Number, + model_size: Union[Number, np.ndarray], + reg_type: str="damping", + ref_model: np.ndarray=None, + byo_matrix: np.ndarray=None, + ): + super().__init__() + self._factor = self._validate_factor(factor) + self._model_size = model_size + self._reg_type = self._validate_reg_type(reg_type) + self._ref_model = ref_model + self._byo_matrix = byo_matrix + self._generate_matrix() + + @property + def model_size(self) -> Number: + """the number of model parameters + + This is always a number describing number of unknowns + """ + if self._reg_type == "damping": + return self._model_size + else: + return self._model_size[0] * self._model_size[1] + + @property + def matrix(self) -> np.ndarray: + """the regularisation matrix + + This is either an identity matrix, or first/second order difference matrix + (generated by Python package ``findiff``), or a custom matrix brought on your + own. + """ + return self._D + + def reg(self, model: np.ndarray) -> Number: + flat_m = self._validate_model(model) + if self._reg_type == "damping": + if self._ref_model is None: + return self._factor * flat_m.T @ flat_m + diff_ref = flat_m - self._ref_model + return self._factor * diff_ref.T @ diff_ref + else: + flat_m = self._validate_model(model) + weighted_m = self.matrix @ flat_m + return self._factor * weighted_m.T @ weighted_m + + def gradient(self, model: np.ndarray) -> np.ndarray: + flat_m = self._validate_model(model) + if self._reg_type == "damping": + if self._ref_model is None: + return 2 * self._factor * flat_m + return 2 * self._factor * (flat_m - self._ref_model) + else: + return 2 * self._factor * self.matrix.T @ self.matrix @ flat_m + + def hessian(self, model: np.ndarray) -> np.ndarray: + if self._reg_type == "damping": + return 2 * self._factor * np.eye(self._model_size) + else: + return 2 * self._factor * self.matrix.T @ self.matrix + + @staticmethod + def _validate_factor(factor): + if not isinstance(factor, Number): + raise ValueError("the regularisation factor must be a number") + elif factor < 0: + raise ValueError("the regularisation factor must be non-negative") + return factor + + @staticmethod + def _validate_reg_type(reg_type): + if reg_type is not None and (not isinstance(reg_type, str) or \ + reg_type not in REG_TYPES): + raise ValueError( + "Please choose a valid regularisation type. `damping`, " + "`flattening` and `smoothing` are supported." + ) + return reg_type + + def _generate_matrix(self): + if self._reg_type == "damping": + if not isinstance(self._model_size, Number): + raise ValueError( + "model_size must be a number when damping is selected" + ) + self._D = np.identity(self._model_size) + elif self._reg_type in REG_TYPES: # 1st/2nd order Tikhonov + if np.size(self._model_size) == 1: + order = REG_TYPES[self._reg_type] + if self._model_size < order+2: + raise ValueError( + f"the model_size needs to be at least >={order+2} " + f"for regularisation type '{self._reg_type}'" + ) + d_dx = findiff.FinDiff(0, 1, order) + self._D = d_dx.matrix((self._model_size,)).toarray() + elif np.size(self._model_size) == 2 and np.ndim(self._model_size) == 1: + nx = self._model_size[0] + ny = self._model_size[1] + order = REG_TYPES[self._reg_type] + if nx < order+2 or ny < order+2: + raise ValueError( + f"the model_size needs to be at least (>={order+2}, >={order+2}) " + f"for regularisation type '{self._reg_type}'" + ) + d_dx = findiff.FinDiff(0, 1, order) # x direction + d_dy = findiff.FinDiff(1, 1, order) # y direction + matx = d_dx.matrix((nx, ny)) # scipy sparse matrix + maty = d_dy.matrix((nx, ny)) # scipy sparse matrix + self._D = np.vstack((matx.toarray(), maty.toarray())) # combine above + else: + raise NotImplementedError( + "only 2D derivative operators implemented" + ) + elif self._reg_type is None: + if not isinstance(self._model_size, Number): + raise ValueError( + "please provide a number for 'model_size' when bringing your " + "own weighting matrix" + ) + if self._byo_matrix is None: + self._D = np.identity(self._model_size) + else: + self._D = self._byo_matrix + if len(self._D.shape) != 2: + raise ValueError( + "the bring-your-own regularisation matrix must be 2-dimensional" + ) + elif self._D.shape[1] != self._model_size: + raise ValueError( + "the bring-your-own regularisation matrix must be in shape (_, M) " + "where M is the model size" + ) + + def _validate_model(self, model): + flat_m = np.ravel(model) + if self._reg_type == "damping": + if flat_m.size != self._model_size: + raise DimensionMismatchError( + entered_name="model", + entered_dimension=model.shape, + expected_source="model", + expected_dimension=self._model_size + ) + else: + if flat_m.shape[0] != self._model_size[0] * self._model_size[1]: + raise DimensionMismatchError( + entered_name="model", + entered_dimension=model.shape, + expected_source="model", + expected_dimension=self._model_size + ) + return flat_m diff --git a/tests/cofi_utils/test_reg_base.py b/tests/cofi_utils/test_reg_base.py new file mode 100644 index 00000000..6a20887c --- /dev/null +++ b/tests/cofi_utils/test_reg_base.py @@ -0,0 +1,48 @@ +import pytest +import numpy as np + +from cofi.utils import BaseRegularisation, QuadraticReg +from cofi.exceptions import DimensionMismatchError + + +def test_base_reg(): + class subclass_reg(BaseRegularisation): + def __init__(self): + super().__init__() + @property + def model_size(self): + return super().model_size + def reg(self, model): + if isinstance(model, np.ndarray): + return model + else: + return super().reg(model) + def gradient(self, model): + return super().gradient(model) + def hessian(self, model): + return super().hessian(model) + test_reg = subclass_reg() + assert test_reg(np.array([1]))[0] == 1 + with pytest.raises(NotImplementedError): test_reg.model_size + with pytest.raises(NotImplementedError): test_reg(1) + with pytest.raises(NotImplementedError): test_reg.gradient(1) + with pytest.raises(NotImplementedError): test_reg.hessian(1) + +def test_add_regs(): + reg1 = QuadraticReg(1, 9) + reg2 = QuadraticReg(1, (3,3), reg_type="flattening") + reg3 = reg1 + reg2 + test_model = np.array([[1,2,3],[1,2,3],[2,3,4]]) + assert reg1(test_model) + reg2(test_model) == reg3(test_model) + assert np.array_equal(reg1.gradient(test_model) + reg2.gradient(test_model), reg3.gradient(test_model)) + assert np.array_equal(reg1.hessian(test_model) + reg2.hessian(test_model), reg3.hessian(test_model)) + assert reg1.model_size == reg2.model_size + assert reg1.model_size == reg3.model_size + +def test_add_regs_invalid(): + reg1 = QuadraticReg(1, 3) + reg2 = QuadraticReg(1, (3,3), reg_type="flattening") + with pytest.raises(DimensionMismatchError): + reg1 + reg2 + with pytest.raises(TypeError): + reg1 + 1 diff --git a/tests/cofi_utils/test_reg_quadratic.py b/tests/cofi_utils/test_reg_quadratic.py new file mode 100644 index 00000000..ba4ad518 --- /dev/null +++ b/tests/cofi_utils/test_reg_quadratic.py @@ -0,0 +1,134 @@ +import pytest +import numpy as np + +from cofi.utils import QuadraticReg +from cofi.exceptions import DimensionMismatchError + + +def test_damping(): + # no ref model + reg = QuadraticReg(1, 3) + reg_mat = reg.matrix + assert reg_mat.shape[0] == 3 + assert reg_mat.shape[0] == reg_mat.shape[1] + assert (reg_mat == np.eye(reg_mat.shape[0])).all() + reg_val = reg(np.array([1,2,3])) + assert reg_val == 14 + grad = reg.gradient(np.array([1,2,3])) + assert grad[0] == 2 + assert grad[1] == 4 + assert grad[2] == 6 + hess = reg.hessian(np.array([1,2,3])) + assert hess[0,0] == 2 + assert hess[1,1] == 2 + assert hess[2,2] == 2 + # ref model + reg = QuadraticReg(1, 3, ref_model=np.array([1,1,1])) + reg_val = reg(np.array([1,2,3])) + assert reg_val == 5 + grad = reg.gradient(np.array([1,2,3])) + assert grad[0] == 0 + assert grad[1] == 2 + assert grad[2] == 4 + hess = reg.hessian(np.array([1,2,3])) + assert hess[0,0] == 2 + assert hess[1,1] == 2 + assert hess[2,2] == 2 + +def test_damping_invalid(): + with pytest.raises(ValueError): QuadraticReg(1, 10, "dampingnf") + with pytest.raises(ValueError): QuadraticReg(-1, 3) + with pytest.raises(ValueError): QuadraticReg(None, 3) + with pytest.raises(ValueError): QuadraticReg("hello", 3) + with pytest.raises(ValueError): QuadraticReg(1, (1,2), "damping") + reg = QuadraticReg(1, 3) + with pytest.raises(DimensionMismatchError): + reg(np.array([1,1])) + +def test_flattening(): + reg = QuadraticReg(factor=1, model_size=(3,3), reg_type="flattening") + reg = QuadraticReg(factor=1, model_size=(3,3), reg_type="roughening") + # 1 + assert reg(np.zeros((3,3))) == 0 + # 2 + assert reg(np.ones((3,3))) == 0 + # 3 + reg_val1 = reg(np.array([[1,2,3],[1,2,3],[2,3,4]])) + reg_val2 = reg(np.array([1,2,3,1,2,3,2,3,4])) + assert reg_val1 == reg_val2 + # 4 + test_model = np.array([[1,1,1],[1,1,1],[1,1,2]]) + reg_val = reg(test_model) + assert reg_val == 5.5 + reg_grad = reg.gradient(test_model) + assert reg_grad.shape == (9,) + reg_hess = reg.hessian(test_model) + assert reg_hess.shape == (9,9) + # 5 - test matrix + reg = QuadraticReg(1, 3, reg_type="flattening") + assert np.allclose(reg.matrix, np.array([[-1.5, 2, -0.5], [-0.5, 0, 0.5], [0.5, -2, 1.5]])) + +def test_flattening_invalid(): + # 1 + with pytest.raises(NotImplementedError, match=r".*only 2D derivative.*"): + QuadraticReg(factor=1, model_size=(1,2,3), reg_type="roughening") + # 2 + with pytest.raises(ValueError, match=r".*at least \(>=3, >=3\).*'roughening'.*"): + QuadraticReg(factor=1, model_size=(2,3), reg_type="roughening") + # 3 + with pytest.raises(ValueError, match=r".*at least >=3 for.*flattening.*"): + QuadraticReg(factor=1, model_size=2, reg_type="flattening") + +def test_smoothing(): + reg = QuadraticReg(factor=1, model_size=(4,4), reg_type="smoothing") + # 1 + assert pytest.approx(reg(np.zeros((4,4)))) == 0 + # 2 + assert pytest.approx(reg(np.ones((4,4)))) == 0 + # 3 + reg_val1 = reg(np.array([[1,2,3,4],[1,2,3,4],[1,2,3,4],[2,3,4,5]])) + reg_val2 = reg(np.array([1,2,3,4,1,2,3,4,1,2,3,4,2,3,4,5])) + assert pytest.approx(reg_val1) == reg_val2 + # 4 + test_model = np.array([[1,1,1,1],[1,1,1,1],[1,1,1,1],[1,1,1,2]]) + reg_val = reg(test_model) + assert pytest.approx(reg_val) == 12 + reg_grad = reg.gradient(test_model) + assert reg_grad.shape == (16,) + reg_hess = reg.hessian(test_model) + assert reg_hess.shape == (16,16) + # 5 - test matrix + reg = QuadraticReg(1, 4, reg_type="smoothing") + assert np.allclose(reg.matrix, np.array([[2.0,-5,4,-1],[1,-2,1,0],[0,1,-2,1],[-1,4,-5,2]])) + +def test_smoothing_invalid(): + # 1 + with pytest.raises(NotImplementedError, match=r".*only 2D derivative.*"): + QuadraticReg(1, (1,2,3), reg_type="smoothing") + # 2 + with pytest.raises(ValueError, match=r".*at least \(>=4, >=4\).*'smoothing'.*"): + QuadraticReg(factor=1, model_size=(2,3), reg_type="smoothing") + # 3 + reg = QuadraticReg(factor=1, model_size=(4,4), reg_type="smoothing") + with pytest.raises(ValueError): reg(np.zeros((3,3))) + # 4 + with pytest.raises(ValueError, match=r".*at least >=4 for.*smoothing.*"): + QuadraticReg(factor=1, model_size=3, reg_type="smoothing") + +def test_byo(): + # 1 + reg = QuadraticReg(1, 3, None) + reg_mat = reg.matrix + assert reg_mat.shape[0] == 3 + assert reg_mat.shape[0] == reg_mat.shape[1] + assert (reg_mat == np.eye(reg_mat.shape[0])).all() + # 2 + reg = QuadraticReg(1, 2, None, byo_matrix=np.array([[1,2],[3,4]])) + +def test_byo_invalid(): + with pytest.raises(ValueError, match=r".*must be 2-dimensional*"): + QuadraticReg(1, 3, None, byo_matrix=np.array([1,2,3])) + with pytest.raises(ValueError, match=r".*must be in shape (_, M)*"): + QuadraticReg(1, 3, None, byo_matrix=np.array([[1,2],[3,4]])) + with pytest.raises(ValueError, match=r".*provide a number for 'model_size'.*"): + QuadraticReg(1, (1,2), None, byo_matrix=np.array([[1,2],[3,4]]))