diff --git a/src/harmonica/_equivalent_sources/cartesian.py b/src/harmonica/_equivalent_sources/cartesian.py index 1e9b982d5..1340f7baf 100644 --- a/src/harmonica/_equivalent_sources/cartesian.py +++ b/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. diff --git a/src/harmonica/_equivalent_sources/spherical.py b/src/harmonica/_equivalent_sources/spherical.py index 52368daff..606311b1c 100644 --- a/src/harmonica/_equivalent_sources/spherical.py +++ b/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. diff --git a/test/test_eq_sources_cartesian.py b/test/test_eq_sources_cartesian.py index b035733a3..d4d736325 100644 --- a/test/test_eq_sources_cartesian.py +++ b/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 diff --git a/test/test_eq_sources_spherical.py b/test/test_eq_sources_spherical.py index 2868d0c88..67c4ceb14 100644 --- a/test/test_eq_sources_spherical.py +++ b/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