From 8e24d0f6890c91493937c76859d3682014b7c3c5 Mon Sep 17 00:00:00 2001
From: Aaron David Schneider <aaron.schneider@nbi.ku.dk>
Date: Fri, 9 Sep 2022 15:51:26 +0200
Subject: [PATCH] add partial energies to calculation

---
 gcm_toolkit/gcmtools.py            | 21 ++++++++++--
 gcm_toolkit/utils/manipulations.py | 52 +++++++++++++++++++++++-------
 setup.py                           |  2 +-
 3 files changed, 61 insertions(+), 14 deletions(-)

diff --git a/gcm_toolkit/gcmtools.py b/gcm_toolkit/gcmtools.py
index afc81b7..c25fdd2 100644
--- a/gcm_toolkit/gcmtools.py
+++ b/gcm_toolkit/gcmtools.py
@@ -350,7 +350,12 @@ def add_theta(self, var_key_out=None, temp_key="T", tag=None):
         )
 
     def add_total_energy(
-        self, var_key_out=None, area_key="area_c", temp_key="T", tag=None
+        self,
+        var_key_out=None,
+        area_key="area_c",
+        temp_key="T",
+        return_all=False,
+        tag=None,
     ):
         """
         Calculate the total Energy of the GCM. See e.g.,
@@ -366,6 +371,8 @@ def add_total_energy(
             Variable key in the dataset for the area of grid cells
         temp_key: str, optional
             The key to look up the temperature
+        return_all: bool, optional
+            Also return the partial energies
         tag : str, optional
             The tag of the dataset that should be used.
             If no tag is provided,
@@ -375,11 +382,21 @@ def add_total_energy(
         -------
         energy : xarray.DataArray
             A dataArray with reduced dimensionality, containing the total energy.
+        therm_energy : xarray.DataArray, optional
+            A dataArray with reduced dimensionality, containing the thermal energy.
+        pot_energy : xarray.DataArray, optional
+            A dataArray with reduced dimensionality, containing the potential energy.
+        kin_energy : xarray.DataArray, optional
+            A dataArray with reduced dimensionality, containing the kinetic energy.
         """
 
         dsi = self.get_one_model(tag)
         return mani.m_add_total_energy(
-            dsi, var_key_out=var_key_out, area_key=area_key, temp_key=temp_key
+            dsi,
+            var_key_out=var_key_out,
+            area_key=area_key,
+            temp_key=temp_key,
+            return_all=return_all,
         )
 
     def add_total_momentum(
diff --git a/gcm_toolkit/utils/manipulations.py b/gcm_toolkit/utils/manipulations.py
index d5934e5..7c8e499 100644
--- a/gcm_toolkit/utils/manipulations.py
+++ b/gcm_toolkit/utils/manipulations.py
@@ -51,7 +51,9 @@ def m_add_horizontal_average(
     return avg
 
 
-def m_add_total_energy(dsi, var_key_out=None, area_key="area_c", temp_key="T"):
+def m_add_total_energy(
+    dsi, var_key_out=None, area_key="area_c", temp_key="T", return_all=False
+):
     """
     Calculate the total Energy of the GCM. See e.g.,
     https://ui.adsabs.harvard.edu/abs/2014Icar..229..355P, Eq. 16
@@ -67,11 +69,20 @@ def m_add_total_energy(dsi, var_key_out=None, area_key="area_c", temp_key="T"):
         Variable key in the dataset for the area of grid cells
     temp_key: str, optional
         The key to look up the temperature
+    return_all: bool, optional
+        Also return the partial energies
 
     Returns
     -------
-    energy : xarray.DataArray
+    tot_energy : xarray.DataArray
         A dataArray with reduced dimensionality, containing the total energy.
+    therm_energy : xarray.DataArray, optional
+        A dataArray with reduced dimensionality, containing the thermal energy.
+    pot_energy : xarray.DataArray, optional
+        A dataArray with reduced dimensionality, containing the potential energy.
+    kin_energy : xarray.DataArray, optional
+        A dataArray with reduced dimensionality, containing the kinetic energy.
+
     """
     # print information
     wrt.write_status("STAT", "Calculate total energy")
@@ -85,19 +96,38 @@ def m_add_total_energy(dsi, var_key_out=None, area_key="area_c", temp_key="T"):
 
     dzdp, rho = _calc_hydrostat_eq(dsi_calc, temp_key)
     z_geo = dzdp.cumulative_integrate(coord=c["Z"])
-    d_energy = (
-        z_geo * dsi_calc.attrs[c["g"]]
-        + dsi_calc.attrs[c["cp"]] * dsi_calc[temp_key]
-        + 0.5 * (dsi_calc[c["U"]] ** 2 + dsi_calc[c["V"]] ** 2)
+
+    d_pot_energy = z_geo * dsi_calc.attrs[c["g"]]
+
+    d_kin_energy = 0.5 * (dsi_calc[c["U"]] ** 2 + dsi_calc[c["V"]] ** 2)
+
+    d_therm_energy = dsi_calc.attrs[c["cp"]] * dsi_calc[temp_key]
+
+    kin_energy = _integrate_over_mass(
+        quant_to_int=d_kin_energy, area=dsi_calc[area_key], dzdp=dzdp, rho=rho
+    )
+    therm_energy = _integrate_over_mass(
+        quant_to_int=d_therm_energy,
+        area=dsi_calc[area_key],
+        dzdp=dzdp,
+        rho=rho,
     )
-    energy = _integrate_over_mass(
-        quant_to_int=d_energy, area=dsi_calc[area_key], dzdp=dzdp, rho=rho
+    pot_energy = _integrate_over_mass(
+        quant_to_int=d_pot_energy, area=dsi_calc[area_key], dzdp=dzdp, rho=rho
     )
 
-    if var_key_out is not None:
-        dsi.update({var_key_out: energy})
+    tot_energy = kin_energy + therm_energy + pot_energy
 
-    return energy
+    if var_key_out is not None:
+        dsi.update({var_key_out: tot_energy})
+        dsi.update({var_key_out + "_th": therm_energy})
+        dsi.update({var_key_out + "_pot": pot_energy})
+        dsi.update({var_key_out + "_kin": kin_energy})
+
+    if return_all:
+        return tot_energy, therm_energy, pot_energy, kin_energy
+    else:
+        return tot_energy
 
 
 def _integrate_over_mass(quant_to_int, area, dzdp, rho):
diff --git a/setup.py b/setup.py
index 4902864..96e4f78 100644
--- a/setup.py
+++ b/setup.py
@@ -18,7 +18,7 @@
 
 setup(
     name="gcm_toolkit",
-    version="v0.2.0",
+    version="v0.2.1",
     packages=find_packages(),
     include_package_data=True,
     scripts=["bin/convert_to_gcmt"],