Beginning of fig 2a validation and fixed api

Author: mdmeeker

examples-unrendered/2d_plus_3d_prototype.py

@@ -1,4 +1,5 @@
 import numpy as np
+import scipy.special as sp
 from low_freq_prototype import LowFreq2DFieldGenerator
 
 from drdmannturb import FluctuationFieldGenerator
@@ -69,7 +70,7 @@
 
 ## New API
 
-generator = LowFreq2DFieldGenerator(grid_dimensions, grid_levels, L2D=L_2D, sigma2=sigma2, z_i=z_i, psi_degs=psi_degs)
+generator = LowFreq2DFieldGenerator(grid_dimensions, grid_levels, L_2D=L_2D, sigma2=sigma2, z_i=z_i, psi_degs=psi_degs)
 
 _, _, u_field = generator.generate()
 _, _, v_field = generator.generate()
@@ -78,29 +79,92 @@
 fluctuation_field_2d3d[:, :, :, 0] += u_field[:, :, np.newaxis]
 fluctuation_field_2d3d[:, :, :, 1] += v_field[:, :, np.newaxis]
 
+#########################################################################################
+# Test Mean, Std, Dev, Skewness, and Kurtosis
+
 
 #########################################################################################
 ## TESTS
 # Helper functions for tests
 
 
-def analytic_F11_2D(k_1):
-    # first_num = (sp.gamma(11/6) * (L_2D**(11/3)))
+def analytic_F11_2D(k1):
+    """
+    Analytic solution for F11(k1) in 2d.
+    """
+    a = 1 + 2 * (k1 * generator.L_2D * np.cos(generator.psi_rad)) ** 2
+    b = 1 + 2 * (k1 * generator.z_i * np.cos(generator.psi_rad)) ** 2
+    p = (L_2D**2 * b) / (z_i**2 * a)
 
-    pass
+    d = 1.0
 
+    first_term_numerator = (sp.gamma(11 / 6) * (L_2D ** (11 / 3))) * (
+        -p * sp.hyp2f1(5 / 6, 1, 1 / 2, p) - 7 * sp.hyp2f1(5 / 6) + 2 * sp.hyp2f1(-1 / 6, 1, 1 / 2, p)
+    )
+    first_term_denominator = (2 * np.pi) ** (11 / 3) * (a ** (11 / 6))
 
-def analytic_F22_2D(k_1):
-    pass
+    second_term_numerator = L_2D ** (14 / 3) * np.sqrt(b)
+    second_term_denominator = 2 * np.sqrt(2) * d ** (7 / 3) * (generator.z_i * np.sin(generator.psi_rad))
 
