Add comments

jmcvey3 · jmcvey3 · commit 8fcf7cf23f8e · 2025-04-25T14:16:08.000-07:00
diff --git a/mhkit/dolfyn/adp/discharge.py b/mhkit/dolfyn/adp/discharge.py
@@ -10,7 +10,9 @@ def discharge(ds, water_depth, rho, mu=None, surface_offset=0, utm_zone=10):
     natively handle ADCP datasets read in using the `dolfyn` module.
 
     Dataset velocity should already be corrected using ADCP-measured
-    bottom track or GPS-measured velocity.
+    bottom track or GPS-measured velocity. The first velocity direction
+    is assumed to be the primary flow axis.
+
     This function linearly interpolates the lowest ADCP depth bin to
     the seafloor, and applies a constant extrapolation from the first
     ADCP bin to the surface.
@@ -30,14 +32,15 @@ def discharge(ds, water_depth, rho, mu=None, surface_offset=0, utm_zone=10):
     rho: float
         Water density in kg/m^3
     mu: float
-        Kinematic visocity based on water temperature and salinity, in Ns/m^2.
-        Required for Reynolds Number.
+        Dynamic visocity based on water temperature and salinity, in Ns/m^2.
+        If not provided, Reynolds Number will not be calculated.
+        Default: None.
     surface_offset: float
         Surface level offset due to changes in tidal level, in meters.
-        Default: 0 m.
+        Positive is down. Default: 0 m.
     utm_zone: int
-        UTM zone that measurements were acquired in. Map of UTM zones for the
-        contiguous US:
+        UTM zone for coordinate transformations (e.g., to compute cross-sectional
+        distances from GPS lat/lon data). Map of UTM zones for the contiguous US:
         https://www.usgs.gov/media/images/mapping-utm-grid-conterminous-48-united-states.
         Default: 10 (the US west coast).
 
