Refactor

fronts.py

@@ -1,41 +1,95 @@
 # module imports
 import numpy as np
+from numpy.lib.utils import _lookfor_generate_cache
 import scipy.signal
+import scipy.interpolate as interp
 import xarray as xr
 import scipy.spatial.distance as sp_dist
 import geopy.distance as gp_dist
 
 
-def wetbulb(ta, hus, plev, steps=100, ta_units=None):
-    # calculates wetbulb temperature from pressure-level data
-    # Inputs: ta - temperature field (xarray)
-    #         hus - specific humidity field (xarray)
-    #         plev - the level of the data (in hPa, scalar)
-    #         steps -  the number of steps in the numerical calculation
-    #         ta_units - the units of the temperature field (if not provided, read from ta)
-    if ta_units == None:
+def calculate_saturation_pressure(ta):
+    """
+    (xarray.DataArray) -> xarray.DataArray
+
+    Calculates the saturation pressure.
+
+    Argument:
+    ta:       Temperature Array in Kelvin
+    """
+    return 6.1094 * np.exp((17.625 * (ta - 273.15)) / (ta - 30.11))
+
+
+def calculate_relative_humidity(hus, plev, es):
+    """
+    Calculate the relative humidity
+
+    Arguments:
+    hus:        specific humidity
+    plev:       level in hPa
+    es:         saturation pressure
+    """
+    return (hus * (plev - es)) / (0.622 * (1 - hus) * es)
+
+
+def calculate_dewpoint_temperature(e):
+    """
+    Calculates the dewpoint from the vapour pressure
+
+    Input:
+    e:      Vapour Pressure
+    """
+    return ((243.5 * np.log(e / 6.112)) / (17.67 - np.log(e / 6.112))) + 273.15
+
+
+def calculate_atmospheric_fields(ta, hus, plev, ta_units=None):
+    """
+    (DataArray, DataArray, Numeric, Str) -> DataArray, DataArray, DataArray, DataArray
+
+    Calculates the Saturation Pressure, the Vapour Pressure, 
+    the Relative Humidity, and the Dewpoint Temperature.
+
+    Inputs:
+    ta:         Temperature
+    hus:        Specific Humidity
+    plev:       Level of Data (hPa)
+    ta_units:   (Optional) units of ta. Defaults to ta.units.
+
+    Output:
+    Tuple(Saturation Pressure, Vapour Pressure, Relative Humidity, Dewpoint Temperature)
+    """
+    if ta_units is None:
         ta_units = ta.units
     # saturation vapor pressure
-    if ta_units == "K" or ta_units == "Kelvin" or ta_units == "kelvin":
-        es = 6.1094 * np.exp((17.625 * (ta - 273.15)) / (ta - 30.11))
-    elif (
-        ta_units == "C"
-        or ta_units == "degC"
-        or ta_units == "deg_C"
-        or ta_units == "Celcius"
-        or ta_units == "celcius"
-    ):
-        es = 6.1094 * np.exp((17.625 * (ta)) / (ta + 243.04))
+    if ta_units.lower() in ["k", "kelvin"]:
+        es = calculate_saturation_pressure(ta)
+    elif ta_units.lower() in ["c", "degc", "deg_c", "celsius"]:
+        es = calculate_saturation_pressure(ta+273.15)
     else:
         raise ValueError(
             "Input temperature unit not recognised, use Kelvin (K) or Celcius (C, degC, deg_C)"
         )
     # relative humidity from specific humidity and sat. vap. pres.
-    rh = (hus * (plev - es)) / (0.622 * (1 - hus) * es)
+    rh = calculate_relative_humidity(hus, plev, es)
     # vapor pressure
     e = es * rh
     # dewpoint temperature
