Fix v022 imports (#343)

* move troposphere-specific interpolation to its subpackage

* re-add init imports

* change interp references

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

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,27 +467,27 @@ 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)
         self._ys = np.unique(self._ys)
 
     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."""

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