From d7a5358199be5346c89ce1225d3c839db92e54f4 Mon Sep 17 00:00:00 2001 From: Scott Staniewicz Date: Wed, 17 Jul 2024 06:46:02 -0400 Subject: [PATCH] Fix v022 imports (#343) * move troposphere-specific interpolation to its subpackage * re-add init imports * change interp references --- src/dolphin/atmosphere/_utils.py | 78 ++++++++++++++++++++++++- src/dolphin/atmosphere/weather_model.py | 25 ++++---- src/dolphin/interpolation.py | 77 ------------------------ 3 files changed, 89 insertions(+), 91 deletions(-) diff --git a/src/dolphin/atmosphere/_utils.py b/src/dolphin/atmosphere/_utils.py index c5abf71a..b08870fe 100644 --- a/src/dolphin/atmosphere/_utils.py +++ b/src/dolphin/atmosphere/_utils.py @@ -1,4 +1,6 @@ import numpy as np +import pandas as pd +from scipy.interpolate import interp1d g0 = np.float64(9.80665) # Standard gravitational constant G1 = np.float64( @@ -105,7 +107,6 @@ def calcgeoh( # Integrate up into the atmosphere from *lowest level* z_h = 0 # initial value for lev, t_level, q_level in zip(range(num_levels, 0, -1), t[::-1], q[::-1]): - # lev is the level number 1-60, we need a corresponding index # into ts and qs # ilevel = num_levels - lev # << this was Ray's original, but is a typo @@ -239,3 +240,78 @@ def _least_nonzero(a): """ mgrid_index = tuple(slice(None, d) for d in a.shape[:-1]) return a[(*tuple(np.mgrid[mgrid_index]), (~np.isnan(a)).argmax(-1))] + + +def interpolate_along_axis(oldCoord, newCoord, data, axis=2): + """Interpolate an array of 3-D data along one axis. This function. + + assumes that the x-coordinate increases monotonically. + """ + if oldCoord.ndim > 1: + stackedData = np.concatenate([oldCoord, data, newCoord], axis=axis) + out = np.apply_along_axis( + interp_vector, axis=axis, arr=stackedData, Nx=oldCoord.shape[axis] + ) + else: + out = np.apply_along_axis( + interp_v, + axis=axis, + arr=data, + old_x=oldCoord, + new_x=newCoord, + left=np.nan, + right=np.nan, + ) + + return out + + +def interp_vector(vec, Nx): + """Interpolate data from a single vector containing the original. + + x, the original y, and the new x, in that order. Nx tells the + number of original x-points. + """ + x = vec[:Nx] + y = vec[Nx : 2 * Nx] + xnew = vec[2 * Nx :] + f = interp1d(x, y, bounds_error=False, copy=False, assume_sorted=True) + return f(xnew) + + +def interp_v(y, old_x, new_x, left=None, right=None, period=None): + """Rearrange np.interp's arguments.""" + return np.interp(new_x, old_x, y, left=left, right=right, period=period) + + +def fillna3d(array: np.ndarray, axis: int = -1, fill_value: float = 0.0) -> np.ndarray: + """Fill in NaNs in 3D arrays using the nearest non-NaN value for "low" NaNs. + + and a specified fill value (default 0.0) for "high" NaNs. + + Parameters + ---------- + array : np.ndarray + 3D array, where the last axis is the "z" dimension. + axis : int, optional + The axis along which to fill values. Default is -1 (the last axis). + fill_value : float, optional + The value used for filling NaNs. Default is 0.0. + + Returns + ------- + np.ndarray + 3D array with low NaNs filled using nearest neighbors and high NaNs + filled with the specified fill value. + + """ + # fill lower NaNs with nearest neighbor + narr = np.moveaxis(array, axis, -1) + nars = narr.reshape((np.prod(narr.shape[:-1]), narr.shape[-1])) + dfd = pd.DataFrame(data=nars).interpolate(axis=1, limit_direction="backward") + out = dfd.values.reshape(array.shape) + + # fill upper NaNs with 0s + outmat = np.moveaxis(out, -1, axis) + outmat[np.isnan(outmat)] = fill_value + return outmat diff --git a/src/dolphin/atmosphere/weather_model.py b/src/dolphin/atmosphere/weather_model.py index acbaa246..5e01d061 100644 --- a/src/dolphin/atmosphere/weather_model.py +++ b/src/dolphin/atmosphere/weather_model.py @@ -12,7 +12,6 @@ import dolphin.atmosphere._utils as utils import dolphin.atmosphere.model_levels as ml -from dolphin.interpolation import fillna3d, interpolate_along_axis logger = logging.getLogger(__name__) @@ -468,15 +467,15 @@ def _uniform_in_z(self, _zlevels=None): new_zs = np.tile(_zlevels, (nx, ny, 1)) # re-assign values to the uniform z - self._t = interpolate_along_axis(self._zs, self._t, new_zs, axis=2).astype( - np.float32 - ) - self._p = interpolate_along_axis(self._zs, self._p, new_zs, axis=2).astype( - np.float32 - ) - self._e = interpolate_along_axis(self._zs, self._e, new_zs, axis=2).astype( - np.float32 - ) + self._t = utils.interpolate_along_axis( + self._zs, self._t, new_zs, axis=2 + ).astype(np.float32) + self._p = utils.interpolate_along_axis( + self._zs, self._p, new_zs, axis=2 + ).astype(np.float32) + self._e = utils.interpolate_along_axis( + self._zs, self._e, new_zs, axis=2 + ).astype(np.float32) self._zs = _zlevels self._xs = np.unique(self._xs) @@ -484,11 +483,11 @@ def _uniform_in_z(self, _zlevels=None): def _check_for_nans(self): """Fill in NaN-values.""" - self._p = fillna3d(self._p) - self._t = fillna3d( + self._p = utils.fillna3d(self._p) + self._t = utils.fillna3d( self._t, fill_value=1e16 ) # to avoid division by zero later on - self._e = fillna3d(self._e) + self._e = utils.fillna3d(self._e) def out_file(self, outLoc): """Path to the output file.""" diff --git a/src/dolphin/interpolation.py b/src/dolphin/interpolation.py index cfccf3a3..f82e3e56 100644 --- a/src/dolphin/interpolation.py +++ b/src/dolphin/interpolation.py @@ -2,9 +2,7 @@ import numba import numpy as np -import pandas as pd from numpy.typing import ArrayLike -from scipy.interpolate import interp1d from .similarity import get_circle_idxs @@ -137,78 +135,3 @@ def _interp_loop( csum += np.exp(-r2[i] / r2_norm) * cphase[i] interpolated_ifg[r0, c0] = np.abs(ifg[r0, c0]) * np.exp(1j * np.angle(csum)) - - -def interpolate_along_axis(oldCoord, newCoord, data, axis=2): - """Interpolate an array of 3-D data along one axis. This function. - - assumes that the x-coordinate increases monotonically. - """ - if oldCoord.ndim > 1: - stackedData = np.concatenate([oldCoord, data, newCoord], axis=axis) - out = np.apply_along_axis( - interp_vector, axis=axis, arr=stackedData, Nx=oldCoord.shape[axis] - ) - else: - out = np.apply_along_axis( - interp_v, - axis=axis, - arr=data, - old_x=oldCoord, - new_x=newCoord, - left=np.nan, - right=np.nan, - ) - - return out - - -def interp_vector(vec, Nx): - """Interpolate data from a single vector containing the original. - - x, the original y, and the new x, in that order. Nx tells the - number of original x-points. - """ - x = vec[:Nx] - y = vec[Nx : 2 * Nx] - xnew = vec[2 * Nx :] - f = interp1d(x, y, bounds_error=False, copy=False, assume_sorted=True) - return f(xnew) - - -def interp_v(y, old_x, new_x, left=None, right=None, period=None): - """Rearrange np.interp's arguments.""" - return np.interp(new_x, old_x, y, left=left, right=right, period=period) - - -def fillna3d(array: np.ndarray, axis: int = -1, fill_value: float = 0.0) -> np.ndarray: - """Fill in NaNs in 3D arrays using the nearest non-NaN value for "low" NaNs. - - and a specified fill value (default 0.0) for "high" NaNs. - - Parameters - ---------- - array : np.ndarray - 3D array, where the last axis is the "z" dimension. - axis : int, optional - The axis along which to fill values. Default is -1 (the last axis). - fill_value : float, optional - The value used for filling NaNs. Default is 0.0. - - Returns - ------- - np.ndarray - 3D array with low NaNs filled using nearest neighbors and high NaNs - filled with the specified fill value. - - """ - # fill lower NaNs with nearest neighbor - narr = np.moveaxis(array, axis, -1) - nars = narr.reshape((np.prod(narr.shape[:-1]), narr.shape[-1])) - dfd = pd.DataFrame(data=nars).interpolate(axis=1, limit_direction="backward") - out = dfd.values.reshape(array.shape) - - # fill upper NaNs with 0s - outmat = np.moveaxis(out, -1, axis) - outmat[np.isnan(outmat)] = fill_value - return outmat