-    t_dewpoint = ((243.5 * np.log(e / 6.112)) / (17.67 - np.log(e / 6.112))) + 273.15
+    t_dewpoint = calculate_dewpoint_temperature(e)
+
+    return (es, rh, e, t_dewpoint)
+
+
+def wetbulb(ta, hus, plev, steps=100, ta_units=None):
+    """
+    calculates wetbulb temperature from pressure-level data
+    Inputs: ta - temperature field (xarray)
+            hus - specific humidity field (xarray)
+            plev - the level of the data (in hPa, scalar)
+            steps -  the number of steps in the numerical calculation
+            ta_units - the units of the temperature field (if not provided, read from ta)
+    """
+
+    es, _, e, t_dewpoint = calculate_atmospheric_fields(ta, hus, plev, ta_units)
 
     # unlike the above, calculating the wetbulb temperature is done numerically
     delta_t = (ta - t_dewpoint) / steps
@@ -44,7 +98,7 @@ def wetbulb(ta, hus, plev, steps=100, ta_units=None):
 
     for i in range(steps):
         cur_t = ta - i * delta_t
-        es_cur_t = 6.1094 * np.exp((17.625 * (cur_t - 273.15)) / (cur_t - 30.11))
+        es_cur_t = calculate_saturation_pressure(cur_t)
         adiabatic_adj = 850 * (ta - cur_t) * (0.00066) * (1 + 0.00115 * cur_t)
         diff = np.abs(es_cur_t - adiabatic_adj - e)
         t_wet.data[diff < cur_diff] = cur_t.data[diff < cur_diff]
@@ -54,90 +108,50 @@ def wetbulb(ta, hus, plev, steps=100, ta_units=None):
 
 
 def dewpoint(ta, hus, plev, ta_units=None):
-    # calculates depoint temperature from pressure-level data
-    # Inputs: ta - temperature field (xarray)
-    #         hus - specific humidity field (xarray)
-    #         plev - the level of the data (in hPa, scalar)
-    #         ta_units - the units of the temperature field (if not provided, read from ta)
-    if ta_units == None:
-        ta_units = ta.units
-    if ta_units == "K" or ta_units == "Kelvin" or ta_units == "kelvin":
-        es = 6.1094 * np.exp((17.625 * (ta - 273.15)) / (ta - 30.11))
-    elif (
-        ta_units == "C"
-        or ta_units == "degC"
-        or ta_units == "deg_C"
-        or ta_units == "Celcius"
-        or ta_units == "celcius"
-    ):
-        es = 6.1094 * np.exp((17.625 * (ta)) / (ta + 243.04))
-    else:
-        raise ValueError(
-            "Input temperature unit not recognised, use Kelvin (K) or Celcius (C, degC, deg_C)"
-        )
-    rh = (hus * (plev - es)) / (0.622 * (1 - hus) * es)
-    e = es * rh
-    t_dewpoint = ((243.5 * np.log(e / 6.112)) / (17.67 - np.log(e / 6.112))) + 273.15
-    return t_dewpoint
+    """
+    calculates depoint temperature from pressure-level data
+    Inputs: ta - temperature field (xarray)
+            hus - specific humidity field (xarray)
+            plev - the level of the data (in hPa, scalar)
+            ta_units - the units of the temperature field (if not provided, read from ta)
+    """
+
+    _, _, _, t_dewpoint = calculate_atmospheric_fields(ta, hus, plev, ta_units)
 
+    return t_dewpoint
 
 def zeropoints(data, dim1, dim2):
