Skip to content

Commit

Permalink
Minor refactors to Fluctuation Generation
Browse files Browse the repository at this point in the history
  • Loading branch information
mdmeeker committed Mar 4, 2025
1 parent 8c157c1 commit ab1533d
Show file tree
Hide file tree
Showing 5 changed files with 328 additions and 25 deletions.
28 changes: 11 additions & 17 deletions drdmannturb/fluctuation_generation/fluctuation_field_generator.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
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 @@ -80,8 +79,8 @@ 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,
lowfreq_params: Optional[dict] = None,
seed: Optional[int] = None,
blend_num=10,
):
Expand Down Expand Up @@ -167,21 +166,16 @@ def __init__(
time_buffer = 3 * Gamma * 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)]
hx, hy, hz = grid_spacing
# TODO: this should not be needed
# grid_levels = [level.GetInt() for level in grid_levels]

# Expand on grid_levels parameter to get grid node counts in each direction
grid_node_counts = 2**grid_levels + 1
Nx, Ny, Nz = grid_node_counts

# Use grid_dimensions and grid_node_counts to get grid spacings in each direction
grid_spacings = [dim / size for dim, size in zip(grid_dimensions, grid_node_counts)]
hx, hy, hz = grid_spacings

# Calculate buffer and margin sizes
n_buffer = ceil(time_buffer / hx)
Expand Down
12 changes: 4 additions & 8 deletions drdmannturb/fluctuation_generation/gaussian_random_fields.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ class GaussianRandomField:
Generator for a discrete Gaussian random field: a random Gaussian variable at each point in a domain. Several
sampling methods are provided by default:
#. fft - usual Fast Fourier Transform
#. fft - SciPy Fast Fourier Transform (default)
#. fftw - "Fastest Fourier Transform in the West"
#. vf_fftw - vector field version of Fastest Fourier Transform in the West
#. dst - Discrete Sine Transform
Expand Down Expand Up @@ -89,6 +89,7 @@ def __init__(
self.ndim = ndim # dimension 2D or 3D
self.all_axes = np.arange(self.ndim)

# TODO: This is a strange block of code
if np.isscalar(grid_level):
if not np.isscalar(grid_shape):
raise ValueError("grid_level and grid_shape must have the same dimensions.")
Expand All @@ -98,17 +99,12 @@ def __init__(
assert len(grid_dimensions) == 3
assert len(grid_level) == 3

h = np.array(
[
grid_dimensions[0] / (2 ** grid_level[0] + 1),
grid_dimensions[1] / (2 ** grid_level[1] + 1),
grid_dimensions[2] / (2 ** grid_level[2] + 1),
]
)
h = grid_dimensions / (2**grid_level + 1)
self.grid_shape = np.array(grid_shape[:ndim])
self.L = h * self.grid_shape

### Extended window (NOTE: extension is done outside)
# TODO: This is also a strange block of code. Vestigial?
N_margin = 0
self.ext_grid_shape = self.grid_shape
self.nvoxels = self.ext_grid_shape.prod()
Expand Down
192 changes: 192 additions & 0 deletions low_freq_dev/usingDRD.ipynb
Original file line number Diff line number Diff line change
@@ -0,0 +1,192 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"\n",
"from drdmannturb.fluctuation_generation.covariance_kernels import Covariance\n",
"from drdmannturb.fluctuation_generation.gaussian_random_fields import GaussianRandomField\n",
"from drdmannturb.fluctuation_generation.sampling_methods import Sampling_FFT\n"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"param_sets = {\n",
" \"figure2_a_eq15\": {\n",
" \"sigma2\": 2.0, # m2/s2\n",
" \"L_2d\": 15_000.0, # m\n",
" \"psi\": np.deg2rad(45.0), # degrees\n",
" \"z_i\": 500.0, # m\n",
" \"L1_factor\": 40,\n",
" \"L2_factor\": 5,\n",
" \"N1\": 10,\n",
" \"N2\": 7,\n",
" \"equation\": \"eq15\",\n",
" },\n",
" \"figure2_b_eq16\": {\n",
" \"sigma2\": 2.0, # m2/s2\n",
" \"L_2d\": 15_000.0, # m\n",
" \"psi\": np.deg2rad(45.0), # degrees\n",
" \"z_i\": 500.0, # m\n",
" \"L1_factor\": 1,\n",
" \"L2_factor\": 0.125,\n",
" \"N1\": 10,\n",
" \"N2\": 7,\n",
" \"equation\": \"eq16\",\n",
" },\n",
" \"figure3_standard_eq14\": {\n",
" \"sigma2\": 0.6, # m2/s2\n",
" \"L_2d\": 15_000.0, # m\n",
" \"psi\": np.deg2rad(45.0), # degrees\n",
" \"z_i\": 500.0, # m\n",
" \"L1_factor\": 4,\n",
" \"L2_factor\": 1,\n",
" \"N1\": 10,\n",
" \"N2\": 7,\n",
" \"equation\": \"eq14\",\n",
" },\n",
" \"figure3_standard_eq15\": {\n",
" \"sigma2\": 0.6, # m2/s2\n",
" \"L_2d\": 15_000.0, # m\n",
" \"psi\": np.deg2rad(45.0), # degrees\n",
" \"z_i\": 500.0, # m\n",
" \"L1_factor\": 4,\n",
" \"L2_factor\": 1,\n",
" \"N1\": 10,\n",
" \"N2\": 7,\n",
" \"equation\": \"eq15\",\n",
" },\n",
"}"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"config = param_sets[\"figure2_a_eq15\"]\n",
"\n",
"_sigma2 = config[\"sigma2\"]\n",
"_L_2d = config[\"L_2d\"]\n",
"_psi = config[\"psi\"]\n",
"_z_i = config[\"z_i\"]\n",
"_L1_factor = config[\"L1_factor\"]\n",
"_L2_factor = config[\"L2_factor\"]\n",
"_N1 = config[\"N1\"]\n",
"_N2 = config[\"N2\"]\n",
"\n",
"# Calculate \n",
"grid_dimensions = np.array([\n",
" _L_2d * _L1_factor, \n",
" _L_2d * _L2_factor, \n",
" _z_i\n",
"])\n",
"grid_levels = np.array([_N1, _N2, 1])"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"class LowFreqCovariance(Covariance):\n",
" def __init__(self, sigma2, L2d, psi, z_i):\n",
" super().__init__()\n",
"\n",
" if self.ndim != 2:\n",
" raise ValueError(\"ndim must be 2 for MannSyed LowFreq covariance.\")\n",
" self.ndim = 2\n",
"\n",
" self.sigma2 = sigma2\n",
" self.L2d = L2d\n",
" self.psi = psi\n",
" self.z_i = z_i\n",
"\n",
" def precompute_Spectrum(Frequencies):\n",
"\n",
" return np.zeros((2, 2, Frequencies[0].size, Frequencies[1].size), dtype=np.complex128)"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"lfc = LowFreqCovariance(\n",
" _sigma2,\n",
" _L_2d,\n",
" _psi,\n",
" _z_i,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"ename": "ValueError",
"evalue": "operands could not be broadcast together with shapes (3,) (2,) ",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mValueError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[10], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m grf \u001b[38;5;241m=\u001b[39m \u001b[43mGaussianRandomField\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 2\u001b[0m \u001b[43m \u001b[49m\u001b[43mlfc\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 3\u001b[0m \u001b[43m \u001b[49m\u001b[43mgrid_level\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mgrid_levels\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 4\u001b[0m \u001b[43m \u001b[49m\u001b[43mgrid_dimensions\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mgrid_dimensions\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 5\u001b[0m \u001b[43m \u001b[49m\u001b[43mgrid_shape\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mgrid_levels\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 6\u001b[0m \u001b[43m \u001b[49m\u001b[43mndim\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 7\u001b[0m \u001b[43m \u001b[49m\u001b[43msampling_method\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mfft\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 8\u001b[0m \u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/OPEN_SOURCE/DRDMannTurb/drdmannturb/fluctuation_generation/gaussian_random_fields.py:109\u001b[0m, in \u001b[0;36mGaussianRandomField.__init__\u001b[0;34m(self, Covariance, grid_level, grid_shape, grid_dimensions, ndim, sampling_method)\u001b[0m\n\u001b[1;32m 101\u001b[0m h \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39marray(\n\u001b[1;32m 102\u001b[0m [\n\u001b[1;32m 103\u001b[0m grid_dimensions[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m/\u001b[39m (\u001b[38;5;241m2\u001b[39m \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m grid_level[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m+\u001b[39m \u001b[38;5;241m1\u001b[39m),\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 106\u001b[0m ]\n\u001b[1;32m 107\u001b[0m )\n\u001b[1;32m 108\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgrid_shape \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39marray(grid_shape[:ndim])\n\u001b[0;32m--> 109\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mL \u001b[38;5;241m=\u001b[39m \u001b[43mh\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgrid_shape\u001b[49m\n\u001b[1;32m 111\u001b[0m \u001b[38;5;66;03m### Extended window (NOTE: extension is done outside)\u001b[39;00m\n\u001b[1;32m 112\u001b[0m N_margin \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n",
"\u001b[0;31mValueError\u001b[0m: operands could not be broadcast together with shapes (3,) (2,) "
]
}
],
"source": [
"grf = GaussianRandomField(\n",
" lfc,\n",
" grid_level=grid_levels,\n",
" grid_dimensions=grid_dimensions,\n",
" grid_shape=2**grid_levels,\n",
" ndim=2,\n",
" sampling_method=\"fft\",\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "DRDMT",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.16"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
94 changes: 94 additions & 0 deletions low_freq_dev/usingDRD.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
import numpy as np

from drdmannturb.fluctuation_generation.covariance_kernels import Covariance

param_sets = {
"figure2_a_eq15": {
"sigma2": 2.0, # m2/s2
"L_2d": 15_000.0, # m
"psi": np.deg2rad(45.0), # degrees
"z_i": 500.0, # m
"L1_factor": 40,
"L2_factor": 5,
"N1": 2**10,
"N2": 2**7,
"equation": "eq15",
},
"figure2_b_eq16": {
"sigma2": 2.0, # m2/s2
"L_2d": 15_000.0, # m
"psi": np.deg2rad(45.0), # degrees
"z_i": 500.0, # m
"L1_factor": 1,
"L2_factor": 0.125,
"N1": 2**10,
"N2": 2**7,
"equation": "eq16",
},
"figure3_standard_eq14": {
"sigma2": 0.6, # m2/s2
"L_2d": 15_000.0, # m
"psi": np.deg2rad(45.0), # degrees
"z_i": 500.0, # m
"L1_factor": 4,
"L2_factor": 1,
"N1": 2**10,
"N2": 2**7,
"equation": "eq14",
},
"figure3_standard_eq15": {
"sigma2": 0.6, # m2/s2
"L_2d": 15_000.0, # m
"psi": np.deg2rad(45.0), # degrees
"z_i": 500.0, # m
"L1_factor": 4,
"L2_factor": 1,
"N1": 2**10,
"N2": 2**7,
"equation": "eq15",
},
}
config = param_sets["figure2_a_eq15"]

_sigma2 = config["sigma2"]
_L_2d = config["L_2d"]
_psi = config["psi"]
_z_i = config["z_i"]
_L1_factor = config["L1_factor"]
_L2_factor = config["L2_factor"]
_N1 = config["N1"]
_N2 = config["N2"]

grid_dimensions = np.array([_L_2d * _L1_factor, _L_2d * _L2_factor, _z_i])
grid_levels = np.array([_N1, _N2, 1])


#########################################################################################################
# Implementation of Covariance


class LowFreqCovariance(Covariance):
def __init__(self, L_2d: float, psi: float, z_i: float, sigma2: float, L1: float, L2: float):
super().__init__()

self.ndim = 2
self.L_2d = L_2d
self.psi = psi
self.z_i = z_i
self.sigma2 = sigma2
self.L1 = L1
self.L2 = L2

def precompute_Spectrum(self, Frequencies: np.ndarray) -> np.ndarray:
sqrtSpectralTensor = np.zeros((2, 2, Frequencies[0].size, Frequencies[1].size), dtype=np.complex128)

# k = np.array(list(np.meshgrid(*Frequencies, indexing="ij")))

# kk = np.sum(k**2, axis=0)
# kappa = np.sqrt(2 * ((k[0, ...] * np.cos(self.psi)) ** 2 + (k[1, ...] * np.sin(self.psi)) ** 2))

return sqrtSpectralTensor


# Need to create a random field object
# fft_sampler = Sampling_FFT()
Loading

0 comments on commit ab1533d

Please sign in to comment.