Add required memory functions for non-gradient boosted equivalent sources

src/harmonica/_equivalent_sources/cartesian.py

@@ -190,6 +190,49 @@ def __init__(
         # Define Green's function for Cartesian coordinates
         self.greens_function = greens_func_cartesian
 
+    def estimate_required_memory(self, coordinates):
+        """
+        Estimate the memory required for storing the Jacobian matrix.
+
+        Parameters
+        ----------
+        coordinates : tuple of arrays
+            Arrays with the coordinates of each data point. Should be in the
+            following order: (``easting``, ``northing``, ``upward``, ...).
+            Only ``easting``, ``northing``, and ``upward`` will be used, all
+            subsequent coordinates will be ignored.
+
+        Returns
+        -------
+        memory_required : int
+            Amount of memory required to store the Jacobian matrix in
+            bytes.
+
+        Examples
+        --------
+        >>> import bordado as bd
+        >>> coordinates = bd.random_coordinates(
+        ...     region=(-1e3, 3e3, 2e3, 5e3),
+        ...     size=100,
+        ...     non_dimensional_coords=100,
+        ...     random_seed=42,
+        ... )
+        >>> eqs = EquivalentSources()
+        >>> n_bytes = eqs.estimate_required_memory(coordinates)
+        >>> int(n_bytes)
+        80000
+        """
+        # Build the sources and assign the points_ attribute
+        coordinates = vdb.n_1d_arrays(coordinates, 3)
+        points = self._build_points(coordinates)
+        # Get the number of sources and data
+        n_data = coordinates[0].size
+        n_points = points[0].size
+        # Compute the size of the Jacobian matrix
+        jacobian_size = n_data * n_points
+        # Estimate size of the Jacobian matrix in bytes
+        return jacobian_size * np.dtype(self.dtype).itemsize
+
     def fit(self, coordinates, data, weights=None):
         """
         Fit the coefficients of the equivalent sources.

src/harmonica/_equivalent_sources/spherical.py

@@ -106,7 +106,7 @@ class EquivalentSourcesSph(vdb.BaseGridder):
     # as xr.Dataset.
     dims = ("spherical_latitude", "longitude")
 
-    # Overwrite the defalt name for the upward coordinate.
+    # Overwrite the default name for the upward coordinate.
     extra_coords_name = "radius"
 
     def __init__(
@@ -123,6 +123,56 @@ def __init__(
         # Define Green's function for spherical coordinates
         self.greens_function = greens_func_spherical
 
+    def estimate_required_memory(self, coordinates):
+        """
+        Estimate the memory required for storing the Jacobian matrix.
+
+        Parameters
+        ----------
+        coordinates : tuple of arrays
+            Arrays with the coordinates of each data point. Should be in the
+            following order: (``longitude``, ``latitude``, ``radius``, ...).
+            Only ``longitude``, ``latitude``, and ``radius`` will be used, all
+            subsequent coordinates will be ignored.
+
+        Returns
+        -------
+        memory_required : int
+            Amount of memory required to store the Jacobian matrix in
+            bytes.
+
+        Examples
+        --------
+        >>> import bordado as bd
+        >>> coordinates = bd.random_coordinates(
+        ...     region=(-70, -60, -40, -30),
+        ...     size=100,
+        ...     non_dimensional_coords=100,
+        ...     random_seed=42,
+        ... )
+        >>> eqs = EquivalentSourcesSph()
+        >>> n_bytes = eqs.estimate_required_memory(coordinates)
+        >>> int(n_bytes)
+        80000
+        """
+        # Build the sources and assign the points_ attribute
+        coordinates = vdb.n_1d_arrays(coordinates, 3)
+        if self.points is None:
+            points = (
+                coordinates[0],
+                coordinates[1],
+                coordinates[2] - self.relative_depth,
+            )
+        else:
+            points = vdb.n_1d_arrays(self.points, 3)
+        # Get the number of sources and data
+        n_data = coordinates[0].size
+        n_points = points[0].size
+        # Compute the size of the Jacobian matrix
+        jacobian_size = n_data * n_points
+        # Estimate size of a single element of the Jacobian matrix in bytes
+        return jacobian_size * np.dtype(coordinates[0].dtype).itemsize
+
     def fit(self, coordinates, data, weights=None):
         """
         Fit the coefficients of the equivalent sources.

test/test_eq_sources_cartesian.py

@@ -477,3 +477,22 @@ def test_zero_depth():
     msg = "Depth value cannot be zero. It should be a non-zero numeric value."
     with pytest.raises(ValueError, match=msg):
         EquivalentSources(depth=zero_depth)
+
+
+@pytest.mark.parametrize("spacing", [100, 500, 1e3])
+@pytest.mark.parametrize(("dtype", "itemsize"), [("float32", 4), ("float64", 8)])
+def test_memory_estimation(spacing, dtype, itemsize):
+    """
+    Test the estimate_required_memory class method.
+    """
+    region = (-1e3, 5e3, 2e3, 8e3)
+    coordinates = bd.grid_coordinates(
+        region=region, spacing=spacing, non_dimensional_coords=0
+    )
+    # Compute expected required memory
+    n_data = coordinates[0].size
+    expected_required_memory = n_data * n_data * itemsize
+    # Estimate required memory
+    eqs = EquivalentSources(dtype=dtype)
+    required_memory = eqs.estimate_required_memory(coordinates)
+    assert required_memory == expected_required_memory

test/test_eq_sources_spherical.py

@@ -300,3 +300,43 @@ def test_error_ignored_args(coordinates_small, data_small, region):
     msg = "The 'bla' arguments are being ignored."
     with pytest.warns(FutureWarning, match=msg):
         eqs.grid(coordinates=grid_coords, bla="bla")
+
+
+@pytest.mark.parametrize("spacing", [100, 500, 1e3])
+def test_memory_estimation(spacing):
+    """
+    Test the estimate_required_memory class method.
+    """
+    region = (-20, 20, -20, 20)
+    coordinates = bd.grid_coordinates(
+        region=region, spacing=spacing, non_dimensional_coords=100.0
+    )
+    # Compute expected required memory
+    n_data = coordinates[0].size
+    expected_required_memory = n_data * n_data * 8
+    # Estimate required memory
+    eqs = EquivalentSourcesSph()
+    required_memory = eqs.estimate_required_memory(coordinates)
+    assert required_memory == expected_required_memory
+
+
+def test_memory_estimation_supplied_points():
+    """
+    Test the estimate_required_memory class method with user-supplied points.
+    """
+    region = (-20, 20, -20, 20)
+    coordinates = bd.grid_coordinates(
+        region=region, shape=(6, 6), non_dimensional_coords=100.0
+    )
+    # instead of 1 point beneath each datapoint, give custom points
+    points = bd.grid_coordinates(
+        region=region, shape=(3, 3), non_dimensional_coords=50.0
+    )
+    # Compute expected required memory
+    n_data = coordinates[0].size
+    n_coords = points[0].size
+    expected_required_memory = n_coords * n_data * 8
+    # Estimate required memory
+    eqs = EquivalentSourcesSph(points=points)
+    required_memory = eqs.estimate_required_memory(coordinates)
+    assert required_memory == expected_required_memory