diff --git a/src/inversion_ideas/regularization/_mesh_based.py b/src/inversion_ideas/regularization/_mesh_based.py index bda8c78..5decaf9 100644 --- a/src/inversion_ideas/regularization/_mesh_based.py +++ b/src/inversion_ideas/regularization/_mesh_based.py @@ -31,6 +31,10 @@ class _MeshBasedRegularization(Objective): @property def n_params(self) -> int: + return self.n_active + + @property + def n_active(self) -> int: return int(np.sum(self.active_cells)) @property @@ -55,20 +59,20 @@ def cell_weights( "It must be an array or a dictionary." ) raise TypeError(msg) - if isinstance(value, np.ndarray) and value.size != self.n_params: + if isinstance(value, np.ndarray) and value.size != self.n_active: msg = ( f"Invalid cell_weights array with '{value.size}' elements. " - f"It must have '{self.n_params}' elements, " + f"It must have '{self.n_active}' elements, " "equal to the number of active cells." ) raise ValueError(msg) if isinstance(value, dict): for key, array in value.items(): - if array.size != self.n_params: + if array.size != self.n_active: msg = ( f"Invalid cell_weights array '{key}' with " f"'{array.size}' elements. " - f"It must have '{self.n_params}' elements, " + f"It must have '{self.n_active}' elements, " "equal to the number of active cells." ) raise ValueError(msg) @@ -88,13 +92,14 @@ class Smallness(_MeshBasedRegularization): active_cells : (n_cells) array or None, optional Array full of bools that indicate the active cells in the mesh. It must have the same amount of elements as cells in the mesh. - cell_weights : (n_params) array or dict of (n_params) arrays or None, optional + cell_weights : (n_active) array or dict of (n_active) arrays or None, optional Array with cell weights. For multiple cell weights, pass a dictionary where keys are strings and values are the different weights arrays. If None, no cell weights are going to be used. - reference_model : (n_params) array or None, optional - Array with values for the reference model. + reference_model : (n_active) array or None, optional + Array with values for the reference model. It must have the same number of + elements as active cells in the mesh. Notes ----- @@ -157,13 +162,13 @@ def __init__( self.cell_weights = ( cell_weights if cell_weights is not None - else np.ones(self.n_params, dtype=np.float64) + else np.ones(self.n_active, dtype=np.float64) ) self.reference_model = ( reference_model if reference_model is not None - else np.zeros(self.n_params, dtype=np.float64) + else np.zeros(self.n_active, dtype=np.float64) ) self.set_name("s") @@ -229,7 +234,7 @@ def hessian(self, model: Model): # noqa: ARG002 ) @property - def weights_matrix(self) -> dia_array: + def weights_matrix(self) -> dia_array[np.float64]: """ Diagonal matrix with the square root of regularization weights on cells. """ @@ -243,7 +248,7 @@ def weights_matrix(self) -> dia_array: return diags_array(np.sqrt(cell_weights)) @property - def _volumes_sqrt_matrix(self) -> dia_array: + def _volumes_sqrt_matrix(self) -> dia_array[np.float64]: """ Diagonal matrix with the square root of cell volumes. """ @@ -266,12 +271,12 @@ class Flatness(_MeshBasedRegularization): active_cells : (n_cells) array or None, optional Array full of bools that indicate the active cells in the mesh. It must have the same amount of elements as cells in the mesh. - cell_weights : (n_params) array or dict of (n_params) arrays or None, optional + cell_weights : (n_active) array or dict of (n_active) arrays or None, optional Array with cell weights. For multiple cell weights, pass a dictionary where keys are strings and values are the different weights arrays. If None, no cell weights are going to be used. - reference_model : (n_params) array or None, optional + reference_model : (n_active) array or None, optional Array with values for the reference model. Notes @@ -346,13 +351,13 @@ def __init__( self.cell_weights = ( cell_weights if cell_weights is not None - else np.ones(self.n_params, dtype=np.float64) + else np.ones(self.n_active, dtype=np.float64) ) self.reference_model = ( reference_model if reference_model is not None - else np.zeros(self.n_params, dtype=np.float64) + else np.zeros(self.n_active, dtype=np.float64) ) self.set_name(direction) @@ -427,7 +432,7 @@ def hessian(self, model: Model): # noqa: ARG002 ) @property - def weights_matrix(self) -> dia_array: + def weights_matrix(self) -> dia_array[np.float64]: """ Diagonal matrix with the square root of cell weights averaged on faces. """ @@ -441,7 +446,7 @@ def weights_matrix(self) -> dia_array: return diags_array(self._average_cells_to_faces @ np.sqrt(cell_weights)) @property - def _volumes_sqrt_matrix(self) -> dia_array: + def _volumes_sqrt_matrix(self) -> dia_array[np.float64]: """ Diagonal matrix with the square root of cell volumes averaged on faces. """ @@ -497,21 +502,22 @@ class SparseSmallness(_MeshBasedRegularization): active_cells : (n_cells) array or None, optional Array full of bools that indicate the active cells in the mesh. It must have the same amount of elements as cells in the mesh. - cell_weights : (n_params) array or dict of (n_params) arrays or None, optional + cell_weights : (n_active) array or dict of (n_active) arrays or None, optional Array with cell weights. For multiple cell weights, pass a dictionary where keys are strings and values are the different weights arrays. If None, no cell weights are going to be used. - reference_model : (n_params) array, optional - Reference model used in the regularization. + reference_model : (n_active) array or None, optional + Array with values for the reference model. It must have the same number of + elements as active cells in the mesh. threshold : float, optional IRLS threshold. Symbolized with :math:`\epsilon` in Fournier and Oldenburg (2019). cooling_factor : float, optional Factor used to cool down the ``threshold`` when updating the IRLS. - model_previous : (n_params) array + model_previous : (n_params) array or None, optional Array with previous model in the iterations. This model is used to build the - ``R`` matrix. + ``R`` matrix. If None, an array full of zeros will be assigned. irls : bool, optional Flag to activate or deactivate IRLS. If False, the class would work as an L2 smallness term. If True, the R matrix will be built using the @@ -547,13 +553,13 @@ def __init__( self.cell_weights = ( cell_weights if cell_weights is not None - else np.ones(self.n_params, dtype=np.float64) + else np.ones(self.n_active, dtype=np.float64) ) self.reference_model = ( reference_model if reference_model is not None - else np.zeros(self.n_params, dtype=np.float64) + else np.zeros(self.n_active, dtype=np.float64) ) self.threshold = threshold self.cooling_factor = cooling_factor @@ -602,7 +608,7 @@ def R(self) -> dia_array: R matrix to approximate lp norm using Lawson's algorithm. """ if not self.irls: - return eye_array(self.n_params) + return eye_array(self.n_active) power = self.norm / 4 - 0.5 diagonal = (self.model_previous**2 + self.threshold**2) ** power return diags_array(diagonal) diff --git a/tests/test_regularizations.py b/tests/test_regularizations.py index 0bc2642..5fbce88 100644 --- a/tests/test_regularizations.py +++ b/tests/test_regularizations.py @@ -2,11 +2,12 @@ Test regularization classes. """ +import numpy as np import pytest from discretize.tensor_mesh import TensorMesh -from scipy.sparse import sparray +from scipy.sparse import dia_array, sparray -from inversion_ideas import Flatness +from inversion_ideas import Flatness, Smallness class TestBugfixFlatness: @@ -24,3 +25,65 @@ def mesh(self): def test_cell_gradient_type(self, mesh, direction): flatness = Flatness(mesh, direction=direction) assert isinstance(flatness._cell_gradient, sparray) + + +class TestSmallness: + """ + Test the :class:`inversion_ideas.Smallness` regularization class. + """ + + @pytest.fixture + def mesh(self): + hx = [(1.0, 10)] + h = [hx, hx, hx] + return TensorMesh(h, origin="CCN") + + @pytest.fixture + def active_cells(self, mesh: TensorMesh): + active_cells = np.ones(mesh.n_cells, dtype=bool) + _, _, z = mesh.cell_centers.T + active_cells[z > -1.0] = False + assert not active_cells.all() + return active_cells + + def test_smallness(self, mesh, active_cells): + n_active = active_cells.sum() + cell_weights = np.full(n_active, fill_value=0.1) + reference_model = np.full(n_active, 1e-8) + smallness = Smallness( + mesh, + active_cells=active_cells, + cell_weights=cell_weights, + reference_model=reference_model, + ) + + model = np.random.default_rng(seed=12312).uniform(size=n_active) + + # Test call + result = smallness(model) + assert np.isscalar(result) + expected = np.sum( + mesh.cell_volumes[active_cells] + * cell_weights + * (model - reference_model) ** 2 + ) + np.testing.assert_allclose(result, expected) + + # Test gradient + gradient = smallness.gradient(model) + expected = ( + 2 + * mesh.cell_volumes[active_cells] + * cell_weights + * (model - reference_model) + ) + assert gradient.size == model.size + np.testing.assert_allclose(gradient, expected) + + # Test hessian + hessian = smallness.hessian(model) + assert hessian.shape == (model.size, model.size) + assert isinstance(hessian, dia_array) + assert hessian.offsets == 0 # should be a diagonal matrix (only main diag) + expected_diagonal = 2 * mesh.cell_volumes[active_cells] * cell_weights + np.testing.assert_allclose(hessian.diagonal(), expected_diagonal)