Merge pull request #8 from cassidymwagner/ltt_sfs

New traditional structure functions

Author: cassidymwagner

oceans_sf/calculate_advection.py

@@ -44,17 +44,7 @@ def calculate_advection(  # noqa: D417
             northward_advection) if scalar is not provided, otherwise returns an ndarray
             of scalar advection.
     """
-    if grid_type == "latlon":
-        xcoords = dx.cumsum()
-        ycoords = dy.cumsum()
-
-        if scalar is not None:
-            dsdy, dsdx = np.gradient(scalar, xcoords, ycoords, axis=(0, 1))
-        else:
-            dudy, dudx = np.gradient(u, xcoords, ycoords, axis=(0, 1))
-            dvdy, dvdx = np.gradient(v, xcoords, ycoords, axis=(0, 1))
-
-    else:
+    if grid_type == "uniform":
         dx = np.abs(x[0] - x[1])
         dy = np.abs(y[0] - y[1])
 
@@ -64,6 +54,16 @@ def calculate_advection(  # noqa: D417
             dudy, dudx = np.gradient(u, dx, dy, axis=(0, 1))
             dvdy, dvdx = np.gradient(v, dx, dy, axis=(0, 1))
 
+    else:
+        xcoords = dx.cumsum()
+        ycoords = dy.cumsum()
+
+        if scalar is not None:
+            dsdy, dsdx = np.gradient(scalar, xcoords, ycoords, axis=(0, 1))
+        else:
+            dudy, dudx = np.gradient(u, xcoords, ycoords, axis=(0, 1))
+            dvdy, dvdx = np.gradient(v, xcoords, ycoords, axis=(0, 1))
+
     if scalar is not None:
         advection = u * dsdx + v * dsdy
     else:

oceans_sf/calculate_structure_function.py

@@ -3,7 +3,7 @@
 from .shift_array2d import shift_array2d
 
 
-def calculate_structure_function(  # noqa: D417
+def calculate_structure_function(  # noqa: D417, C901
     u,
     v,
     adv_e,
@@ -13,8 +13,8 @@ def calculate_structure_function(  # noqa: D417
     skip_velocity_sf=False,
     scalar=None,
     adv_scalar=None,
-    traditional_order=0,
-    boundary="Periodic",
+    traditional_type=None,
+    boundary="periodic-all",
 ):
     """
     Calculate structure function, either advective or traditional.
