Add missing rows support

Author: znicholls

src/pandas_openscm/unit_conversion.py

@@ -185,8 +185,9 @@ def convert_unit(
         then the unit conversion will not happen.
 
         If this is a [pd.Series][pandas.Series],
-        then it will be passed directly to [convert_unit_from_target_series][]
-        and the requirements for that function apply.
+        then it will be passed to [convert_unit_from_target_series][]
+        after filling any rows in `df` that are not in `desired_unit`
+        with the unit from `df` (i.e. unspecified rows are not converted).
 
         For further details, see examples
 
@@ -205,7 +206,6 @@ def convert_unit(
     :
         `df` with converted units
 
-
     Examples
     --------
     >>> import pandas as pd
@@ -253,22 +253,22 @@ def convert_unit(
     >>> convert_unit(
     ...     start,
     ...     pd.Series(
-    ...         ["K", "mK", "degF"],
+    ...         ["K", "degF"],
     ...         index=pd.MultiIndex.from_tuples(
     ...             (
     ...                 ("sa", "temperature"),
-    ...                 ("sb", "temperature"),
+    ...                 # ("sb", "temperature") not included therefore not converted
     ...                 ("sb", "body temperature"),
     ...             ),
     ...             names=["scenario", "variable"],
     ...         ),
     ...     ),
     ... )
-                                        2020      2030      2050
+                                      2020     2030     2050
     scenario variable         unit
-    sa       temperature      K        0.001     0.002     0.003
-    sb       temperature      mK    1100.000  1200.000  1300.000
-             body temperature degF    98.600   100.580   100.220
+    sa       temperature      K      0.001    0.002    0.003
+    sb       temperature      K      1.100    1.200    1.300
+             body temperature degF  98.600  100.580  100.220
     """
     df_units_s = df.index.get_level_values(unit_level).to_series(
         index=df.index.droplevel(unit_level), name="df_unit"
@@ -288,7 +288,13 @@ def convert_unit(
         desired_unit_s = df_units_s.replace(desired_unit)
 
     elif isinstance(desired_unit, pd.Series):
-        desired_unit_s = desired_unit
+        missing = df_units_s.index.difference(desired_unit.index)
+        if missing.empty:
+            desired_unit_s = desired_unit
+        else:
+            desired_unit_s = pd.concat(
+                [desired_unit, multi_index_lookup(df_units_s, missing)]
+            )
 
     else:
         raise NotImplementedError(type(desired_unit))

tests/integration/test_unit_conversion.py

@@ -135,9 +135,86 @@ def test_convert_series():
         timepoints=np.array([1850.0, 2000.0, 2050.0, 2100.0]),
     )
 
+    target_units = (
+        start.loc[start.index.get_level_values("variable") != "temperature"]
+        .reset_index("unit")["unit"]
+        .replace({"W / m^2": "ZJ / yr / m^2", "ZJ": "PJ"})
+    )
+
+    res = convert_unit(start, target_units)
+
+    np.testing.assert_allclose(
+        res.loc[res.index.get_level_values("variable") == "temperature", :].values,
+        start.loc[start.index.get_level_values("variable") == "temperature", :].values,
+    )
+
+    np.testing.assert_allclose(
+        res.loc[res.index.get_level_values("variable") == "erf", :].values,
+        (60.0 * 60.0 * 24.0 * 365.25)
+        * 1e-21
+        * start.loc[start.index.get_level_values("variable") == "erf", :].values,
+    )
+
+    np.testing.assert_allclose(
+        res.loc[res.index.get_level_values("variable") == "ohc", :].values,
+        1e6 * start.loc[start.index.get_level_values("variable") == "ohc", :].values,
+    )
+
+
+def test_convert_series_all_rows():
+    # Check that conversion works if user supplies a Series of target units
+    start = create_test_df(
+        variables=[
+            ("temperature", "K"),
+            ("erf", "W / m^2"),
+            ("ohc", "ZJ"),
+        ],
+        n_scenarios=2,
+        n_runs=2,
+        timepoints=np.array([1850.0, 2000.0, 2050.0, 2100.0]),
+    )
+
+    target_units = start.reset_index("unit")["unit"].replace(
+        {"W / m^2": "ZJ / yr / m^2", "ZJ": "PJ"}
+    )
+
+    res = convert_unit(start, target_units)
+
+    np.testing.assert_allclose(
+        res.loc[res.index.get_level_values("variable") == "temperature", :].values,
+        start.loc[start.index.get_level_values("variable") == "temperature", :].values,
+    )
+
+    np.testing.assert_allclose(
+        res.loc[res.index.get_level_values("variable") == "erf", :].values,
+        (60.0 * 60.0 * 24.0 * 365.25)
+        * 1e-21
+        * start.loc[start.index.get_level_values("variable") == "erf", :].values,
+    )
+
+    np.testing.assert_allclose(
+        res.loc[res.index.get_level_values("variable") == "ohc", :].values,
+        1e6 * start.loc[start.index.get_level_values("variable") == "ohc", :].values,
+    )
+
+
+def test_convert_series_extra_rows():
+    # Check that conversion works if user supplies a Series of target units
+    start = create_test_df(
+        variables=[
+            ("temperature", "K"),
+            ("erf", "W / m^2"),
+            ("ohc", "ZJ"),
+        ],
+        n_scenarios=2,
+        n_runs=2,
+        timepoints=np.array([1850.0, 2000.0, 2050.0, 2100.0]),
+    )
+
     target_units = start.reset_index("unit")["unit"].replace(
         {"W / m^2": "ZJ / yr / m^2", "ZJ": "PJ"}
     )
+    target_units.loc[("scenario_2", "temperature", 0)] = "kK"
 
     res = convert_unit(start, target_units)