From 1afb31ab174e88d928109f3a50754b7533115ed5 Mon Sep 17 00:00:00 2001 From: James Mineau Date: Tue, 13 Aug 2024 14:20:52 -0600 Subject: [PATCH] vhd using pint --- lair/air/meteorology.py | 10 +++++++-- lair/air/soundings.py | 48 ++++++++++++++++++++++++++--------------- 2 files changed, 39 insertions(+), 19 deletions(-) diff --git a/lair/air/meteorology.py b/lair/air/meteorology.py index 2c6de6b..1d0214f 100644 --- a/lair/air/meteorology.py +++ b/lair/air/meteorology.py @@ -2,6 +2,12 @@ Meteorological calculations. Inspired by AOS 330 at UW-Madison with Grant Petty. + +.. note:: + It would be nice to be able to wrap these funtions with `pint` - however, + because I input both numpy and xarray arrays, this will not work. + Waiting for https://github.com/xarray-contrib/pint-xarray/pull/143 + For now, we will assume all inputs are in SI units. """ import numpy as np @@ -153,7 +159,7 @@ def virt_T(T: float, q: float) -> float: return T * (1 + 0.61 * q) -def poisson(T: float, p: float, p0: float = 1000 * units('hPa')) -> float: +def poisson(T: float, p: float, p0: float = 1e5) -> float: """ Calculate the potential temperature. (Poission's equation) @@ -174,7 +180,7 @@ def poisson(T: float, p: float, p0: float = 1000 * units('hPa')) -> float: return T * (p0/p)**(Rd/cp) -def inv_poisson(p: float, theta: float, p0: float = 1000 * units('hPa')) -> float: +def inv_poisson(p: float, theta: float, p0: float = 1e5) -> float: """ Calculate the temperature from potential temperature. (Inverse Poission's equation) diff --git a/lair/air/soundings.py b/lair/air/soundings.py index 53f9d87..6fe2222 100644 --- a/lair/air/soundings.py +++ b/lair/air/soundings.py @@ -9,7 +9,6 @@ from collections import deque import datetime as dt -import numpy as np import os import pandas as pd import requests @@ -17,9 +16,9 @@ import xarray as xr from lair.config import vprint, GROUP_DIR -from lair.constants import Rd +from lair import units +from lair.constants import Rd, cp from lair.air.meteorology import ideal_gas_law, hypsometric, poisson -from lair.units import C2K, hPa, J, kg, K, m #: Sounding data directory @@ -260,8 +259,17 @@ def get_soundings(station='SLC', start=None, end=None, sounding_dir=None, months download_soundings(station, start, end, sounding_dir, months) soundings = [] - # TODO filter by start/end for file in files: + # Skip files that don't match the date range + date = file.split('_')[1].split('.')[0] + date = dt.datetime.strptime(date, '%Y%m%d%H') + if start: + if date <= start: + continue + if end: + if date > end: + continue + path = os.path.join(sounding_dir, file) try: soundings.append(Sounding(path)) @@ -280,7 +288,7 @@ def get_soundings(station='SLC', start=None, end=None, sounding_dir=None, months return data -def valleyheatdeficit(data, integration_height=2200) -> xr.DataArray: +def valleyheatdeficit(data: xr.Dataset, integration_height=2200) -> xr.DataArray: """ Calculate the valley heat deficit. @@ -294,38 +302,44 @@ def valleyheatdeficit(data, integration_height=2200) -> xr.DataArray: data : xr.DataSet The sounding data. integration_height : int - The height to integrate to. + The height to integrate to [m]. Returns ------- xr.DataArray - The valley heat deficit. + The valley heat deficit [MJ/m^2]. """ - h0 = data.elevation * m - cp = 1005 * J/kg/K - + h0 = data.elevation + # Subset to the heights between the surface and the integration height data = data.sel(height=slice(h0, integration_height)) - T = C2K(data.temperature) - p = data.pressure * hPa + T = data.temperature.pint.quantify('degC').pint.to('degK') + p = data.pressure.pint.quantify('hPa').pint.to('Pa') # Calculate potential temperature using poisson's equation - theta = poisson(T, p) + theta = poisson(T=T, p=p, p0=1e5 * units('Pa')) theta_h = theta.sel(height=integration_height, method='nearest') # Calculate virtual temperature to account for water vapor Tv = hypsometric(p1=p.isel(height=slice(0, -1)).values, p2=p.isel(height=slice(1, None)).values, - deltaz=data.interpolation_interval) + deltaz=data.interpolation_interval * units('m')) layer_heights = T.height.values[:-1] + data.interpolation_interval / 2 Tv = xr.DataArray(Tv, coords=[data.time, layer_heights], - dims=['time', 'height']).interp_like(T, method='linear') + dims=['time', 'height'])\ + .interp_like(T, method='linear')\ + .pint.quantify('degK') # Calculate the density using the ideal gas law rho = ideal_gas_law(solve_for='rho', p=p, T=Tv, R=Rd) + # Set pint units - Setting units to kg/m3 doesnt change the numbers + # pint-xarray hasnt implemented .to_base_units() yet + # when they do, we can change this to .pint.to_base_units() + rho = rho.pint.to('kg/m^3') # Calculate the heat deficit by integrating using the trapezoid method - heat_deficit = (cp * rho * (theta_h - theta)).dropna('height', how='all').integrate('height') + heat_deficit = (cp * rho * (theta_h - theta)).dropna('height', how='all')\ + .integrate('height') * (1 * units('m')) # J/m2 - return heat_deficit # J/m2 + return heat_deficit.pint.to('MJ/m^2').pint.dequantify()