-    # finds zero-crossing points in a gridded data set along the lines of each dimension
-    # inputs: data - 2d data field (numpy array)
-    #         dim1 - coords of the first dim of data (np array)
-    #         dim2 - coords of the second dim of data (np array)
-    n1, n2 = data.shape
-    # assuming regularly spaced grid:
-    d_dim2 = dim2[1] - dim2[0]
-    d_dim1 = dim1[1] - dim1[0]
-    tloc_1 = []
-    tloc_2 = []
-    for lonn in range(0, n2 - 1):
-        for latn in range(0, n1 - 1):
-            flag = False
-            if data[latn, lonn] == 0:
-                tloc_1.append([dim1[latn], dim2[lonn]])
-                flag = True
-            else:
-                if (
-                    np.isfinite(data[latn, lonn])
-                    and np.isfinite(data[latn, lonn + 1])
-                    and not flag
-                ):
-                    if (data[latn, lonn] > 0 and data[latn, lonn + 1] < 0) or (
-                        data[latn, lonn] < 0 and data[latn, lonn + 1] > 0
-                    ):
-                        tloc_1.append(
-                            [
-                                dim1[latn],
-                                dim2[lonn]
-                                + d_dim2
-                                * np.abs(
-                                    data[latn, lonn]
-                                    / (data[latn, lonn] - data[latn, lonn + 1])
-                                ),
-                            ]
-                        )
-                if (
-                    np.isfinite(data[latn, lonn])
-                    and np.isfinite(data[latn + 1, lonn])
-                    and not flag
-                ):
-                    if (data[latn, lonn] > 0 and data[latn + 1, lonn] < 0) or (
-                        data[latn, lonn] < 0 and data[latn + 1, lonn] > 0
-                    ):
-                        tloc_2.append(
-                            [
-                                dim1[latn]
-                                + d_dim1
-                                * np.abs(
-                                    data[latn, lonn]
-                                    / (data[latn, lonn] - data[latn + 1, lonn])
-                                ),
-                                dim2[lonn],
-                            ]
-                        )
-    return np.array(tloc_1 + tloc_2)
+    """
+    finds zero-crossing points in a gridded data set along the lines of each dimension
+    inputs: data - 2d data field (numpy array)
+            dim1 - coords of the first dim of data (np array)
+            dim2 - coords of the second dim of data (np array)
+    """
+
+    ## Find points where the value itself is zero:
+    zero_locations = [
+        [dim1[idx1], dim2[idx2]] 
+        for idx1, idx2 in zip(*np.where(data==0))
+    ]
+
+    ## Find zeropoints along latitude
+    for dim1_val, dim2_data in zip(dim1, data):
+        # Multiply each data point with the next. Negative values then indicate change in sign
+        indicator_array = dim2_data[:-1] * dim2_data[1:]
+        zero_locations.extend([
+            [dim1_val, interp.interp1d(dim2_data[i:i+2], dim2[i:i+2])(0)] 
+            for i in np.where(indicator_array < 0)[0]
+        ])
+
+    for dim2_val, dim1_data in zip(dim2, data.T):
+        # Multiply each data point with the next. Negative values then indicate change in sign
+        indicator_array = dim1_data[:-1] * dim1_data[1:]
+        zero_locations.extend([
+            [interp.interp1d(dim1_data[i:i+2], dim1[i:i+2])(0), dim2_val]
+            for i in np.where(indicator_array < 0)[0]
+        ])
+
+    return np.array(zero_locations)
 
 
 def frontfields(data, ua, va, threshold=-0.3e-10):

test_front.py

@@ -3,27 +3,96 @@
 import xarray as xr
 import numpy as np
 
+
+def create_data(data):
+    d = np.array(data)
+    d1 = np.arange(d.shape[0])
+    d2 = np.arange(d.shape[1])
+    return d, d1, d2
+
+
+def test_no_zeropoints():
+    d, d1, d2 = create_data([[1, 1, 1], [1, 1, 1], [1, 1, 1]])
+    assert fronts.zeropoints(d, d1, d2).size == 0
+
+
+def test_zeropoints_on_grid():
+    d, d1, d2 = create_data([[-1, 0, 1], [-1, np.nan, 1], [-1, np.nan, 1]])
+    res = fronts.zeropoints(d, d1, d2)
+    expected = np.array([[0, 1]])
+    np.testing.assert_allclose(res, expected)
+
+
+def test_zeropoints_between_grid():
+    d, d1, d2 = create_data([[-1, 1, 1], [-1, 1, 1], [-1, 1, 1]])
+    res = fronts.zeropoints(d, d1, d2)
+    expected = np.array([[0, 0.5], [1, 0.5], [2, 0.5]])
+    np.testing.assert_allclose(res, expected)
+
+
+def test_zeropoints_on_grid_y():
+    d, d1, d2 = create_data([[-1, -1, -1], [0, 0, 0], [1, 1, 1]])
+    res = fronts.zeropoints(d, d1, d2)
+    expected = np.array([[1, 0], [1, 1], [1, 2]])
+    np.testing.assert_allclose(res, expected)
+
+
+def test_zeropoints_between_grid_y():
+    d, d1, d2 = create_data([[-1, -1, -1], [1, 1, 1], [1, 1, 1]])
+    res = fronts.zeropoints(d, d1, d2)
+    expected = np.array([[0.5, 0], [0.5, 1], [0.5, 2]])
+    np.testing.assert_allclose(res, expected)
+
+
+def test_zeropoints_fractional():
+    d, d1, d2 = create_data([[-1, -1, -1], [3, 3, 3], [1, 1, 1]])
+    res = fronts.zeropoints(d, d1, d2)
+    expected = np.array([[0.25, 0], [0.25, 1], [0.25, 2]])
+    np.testing.assert_allclose(res, expected)
+
+
+def test_zeropoints_diagonal():
+    d, d1, d2 = create_data([[-3, -1, 2], [-2, 0, 1], [1, 1, 4]])
+    res = fronts.zeropoints(d, d1, d2)
+    expected = np.array([[1, 1], [0, 4 / 3], [5 / 3, 0]])
+    np.testing.assert_allclose(res, expected)
+
+
 def test_era5():
 
