fix pre-commit and tests from #262 (#294)

* fix pre-commit errors

* fix tests

Author: chengzhuzhang

e3sm_to_cmip/cmor_handlers/_formulas.py

@@ -134,10 +134,18 @@ def mmrbc(ds: xr.Dataset) -> xr.DataArray:
     mmrbc = Mass_bc or bc_a1 + bc_a3 + bc_a4 + bc_c1 + bc_c3 + bc_c4
     """
 
-    if all(key in ds.data_vars for key in ["bc_a1", "bc_a3", "bc_a4",
-                                           "bc_c1", "bc_c3", "bc_c4"]):
-        result = ( ds["bc_a1"] + ds["bc_a3"] + ds["bc_a4"] +
-                   ds["bc_c1"] + ds["bc_c3"] + ds["bc_c4"] )
+    if all(
+        key in ds.data_vars
+        for key in ["bc_a1", "bc_a3", "bc_a4", "bc_c1", "bc_c3", "bc_c4"]
+    ):
+        result = (
+            ds["bc_a1"]
+            + ds["bc_a3"]
+            + ds["bc_a4"]
+            + ds["bc_c1"]
+            + ds["bc_c3"]
+            + ds["bc_c4"]
+        )
     elif "Mass_bc" in ds:
         result = ds["Mass_bc"]
     else:
@@ -155,10 +163,8 @@ def mmrdust(ds: xr.Dataset) -> xr.DataArray:
     mmrdust = Mass_dst or dst_a1 + dst_a3 + dst_c1 + dst_c3
     """
 
-    if all(key in ds.data_vars for key in ["dst_a1", "dst_a3",
-                                           "dst_c1", "dst_c3"]):
-        result = ( ds["dst_a1"] + ds["dst_a3"] + 
-                   ds["dst_c1"] + ds["dst_c3"] )
+    if all(key in ds.data_vars for key in ["dst_a1", "dst_a3", "dst_c1", "dst_c3"]):
+        result = ds["dst_a1"] + ds["dst_a3"] + ds["dst_c1"] + ds["dst_c3"]
     elif "Mass_dst" in ds:
         result = ds["Mass_dst"]
     else:
@@ -177,21 +183,44 @@ def mmroa(ds: xr.Dataset) -> xr.DataArray:
                                    soa_a1 + soa_a2 + soa_a3 + soa_c1 + soa_c2 + soa_c3
     """
 
-    if all(key in ds.data_vars for key in ["pom_a1", "pom_a3", "pom_a4",
-                                           "pom_c1", "pom_c3", "pom_c4",
-                                           "soa_a1", "soa_a2", "soa_a3",
-                                           "soa_c1", "soa_c2", "soa_c3"]):
-        result = ( ds["pom_a1"] + ds["pom_a3"] + ds["pom_a4"] +
-                   ds["pom_c1"] + ds["pom_c3"] + ds["pom_c4"] +
-                   ds["soa_a1"] + ds["soa_a2"] + ds["soa_a3"] +
-                   ds["soa_c1"] + ds["soa_c2"] + ds["soa_c3"] )
+    if all(
+        key in ds.data_vars
+        for key in [
+            "pom_a1",
+            "pom_a3",
+            "pom_a4",
+            "pom_c1",
+            "pom_c3",
+            "pom_c4",
+            "soa_a1",
+            "soa_a2",
+            "soa_a3",
+            "soa_c1",
+            "soa_c2",
+            "soa_c3",
+        ]
+    ):
+        result = (
+            ds["pom_a1"]
+            + ds["pom_a3"]
+            + ds["pom_a4"]
+            + ds["pom_c1"]
+            + ds["pom_c3"]
+            + ds["pom_c4"]
+            + ds["soa_a1"]
+            + ds["soa_a2"]
+            + ds["soa_a3"]
+            + ds["soa_c1"]
+            + ds["soa_c2"]
+            + ds["soa_c3"]
+        )
     elif all(key in ds.data_vars for key in ["Mass_pom", "Mass_soa"]):
         result = ds["Mass_pom"] + ds["Mass_soa"]
     else:
         raise KeyError(
             "No formula could be applied for 'mmroa'. Check the handler entry for 'mmroa' "
             "and input file(s) contain either 'pom_a1', 'pom_a3', 'pom_a4', "
-            "'pom_c1', 'pom_c3', 'pom_c4', 'soa_a1', 'soa_a2', 'soa_a3', " 
+            "'pom_c1', 'pom_c3', 'pom_c4', 'soa_a1', 'soa_a2', 'soa_a3', "
             "'soa_c1', 'soa_c2', 'soa_c3', or 'Mass_pom' and 'Mass_soa'."
         )
 
@@ -203,10 +232,18 @@ def mmrsoa(ds: xr.Dataset) -> xr.DataArray:
     mmrsoa = Mass_soa or soa_a1 + soa_a2 + soa_a3 + soa_c1 + soa_c2 + soa_c3
     """
 
