[PR]: Improving regrid2 performance

tests/test_regrid.py

@@ -384,23 +384,6 @@ def test_vertical_placeholder(self):
         with pytest.raises(NotImplementedError, match=""):
             regridder.vertical("so", ds)
 
-    def test_missing_dimension(self):
-        ds = fixtures.generate_dataset(
-            decode_times=True, cf_compliant=False, has_bounds=True
-        )
-
-        del ds.lat.attrs["axis"]
-
-        output_grid = grid.create_gaussian_grid(32)
-
-        regridder = regrid2.Regrid2Regridder(ds, output_grid)
-
-        with pytest.raises(
-            RuntimeError,
-            match="Could not find axis 'lat', ensure 'lat' exists and the attributes are correct.",
-        ):
-            regridder.horizontal("ts", ds)
-
     @pytest.mark.filterwarnings("ignore:.*invalid value.*divide.*:RuntimeWarning")
     def test_output_bounds(self):
         ds = fixtures.generate_dataset(
@@ -499,45 +482,15 @@ def test_regrid_input_mask(self):
 
         expected_output = np.array(
             [
-                [0.0, 0.0, 0.0, 0.0],
-                [0.70710677, 0.70710677, 0.70710677, 0.70710677],
-                [0.70710677, 0.70710677, 0.70710677, 0.70710677],
-                [0.0, 0.0, 0.0, 0.0],
-            ],
-            dtype=np.float32,
+                [0.0] * 4,
+                [0.7071067811865476] * 4,
+                [0.7071067811865476] * 4,
+                [0.0] * 4,
+            ]
         )
 
         assert np.all(output_data.ts.values == expected_output)
 
-    def test_regrid_output_mask(self):
-        output_mask = [
-            [0, 0, 0, 0],
-            [1, 1, 1, 1],
-            [1, 1, 1, 1],
-            [0, 0, 0, 0],
-        ]
-
-        self.fine_2d_ds["mask"] = (("lat", "lon"), output_mask)
-
-        regridder = regrid2.Regrid2Regridder(self.coarse_2d_ds, self.fine_2d_ds)
-
-        output_data = regridder.horizontal("ts", self.coarse_2d_ds)
-
-        expected_output = np.array(
-            [
-                [1.0, 1.0, 1.0, 1.0],
-                [1e20, 1e20, 1e20, 1e20],
-                [1e20, 1e20, 1e20, 1e20],
-                [1.0, 1.0, 1.0, 1.0],
-            ],
-            dtype=np.float32,
-        )
-
-        # need to replace nans since nan != nan
-        output_data["ts"] = output_data.ts.fillna(1e20)
-
-        assert np.all(output_data.ts.values == expected_output)
-
     def test_preserve_attrs(self):
         regridder = regrid2.Regrid2Regridder(self.coarse_2d_ds, self.fine_2d_ds)
 
@@ -547,7 +500,7 @@ def test_preserve_attrs(self):
         assert output_data["ts"].attrs == self.da_attrs
 
         for x in output_data.coords:
-            assert output_data[x].attrs == self.coarse_2d_ds[x].attrs
+            assert output_data[x].attrs == self.coarse_2d_ds[x].attrs, f"{x}"
 
     def test_regrid_2d(self):
         regridder = regrid2.Regrid2Regridder(self.coarse_2d_ds, self.fine_2d_ds)
@@ -582,7 +535,7 @@ def test_regrid_4d(self):
 
     def test_map_longitude_coarse_to_fine(self):
         mapping, weights = regrid2._map_longitude(
-            self.coarse_lon_bnds, self.fine_lon_bnds
+            self.coarse_lon_bnds.values, self.fine_lon_bnds.values
         )
 
         expected_mapping = [
@@ -593,33 +546,33 @@ def test_map_longitude_coarse_to_fine(self):
         ]
 
         expected_weigths = [
-            [[90]],
-            [[90]],
-            [[90]],
-            [[90]],
+            [90],
+            [90],
+            [90],
+            [90],
         ]
 
         np.testing.assert_allclose(mapping, expected_mapping)
         np.testing.assert_allclose(weights, expected_weigths)
 
     def test_map_longitude_fine_to_coarse(self):
         mapping, weights = regrid2._map_longitude(
-            self.fine_lon_bnds, self.coarse_lon_bnds
+            self.fine_lon_bnds.values, self.coarse_lon_bnds.values
         )
 
         expected_mapping = [
             [0, 1],
             [2, 3],
         ]
 
-        expected_weigths = [[[90, 90]], [[90, 90]]]
+        expected_weigths = [[90, 90], [90, 90]]
 
         np.testing.assert_allclose(mapping, expected_mapping)
         np.testing.assert_allclose(weights, expected_weigths)
 
     def test_map_latitude_coarse_to_fine(self):
         mapping, weights = regrid2._map_latitude(
-            self.coarse_lat_bnds, self.fine_lat_bnds
+            self.coarse_lat_bnds.values, self.fine_lat_bnds.values
         )
 
         expected_mapping = [
@@ -648,7 +601,7 @@ def test_map_latitude_coarse_to_fine(self):
 
     def test_map_latitude_fine_to_coarse(self):
         mapping, weights = regrid2._map_latitude(
-            self.fine_lat_bnds, self.coarse_lat_bnds
+            self.fine_lat_bnds.values, self.coarse_lat_bnds.values
         )
 
         expected_mapping = [
@@ -658,9 +611,9 @@ def test_map_latitude_fine_to_coarse(self):
         ]
 
         expected_weigths = [
-            [[0.29289322], [0.20710678]],
-            [[0.5], [0.5]],
-            [[0.20710678], [0.29289322]],
+            [0.29289322, 0.20710678],
+            [0.5, 0.5],
+            [0.20710678, 0.29289322],
         ]
 
         np.testing.assert_allclose(mapping, expected_mapping)