diff --git a/fronts.py b/fronts.py index 06717a8..67226c0 100644 --- a/fronts.py +++ b/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): diff --git a/test_front.py b/test_front.py index f434fa2..e7c8649 100644 --- a/test_front.py +++ b/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 diff --git a/unittest_fronts.py b/unittest_fronts.py new file mode 100644 index 0000000..f9f1857 --- /dev/null +++ b/unittest_fronts.py @@ -0,0 +1,153 @@ +import unittest +import numpy as np +from scipy.signal.ltisys import ZerosPolesGain +import xarray as xr +from fronts import * + +class DataGenerator: + def get_latitude(self, nlat=None, latitudes=None): + if latitudes is not None: + nlat = len(latitudes) + elif nlat is None: + nlat = 21 + latitudes = np.linspace(-90, 90, nlat, endpoint=True) + else: + latitudes = np.linspace(-90, 90, nlat, endpoint=True) + return xr.DataArray( + latitudes, dims=('latitude',), coords={'latitude': latitudes}, + attrs={'units': 'degree_north', 'name': 'Latitude'} + ) + + def get_longitude(self, nlon=None, longitudes=None): + if longitudes is not None: + nlon = len(longitudes) + elif nlon is None: + nlon = 40 + longitudes = np.linspace(0, 360, nlon, endpoint=False) + else: + longitudes = np.linspace(0, 360, nlon, endpoint=False) + return xr.DataArray( + longitudes, dims=('longitude',), coords={'longitude': longitudes}, + attrs={'units': 'degree_east', 'name': 'Longitude'} + ) + + def get_temp(self, data=None, units="K"): + if data is None: + nlat = 21 + nlon = 40 + data = np.ones(nlat, nlon, dtype=np.float32)*285.0 + else: + nlat = len(data) + nlon = len(data[0]) + data = np.ndarray(data) + lat = self.get_latitude(nlat=nlat) + lon = self.get_longitude(nlon=nlon) + return xr.DataArray( + data, dims=('latitude', 'longitude'), + coords={'latitude': lat, 'longitude': lon}, + attrs={'name': 'temperature', 'units': units} + ) + + +class Test_Saturation_Pressure(unittest.TestCase): + def test_saturation_pressure_negative(self): + t_in = -12.0 + 273.15 + self.assertAlmostEqual(calculate_saturation_pressure(t_in), 2.4459, places=4) + + def test_saturation_pressure_zero(self): + t_in = 0.0 + 273.15 + self.assertAlmostEqual(calculate_saturation_pressure(t_in), 6.1094, places=4) + + def test_saturation_pressure_positive(self): + t_in = 21.0 + 273.15 + self.assertAlmostEqual(calculate_saturation_pressure(t_in), 24.81888, places=4) + + +# class Test_Dewpoint_Temperature(unittest.TestCase): +# def test_dewpoint_temperature(self): +# self.assertFalse(True, msg="Dewpoint Temperature Tests not yet implemented") + + +# class Test_Relative_Humidity(unittest.TestCase): +# def test_relative_humidity(self): +# self.assertFalse(True, msg="Relative Humidity Test not yet implemented") + +class Test_Zeropoints(unittest.TestCase): + + def create_data(self, 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(self): + d, d1, d2 = self.create_data( + [[1, 1, 1], + [1, 1, 1], + [1, 1, 1]]) + self.assertEqual(zeropoints(d, d1, d2).size, 0) + + def test_zeropoints_on_grid(self): + d, d1, d2 = self.create_data( + [[-1, 0, 1], + [-1, np.nan, 1], + [-1, np.nan, 1]] + ) + res = zeropoints(d, d1, d2) + expected = np.array([[0, 1]]) + np.testing.assert_allclose(res, expected) + + def test_zeropoints_between_grid(self): + d, d1, d2 = self.create_data( + [[-1, 1, 1], + [-1, 1, 1], + [-1, 1, 1]] + ) + res = 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(self): + d, d1, d2 = self.create_data( + [[-1, -1, -1], + [ 0, 0, 0], + [ 1, 1, 1]] + ) + res = 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(self): + d, d1, d2 = self.create_data( + [[-1, -1, -1], + [1, 1, 1], + [1, 1, 1]] + ) + res = 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(self): + d, d1, d2 = self.create_data( + [[-1, -1, -1], + [3, 3, 3], + [1, 1, 1]] + ) + res = 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(self): + d, d1, d2 = self.create_data( + [[-3, -1, 2], + [-2, 0, 1], + [1, 1, 4]] + ) + res = zeropoints(d, d1, d2) + expected = np.array([[1, 1], [0, 4/3], [5/3, 0]]) + np.testing.assert_allclose(res, expected) + + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file