ENH for p-data, dp calc that incorporates ps value

aospy/calc.py

@@ -16,9 +16,9 @@
 from .io import (_data_in_label, _data_out_label, _ens_label, _yr_label, dmget,
                  data_in_name_gfdl)
 from .timedate import TimeManager, _get_time
-from .utils import (get_parent_attr, level_thickness, apply_time_offset,
-                    monthly_mean_ts, monthly_mean_at_each_ind,
-                    pfull_from_ps, to_pfull_from_phalf, dp_from_ps, int_dp_g)
+from .utils import (get_parent_attr, apply_time_offset, monthly_mean_ts,
+                    monthly_mean_at_each_ind, pfull_from_ps,
+                    to_pfull_from_phalf, dp_from_ps, dp_from_p, int_dp_g)
 
 
 dp = Var(
@@ -438,6 +438,7 @@ def _create_input_data_obj(self, var, start_date=False,
 
     def _get_pressure_vals(self, var, start_date, end_date, n=0):
         """Get pressure array, whether sigma or standard levels."""
+        ps = self._create_input_data_obj(self.ps, start_date, end_date)
 
         if self.dtype_in_vert == 'pressure':
             if np.any(self.pressure):
@@ -447,12 +448,11 @@ def _get_pressure_vals(self, var, start_date, end_date, n=0):
             if var.name == 'p':
                 data = pressure
             elif var.name == 'dp':
-                data = level_thickness(pressure)
+                data = dp_from_p(pressure, ps)
 
         elif self.dtype_in_vert == 'sigma':
             bk = self.model[n].bk
             pk = self.model[n].pk
-            ps = self._create_input_data_obj(self.ps, start_date, end_date)
             pfull_coord = self.model[n].pfull
             if var.name == 'p':
                 data = pfull_from_ps(bk, pk, ps, pfull_coord)

aospy/test/test_utils.py

@@ -1,18 +1,19 @@
 #!/usr/bin/env python
 """Test suite for aospy.utils module."""
-
 import sys
 import unittest
 
 import numpy as np
+import xray
 
 import aospy.utils as au
 
 
 class AospyUtilsTestCase(unittest.TestCase):
     def setUp(self):
         self.p_in_hpa = np.array([1000, 925, 850, 775, 700, 600, 500, 400, 300,
-                                  200, 150, 100, 70, 50, 30, 20, 10])
+                                  200, 150, 100, 70, 50, 30, 20, 10],
+                                 dtype=np.float64)
         self.p_in_pa = self.p_in_hpa*1e2
         self.p_top = 0
         self.p_bot = 1.1e5
@@ -37,15 +38,30 @@ def test_to_pascal_array(self):
                                       self.p_in_pa)
         np.testing.assert_array_equal(au.to_pascal(self.p_in_pa), self.p_in_pa)
 
-    def test_phalf_from_pfull(self):
+    def test_to_phalf_from_pfull(self):
+        # S. Hill 2015-11-10: This needs to be rewritten.
         np.testing.assert_array_equal(
-            au.phalf_from_pfull(self.p_in_pa, self.p_top, self.p_bot),
+            au.to_phalf_from_pfull(self.p_in_pa, self.p_top, self.p_bot),
             self.phalf
         )
 
-    # def test_dp_from_p(self):
-        # ps = 9e4
-        # dp_true = self.p_in_pa
+
+def test_dp_from_p():
+    path = (
+        '/archive/Spencer.Hill/am2/am2clim_reyoi/gfdl.ncrc2-default-prod/pp/'
+        'atmos/ts/monthly/30yr/atmos.198301-201212.ucomp.nc'
+    )
+    ds = xray.open_dataset(path)
+    p = ds.level
+    path = (
+        '/archive/Spencer.Hill/am2/am2clim_reyoi/gfdl.ncrc2-default-prod/pp/'
+        'atmos/ts/monthly/30yr/atmos.198301-201212.ps.nc'
+    )
+    ps = xray.open_dataset(path).ps
+    dp = au.dp_from_p(p, ps)
+    np.testing.assert_array_equal(p.level, dp.level)
+    # TODO: More tests
+
 
 if __name__ == '__main__':
     sys.exit(unittest.main())

aospy/utils.py

@@ -1,12 +1,14 @@
 """aospy.utils: utility functions for the aospy module."""
+import logging
 import warnings
 
 from infinite_diff import FiniteDiff
 import numpy as np
 import pandas as pd
 import xray
 
-from .__config__ import PHALF_STR, PFULL_STR, PLEVEL_STR, TIME_STR, user_path
+from .__config__ import (PHALF_STR, PFULL_STR, PLEVEL_STR, TIME_STR,
+                         LAT_STR, LON_STR, user_path)
 from .constants import grav, Constant
 
 
