GEOPY-1912: Add custom Error for problem too large

simpeg_drivers/driver.py

@@ -27,9 +27,8 @@
 
 from dask.distributed import get_client, Client, LocalCluster, performance_report
 from geoapps_utils.driver.driver import BaseDriver
-from geoapps_utils.driver.data import BaseData
+from geoapps_utils.utils.importing import GeoAppsError
 
-from geoh5py.data import Data
 from geoh5py.groups import SimPEGGroup
 from geoh5py.shared.utils import fetch_active_workspace
 from geoh5py.ui_json import InputFile
@@ -345,13 +344,23 @@ def run(self):
                 self.out_group.add_file(self.params.input_file.path_name)
 
         predicted = None
-        if self.params.forward_only:
-            print("Running the forward simulation ...")
-            predicted = simpeg_inversion.invProb.get_dpred(self.models.starting, None)
-        else:
-            # Run the inversion
-            self.start_inversion_message()
-            simpeg_inversion.run(self.models.starting)
+        try:
+            if self.params.forward_only:
+                print("Running the forward simulation ...")
+                predicted = simpeg_inversion.invProb.get_dpred(
+                    self.models.starting, None
+                )
+            else:
+                # Run the inversion
+                self.start_inversion_message()
+                simpeg_inversion.run(self.models.starting)
+
+        except np.core._exceptions._ArrayMemoryError as error:  # pylint: disable=protected-access
+            raise GeoAppsError(
+                "Memory Error: Sensitivities too large for system. \n"
+                "Try reducing the number of data, reducing the number of cells in the mesh\n"
+                "or increase the number of tiles."
+            ) from error
 
         self.logger.end()
         sys.stdout = self.logger.terminal

tests/run_tests/driver_grav_test.py

@@ -11,9 +11,12 @@
 from __future__ import annotations
 
 from pathlib import Path
+from unittest.mock import patch
 
 import numpy as np
+from geoapps_utils.utils.importing import GeoAppsError
 from geoh5py.workspace import Workspace
+from pytest import raises
 
 from simpeg_drivers.params import ActiveCellsOptions
 from simpeg_drivers.potential_fields import (
@@ -64,6 +67,37 @@ def test_gravity_fwr_run(
     fwr_driver.run()
 
 
+def test_array_too_large_run(
+    tmp_path: Path,
+):
+    workpath = tmp_path.parent / "test_gravity_fwr_run0" / "inversion_test.ui.geoh5"
+
+    with Workspace(workpath) as geoh5:
+        gz = geoh5.get_entity("Iteration_0_gz")[0]
+        mesh = geoh5.get_entity("mesh")[0]
+        topography = geoh5.get_entity("topography")[0]
+
+        # Run the inverse
+        active_cells = ActiveCellsOptions(topography_object=topography)
+        params = GravityInversionOptions(
+            geoh5=geoh5,
+            mesh=mesh,
+            active_cells=active_cells,
+            data_object=gz.parent,
+            gz_channel=gz,
+            gz_uncertainty=1e-4,
+            starting_model=1e-4,
+        )
+
+    with patch(
+        "simpeg.inversion.BaseInversion.run",
+        side_effect=np.core._exceptions._ArrayMemoryError((0,), np.dtype("float64")),  # pylint: disable=protected-access
+    ):
+        with raises(GeoAppsError, match="Memory Error"):
+            driver = GravityInversionDriver(params)
+            driver.run()
+
+
 def test_gravity_run(
     tmp_path: Path,
     max_iterations=1,
@@ -111,7 +145,6 @@ def test_gravity_run(
             y_norm=gradient_norms,
             z_norm=gradient_norms,
             gradient_type="components",
-            gz_channel_bool=True,
             gz_channel=gz,
             gz_uncertainty=2e-3,
             lower_bound=0.0,