Theta, rcb and angular momentum (#44)

doc/source/usage.rst

@@ -32,7 +32,32 @@ The GCMT class can also be used for quick plotting of data, since it wraps a few
 
 
 .. autoclass:: gcm_toolkit.GCMT
-    :members: __init__, get, get_models, models, add_horizontal_average, add_total_energy, add_meridional_overturning, read_raw, read_reduced, load, save
+    :members: __init__, get, get_models, models, read_raw, read_reduced, load, save
+
+
+Postprocessing
+--------------
+The GCMT class has quite a few methods that can calculate diagnostics such as total energy, location of the RCB and so on.
+You find their documentation here.
+
+Example for the total energy:
+
+.. code-block:: python
+
+        from gcm_toolkit import GCMT
+        import gcm_toolkit.gcm_plotting as gcmp
+
+        # Load data
+        tools = GCMT()
+        tools.read_raw(..., tag='model_name')
+        E = tools.add_total_energy(var_key_out='E', tag='model_name')
+
+        # You will also have the energy stored in the dataset if var_key_out is not None:
+        ds = tools['model_name']
+        E == ds.E
+
+.. autoclass:: gcm_toolkit.GCMT
+    :members: add_theta, add_horizontal_average, add_total_energy, add_meridional_overturning, add_rcb, add_total_momentum
 
 
 Plotting

gcm_toolkit/gcmtools.py

@@ -274,6 +274,81 @@ def add_horizontal_average(
             dsi, var_key, var_key_out=var_key_out, area_key=area_key
         )
 
+    def add_rcb(
+        self,
+        tol=0.01,
+        var_key_out=None,
+        area_key="area_c",
+        temp_key="T",
+        tag=None,
+    ):
+        """
+        Calculate the radiative convective boundary (rcb) by searching
+        (from the bottom upwards) for the first occurance of a deviation
+        from an adiabatic temperature profile.
+
+        Operates on and calculates the horizontal average.
+
+        Parameters
+        ----------
+        tol: float
+            tolerance for the relative deviation from adiabat
+        var_key_out: str, optional
+            variable name used to store the outcome.
+            If not provided, this script will just
+            return the averages and not change the dataset inplace.
+        area_key: str, optional
+            Variable key in the dataset for the area of grid cells
+        temp_key: str, optional
+            The key to look up the temperature (needed for density calculation)
+        tag : str, optional
+            The tag of the dataset that should be used.
+            If no tag is provided, and multiple datasets are available,
+            an error is raised.
+
+        Returns
+        -------
+        rcb : xarray.DataArray
+            A dataArray with reduced dimensionality,
+            containing the pressure of the rcb location.
+        """
+        dsi = self.get_one_model(tag)
+        return mani.m_add_rcb(
+            dsi,
+            tol=tol,
+            var_key_out=var_key_out,
+            area_key=area_key,
+            temp_key=temp_key,
+        )
+
+    def add_theta(self, var_key_out=None, temp_key="T", tag=None):
+        """
+        Convert temperature to potential temperature with respect to model boundary.
+
+        Parameters
+        ----------
+        var_key_out: str, optional
+            variable name used to store the outcome.
+            If not provided, this script will just
+            return theta and not change the dataset inplace.
+        temp_key: str, optional
+            The key to look up the temperature
+        tag : str, optional
+            The tag of the dataset that should be used.
+            If no tag is provided, and multiple datasets are available,
+            an error is raised.
+
+        Returns
+        -------
+        theta : xarray.DataArray
+            A dataArray with reduced dimensionality,
+            containing the potential temperature
+        """
+        dsi = self.get_one_model(tag)
+        return mani.m_add_theta(
+            dsi, var_key_out=var_key_out, temp_key=temp_key
+        )
+
     def add_total_energy(
         self, var_key_out=None, area_key="area_c", temp_key="T", tag=None
     ):
@@ -283,8 +358,6 @@ def add_total_energy(
 
         Parameters
         ----------
-        ds: xarray.Dataset
-            The dataset for which the calculation should be performed
         var_key_out: str, optional
             variable name used to store the outcome. If not provided,
             this script will just return the averages and not change
@@ -309,6 +382,38 @@ def add_total_energy(
             dsi, var_key_out=var_key_out, area_key=area_key, temp_key=temp_key
         )
 
+    def add_total_momentum(
+        self, var_key_out=None, area_key="area_c", temp_key="T", tag=None
+    ):
+        """
+        Calculate the total angular momentum of the GCM. See e.g.,
+        https://ui.adsabs.harvard.edu/abs/2014Icar..229..355P, Eq. 17
+
+        Parameters
+        ----------
+        var_key_out: str, optional
+            variable name used to store the outcome.
+            If not provided, this script will just
+            return the averages and not change the dataset inplace.
+        area_key: str, optional
+            Variable key in the dataset for the area of grid cells
+        temp_key: str, optional
+            The key to look up the temperature (needed for density calculation)
+        tag : str, optional
+            The tag of the dataset that should be used.
+            If no tag is provided,
+            and multiple datasets are available, an error is raised.
+
+        Returns
+        -------
+        momentum : xarray.DataArray
+            A dataArray with reduced dimensionality, containing the total momentum.
+        """
+        dsi = self.get_one_model(tag)
+        return mani.m_add_total_momentum(
+            dsi, var_key_out=var_key_out, area_key=area_key, temp_key=temp_key
+        )
+
     def add_meridional_overturning(
         self, v_data="V", var_key_out=None, tag=None
     ):

gcm_toolkit/tests/test_manipulations.py

@@ -46,6 +46,63 @@ def test_total_energy(all_nc_testdata):
     assert set(dsi.E_g.dims) == {"time"}
 
 
+def test_total_momentum(all_nc_testdata):
+    """Create a minimal gcm_toolkit object and do simple tests on the total momentum.
+    """
+    dirname, expected = all_nc_testdata
+
+    tools = GCMT(write="off")
+    tools.read_reduced(data_path=dirname)
+
+    dsi = tools.get_models()
+
+    area_key = expected.get("area_key", "area_c")
+    TM = tools.add_total_momentum(
+        var_key_out="m_g", area_key=area_key, temp_key="T"
+    )
+
+    assert hasattr(dsi, "m_g")
+    assert (TM == dsi.m_g).all()
+    assert set(dsi.m_g.dims) == {"time"}
+
+
+def test_rcb(all_nc_testdata):
+    """Create a minimal gcm_toolkit object and do simple tests on the rcb location.
+    """
+    dirname, expected = all_nc_testdata
+
+    tools = GCMT(write="off")
+    tools.read_reduced(data_path=dirname)
+
+    dsi = tools.get_models()
+
+    area_key = expected.get("area_key", "area_c")
+    rcb = tools.add_rcb(
+        tol=0.01, var_key_out="rcb", area_key=area_key, temp_key="T"
+    )
+
+    assert hasattr(dsi, "rcb")
+    assert (rcb == dsi.rcb).all()
+    assert set(dsi.rcb.dims) == {"time"}
+
+
+def test_theta(all_nc_testdata):
+    """Create a minimal gcm_toolkit object and do simple tests on the calculation of theta.
+    """
+    dirname, expected = all_nc_testdata
+
+    tools = GCMT(write="off")
+    tools.read_reduced(data_path=dirname)
+
+    dsi = tools.get_models()
+
+    theta = tools.add_theta(var_key_out="theta", temp_key="T")
+
+    assert hasattr(dsi, "theta")
+    assert (theta == dsi.theta).all()
+    assert set(dsi.theta.dims) == {"Z", "time", "lat", "lon"}
+
+
 def test_horizontal_overturning(all_nc_testdata):
     """Create a minimal gcm_toolkit object and do simple tests on it."""
 

gcm_toolkit/tests/test_utils.py

@@ -34,6 +34,11 @@ def test_data_testing_utils(all_nc_testdata):
     assert not is_the_data_cloudy(ds_new)
     assert not is_the_data_basic(ds_new)
 
+    ds_new = ds.copy()
+    ds_new["Z"] = ds_new.Z[::-1]
+    assert not is_the_data_cloudy(ds_new)
+    assert not is_the_data_basic(ds_new)
+
 
 def test_convert_pressure_failures():
     with pytest.raises(ValueError):