+    return generator.c * (
+        (first_term_numerator / first_term_denominator) + (second_term_numerator / second_term_denominator)
+    )
+
+
+def analytic_F22_2D(k1):
+    a = 1 + 2 * (k1 * generator.L_2D * np.cos(generator.psi_rad)) ** 2
+    b = 1 + 2 * (k1 * generator.z_i * np.cos(generator.psi_rad)) ** 2
+    p = (L_2D**2 * b) / (z_i**2 * a)
+
+    d = 1.0
+
+    leading_factor_num = generator.z_i**4 * a ** (1 / 6) ** generator.L_2D * sp.gamma(17 / 6)
+    leading_factor_denom = (
+        55
+        * np.sqrt(2 * np.pi)
+        * (generator.L_2D**2 - generator.z_i**2) ** 2
+        * b
+        * sp.gamma(7 / 3)
+        * np.sin(generator.psi_rad)
+    )
+    leading_factor = -leading_factor_num / leading_factor_denom
+
+    line_1 = -9 - 25 * sp.hyp2f1(-1 / 6, 1, 1 / 2, p)
+    line_2 = p**2 * (15 - 30 * sp.hyp2f1(-1 / 6, 1, 1 / 2, p) - 59 * sp.hyp2f1(5 / 6, 1, 1 / 2, p))
+    line_3 = 35 * sp.hyp2f1(5 / 6, 1, 1 / 2, p) + 15 * p**3 * sp.hyp2f1(5 / 6, 1, 1 / 2, p)
+    line_4 = p * (-54 + 88 * sp.hyp2f1(-1 / 6, 1, 1 / 2, p) + 9 * sp.hyp2f1(5 / 6, 1, 1 / 2, p))
+
+    term_1 = leading_factor * (line_1 + line_2 + line_3 + line_4)
+
+    term_2 = (L_2D ** (14 / 3)) / (np.sqrt(2 * b) * d ** (7 / 3) * z_i * np.sin(generator.psi_rad))
+
+    paren = term_1 - term_2
+
+    return generator.c * k1**2 * paren
+
+
+def estimate_spectra_2d(padded_u_field, padded_v_field) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Estimate the spectra of the 2d fields.
+    """
+
+    u_fft = np.fft.fft2(padded_u_field)
+    v_fft = np.fft.fft2(padded_v_field)
+
+    F_11 = np.abs(u_fft) ** 2
+    F_22 = np.abs(v_fft) ** 2
+
+    N = padded_u_field.shape[0]
+    F_11 = F_11 / N
+    F_22 = F_22 / N
+
+    F_11
+
+    return np.zeros(N), np.zeros(N)
 
-#########################################################################################
-# Test Mean, Std, Dev, Skewness, and Kurtosis
 
 #########################################################################################
 # Test energy spectrum
-#
-# ie,
 
 
 #########################################################################################

examples-unrendered/low_freq_prototype.py

@@ -7,16 +7,16 @@
 class LowFreq2DFieldGenerator:
     def __init__(
         self,
-        grid_dimensions,
-        grid_levels,
-        L2D=15_000.0,
-        sigma2=0.6,
-        z_i=500.0,
-        psi_degs=43.0,
+        grid_dimensions: np.ndarray,
+        grid_levels: np.ndarray,
+        L_2D: float = 15_000.0,
+        sigma2: float = 0.6,
+        z_i: float = 500.0,
+        psi_degs: float = 43.0,
         c: Optional[float] = None,
     ):
         # Field parameters
-        self.L2D = L2D
+        self.L_2D = L_2D
         self.sigma2 = sigma2
         self.z_i = z_i
         self.psi_degs = psi_degs
@@ -25,11 +25,10 @@ def __init__(
         self.L1, self.L2 = grid_dimensions[:2]  # Default was 60k x 15k
         self.Nx, self.Ny = 2 ** grid_levels[:2]  # Default was 1024 x 256
 
-        self.c = self._solve_for_c(sigma2, L2D, z_i)
         self.psi_rad = np.deg2rad(self.psi_degs)
 
         if c is None:
-            self.c = self._solve_for_c(self.sigma2, self.L2D, self.z_i)
+            self.c = self._solve_for_c()
         else:
             self.c = c
 
@@ -42,37 +41,33 @@ def _compute_kappa(self, kx, ky):
 
         return np.sqrt(2.0 * ((kx**2) * cos2 + (ky**2) * sin2))
 
-    def _compute_E(self, kappa):
+    def _compute_E(self, kappa: float) -> float:
+        """
+        Compute the energy spectrum E(kappa) for a given kappa.
+        """
         if kappa < 1e-12:
             return 0.0
-        denom = (1.0 / (self.L2D**2) + kappa**2) ** (7.0 / 3.0)
+        denom = (1.0 / (self.L_2D**2) + kappa**2) ** (7.0 / 3.0)
         atten = 1.0 / (1.0 + (kappa * self.z_i) ** 2)
         return self.c * (kappa**3) / denom * atten
 
-    def _solve_for_c(self, sigma2, L2D, z_i):
+    def _solve_for_c(self):
         """
-        Solve for c so that integral of E(k) from k=0..inf = sigma2.
-        (1D version for demonstration.)
+        Solve for scaling constant c so that integral of E(k) from k=0..inf = sigma2.
         """
 
-        def integrand(k):
-            return (k**3 / ((1.0 / (L2D**2) + k**2) ** (7.0 / 3.0))) * (1.0 / (1.0 + (k * z_i) ** 2))
+        def integrand(k: float) -> float:
+            return (k**3 / ((1.0 / (self.L_2D**2) + k**2) ** (7.0 / 3.0))) * (1.0 / (1.0 + (k * self.z_i) ** 2))
 
         val, _ = integrate.quad(integrand, 0, np.inf)
-        self.c = sigma2 / val
+        return self.sigma2 / val
 
     def generate(
         self,
+        pad: bool = True,
     ):
         """