-    if all(key in ds.data_vars for key in ["soa_a1", "soa_a2", "soa_a3",
-                                           "soa_c1", "soa_c2", "soa_c3"]):
-        result = ( ds["soa_a1"] + ds["soa_a2"] + ds["soa_a3"] +
-                   ds["soa_c1"] + ds["soa_c2"] + ds["soa_c3"] )
+    if all(
+        key in ds.data_vars
+        for key in ["soa_a1", "soa_a2", "soa_a3", "soa_c1", "soa_c2", "soa_c3"]
+    ):
+        result = (
+            ds["soa_a1"]
+            + ds["soa_a2"]
+            + ds["soa_a3"]
+            + ds["soa_c1"]
+            + ds["soa_c2"]
+            + ds["soa_c3"]
+        )
     elif "Mass_soa" in ds:
         result = ds["Mass_soa"]
     else:
@@ -224,10 +261,18 @@ def mmrss(ds: xr.Dataset) -> xr.DataArray:
     mmrss = Mass_ncl or ncl_a1 + ncl_a2 + ncl_a3 + ncl_c1 + ncl_c2 + ncl_c3
     """
 
-    if all(key in ds.data_vars for key in ["ncl_a1", "ncl_a2", "ncl_a3",
-                                           "ncl_c1", "ncl_c2", "ncl_c3"]):
-        result = ( ds["ncl_a1"] + ds["ncl_a2"] + ds["ncl_a3"] +
-                   ds["ncl_c1"] + ds["ncl_c2"] + ds["ncl_c3"] )
+    if all(
+        key in ds.data_vars
+        for key in ["ncl_a1", "ncl_a2", "ncl_a3", "ncl_c1", "ncl_c2", "ncl_c3"]
+    ):
+        result = (
+            ds["ncl_a1"]
+            + ds["ncl_a2"]
+            + ds["ncl_a3"]
+            + ds["ncl_c1"]
+            + ds["ncl_c2"]
+            + ds["ncl_c3"]
+        )
     elif "Mass_ncl" in ds:
         result = ds["Mass_ncl"]
     else:
@@ -245,15 +290,50 @@ def mmrso4(ds: xr.Dataset) -> xr.DataArray:
     mmrso4 = Mass_so4 or so4_a1+so4_c1+so4_a2+so4_c2+so4_a3+so4_c3 for MAM5
     mmrso4 = Mass_so4 or so4_a1+so4_c1+so4_a2+so4_c2+so4_a3+so4_c3 for MAM4
     """