@@ -51,7 +54,38 @@ def discharge(ds, water_depth, rho, mu=None, surface_offset=0, utm_zone=10):
          - `reynolds_number`: (1) Reynolds number, unitless
     """
 
-    def extrap2bottom(vel, bottom, rng):
+    def _extrapolate_to_bottom(vel, bottom, rng):
+        """
+        Linearly extrapolate velocity values from the deepest valid bin down to zero at the seafloor.
+
+        This function sets velocity to zero at the seafloor and linearly interpolates
+        between the last valid velocity bin and this zero-velocity boundary. If no valid
+        velocity is found in a particular profile, no update is performed for that profile.
+        This function assumes `rng` extends at least to (or below) the deepest seafloor depth
+        specified in `bottom`.
+
+        Parameters
+        ----------
+        vel : numpy.ndarray
+            A velocity array of shape (dir, range, time), typically containing:
+                - `dir` : velocity component dimension (e.g., 2 or 3 for 2D or 3D flow).
+                - `range` : vertical/bin dimension (positive downward).
+                - `time` : time dimension corresponding to each profile.
+            The array is modified in-place (the updated values are also returned).
+        bottom : array-like
+            Array of length equal to the time dimension in `vel`, specifying the seafloor
+            depth (in the same coordinate system as `rng`) at each time step.
+        rng : array-like
+            The vertical/bin positions corresponding to `vel` along the `range` dimension,
+            sorted in ascending order (e.g., depth from the water surface downward).
+
+        Returns
+        -------
+        vel : numpy.ndarray
+            The same array passed in, with updated values below the last valid velocity bin
+            for each time step (linear extrapolation to zero at the seafloor).
+        """
+
         for idx in range(vel.shape[-1]):
             z_bot = bottom[idx]
             # Fetch lowest range index
@@ -71,7 +105,47 @@ def extrap2bottom(vel, bottom, rng):
 
         return vel
 
-    def latlon2utm(ds, proj):
+    def _convert_latlon_to_utm(ds, proj):
+        """
+        Convert latitude/longitude coordinates to UTM coordinates.
+
+        This function uses the Cartopy `transform_point` and `transform_points` methods to
+        project GPS latitude/longitude data into the specified UTM coordinate reference
+        system. The resulting (x, y) coordinates are stored in an xarray DataArray that is
+        interpolated onto the main time axis of `ds`.
+
+        The function sets `proj.x0` and `proj.y0` to the UTM coordinates of the mean
+        longitude and latitude from `ds`. This can be used as a reference origin.
+        Missing or NaN lat/lon values are handled via interpolation and extrapolation
+        onto the `ds["time"]` axis.
+        This function modifies the `proj` object by adding `x0` and `y0` attributes,
+        which may be used for subsequent coordinate transformations or offsets.
+
+        Parameters
+        ----------
+        ds : xarray.Dataset
+            A dataset that must contain at least the following variables:
+            - "latitude_gps"  : (time_gps) latitude values in degrees North.
+            - "longitude_gps" : (time_gps) longitude values in degrees East.
+            - "time"          : time axis onto which the projected coordinates will be
+                                interpolated.
+        proj : cartopy.crs.Projection
+            A Cartopy UTM projection or similar projection object. This is used both to
+            store the reference origin (`x0`, `y0`) and to transform lat/lon coordinates
+            into UTM.
+
+        Returns
+        -------
+        xy : xarray.DataArray
+            A DataArray of shape (gps=2, time), where:
+            - The first dimension (indexed by "gps") corresponds to ["x", "y"] UTM
+                coordinates.
+            - The second dimension ("time") matches `ds["time"]`.
+            The returned coordinates are interpolated in time using `ds["longitude_gps"]`
+            and `ds["latitude_gps"]`, with values extrapolated if necessary.
+
+        """
+
         PlateC = ccrs.PlateCarree()
         proj.x0, proj.y0 = proj.transform_point(
             ds["longitude_gps"].mean(), ds["latitude_gps"].mean(), PlateC
@@ -90,11 +164,59 @@ def latlon2utm(ds, proj):
         return xy
 
     def _distance(proj, x, y):
-        # GPS distance traveled in meters
+        """
+        Compute the planar distance from the projection's reference origin.
+
+        Parameters
+        ----------
+        proj : cartopy.crs.Projection
+            A projection object with attributes `x0` and `y0`, which define the
+            reference origin in the projected coordinate system.
+        x : float or array-like
+            One or more x-coordinates in the same units (m) as `proj.x0`.
+        y : float or array-like
+            One or more y-coordinates in the same units (m) as `proj.y0`.
+
+        Returns
+        -------
+        dist : float or numpy.ndarray
+            The distance(s) in m from the point(s) `(x, y)` to `(proj.x0, proj.y0)`.
+            If `x` and `y` are arrays, the output is an array of the same shape.
+        """
+
         return np.sqrt((proj.x0 - x) ** 2 + (proj.y0 - y) ** 2)
 
-    def calc_Q(vel, x, depth, surface_zoff=None):
-        # depth and surface_zoff should be positive in down direction
+    def _calc_discharge(vel, x, depth, surface_zoff=None):
+        """
+        Calculate the integrated flux (e.g., discharge) by double integration of velocity
+        over the cross-sectional area: depth and lateral distance.
+
+        Missing (NaN) velocities are treated as zero.
+        Ensure `depth` and `surface_zoff` are both positive downward.
+
+        Parameters
+        ----------
+        vel : numpy.ndarray or xarray.DataArray
+            A 2D array of shape (nz, nx) corresponding to velocity values (m/s).
+            - `nz` is the number of vertical bins (downward).
+            - `nx` is the number of horizontal points.
+        x : array-like
+            Horizontal positions (m) of length `nx`. If `x` is in descending order
+            (i.e., `x[0] > x[-1]`), the resulting flux is assigned a negative sign to
+            indicate reverse orientation.
+        depth : array-like
+            Vertical positions (m) of length `nz`, positive downward. This is used
+            for integration along the vertical dimension.
+        surface_zoff : float, optional
+            Surface level offset due to changes in tidal level, in meters.
+            Positive is down.
+
+        Returns
+        -------
+        Q : float
+            The integrated flux (e.g., discharge) in units of m^3/s
+
+        """
         vel = vel.copy()
         vel = vel.fillna(0)
         if surface_zoff is not None:
@@ -109,37 +231,44 @@ def calc_Q(vel, x, depth, surface_zoff=None):
 
     # Extrapolate to bed
     vel = ds["vel"].copy()
-    vel.values = extrap2bottom(ds["vel"].values, water_depth, ds["range"].values)
+    vel.values = _extrapolate_to_bottom(
+        ds["vel"].values, water_depth, ds["range"].values
+    )
     vel_x = vel[0]
     # Get position at each timestep in UTM grid
     proj = ccrs.UTM(utm_zone)
-    xy = latlon2utm(ds, proj)
+    xy = _convert_latlon_to_utm(ds, proj)
     # Distance from UTM grid origin (mean of GPS points)
     _x = _distance(proj, xy[0], xy[1])
     # Set distance range for entire transect
-    Q_x_range = [_x.min(), _x.max()]  # meters
+    q_x_range = [_x.min(), _x.max()]  # meters
 
     # Calculate discharge, power, kinetic energy, and reynolds number
-    _xinds = (Q_x_range[0] < _x) & (_x < Q_x_range[1])
+    _xinds = (q_x_range[0] < _x) & (_x < q_x_range[1])
     out = {}
     if _xinds.any():
-        U = vel_x[:, _xinds]  # m/s
+        speed = vel_x[:, _xinds]  # m/s
         # Volume Flux, aka Discharge
-        out["Q"] = calc_Q(
-            U, xy[0][_xinds], ds["range"], surface_offset
+        out["Q"] = _calc_discharge(
+            speed, xy[0][_xinds], ds["range"], surface_offset
         )  # m/s * m * m = m^3/s
         # Kinetic Energy Flux, aka Power
         out["P"] = (
-            0.5 * rho * calc_Q(U**3, xy[0][_xinds], ds["range"], surface_offset)
+            0.5
+            * rho
+            * _calc_discharge(speed**3, xy[0][_xinds], ds["range"], surface_offset)
         )  # kg/m^3 * m^3/s^3 * m * m = kg*m^2/s = W
         # Power Density
-        out["J"] = (0.5 * rho * U**3).mean().item()  # kg/m^3 * m^3/s^3 = kg/s^3 = W/m^2
-        # Hydraulic Depth
-        L = abs(np.trapz((water_depth - surface_offset)[_xinds], xy[0][_xinds])) / (
+        out["J"] = (
+            (0.5 * rho * speed**3).mean().item()
+        )  # kg/m^3 * m^3/s^3 = kg/s^3 = W/m^2
+        hydraulic_depth = abs(
+            np.trapz((water_depth - surface_offset)[_xinds], xy[0][_xinds])
+        ) / (
             xy[0][_xinds].max() - xy[0][_xinds].min()
         )  # area / surface-width
         # Reynolds Number
-        out["Re"] = ((rho * ds.velds.U_mag.mean() * L) / mu).item()
+        out["Re"] = ((rho * ds.velds.U_mag.mean() * hydraulic_depth) / mu).item()
     else:
         out["Q"] = np.nan
         out["P"] = np.nan