-        Approx approach:
-        1) compute c from sigma2
-        2) for each k1, compute phi_11(k1) using e.g. E(kappa)/(pi*kappa)
-            or the simpler "2D swirl" formula
-        3) multiply by factor ~ 2 * pi^2 / L1
-        4) randomize phases in k2
-        5) iFFT => real field
-        We'll do just the 'u' field for demonstration, but you can do 'v' similarly.
+        Generate a 2D low-frequency field.
         """
         L1, L2 = self.L1, self.L2
         Nx, Ny = self.Nx, self.Ny
@@ -120,48 +115,57 @@ def generate(
         if var_now > 1e-12:
             u_field *= np.sqrt(self.sigma2 / var_now)
 
-        u_field = np.pad(u_field, ((0, 1), (0, 1)), mode="wrap")
+        if pad:
+            u_field = np.pad(u_field, ((0, 1), (0, 1)), mode="wrap")
+
+        return np.linspace(0, L1 / 1000, Nx), np.linspace(0, L2 / 1000, Ny), u_field
+
 
-        return np.linspace(0, L1, Nx), np.linspace(0, L2, Ny), u_field
+if __name__ == "__main__":
+    import matplotlib.pyplot as plt
 
+    # Domain: 60 km x 15 km
+    L1 = 60_000.0
+    L2 = 15_000.0
+    # Nx = 1024 # = 2^10
+    # Ny = 256 # = 2^8
 
-# TODO: Update this
-# if __name__ == "__main__":
-#     # Domain: 60 km x 15 km
-#     L1 = 60_000.0
-#     L2 = 15_000.0
-#     Nx = 1024
-#     Ny = 256
+    # Figure 3 parameters
+    L2D = 15000.0  # [m]
+    sigma2 = 0.6  # [m^2/s^2]
+    z_i = 500.0  # [m]
+    psi_degs = 43.0  # anisotropy angle
 
-#     # Figure 3 parameters
-#     L2D = 15000.0  # [m]
-#     sigma2 = 0.6  # [m^2/s^2]
-#     z_i = 500.0  # [m]
-#     psi_degs = 43.0  # anisotropy angle
+    generator = LowFreq2DFieldGenerator(
+        grid_dimensions=np.array([L1, L2]),
+        grid_levels=np.array([10, 8]),
+        L_2D=L2D,
+        sigma2=sigma2,
+        z_i=z_i,
+        psi_degs=psi_degs,
+    )
 
-#     # Generate large-scale u-component
-#     u_field = generate_2D_lowfreq_approx(Nx, Ny, L1, L2, psi_degs, sigma2, L2D, z_i)
+    # Generate large-scale u-component
+    _, _, u_field = generator.generate(pad=False)
 
-#     # Generate large-scale v-component similarly
-#     # (Here, we assume same sigma^2 and same approach.)
-#     v_field = generate_2D_lowfreq_approx(Nx, Ny, L1, L2, psi_degs, sigma2, L2D, z_i)
+    # Generate large-scale v-component similarly
+    # (Here, we assume same sigma^2 and same approach.)
+    x, y, v_field = generator.generate(pad=False)
 
-#     x = np.linspace(0, L1 / 1000, Nx)
-#     y = np.linspace(0, L2 / 1000, Ny)
-#     X, Y = np.meshgrid(x, y, indexing="ij")
+    X, Y = np.meshgrid(x, y, indexing="ij")
 
-#     fig, axs = plt.subplots(2, 1, figsize=(10, 8))
-#     im1 = axs[0].pcolormesh(X, Y, u_field, shading="auto", cmap="RdBu_r")
-#     cb1 = plt.colorbar(im1, ax=axs[0], label="m/s")
-#     axs[0].set_title("(a) u")
-#     axs[0].set_xlabel("x [km]")
-#     axs[0].set_ylabel("y [km]")
+    fig, axs = plt.subplots(2, 1, figsize=(10, 8))
+    im1 = axs[0].pcolormesh(X, Y, u_field, shading="auto", cmap="RdBu_r")
+    cb1 = plt.colorbar(im1, ax=axs[0], label="m/s")
+    axs[0].set_title("(a) u")
+    axs[0].set_xlabel("x [km]")
+    axs[0].set_ylabel("y [km]")
 
-#     im2 = axs[1].pcolormesh(X, Y, v_field, shading="auto", cmap="RdBu_r")
-#     cb2 = plt.colorbar(im2, ax=axs[1], label="m/s")
-#     axs[1].set_title("(b) v")
-#     axs[1].set_xlabel("x [km]")
-#     axs[1].set_ylabel("y [km]")
+    im2 = axs[1].pcolormesh(X, Y, v_field, shading="auto", cmap="RdBu_r")
+    cb2 = plt.colorbar(im2, ax=axs[1], label="m/s")
+    axs[1].set_title("(b) v")
+    axs[1].set_xlabel("x [km]")
+    axs[1].set_ylabel("y [km]")
 
-#     plt.tight_layout()
-#     plt.show()
+    plt.tight_layout()
+    plt.show()