@@ -226,28 +228,67 @@ def int_dp_g(arr, dp):
                      vert_coord_name(dp)) / grav.value
 
 
-def dp_from_p(p, ps):
-    """Get level thickness of pressure data, incorporating surface pressure."""
-    # Top layer goes to 0 hPa; bottom layer goes to 1100 hPa.
-    p = to_pascal(p)
-    p_top = np.array([0])
-    p_bot = np.array([1.1e5])
+def dp_from_p(p, ps, p_top=0., p_bot=1.1e5):
+    """Get level thickness of pressure data, incorporating surface pressure.
+
+    Level edges are defined as halfway between the levels, as well as the user-
+    specified uppermost and lowermost values.  The dp of levels whose bottom
+    pressure is less than the surface pressure is not changed by ps, since they
+    don't intersect the surface.  If ps is in between a level's top and bottom
+    pressures, then its dp becomes the pressure difference between its top and
+    ps.  If ps is less than a level's top and bottom pressures, then that level
+    is underground and its values are masked.
+
+    Note that postprocessing routines (e.g. at GFDL) typically mask out data
+    wherever the surface pressure is less than the level's given value, not the
+    level's upper edge.  This masks out more levels than the
+
+    """
+    p_vals = to_pascal(p.values)
 
     # Layer edges are halfway between the given pressure levels.
-    p_edges_interior = 0.5*(p.isel(phalf=slice(0, -1)) +
-                            p.isel(phalf=slice(1, None)))
-    p_edges = xray.concat((p_bot, p_edges_interior, p_top), dim=PHALF_STR)
-    p_edge_above = p_edges.isel(phalf=slice(1, None))
-    p_edge_below = p_edges.isel(phalf=slice(0, -1))
-    dp_interior = p_edge_below - p_edge_above
-    dp_interior.rename({PHALF_STR: PFULL_STR})
-
-    ps = to_pascal(ps)
-    # If ps < p_edge_below, then ps becomes the layer's bottom boundary.
-    dp_adj_sfc = ps - p_edge_above
-    dp = np.where(np.sign(ps - p_edge_below) > 0, dp_interior, dp_adj_sfc)
-    # Mask where ps is less than the p.
-    return np.ma.masked_where(ps < p, dp)
+    p_edges_interior = 0.5*(p_vals[:-1] + p_vals[1:])
+    p_edges = np.concatenate(([p_bot], p_edges_interior, [p_top]))
+    p_edge_above = p_edges[1:]
+    p_edge_below = p_edges[:-1]
+    dp = p_edge_below - p_edge_above
+    if not all(np.sign(dp)):
+        raise ValueError("dp array not all > 0 : {}".format(dp))
+    # Pressure difference between ps and the upper edge of each pressure level.
+    p_edge_above_xray = xray.DataArray(p_edge_above, dims=p.dims,
+                                       coords=p.coords)
+    dp_to_sfc = ps - p_edge_above_xray
+    # Find the level adjacent to the masked, under-ground levels.
+    change = xray.DataArray(np.zeros(dp_to_sfc.shape), dims=dp_to_sfc.dims,
+                            coords=dp_to_sfc.coords)
+    change[{PLEVEL_STR: slice(1, None)}] = np.diff(
+        np.sign(ps - to_pascal(p.copy()))
+    )
+    dp_combined = xray.DataArray(np.where(change, dp_to_sfc, dp),
+                                 dims=dp_to_sfc.dims, coords=dp_to_sfc.coords)
+    # Mask levels that are under ground.
+    above_ground = ps > to_pascal(p.copy())
+    above_ground[PLEVEL_STR].values = p.values
+    dp_with_ps = dp_combined.where(above_ground)
+    # Revert to original dim order.
+    possible_dim_orders = [
+        (TIME_STR, PLEVEL_STR, LAT_STR, LON_STR),
+        (TIME_STR, PLEVEL_STR, LAT_STR),
+        (TIME_STR, PLEVEL_STR, LON_STR),
+        (TIME_STR, PLEVEL_STR),
+        (PLEVEL_STR, LAT_STR, LON_STR),
+        (PLEVEL_STR, LAT_STR),
+        (PLEVEL_STR, LON_STR),
+        (PLEVEL_STR,),
+    ]
+    for dim_order in possible_dim_orders:
+        try:
+            return dp_with_ps.transpose(*dim_order)
+        except ValueError:
+            logging.debug("Failed transpose to dims: {}".format(dim_order))
+    else:
+        logging.debug("No transpose was successful.")
+        return dp_with_ps
 
 
 def level_thickness(p):