-    
-    if all(key in ds.data_vars for key in ["so4_a1", "so4_a2", "so4_a3", "so4_a5",
-                                           "so4_c1", "so4_c2", "so4_c3", "so4_c5"]):
-        result = ( ds["so4_a1"] + ds["so4_a2"] + ds["so4_a3"] + ds["so4_a5"] +
-                   ds["so4_c1"] + ds["so4_c2"] + ds["so4_c3"] + ds["so4_c5"] ) * 96.0636 / 115.10734
-    elif all(key in ds.data_vars for key in ["so4_a1", "so4_a2", "so4_a3",
-                                             "so4_c1", "so4_c2", "so4_c3"]):
-        result = ( ds["so4_a1"] + ds["so4_a2"] + ds["so4_a3"] + 
-                   ds["so4_c1"] + ds["so4_c2"] + ds["so4_c3"] ) * 96.0636 / 115.10734
+
+    if all(
+        key in ds.data_vars
+        for key in [
+            "so4_a1",
+            "so4_a2",
+            "so4_a3",
+            "so4_a5",
+            "so4_c1",
+            "so4_c2",
+            "so4_c3",
+            "so4_c5",
+        ]
+    ):
+        result = (
+            (
+                ds["so4_a1"]
+                + ds["so4_a2"]
+                + ds["so4_a3"]
+                + ds["so4_a5"]
+                + ds["so4_c1"]
+                + ds["so4_c2"]
+                + ds["so4_c3"]
+                + ds["so4_c5"]
+            )
+            * 96.0636
+            / 115.10734
+        )
+    elif all(
+        key in ds.data_vars
+        for key in ["so4_a1", "so4_a2", "so4_a3", "so4_c1", "so4_c2", "so4_c3"]
+    ):
+        result = (
+            (
+                ds["so4_a1"]
+                + ds["so4_a2"]
+                + ds["so4_a3"]
+                + ds["so4_c1"]
+                + ds["so4_c2"]
+                + ds["so4_c3"]
+            )
+            * 96.0636
+            / 115.10734
+        )
     elif "Mass_so4" in ds:
         result = ds["Mass_so4"] * 96.0636 / 115.10734
     else:

e3sm_to_cmip/mpas.py

@@ -728,7 +728,7 @@ def _compute_moc_time_series(
     lat_bnds = np.zeros((len(lat) - 1, 2))
     lat_bnds[:, 0] = lat[0:-1]
     lat_bnds[:, 1] = lat[1:]
-    lat = 0.5 * (lat_bnds[:, 0] + lat_bnds[:, 1])
+    lat = 0.5 * (lat_bnds[:, 0] + lat_bnds[:, 1])  # type: ignore
 
     lat_bnds = xarray.DataArray(lat_bnds, dims=("lat", "nbnd"))  # type: ignore
     lat = xarray.DataArray(lat, dims=("lat",))  # type: ignore

tests/cmor_handlers/test__formulas.py

@@ -142,15 +142,17 @@ def test_mmrbc():
     ds = xr.Dataset(
         data_vars={
             "bc_a1": _dummy_dataarray(),
+            "bc_a3": _dummy_dataarray(),
             "bc_a4": _dummy_dataarray(),
             "bc_c1": _dummy_dataarray(),
+            "bc_c3": _dummy_dataarray(),
             "bc_c4": _dummy_dataarray(),
         }
     )
 
     result = mmrbc(ds)
     expected = xr.DataArray(
-        dims=["lat", "lon"], data=np.array([[0, 4, 8], [0, 4, 8], [0, 4, 8]])
+        dims=["lat", "lon"], data=np.array([[0, 6, 12], [0, 6, 12], [0, 6, 12]])
     )
     xr.testing.assert_allclose(result, expected)
 
@@ -173,7 +175,10 @@ def test_mmrso4():
 
     result = mmrso4(ds)
     expected = xr.DataArray(
-        dims=["lat", "lon"], data=np.array([[0, 6, 12], [0, 6, 12], [0, 6, 12]])
+        dims=["lat", "lon"],
+        data=np.array(
+            [[0, 5.00734, 10.01468], [0, 5.00734, 10.01468], [0, 5.00734, 10.01468]]
+        ),
     )
     xr.testing.assert_allclose(result, expected)