Checkpointing

Author: mdmeeker

drdmannturb/fluctuation_generation/fluctuation_field_generator.py

@@ -12,6 +12,7 @@
 import torch
 
 from ..common import CPU_Unpickler
+from ..parameters import LowFreqParameters
 from ..spectra_fitting import CalibrationProblem
 from .covariance_kernels import MannCovariance, VonKarmanCovariance
 from .gaussian_random_fields import VectorGaussianRandomField
@@ -79,6 +80,7 @@ def __init__(
         length_scale: Optional[float] = None,
         time_scale: Optional[float] = None,
         energy_spectrum_scale: Optional[float] = None,
+        lowfreq_params: Optional[LowFreqParameters] = None,
         path_to_parameters: Optional[Union[str, PathLike]] = None,
         seed: int = None,
         blend_num=10,
@@ -89,7 +91,7 @@ def __init__(
         friction_velocity : float
             The reference wind friction velocity :math:`u_*`
         reference_height : float
-            Reference height :math:`L`
+            Reference height :math:`z_{\text{ref}}`
         grid_dimensions : np.ndarray
             Numpy array denoting the grid size; the real dimensions of the domain of interest.
         grid_levels : np.ndarray
@@ -162,17 +164,19 @@ def __init__(
 
         # define margins and buffer
         time_buffer = 3 * Gamma * L
-        spatial_margin = 1 * L
+        spatial_margin = 1 * L  # TODO: where is this used?
 
         # Grid calculations
         def calc_grid_size(level):
             return 2**level + 1
 
+        # TODO: why is this needed?
         try:
             grid_levels = [level.GetInt() for level in grid_levels]
         except AttributeError:
             pass
 
+        # TODO: Same as above...
         grid_sizes = [calc_grid_size(level) for level in grid_levels]
         Nx, Ny, Nz = grid_sizes
         grid_spacing = [dim / size for dim, size in zip(grid_dimensions, grid_sizes)]
@@ -334,10 +338,17 @@ def _normalize_block(
         Raises
         ------
         ValueError
+            In the case that curr_block does not satisfy not np.any (ie, it is empty, or all zeros).
             "No fluctuation field has been generated, call the .generate() method first."
+
+        ValueError
+            If windprofiletype is not one of "LOG" or "PL".
+
+        ValueError
+            In the case that any of the parameters zref, uref, or z0 are not positive.
         """
         if not np.any(curr_block):
-            raise ValueError("No fluctuation field has been generated, call the .generate() method first.")
+            raise ValueError("No fluctuation field has been generated. The .generate() method must be called first.")
 
         if windprofiletype not in ["LOG", "PL"]:
             raise ValueError('windprofiletype must be either "LOG" or "PL"')

drdmannturb/fluctuation_generation/lowfreq_syed_mann.py

@@ -9,7 +9,7 @@
 
 
 def _compute_kappa(k1: float, k2: float, psi: float) -> float:
-    """
+    r"""
     Subroutine to compute the horizontal wavevector :math:`\kappa`, defined by
 
     .. math::
@@ -36,7 +36,7 @@ def _compute_kappa(k1: float, k2: float, psi: float) -> float:
 
 
 def _compute_E(kappa: float, c: float, L2D: float, z_i: float) -> float:
-    """
+    r"""
     Subroutine to compute the energy spectrum :math:`E(\kappa)` with the attenuation factor,
     defined by
 
@@ -64,7 +64,7 @@ def _compute_E(kappa: float, c: float, L2D: float, z_i: float) -> float:
 
 
 def _estimate_c(sigma2: float, L2D: float, z_i: float) -> float:
-    """
+    r"""
     Subroutine to estimate the scaling factor :math:`c` from the target variance :math:`\sigma^2`.
 
     This is achieved by approximating the integral of :math:`E(\kappa)` from :math:`\kappa=0` to
@@ -105,8 +105,9 @@ def generate_2D_lowfreq(
     L2D: float,
     z_i: float,
     c: Optional[float] = None,
+    seed: Optional[int] = None,
 ) -> np.ndarray:
-    """
+    r"""
     Generates the 2D low-frequency wind fluctuation component of the Syed-Mann (2024) 2D+3D model.
 
     Parameters

drdmannturb/fluctuation_generation/sampling_methods.py

@@ -3,7 +3,6 @@
 """
 
 import os
-from math import *
 
 import numpy as np
 import scipy.fftpack as fft
@@ -16,14 +15,16 @@
 
 
 class Sampling_method_base:
-    """Meta class for different sampling methods. Each of these requires a ``RandomField`` object, which is a subclass of :py:class:``GaussianRandomField``."""
+    """Meta class for different sampling methods. Each of these requires a ``RandomField`` object, which is
+    a subclass of :py:class:``GaussianRandomField``."""
 
     def __init__(self, RandomField):
         """
         Parameters
         ----------
         RandomField : GaussianRandomField
-            The random field from which to sample from. This object also determines all of the physical quantities and domain partitioning.
+            The random field from which to sample from. This object also determines all of the physical quantities
+            and domain partitioning.
         """
         self.L, self.Nd, self.ndim = (
             RandomField.L,
@@ -34,30 +35,25 @@ def __init__(self, RandomField):
 
 
 class Sampling_method_freq(Sampling_method_base):
-    """Sampling method specifically in the frequency domain. This metaclass involves a single precomputation of the covariance spectrum of the underlying ``GaussianRandomField``. Refer to specific subclasses for details on what each of these entails, but generally, the approximate square-root of each associated spectral tensor is computed and transformed into the frequency domain.
+    """Sampling method specifically in the frequency domain. This metaclass involves a single precomputation of the
+    covariance spectrum of the underlying ``GaussianRandomField``. Refer to specific subclasses for details on what
+    each of these entails, but generally, the approximate square-root of each associated spectral tensor is computed
+    and transformed into the frequency domain.
 
     The norm of the transform is defined as the square-root of the length-scale.
     """
 
     def __init__(self, RandomField):
         super().__init__(RandomField)
         L, Nd, d = self.L, self.Nd, self.ndim
-        self.Frequencies = [
-            (2 * pi / L[j]) * (Nd[j] * fft.fftfreq(Nd[j])) for j in range(d)
-        ]
+        self.Frequencies = [(2 * np.pi / L[j]) * (Nd[j] * fft.fftfreq(Nd[j])) for j in range(d)]
         self.TransformNorm = np.sqrt(L.prod())
         self.Spectrum = RandomField.Covariance.precompute_Spectrum(self.Frequencies)
 
 
-#######################################################################################################
-# 	Fourier Transform (FFTW)
-#######################################################################################################
-### - Only stationary covariance
-### - Uses the Fastest Fourier Transform in the West
-
-
 class Sampling_FFTW(Sampling_method_freq):
-    """Sampling with FFTW. Two stencils for the forward and inverse FFTs are generated using the following FFTW flags: ``"FFTW_MEASURE", "FFTW_DESTROY_INPUT", "FFTW_UNALIGNED"``.
+    """Sampling with FFTW. Two stencils for the forward and inverse FFTs are generated using the following FFTW flags:
+    ``"FFTW_MEASURE", "FFTW_DESTROY_INPUT", "FFTW_UNALIGNED"``.
 
     Due to properties of the FFT, only stationary covariances are admissible.
     """
@@ -74,12 +70,8 @@ def __init__(self, RandomField):
         flags = ("FFTW_MEASURE", "FFTW_DESTROY_INPUT", "FFTW_UNALIGNED")
         self.fft_x = pyfftw.empty_aligned(shpR, dtype="float64")
         self.fft_y = pyfftw.empty_aligned(shpC, dtype="complex128")
-        self.fft_plan = pyfftw.FFTW(
-            self.fft_x, self.fft_y, axes=axes, direction="FFTW_FORWARD", flags=flags
-        )
-        self.ifft_plan = pyfftw.FFTW(
-            self.fft_y, self.fft_x, axes=axes, direction="FFTW_BACKWARD", flags=flags
-        )
+        self.fft_plan = pyfftw.FFTW(self.fft_x, self.fft_y, axes=axes, direction="FFTW_FORWARD", flags=flags)
+        self.ifft_plan = pyfftw.FFTW(self.fft_y, self.fft_x, axes=axes, direction="FFTW_BACKWARD", flags=flags)
         self.Spectrum_half = self.Spectrum[..., : shpC[-1]] * np.sqrt(self.Nd.prod())
 
     def __call__(self, noise):
@@ -90,27 +82,25 @@ def __call__(self, noise):
         return self.fft_x[self.DomainSlice] / self.TransformNorm
 
 
-#######################################################################################################
-# 	Vector Field Fourier Transform (VF_FFTW)
-#######################################################################################################
-### - Random vector fields
-### - Only stationary covariance
-### - Uses the Fastest Fourier Transform in the West
+class Sampling_VF_FFTW(Sampling_method_freq):
+    """Random vector fields using
 
+    FFTW applied to a vector field. This should be used in conjunction with :py:class:`VectorGaussianRandomField`.
+    This sampling method is also multi-threaded across 4 threads, or else the maximum allowed by the environment. As in
+    :py:class:`Sampling_FFTW`, the following FFTW flags are used: ``"FFTW_MEASURE", "FFTW_DESTROY_INPUT",
+    "FFTW_UNALIGNED"``.
 
-class Sampling_VF_FFTW(Sampling_method_freq):
-    """FFTW applied to a vector field. This should be used in conjunction with :py:class:`VectorGaussianRandomField`. This sampling method is also multi-threaded across 4 threads, or else the maximum allowed by the environment. As in :py:class:`Sampling_FFTW`, the following FFTW flags are used: ``"FFTW_MEASURE", "FFTW_DESTROY_INPUT", "FFTW_UNALIGNED"``."""
+    Due to properties of the FFT, only stationary covariances are admissible.
+    """
 
     def __init__(self, RandomField):
         super().__init__(RandomField)
 
         import pyfftw
 
-        n_cpu = 4
-        try:
-            n_cpu = int(os.environ["OMP_NUM_THREADS"])
-        except:
-            pass
+        # WARN: User might have OMP_NUM_THREADS set to something invalid here
+        n_cpu = int(os.environ.get("OMP_NUM_THREADS", 4))
+
         shpR = RandomField.ext_grid_shape
         shpC = shpR.copy()
         shpC[-1] = int(shpC[-1] // 2) + 1
@@ -135,9 +125,7 @@ def __init__(self, RandomField):
             threads=n_cpu,
         )
         self.Spectrum_half = self.Spectrum[..., : shpC[-1]] * np.sqrt(self.Nd.prod())
-        self.hat_noise = np.stack(
-            [np.zeros(shpC, dtype="complex64") for _ in range(3)], axis=-1
-        )
+        self.hat_noise = np.stack([np.zeros(shpC, dtype="complex64") for _ in range(3)], axis=-1)
         self.shpC = shpC
 
     def __call__(self, noise):
@@ -146,25 +134,18 @@ def __call__(self, noise):
             self.fft_x[:] = noise[..., i]
             self.fft_plan()
             self.hat_noise[..., i] = self.fft_y[:]
-        self.hat_noise = np.einsum(
-            "kl...,...l->...k", self.Spectrum_half, self.hat_noise
-        )
+        self.hat_noise = np.einsum("kl...,...l->...k", self.Spectrum_half, self.hat_noise)
         for i in range(noise.shape[-1]):
             self.fft_y[:] = self.hat_noise[..., i]
             self.ifft_plan()
             tmp[..., i] = self.fft_x[:]
         return tmp[self.DomainSlice] / self.TransformNorm
 
 
-#######################################################################################################
-# 	Fourier Transform
-#######################################################################################################
-### - Only stationary covariance
-### - Uses sscipy.fftpack (slower solution)
-
-
 class Sampling_FFT(Sampling_method_freq):
-    """Sampling using ``scipy.fftpack``, which is considerably slower than with FFTW but is a simpler interface."""
+    """Sampling using ``scipy.fftpack``, which is considerably slower than with FFTW but is a simpler interface.
+
+    Due to properties of the FFT, only stationary covariances are admissible."""
 
     def __init__(self, RandomField):
         super().__init__(RandomField)
@@ -176,15 +157,11 @@ def __call__(self, noise):
         return y.real[self.DomainSlice] / self.TransformNorm
 
 
-#######################################################################################################
-# 	Sine Transform
-#######################################################################################################
-### - Only stationary covariance
-### - Uses sscipy.fftpack (non the fastest solution)
-
-
 class Sampling_DST(Sampling_method_freq):
-    """Sampling using the discrete sine transform from ``scipy.fftpack``, with all other operations being identical as other sampling methods."""
+    """Sampling using the discrete sine transform from ``scipy.fftpack``, with all other operations being identical as
+    other sampling methods. Should only be used for stationary covariances
+
+    Due to properties of the FFT, only stationary covariances are admissible."""
 
     def __init__(self, RandomField):
         super().__init__(RandomField)
@@ -196,15 +173,11 @@ def __call__(self, noise):
         return y[self.DomainSlice] / self.TransformNorm
 
 
-#######################################################################################################
-# 	Cosine Transform
-#######################################################################################################
-### - Only stationary covariance
-### - Uses sscipy.fftpack (non the fastest solution)
-
-
 class Sampling_DCT(Sampling_method_freq):
-    """Sampling using the discrete cosine transform from ``scipy.fftpack``, with all other operations being identical as other sampling methods."""
+    """Sampling using the discrete cosine transform from ``scipy.fftpack``, with all other operations being identical
+    to other sampling methods.
+
+    Due to properties of the FFT, only stationary covariances are admissible."""
 
     def __init__(self, RandomField):
         super().__init__(RandomField)