From 56aa5fa5d16af2b0ac4e5908b816589b8fd9cfb2 Mon Sep 17 00:00:00 2001 From: Tyler Sutterley Date: Thu, 26 Oct 2023 19:51:11 -0700 Subject: [PATCH] feat: add generic wrapper function for reading ATLAS constituents (#250) --- doc/source/api_reference/io/ATLAS.rst | 2 + pyTMD/io/ATLAS.py | 99 ++++++++++++++++++--------- pyTMD/io/FES.py | 2 +- pyTMD/io/GOT.py | 2 +- pyTMD/io/OTIS.py | 7 +- 5 files changed, 73 insertions(+), 39 deletions(-) diff --git a/doc/source/api_reference/io/ATLAS.rst b/doc/source/api_reference/io/ATLAS.rst index aa77b124..29c3fd85 100644 --- a/doc/source/api_reference/io/ATLAS.rst +++ b/doc/source/api_reference/io/ATLAS.rst @@ -26,6 +26,8 @@ Calling Sequence .. autofunction:: pyTMD.io.ATLAS.read_netcdf_grid +.. autofunction:: pyTMD.io.ATLAS.read_netcdf_file + .. autofunction:: pyTMD.io.ATLAS.read_netcdf_elevation .. autofunction:: pyTMD.io.ATLAS.read_netcdf_transport diff --git a/pyTMD/io/ATLAS.py b/pyTMD/io/ATLAS.py index ba95322f..0ef41f42 100644 --- a/pyTMD/io/ATLAS.py +++ b/pyTMD/io/ATLAS.py @@ -1,7 +1,7 @@ #!/usr/bin/env python u""" ATLAS.py -Written by Tyler Sutterley (04/2023) +Written by Tyler Sutterley (10/2023) Reads files for a tidal model and makes initial calculations to run tide program Includes functions to extract tidal harmonic constants from OTIS tide models for @@ -55,6 +55,7 @@ interpolate.py: interpolation routines for spatial data UPDATE HISTORY: + Updated 10/2023: add generic wrapper function for reading constituents Updated 04/2023: using pathlib to define and expand tide model paths Updated 03/2023: add basic variable typing to function inputs Updated 12/2022: refactor tide read programs under io @@ -275,14 +276,9 @@ def extract_constants( model_file = pathlib.Path(model_file).expanduser() if not model_file.exists(): raise FileNotFoundError(str(model_file)) - if (kwargs['type'] == 'z'): - # read constituent from elevation file - hc, cons = read_netcdf_elevation(model_file, - compressed=kwargs['compressed']) - elif kwargs['type'] in ('U','u','V','v'): - # read constituent from transport file - hc, cons = read_netcdf_transport(model_file, kwargs['type'], - compressed=kwargs['compressed']) + # read constituent from netCDF4 file + hc, cons = read_netcdf_file(model_file, kwargs['type'], + compressed=kwargs['compressed']) # append constituent to list constituents.append(cons) # replace original values with extend matrices @@ -414,14 +410,9 @@ def read_constants( model_file = pathlib.Path(model_file).expanduser() if not model_file.exists(): raise FileNotFoundError(str(model_file)) - if (kwargs['type'] == 'z'): - # read constituent from elevation file - hc, cons = read_netcdf_elevation(model_file, - compressed=kwargs['compressed']) - elif kwargs['type'] in ('U','u','V','v'): - # read constituent from transport file - hc, cons = read_netcdf_transport(model_file, kwargs['type'], - compressed=kwargs['compressed']) + # read constituent from netCDF4 file + hc, cons = read_netcdf_file(model_file, kwargs['type'], + compressed=kwargs['compressed']) # replace original values with extend matrices hc = extend_matrix(hc) hc.mask[:] |= bathymetry.mask[:] @@ -743,6 +734,46 @@ def read_netcdf_grid( f.close() if kwargs['compressed'] else None return (lon, lat, bathymetry) +# PURPOSE: wrapper function for reading netCDF4 constituent files +def read_netcdf_file( + input_file: str | pathlib.Path, + variable: str, + **kwargs + ): + """ + Wrapper function for reading netCDF4 files to extract + real and imaginary components for constituent + + Parameters + ---------- + input_file: str or pathlib.Path + input transport file + variable: str + Tidal variable to read + + - ``'z'``: heights + - ``'u'``: horizontal transport velocities + - ``'U'``: horizontal depth-averaged transport + - ``'v'``: vertical transport velocities + - ``'V'``: vertical depth-averaged transport + + compressed: bool, default False + Input file is gzip compressed + + Returns + ------- + hc: np.ndarray + tidal constituent + con: str + tidal constituent ID + """ + if (variable == 'z'): + # read constituent from elevation file + return read_netcdf_elevation(input_file, **kwargs) + elif variable in ('U','u','V','v'): + # read constituent from transport file + return read_netcdf_transport(input_file, variable, **kwargs) + # PURPOSE: read elevation file to extract real and imaginary components for # constituent def read_netcdf_elevation( @@ -761,7 +792,7 @@ def read_netcdf_elevation( Returns ------- - h: np.ndarray + hc: np.ndarray tidal elevation con: str tidal constituent ID @@ -783,15 +814,15 @@ def read_netcdf_elevation( nx = fileID.dimensions['nx'].size ny = fileID.dimensions['ny'].size # real and imaginary components of elevation - h = np.ma.zeros((ny,nx), dtype=np.complex64) - h.mask = np.zeros((ny,nx), dtype=bool) - h.data.real[:,:] = fileID.variables['hRe'][:,:].T - h.data.imag[:,:] = fileID.variables['hIm'][:,:].T + hc = np.ma.zeros((ny,nx), dtype=np.complex64) + hc.mask = np.zeros((ny,nx), dtype=bool) + hc.data.real[:,:] = fileID.variables['hRe'][:,:].T + hc.data.imag[:,:] = fileID.variables['hIm'][:,:].T # close the file fileID.close() f.close() if kwargs['compressed'] else None # return the elevation and constituent - return (h, con.strip()) + return (hc, con.strip()) # PURPOSE: read transport file to extract real and imaginary components for # constituent @@ -820,7 +851,7 @@ def read_netcdf_transport( Returns ------- - tr: np.ndarray + hc: np.ndarray tidal transport con: str tidal constituent ID @@ -842,19 +873,19 @@ def read_netcdf_transport( nx = fileID.dimensions['nx'].size ny = fileID.dimensions['ny'].size # real and imaginary components of transport - tr = np.ma.zeros((ny,nx), dtype=np.complex64) - tr.mask = np.zeros((ny,nx), dtype=bool) + hc = np.ma.zeros((ny,nx), dtype=np.complex64) + hc.mask = np.zeros((ny,nx), dtype=bool) if variable in ('U','u'): - tr.data.real[:,:] = fileID.variables['uRe'][:,:].T - tr.data.imag[:,:] = fileID.variables['uIm'][:,:].T + hc.data.real[:,:] = fileID.variables['uRe'][:,:].T + hc.data.imag[:,:] = fileID.variables['uIm'][:,:].T elif variable in ('V','v'): - tr.data.real[:,:] = fileID.variables['vRe'][:,:].T - tr.data.imag[:,:] = fileID.variables['vIm'][:,:].T + hc.data.real[:,:] = fileID.variables['vRe'][:,:].T + hc.data.imag[:,:] = fileID.variables['vIm'][:,:].T # close the file fileID.close() f.close() if kwargs['compressed'] else None # return the transport components and constituent - return (tr, con.strip()) + return (hc, con.strip()) # PURPOSE: output grid file in ATLAS netCDF format def output_netcdf_grid( @@ -870,7 +901,7 @@ def output_netcdf_grid( lat_v: np.ndarray ): """ - Writes grid files in ATLAS netCDF format + Writes grid parameters to netCDF4 files in ATLAS format Parameters ---------- @@ -964,7 +995,7 @@ def output_netcdf_elevation( constituent: str ): """ - Writes elevation files in ATLAS netCDF format + Writes elevation constituents to netCDF4 files in ATLAS format Parameters ---------- @@ -1051,7 +1082,7 @@ def output_netcdf_transport( constituent: str ): """ - Writes transport files in ATLAS netCDF format + Writes transport constituents to netCDF4 files in ATLAS format Parameters ---------- diff --git a/pyTMD/io/FES.py b/pyTMD/io/FES.py index 6600053e..9e0cf851 100644 --- a/pyTMD/io/FES.py +++ b/pyTMD/io/FES.py @@ -747,7 +747,7 @@ def output_netcdf_file( **kwargs ): """ - Writes tidal constituent files in FES2014 netCDF format + Writes tidal constituents to netCDF4 files in FES2014 format Parameters ---------- diff --git a/pyTMD/io/GOT.py b/pyTMD/io/GOT.py index 734bb2b8..ba253b48 100644 --- a/pyTMD/io/GOT.py +++ b/pyTMD/io/GOT.py @@ -678,7 +678,7 @@ def output_netcdf_file( constituent: str ): """ - Writes tidal constituent files in GOT netCDF format + Writes tidal constituents to netCDF4 files in GOT format Parameters ---------- diff --git a/pyTMD/io/OTIS.py b/pyTMD/io/OTIS.py index 3c3ed504..59dd2d28 100644 --- a/pyTMD/io/OTIS.py +++ b/pyTMD/io/OTIS.py @@ -1,7 +1,7 @@ #!/usr/bin/env python u""" OTIS.py -Written by Tyler Sutterley (09/2023) +Written by Tyler Sutterley (10/2023) Reads files for a tidal model and makes initial calculations to run tide program Includes functions to extract tidal harmonic constants from OTIS tide models for @@ -59,6 +59,7 @@ interpolate.py: interpolation routines for spatial data UPDATE HISTORY: + Updated 10/2023: fix transport variable entry for TMD3 models Updated 09/2023: prevent overwriting ATLAS compact x and y coordinates Updated 08/2023: changed ESR netCDF4 format to TMD3 format Updated 04/2023: using pathlib to define and expand tide model paths @@ -384,7 +385,7 @@ def extract_constants( _,_,hc = combine_atlas_model(x0, y0, v0, pmask, uvlocal, variable='v') elif (kwargs['grid'] == 'TMD3'): - hc = read_netcdf_file(model_file, i, type='v') + hc = read_netcdf_file(model_file, i, variable='v') elif isinstance(model_file,list): u,hc = read_otis_transport(model_file[i], 0) else: @@ -618,7 +619,7 @@ def read_constants( _,_,hc = combine_atlas_model(x0, y0, v0, pmask, uvlocal, variable='v') elif (kwargs['grid'] == 'TMD3'): - hc = read_netcdf_file(model_file, i, type='v') + hc = read_netcdf_file(model_file, i, variable='v') elif isinstance(model_file,list): u,hc = read_otis_transport(model_file[i], 0) else: