Skip to content

Commit

Permalink
Checkpointing
Browse files Browse the repository at this point in the history
  • Loading branch information
@mdmeeker
mdmeeker committed Jan 15, 2025
1 parent 08ad421 commit 2f03693
Show file tree
Hide file tree
Showing 4 changed files with 146 additions and 81 deletions.
17 changes: 14 additions & 3 deletions drdmannturb/fluctuation_generation/fluctuation_field_generator.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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)]
Expand Down Expand Up @@ -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"')
Expand Down
9 changes: 5 additions & 4 deletions drdmannturb/fluctuation_generation/lowfreq_syed_mann.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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::
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
101 changes: 37 additions & 64 deletions drdmannturb/fluctuation_generation/sampling_methods.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
"""

import os
from math import *

import numpy as np
import scipy.fftpack as fft
Expand All @@ -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,
Expand All @@ -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.
"""
Expand All @@ -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):
Expand All @@ -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
Expand All @@ -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):
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand Down
Loading

0 comments on commit 2f03693

Please sign in to comment.