-    tafile=xr.open_dataarray('/g/data/rt52/era5/pressure-levels/reanalysis/t/2020/t_era5_oper_pl_20200101-20200131.nc')
-    uafile=xr.open_dataarray('/g/data/rt52/era5/pressure-levels/reanalysis/u/2020/u_era5_oper_pl_20200101-20200131.nc')
-    vafile=xr.open_dataarray('/g/data/rt52/era5/pressure-levels/reanalysis/v/2020/v_era5_oper_pl_20200101-20200131.nc')
-    husfile=xr.open_dataarray('/g/data/rt52/era5/pressure-levels/reanalysis/q/2020/q_era5_oper_pl_20200101-20200131.nc')
+    tafile = xr.open_dataarray(
+        "/g/data/rt52/era5/pressure-levels/reanalysis/t/2020/t_era5_oper_pl_20200101-20200131.nc"
+    )
+    uafile = xr.open_dataarray(
+        "/g/data/rt52/era5/pressure-levels/reanalysis/u/2020/u_era5_oper_pl_20200101-20200131.nc"
+    )
+    vafile = xr.open_dataarray(
+        "/g/data/rt52/era5/pressure-levels/reanalysis/v/2020/v_era5_oper_pl_20200101-20200131.nc"
+    )
+    husfile = xr.open_dataarray(
+        "/g/data/rt52/era5/pressure-levels/reanalysis/q/2020/q_era5_oper_pl_20200101-20200131.nc"
+    )
 
     n = 0
 
     # 32 index in level --> 900 hPa
-    ta=tafile[n,32,0:100,0:100]
-    ua=uafile[n,32,0:100,0:100]
-    va=vafile[n,32,0:100,0:100]
-    hus=husfile[n,32,0:100,0:100]
-
-    t_wet=fronts.wetbulb(ta,hus,900,steps=120)
-
-    frontdata=fronts.front(t_wet,ua,va,threshold_i=-1e-10,numsmooth=9,minlength=50)
-
-    timestring=np.datetime_as_string(tafile.time.data[n],unit='h')
-    with open(f'tests/900hPa_fronts_{timestring}.json') as sample_file:
+    ta = tafile[n, 32, 0:100, 0:100]
+    ua = uafile[n, 32, 0:100, 0:100]
+    va = vafile[n, 32, 0:100, 0:100]
+    hus = husfile[n, 32, 0:100, 0:100]
+
+    t_wet = fronts.wetbulb(ta, hus, 900, steps=120)
+
+    frontdata = fronts.front(
+        t_wet, ua, va, threshold_i=-1e-10, numsmooth=9, minlength=50
+    )
+
+    timestring = np.datetime_as_string(tafile.time.data[n], unit="h")
+
+    with open("tests/test_fronts.json", "w") as outfile:
+        json.dump(frontdata, outfile)
+
+    with open(f"tests/900hPa_fronts_{timestring}.json") as sample_file:
         sample = json.load(sample_file)
 
         assert frontdata == sample