@@ -45,11 +45,14 @@ def calculate_structure_function(  # noqa: D417
             Array of scalar values. Defaults to None.
         adv_scalar: ndarray, optional
             Array of scalar advection values. Defaults to None.
-        traditional_order: int, optional
-            Order for calculating traditional non-advective structure functions.
-            If 0, no traditional structure functions are calculated. Defaults to 0.
+        traditional_type: list, optional
+            List of traditional structure function types to calculate.
+            Accepted types are: "LL", "LLL", "LTT", "LSS". If None,
+            no traditional structure functions are calculated. Defaults to None.
         boundary: str, optional
-            Boundary condition for shifting arrays. Defaults to "Periodic".
+            Boundary condition for shifting arrays. Accepted strings
+            are "periodic-x", "periodic-y", and "periodic-all".
+            Defaults to "periodic-all".
 
     Returns
     -------
@@ -119,20 +122,37 @@ def calculate_structure_function(  # noqa: D417
                 + (inputs["adv_n_" + direction + "_shift"] - adv_n)
                 * (inputs["v_" + direction + "_shift"] - v)
             )
-            if traditional_order > 0:
-                N = traditional_order
-                SF_dict["SF_trad_velocity_" + direction] = np.nanmean(
-                    (inputs["u_" + direction + "_shift"] - u) ** N
-                )
+            if traditional_type is not None:
+                if any("LL" in t for t in traditional_type):
+                    SF_dict["SF_LL_" + direction] = np.nanmean(
+                        (inputs["u_" + direction + "_shift"] - u) ** 2
+                    )
+                if any("LLL" in t for t in traditional_type):
+                    SF_dict["SF_LLL_" + direction] = np.nanmean(
+                        (inputs["u_" + direction + "_shift"] - u) ** 3
+                    )
+                if any("LTT" in t for t in traditional_type):
+                    if direction == "right":
+                        SF_dict["SF_LTT_" + direction] = np.nanmean(
+                            (inputs["u_" + direction + "_shift"] - u)
+                            * (inputs["v_" + direction + "_shift"] - v) ** 2
+                        )
+
+                    if direction == "down":
+                        SF_dict["SF_LTT_" + direction] = np.nanmean(
+                            (inputs["v_" + direction + "_shift"] - v)
+                            * (inputs["u_" + direction + "_shift"] - u) ** 2
+                        )
 
         if scalar is not None:
             SF_dict["SF_scalar_" + direction] = np.nanmean(
                 (inputs["adv_scalar_" + direction + "_shift"] - adv_scalar)
                 * (inputs["scalar_" + direction + "_shift"] - scalar)
             )
-            if traditional_order > 0:
-                N = traditional_order
-                SF_dict["SF_trad_scalar_" + direction] = np.nanmean(
-                    (inputs["scalar_" + direction + "_shift"] - scalar) ** N
-                )
+            if traditional_type is not None:
+                if any("LSS" in t for t in traditional_type):
+                    SF_dict["SF_LSS_" + direction] = np.nanmean(
+                        (inputs["u_" + direction + "_shift"] - u)
+                        * (inputs["scalar_" + direction + "_shift"] - scalar) ** 2
+                    )
     return SF_dict

oceans_sf/generate_structure_functions.py

@@ -16,10 +16,10 @@ def generate_structure_functions(  # noqa: C901, D417
     y,
     skip_velocity_sf=False,
     scalar=None,
-    traditional_order=0,
+    traditional_type=None,
     dx=None,
     dy=None,
-    boundary="Periodic",
+    boundary="periodic-all",
     even="True",
     grid_type="uniform",
     nbins=10,
@@ -47,15 +47,17 @@ def generate_structure_functions(  # noqa: C901, D417
             Defaults to False.
         scalar: ndarray, optional
             2D array of scalar values. Defaults to None.
-        traditional_order: int, optional
-            Order for calculating traditional non-advective structure functions.
-            If 0, no traditional structure functions are calculated. Defaults to 0.
+        traditional_type: list, optional
+            List of traditional structure function types to calculate.
+            Accepted types are: "LL", "LLL", "LTT", "LSS". If None,
+            no traditional structure functions are calculated. Defaults to None.
         dx: float, optional
             Grid spacing in the x-direction. Defaults to None.
         dy: float, optional
             Grid spacing in the y-direction. Defaults to None.
         boundary: str, optional
-            Boundary condition of the data. Defaults to "Periodic".
+            Boundary condition of the data. Accepted strings are "periodic-x",
+            "periodic-y", and "periodic-all". Defaults to "periodic-all".
         even: bool, optional
             Flag indicating if the grid is evenly spaced. Defaults to True.
         grid_type:str, optional
@@ -75,20 +77,30 @@ def generate_structure_functions(  # noqa: C901, D417
     SF_m = None
     SF_z_scalar = None
     SF_m_scalar = None
-    SF_z_trad_velocity = None
-    SF_m_trad_velocity = None
-    SF_z_trad_scalar = None
-    SF_m_trad_scalar = None
     adv_E = None
     adv_N = None
     adv_scalar = None
+    SF_z_LL = None
+    SF_m_LL = None
+    SF_z_LLL = None
+    SF_m_LLL = None
+    SF_z_LTT = None
+    SF_m_LTT = None
+    SF_z_LSS = None
+    SF_m_LSS = None
 
     # Define a list of separation distances to iterate over.
     # Periodic is half the length since the calculation will wrap the data.
-    if boundary == "Periodic":
+    if boundary == "periodic-all":
         sep_z = range(1, int(len(x) / 2))
         sep_m = range(1, int(len(y) / 2))
-    else:
+    elif boundary == "periodic-x":
+        sep_z = range(1, int(len(x) / 2))
+        sep_m = range(1, int(len(y) - 1))
+    elif boundary == "periodic-y":
+        sep_z = range(1, int(len(x) - 1))
+        sep_m = range(1, int(len(y) / 2))
+    elif boundary is None:
         sep_z = range(1, int(len(x) - 1))
         sep_m = range(1, int(len(y) - 1))
 
@@ -101,22 +113,33 @@ def generate_structure_functions(  # noqa: C901, D417
         SF_z = np.zeros(len(sep_z) + 1)
         SF_m = np.zeros(len(sep_m) + 1)
         adv_E, adv_N = calculate_advection(u, v, x, y, dx, dy, grid_type)
-        if traditional_order > 0:
-            SF_z_trad_velocity = np.zeros(len(sep_z) + 1)
-            SF_m_trad_velocity = np.zeros(len(sep_m) + 1)
+        if traditional_type is not None:
+            if any("LL" in t for t in traditional_type):
+                SF_z_LL = np.zeros(len(sep_z) + 1)
+                SF_m_LL = np.zeros(len(sep_m) + 1)
+            if any("LLL" in t for t in traditional_type):
+                SF_z_LLL = np.zeros(len(sep_z) + 1)
+                SF_m_LLL = np.zeros(len(sep_m) + 1)
+            if any("LTT" in t for t in traditional_type):
+                SF_z_LTT = np.zeros(len(sep_z) + 1)
+                SF_m_LTT = np.zeros(len(sep_m) + 1)
 
     if scalar is not None:
         SF_z_scalar = np.zeros(len(sep_z) + 1)
         SF_m_scalar = np.zeros(len(sep_m) + 1)
         adv_scalar = calculate_advection(u, v, x, y, dx, dy, grid_type, scalar)
-        if traditional_order > 0:
-            SF_z_trad_scalar = np.zeros(len(sep_z) + 1)
-            SF_m_trad_scalar = np.zeros(len(sep_m) + 1)
+        if traditional_type is not None:
+            if any("LSS" in t for t in traditional_type):
+                SF_z_LSS = np.zeros(len(sep_z) + 1)
+                SF_m_LSS = np.zeros(len(sep_m) + 1)
 
     # Iterate over separations right and down
     for down in sep_m:
         right = 1
-        yroll = shift_array1d(y, shift_by=down, boundary=boundary)
+        if boundary == "periodic-all" or boundary == "periodic-y":
+            yroll = shift_array1d(y, shift_by=down, boundary="Periodic")
+        else:
+            yroll = shift_array1d(y, shift_by=down, boundary=None)
 
         SF_dicts = calculate_structure_function(
             u,
@@ -128,18 +151,24 @@ def generate_structure_functions(  # noqa: C901, D417
             skip_velocity_sf,
             scalar,
             adv_scalar,
-            traditional_order,
+            traditional_type,
             boundary,
         )
 
         if skip_velocity_sf is False:
             SF_m[down] = SF_dicts["SF_velocity_down"]
-            if traditional_order > 0:
-                SF_m_trad_velocity[down] = SF_dicts["SF_trad_velocity_down"]
+            if traditional_type is not None:
+                if any("LL" in t for t in traditional_type):
+                    SF_m_LL[down] = SF_dicts["SF_LL_down"]
+                if any("LLL" in t for t in traditional_type):
+                    SF_m_LLL[down] = SF_dicts["SF_LLL_down"]
+                if any("LTT" in t for t in traditional_type):
+                    SF_m_LTT[down] = SF_dicts["SF_LTT_down"]
         if scalar is not None:
             SF_m_scalar[down] = SF_dicts["SF_scalar_down"]
-            if traditional_order > 0:
-                SF_m_trad_scalar[down] = SF_dicts["SF_trad_scalar_down"]
+            if traditional_type is not None:
+                if any("LSS" in t for t in traditional_type):
+                    SF_m_LSS[down] = SF_dicts["SF_LSS_down"]
 
         # Calculate separation distances in y
         tmp, yd[down] = calculate_separation_distances(
@@ -148,7 +177,10 @@ def generate_structure_functions(  # noqa: C901, D417
 
     for right in sep_z:
         down = 1
-        xroll = shift_array1d(x, shift_by=right, boundary=boundary)
+        if boundary == "periodic-all" or boundary == "periodic-x":
+            xroll = shift_array1d(x, shift_by=right, boundary="Periodic")
+        else:
+            xroll = shift_array1d(x, shift_by=right, boundary=None)
 
         SF_dicts = calculate_structure_function(
             u,
@@ -160,18 +192,24 @@ def generate_structure_functions(  # noqa: C901, D417
             skip_velocity_sf,
             scalar,
             adv_scalar,
-            traditional_order,
+            traditional_type,
             boundary,
         )
 
         if skip_velocity_sf is False:
             SF_z[right] = SF_dicts["SF_velocity_right"]
-            if traditional_order > 0:
-                SF_z_trad_velocity[right] = SF_dicts["SF_trad_velocity_right"]
+            if traditional_type is not None:
+                if any("LL" in t for t in traditional_type):
+                    SF_z_LL[right] = SF_dicts["SF_LL_right"]
+                if any("LLL" in t for t in traditional_type):
+                    SF_z_LLL[right] = SF_dicts["SF_LLL_right"]
+                if any("LTT" in t for t in traditional_type):
+                    SF_z_LTT[right] = SF_dicts["SF_LTT_right"]
         if scalar is not None:
             SF_z_scalar[right] = SF_dicts["SF_scalar_right"]
-            if traditional_order > 0:
-                SF_z_trad_scalar[right] = SF_dicts["SF_trad_scalar_right"]
+            if traditional_type is not None:
+                if any("LSS" in t for t in traditional_type):
+                    SF_z_LSS[right] = SF_dicts["SF_LSS_right"]
 
         # Calculate separation distances in x
         xd[right], tmp = calculate_separation_distances(
@@ -183,15 +221,23 @@ def generate_structure_functions(  # noqa: C901, D417
         if skip_velocity_sf is False:
             xd_bin, SF_z = bin_data(xd, SF_z, nbins)
             yd_bin, SF_m = bin_data(yd, SF_m, nbins)
-            if traditional_order > 0:
-                xd_bin, SF_z_trad_velocity = bin_data(xd, SF_z_trad_velocity, nbins)
-                yd_bin, SF_m_trad_velocity = bin_data(yd, SF_m_trad_velocity, nbins)
+            if traditional_type is not None:
+                if any("LL" in t for t in traditional_type):
+                    xd_bin, SF_z_LL = bin_data(xd, SF_z_LL, nbins)
+                    yd_bin, SF_m_LL = bin_data(yd, SF_m_LL, nbins)
+                if any("LLL" in t for t in traditional_type):
+                    xd_bin, SF_z_LLL = bin_data(xd, SF_z_LLL, nbins)
+                    yd_bin, SF_m_LLL = bin_data(yd, SF_m_LLL, nbins)
+                if any("LTT" in t for t in traditional_type):
+                    xd_bin, SF_z_LTT = bin_data(xd, SF_z_LTT, nbins)
+                    yd_bin, SF_m_LTT = bin_data(yd, SF_m_LTT, nbins)
         if scalar is not None:
             xd_bin, SF_z_scalar = bin_data(xd, SF_z_scalar, nbins)
             yd_bin, SF_m_scalar = bin_data(yd, SF_m_scalar, nbins)
-            if traditional_order > 0:
-                xd_bin, SF_z_trad_scalar = bin_data(xd, SF_z_trad_scalar, nbins)
-                yd_bin, SF_m_trad_scalar = bin_data(yd, SF_m_trad_scalar, nbins)
+            if traditional_type is not None:
+                if any("LSS" in t for t in traditional_type):
+                    xd_bin, SF_z_LSS = bin_data(xd, SF_z_LSS, nbins)
+                    yd_bin, SF_m_LSS = bin_data(yd, SF_m_LSS, nbins)
         xd = xd_bin
         yd = yd_bin
 
@@ -200,10 +246,14 @@ def generate_structure_functions(  # noqa: C901, D417
         "SF_advection_velocity_meridional": SF_m,
         "SF_advection_scalar_zonal": SF_z_scalar,
         "SF_advection_scalar_meridional": SF_m_scalar,
-        "SF_traditional_velocity_zonal": SF_z_trad_velocity,
-        "SF_traditional_velocity_meridional": SF_m_trad_velocity,
-        "SF_traditional_scalar_zonal": SF_z_trad_scalar,
-        "SF_traditional_scalar_meridional": SF_m_trad_scalar,
+        "SF_LL_zonal": SF_z_LL,
+        "SF_LL_meridional": SF_m_LL,
+        "SF_LLL_zonal": SF_z_LLL,
+        "SF_LLL_meridional": SF_m_LLL,
+        "SF_LTT_zonal": SF_z_LTT,
+        "SF_LTT_meridional": SF_m_LTT,
+        "SF_LSS_zonal": SF_z_LSS,
+        "SF_LSS_meridional": SF_m_LSS,
         "x-diffs": xd,
         "y-diffs": yd,
     }