GEOPY-1889: Return nan on predicted where data is nan

simpeg_drivers/components/data.py

@@ -332,12 +332,9 @@ def get_normalizations(self) -> list[dict]:
             n_locs = len(self.locations)
             for comp in self.params.active_components:
                 normalizations[chan][comp] = np.ones(n_locs)
-                if comp in ["potential", "chargeability"]:
-                    normalizations[chan][comp] = 1
+
                 if comp in ["gz", "bz", "gxz", "gyz", "bxz", "byz"]:
                     normalizations[chan][comp] = np.full(n_locs, -1)
-                elif self.params.inversion_type in ["magnetotellurics"]:
-                    normalizations[chan][comp] = np.ones(n_locs)
                 elif self.params.inversion_type in ["tipper"]:
                     if "imag" in comp:
                         normalizations[chan][comp] = np.full(n_locs, -1)

simpeg_drivers/components/factories/directives_factory.py

@@ -570,18 +570,20 @@ def reshape(values):
             data[ordering[:, 0], ordering[:, 1], ordering[:, 2]] = values
             return data
 
+        transforms = []
+        for chan in range(len(channels)):
+            for comp in components:
+                transform = 1 / inversion_object.normalizations[chan][comp]
+
+                if inversion_object.observed[comp]:
+                    transform[np.isnan(inversion_object.observed[comp][chan])] = np.nan
+
+                transforms.append(transform)
+
         kwargs = {
             "data_type": inversion_object.observed_data_types,
             "association": "VERTEX",
-            "transforms": [
-                np.hstack(
-                    [
-                        1 / inversion_object.normalizations[chan][comp]
-                        for chan in range(len(channels))
-                        for comp in components
-                    ],
-                ),
-            ],
+            "transforms": [np.hstack(transforms)],
             "channels": channels,
             "components": components,
             "reshape": reshape,

tests/run_tests/driver_grav_test.py

@@ -195,9 +195,7 @@ def test_gravity_run(
             for pred in run_ws.get_entity("Iteration_0_gz")
             if pred.parent.parent.name == "Gravity Inversion"
         )
-        assert not any(np.isnan(predicted.values)), (
-            "Predicted data should not have nans."
-        )
+        assert np.isnan(predicted.values).sum() == 1
 
         l2_data_group = predicted.parent.get_entity("L2 iterations")[0]
         assert isinstance(l2_data_group, PropertyGroup)

tests/run_tests/driver_mt_test.py

@@ -42,7 +42,7 @@
 # To test the full run and validate the inversion.
 # Move this file out of the test directory and run.
 
-target_run = {"data_norm": 3.4932e-02, "phi_d": 1.4260e-01, "phi_m": 3.3270e00}
+target_run = {"data_norm": 3.4932e-02, "phi_d": 1.4110e-01, "phi_m": 3.3300e00}
 
 
 def setup_data(workspace, survey):
@@ -113,7 +113,7 @@ def test_magnetotellurics_fwr_run(
         ),
         model=ModelOptions(background=100.0),
     )
-    with get_workspace(tmp_path / "inversion_test.ui.geoh5") as geoh5:
+    with get_workspace(tmp_path / "inversion_test.geoh5") as geoh5:
         components = SyntheticsComponents(geoh5, options=opts)
 
         # Test for label index supporting ints - bypass setter
@@ -141,15 +141,15 @@ def test_magnetotellurics_fwr_run(
     fwr_driver = MTForwardDriver(params)
     fwr_driver.run()
 
-    with Workspace(tmp_path / "inversion_test.ui.geoh5") as geoh5:
+    with Workspace(tmp_path / "inversion_test.geoh5") as geoh5:
         assert geoh5.get_entity("Iteration_0_zyy_real_[0]")[0] is not None
 
 
 def test_bad_uncertainties(
     tmp_path: Path,
 ):
     workpath = (
-        tmp_path.parent / "test_magnetotellurics_fwr_run0" / "inversion_test.ui.geoh5"
+        tmp_path.parent / "test_magnetotellurics_fwr_run0" / "inversion_test.geoh5"
     )
 
     with Workspace(workpath) as geoh5:
@@ -200,12 +200,10 @@ def test_bad_uncertainties(
 
 def test_magnetotellurics_run(tmp_path: Path, max_iterations=1, pytest=True):
     # pass
-    workpath = tmp_path / "inversion_test.ui.geoh5"
+    workpath = tmp_path / "inversion_test.geoh5"
     if pytest:
         workpath = (
-            tmp_path.parent
-            / "test_magnetotellurics_fwr_run0"
-            / "inversion_test.ui.geoh5"
+            tmp_path.parent / "test_magnetotellurics_fwr_run0" / "inversion_test.geoh5"
         )
 
     with Workspace(workpath) as geoh5:
@@ -217,6 +215,12 @@ def test_magnetotellurics_run(tmp_path: Path, max_iterations=1, pytest=True):
 
         orig_zyy_real_1 = geoh5.get_entity("Iteration_0_zyy_real_[0]")[0].values
 
+        # Also NDV the data for one of them
+        data = geoh5.get_entity(data_kwargs["zyy_real_channel"].properties[0])[0]
+        vals = data.values
+        vals[0] = np.nan
+        data.values = vals
+
         # Run the inverse
         params = MTInversionOptions.build(
             geoh5=geoh5,
@@ -247,6 +251,11 @@ def test_magnetotellurics_run(tmp_path: Path, max_iterations=1, pytest=True):
         output = get_inversion_output(
             driver.params.geoh5.h5file, driver.params.out_group.uid
         )
+
+        # Check the predicted returns a nan
+        data = run_ws.get_entity("Iteration_1_zyy_real_[0]")[0]
+        assert np.isnan(data.values)[0]
+
         output["data"] = orig_zyy_real_1
         assert not run_ws.get_entity("Iteration_0_sensitivities")[0]
         if pytest: