diff --git a/src/proteus/atmos_clim/common.py b/src/proteus/atmos_clim/common.py index b1170670..63e9abae 100644 --- a/src/proteus/atmos_clim/common.py +++ b/src/proteus/atmos_clim/common.py @@ -13,9 +13,28 @@ log = logging.getLogger("fwl."+__name__) -def read_ncdf_profile(nc_fpath:str): - """ - Read temperature, pressure, height data from NetCDF file. +def read_ncdf_profile(nc_fpath:str, extra_keys:list=[]): + """Read data from atmosphere NetCDF output file. + + Automatically reads pressure (p), temperature (t), height (z) arrays with + cell-centre (N) and cell-edge (N+1) values interleaved into a single combined array of + length (2*N+1). + + Extra keys can be read-in using the extra_keys parameter. These will be stored with + the same dimensions as in the NetCDF file. + + Parameters + ---------- + nc_fpath : str + Path to NetCDF file. + + extra_keys : list + List of extra keys (strings) to read from the file. + + Returns + ---------- + out : dict + Dictionary containing numpy arrays of data from the file. """ # open file @@ -30,14 +49,15 @@ def read_ncdf_profile(nc_fpath:str): z = np.array(ds.variables["z"][:]) zl = np.array(ds.variables["zl"][:]) - nlev = len(p) + nlev_c = len(p) + nlev_l = len(pl) # read pressure, temperature, height data into dictionary values out = {} out["p"] = [pl[0]] out["t"] = [tl[0]] out["z"] = [zl[0]] - for i in range(nlev): + for i in range(nlev_c): out["p"].append(p[i]) out["p"].append(pl[i+1]) @@ -47,12 +67,17 @@ def read_ncdf_profile(nc_fpath:str): out["z"].append(z[i]) out["z"].append(zl[i+1]) + # Read extra keys + for key in extra_keys: + if key in ds.variables.keys(): + out[key] = np.array(ds.variables[key][:]) + # close file ds.close() # convert to np arrays - for k in out.keys(): - out[k] = np.array(out[k], dtype=float) + for key in out.keys(): + out[key] = np.array(out[key], dtype=float) return out diff --git a/tests/integration/test_integration_physical.py b/tests/integration/test_integration_physical.py index d1d2e112..9c6bf993 100644 --- a/tests/integration/test_integration_physical.py +++ b/tests/integration/test_integration_physical.py @@ -6,7 +6,7 @@ import pytest from helpers import PROTEUS_ROOT from pandas.testing import assert_frame_equal -from np.testing import assert_array_equal +from numpy.testing import assert_allclose from proteus import Proteus from proteus.utils.coupler import ReadHelpfileFromCSV @@ -47,7 +47,7 @@ def test_physical_atmosphere(physical_run): # Load atmosphere reference atm_ref = read_ncdf_profile(ref_dir / '2002_atm.nc') - # Compare to expected + # Compare to expected array values. # Cannot simply compare the files as black-boxes, because they contain date information - for field in ["tmp","p","fl_U_LW"]: + for field in ["t", "p", "z", "fl_U_LW", "fl_D_LW", "fl_U_SW", "fl_D_SW"]: assert_allclose(atm_out[field], atm_ref[field], rtol=5e-3)