From 5e57beef48711ae5866f5fe4118a03feadfdf527 Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Thu, 25 Sep 2025 09:13:57 -0400 Subject: [PATCH 01/22] First pass and Sine and Cosine bases and MMTs --- dedalus/core/basis.py | 217 +++++++++++++++++++++++-------------- dedalus/core/transforms.py | 86 ++++++++++++++- 2 files changed, 218 insertions(+), 85 deletions(-) diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index e7705f40..aff14b57 100644 --- a/dedalus/core/basis.py +++ b/dedalus/core/basis.py @@ -36,6 +36,8 @@ 'RealFourier', 'ComplexFourier', 'Fourier', + 'Sine', + 'Cosine', 'DiskBasis', 'AnnulusBasis', 'SphereBasis', @@ -855,6 +857,7 @@ def __init__(self, coord, size, bounds, dealias=(1,), library=None): self.dealias = dealias self.library = library # Other attributes + self.grid_params = (coord, bounds, dealias, library) self.constant_mode_value = 1 # No permutations by default self.forward_coeff_permutation = None @@ -1332,29 +1335,64 @@ def _group_matrix(group, input_basis, output_basis): # return 0 -# class Sine(Basis, metaclass=CachedClass): -# """Sine series basis.""" -# space_type = ParityInterval -# const = None -# supported_dtypes = {np.float64, np.complex128} +class Sine(FourierBase, metaclass=CachedClass): + """Sine series basis.""" -# def __add__(self, other): -# space = self.space -# if other is space.Sine: -# return space.Sine -# else: -# return NotImplemented + native_bounds = (0, np.pi) + default_library = "fftw" + group_shape = (1,) + transforms = {} -# def __mul__(self, other): -# space = self.space -# if other is None: -# return space.Sine -# elif other is space.Sine: -# return space.Cosine -# elif other is space.Cosine: -# return space.Sine -# else: -# return NotImplemented + @CachedAttribute + def _native_wavenumbers(self): + # Excludes Nyquist mode + kmax = self.size - 1 + return np.arange(0, kmax+1) + + def valid_elements(self, tensorsig, grid_space, elements): + vshape = tuple(cs.dim for cs in tensorsig) + elements[0].shape + valid = np.ones(shape=vshape, dtype=bool) + if not grid_space[0]: + # Drop sine part of k=0 for all cartesian components and spin scalars + if not isinstance(self.coord, AzimuthalCoordinate) or not tensorsig: + # Drop sine part of k=0 + groups = self.elements_to_groups(grid_space, elements) + allcomps = tuple(slice(None) for cs in tensorsig) + selection = (groups[0] == 0) + valid[allcomps + (selection,)] = False + return valid + + def __add__(self, other): + if other is self: + return self + if isinstance(other, Sine): + if self.grid_params == other.grid_params: + size = max(self.size, other.size) + return self.clone_with(size=size) + return NotImplemented + + def __mul__(self, other): + if other is None: + return self + if other is self: + return Cosine(self.coord, self.size, self.bounds, self.dealias, self.library) + if isinstance(other, Sine): + if self.grid_params == other.grid_params: + size = max(self.size, other.size) + return Cosine(self.coord, size, self.bounds, self.dealias, self.library) + if isinstance(other, Cosine): + if self.grid_params == other.grid_params: + size = max(self.size, other.size) + return Sine(self.coord, size, self.bounds, self.dealias, self.library) + return NotImplemented + + def _native_grid(self, scale): + """Native flat global grid.""" + N, = self.grid_shape((scale,)) + return (np.pi / N) * (1/2 + np.arange(N)) + + def __pow__(self, other): + return NotImplemented # def __pow__(self, other): # space = self.space @@ -1365,44 +1403,59 @@ def _group_matrix(group, input_basis, output_basis): # else: # return NotImplemented -# def include_mode(self, mode): -# # Drop k=0 and Nyquist mode -# k = mode -# return (1 <= k <= self.space.kmax) +class Cosine(FourierBase, metaclass=CachedClass): + """Cosine series basis.""" -# class Cosine(Basis, metaclass=CachedClass): -# """Cosine series basis.""" -# space_type = ParityInterval -# const = 1 + native_bounds = (0, np.pi) + default_library = "fftw" + group_shape = (1,) + transforms = {} -# def __add__(self, other): -# space = self.space -# if other is None: -# return space.Cosine -# elif other is space.Cosine: -# return space.Cosine -# else: -# return NotImplemented + @CachedAttribute + def _native_wavenumbers(self): + # Excludes Nyquist mode + kmax = self.size - 1 + return np.arange(0, kmax+1) -# def __mul__(self, other): -# space = self.space -# if other is None: -# return space.Cosine -# elif other is space.Sine: -# return space.Sine -# elif other is space.Cosine: -# return space.Cosine -# else: -# return NotImplemented + def valid_elements(self, tensorsig, grid_space, elements): + vshape = tuple(cs.dim for cs in tensorsig) + elements[0].shape + valid = np.ones(shape=vshape, dtype=bool) + return valid -# def __pow__(self, other): -# return self.space.Cosine + def __add__(self, other): + if other is None: + return self + if other is self: + return self + if isinstance(other, Cosine): + if self.grid_params == other.grid_params: + size = max(self.size, other.size) + return self.clone_with(size=size) + return NotImplemented + + def __mul__(self, other): + if other is None: + return self + if other is self: + return self + if isinstance(other, Sine): + if self.grid_params == other.grid_params: + size = max(self.size, other.size) + return Sine(self.coord, size, self.bounds, self.dealias, self.library) + if isinstance(other, Cosine): + if self.grid_params == other.grid_params: + size = max(self.size, other.size) + return Cosine(self.coord, size, self.bounds, self.dealias, self.library) + return NotImplemented -# def include_mode(self, mode): -# # Drop Nyquist mode -# k = mode -# return (0 <= k <= self.space.kmax) + def _native_grid(self, scale): + """Native flat global grid.""" + N, = self.grid_shape((scale,)) + return (np.pi / N) * (1/2 + np.arange(N)) + + def __pow__(self, other): + return NotImplemented # class InterpolateSine(operators.Interpolate): @@ -1457,44 +1510,42 @@ def _group_matrix(group, input_basis, output_basis): # return 0 -# class DifferentiateSine(operators.Differentiate): -# """Sine series differentiation.""" +class DifferentiateSine(operators.Differentiate, operators.SpectralOperator1D): + """Sine series differentiation.""" -# input_basis_type = Sine -# bands = [0] -# separable = True + input_basis_type = Sine + subaxis_dependence = [True] + subaxis_coupling = [False] -# @staticmethod -# def output_basis(space, input_basis): -# return space.Cosine + @staticmethod + def _output_basis(input_basis): + return Cosine(input_basis.coord, input_basis.size, input_basis.bounds, input_basis.dealias, input_basis.library) -# @staticmethod -# def _build_subspace_entry(i, j, space, input_basis): -# # dx(sin(n*x)) = n*cos(n*x) -# if i == j: -# return j / space.COV.stretch -# else: -# return 0 + @staticmethod + def _group_matrix(group, input_basis, output_basis): + # Rescale group (native wavenumber) to get physical wavenumber + k = group / input_basis.COV.stretch + # dx sin(k*x) = k * cos(k*x) + return np.array([[k]]) -# class DifferentiateCosine(operators.Differentiate): -# """Cosine series differentiation.""" +class DifferentiateCosine(operators.Differentiate, operators.SpectralOperator1D): + """Cosine series differentiation.""" -# input_basis_type = Cosine -# bands = [0] -# separable = True + input_basis_type = Cosine + subaxis_dependence = [True] + subaxis_coupling = [False] -# @staticmethod -# def output_basis(space, input_basis): -# return space.Sine + @staticmethod + def _output_basis(input_basis): + return Sine(input_basis.coord, input_basis.size, input_basis.bounds, input_basis.dealias, input_basis.library) -# @staticmethod -# def _build_subspace_entry(i, j, space, input_basis): -# # dx(cos(n*x)) = -n*sin(n*x) -# if i == j: -# return (-j) / space.COV.stretch -# else: -# return 0 + @staticmethod + def _group_matrix(group, input_basis, output_basis): + # Rescale group (native wavenumber) to get physical wavenumber + k = group / input_basis.COV.stretch + # dx cos(k*x) = -k * sin(k*x) + return np.array([[-k]]) # class HilbertTransformSine(operators.HilbertTransform): diff --git a/dedalus/core/transforms.py b/dedalus/core/transforms.py index 00758fb2..02e91f65 100644 --- a/dedalus/core/transforms.py +++ b/dedalus/core/transforms.py @@ -629,6 +629,88 @@ def backward(self, cdata, gdata, axis): plan.backward(temp, gdata) +class SineTransform(SeparableTransform): + """ + Abstract base class for sine transforms. + + Parameters + ---------- + grid_size : int + Grid size (N) along transform dimension. + coeff_size : int + Coefficient size (M) along transform dimension. + + Notes + ----- + Let KN = (N - 1) be the maximum (Nyquist) mode on the grid. + Let KM = (M - 1) be the maximum retained mode in coeff space. + Then K = min(KN, KM) is the maximum wavenumber used in the transforms. + A unit-amplitude normalization is used. + + Forward transform: + if k == 0: + a(k) = 0 + elif k <= K: + a(k) = (2/N) \sum_{x=0}^{N-1} f(x) \sin(\pi k x / N) + else: + a(k) = 0 + + Backward transform: + f(x) = \sum_{k=1}^{K} a(k) \sin(\pi k x / N) + + Coefficient ordering: + The sine coefficients are ordered simply as + [0, a(1), a(2), ..., a(KM)] + """ + + def __init__(self, grid_size, coeff_size): + self.N = grid_size + self.M = coeff_size + self.KN = (self.N - 1) + self.KM = (self.M - 1) + self.Kmax = min(self.KN, self.KM) + + @property + def wavenumbers(self): + """One-dimensional global wavenumber array.""" + return np.arange(self.KM + 1) + + +@register_transform(basis.Sine, 'matrix') +class SineMMT(SineTransform, SeparableMatrixTransform): + """Sine MMT.""" + + @CachedAttribute + def forward_matrix(self): + """Build forward transform matrix.""" + N = self.N + M = self.M + Kmax = self.Kmax + K = self.wavenumbers[:, None] + X = np.arange(N)[None, :] + 1/2 + dX = N / np.pi + quadrature = (2 / N) * np.sin(K*X/dX) + # Zero higher modes for transforms with grid_size < coeff_size + quadrature *= (K <= self.Kmax) + # Ensure C ordering for fast dot products + return np.asarray(quadrature, order='C') + + @CachedAttribute + def backward_matrix(self): + """Build backward transform matrix.""" + N = self.N + M = self.M + Kmax = self.Kmax + K = self.wavenumbers[None, :] + X = np.arange(N)[:, None] + 1/2 + dX = N / np.pi + functions = np.sin(K*X/dX) + # Zero higher modes for transforms with grid_size < coeff_size + functions *= (K <= self.Kmax) + # Ensure C ordering for fast dot products + return np.asarray(functions, order='C') + + class CosineTransform(SeparableTransform): """ Abstract base class for cosine transforms. @@ -676,7 +758,7 @@ def wavenumbers(self): return np.arange(self.KM + 1) -#@register_transform(basis.Cosine, 'matrix') +@register_transform(basis.Cosine, 'matrix') class CosineMMT(CosineTransform, SeparableMatrixTransform): """Cosine MMT.""" @@ -687,7 +769,7 @@ def forward_matrix(self): M = self.M Kmax = self.Kmax K = self.wavenumbers[:, None] - X = np.arange(N)[None, :] + X = np.arange(N)[None, :] + 1/2 dX = N / np.pi quadrature = (2 / N) * np.cos(K*X/dX) quadrature[0] = 1 / N From eb0f2320038d6cc6fc9b5be8030cb9c986b4b6f7 Mon Sep 17 00:00:00 2001 From: "J. S. Oishi" Date: Thu, 25 Sep 2025 17:05:28 -0400 Subject: [PATCH 02/22] Added interpolation and integration for sin/cos. Note that sine integration is very inaccurate: only four digits for integrating sin(3pi x) from 0 to 1. --- dedalus/core/basis.py | 109 ++++++++++++++++++++++++++++-------------- 1 file changed, 73 insertions(+), 36 deletions(-) diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index aff14b57..66269557 100644 --- a/dedalus/core/basis.py +++ b/dedalus/core/basis.py @@ -1458,56 +1458,93 @@ def __pow__(self, other): return NotImplemented -# class InterpolateSine(operators.Interpolate): -# """Sine series interpolation.""" +class InterpolateSine(operators.Interpolate, operators.SpectralOperator1D): + """RealFourier interpolation.""" -# input_basis_type = Sine + input_basis_type = Sine + basis_subaxis = 0 + subaxis_dependence = [True] + subaxis_coupling = [True] -# @staticmethod -# def _build_subspace_entry(j, space, input_basis, position): -# # sin(n*x) -# x = space.COV.native_coord(position) -# return math.sin(j*x) + @staticmethod + def _output_basis(input_basis, position): + return None + + @staticmethod + def _full_matrix(input_basis, output_basis, position): + # Build native interpolation vector + x = input_basis.COV.native_coord(position) + k = input_basis.native_wavenumbers + interp_vector = np.zeros(k.size) + interp_vector[1:] = np.sin(k[1:] * x) + # Return with shape (1, N) + return interp_vector[None, :] +class InterpolateCosine(operators.Interpolate, operators.SpectralOperator1D): + """RealFourier interpolation.""" -# class InterpolateCosine(operators.Interpolate): -# """Cosine series interpolation.""" + input_basis_type = Cosine + basis_subaxis = 0 + subaxis_dependence = [True] + subaxis_coupling = [True] -# input_basis_type = Cosine + @staticmethod + def _output_basis(input_basis, position): + return None -# @staticmethod -# def _build_subspace_entry(j, space, input_basis, position): -# # cos(n*x) -# x = space.COV.native_coord(position) -# return math.cos(j*x) + @staticmethod + def _full_matrix(input_basis, output_basis, position): + # Build native interpolation vector + x = input_basis.COV.native_coord(position) + k = input_basis.native_wavenumbers + interp_vector = np.cos(k * x) + # Return with shape (1, N) + return interp_vector[None, :] +class IntegrateSine(operators.Integrate, operators.SpectralOperator1D): + """RealFourier integration.""" -# class IntegrateSine(operators.Integrate): -# """Sine series integration.""" + input_coord_type = Coordinate + input_basis_type = Sine + subaxis_dependence = [True] + subaxis_coupling = [False] -# input_basis_type = Sine + @staticmethod + def _output_basis(input_basis): + return Cosine(input_basis.coord, input_basis.size, input_basis.bounds, input_basis.dealias, input_basis.library) + + @staticmethod + def _group_matrix(group, input_basis, output_basis): + # Rescale group (native wavenumber) to get physical wavenumber + k = group / input_basis.COV.stretch + #integral(sin(n*x), 0, pi) = (2 / n) * (n % 2) + # integ sin(k*x) = 0 + if (group%2): + return np.array([[2/k]]) + else: + return np.array([[0]]) -# @staticmethod -# def _build_subspace_entry(j, space, input_basis): -# # integral(sin(n*x), 0, pi) = (2 / n) * (n % 2) -# if (j % 2): -# return 0 -# else: -# return (2 / j) * space.COV.stretch +class IntegrateCosine(operators.Integrate, operators.SpectralOperator1D): + """RealFourier integration.""" -# class IntegrateCosine(operators.Integrate): -# """Cosine series integration.""" + input_coord_type = Coordinate + input_basis_type = Cosine + subaxis_dependence = [True] + subaxis_coupling = [False] -# input_basis_type = Cosine + @staticmethod + def _output_basis(input_basis): + return None -# @staticmethod -# def _build_subspace_entry(j, space, input_basis): -# # integral(cos(n*x), 0, pi) = pi * δ(n, 0) -# if j == 0: -# return np.pi * space.COV.stretch -# else: -# return 0 + @staticmethod + def _group_matrix(group, input_basis, output_basis): + # integ cos(k*x) = L * δ(k, 0) + if group == 0: + L = input_basis.COV.problem_length + return np.array([[L]]) + else: + return np.array([[0]]) class DifferentiateSine(operators.Differentiate, operators.SpectralOperator1D): From db95469c8d0a19a900975bb9c4ca55a1c8196a49 Mon Sep 17 00:00:00 2001 From: "J. S. Oishi" Date: Fri, 26 Sep 2025 12:14:26 -0400 Subject: [PATCH 03/22] Fixed sine integration. --- dedalus/core/basis.py | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index 66269557..c39a1040 100644 --- a/dedalus/core/basis.py +++ b/dedalus/core/basis.py @@ -1507,22 +1507,20 @@ class IntegrateSine(operators.Integrate, operators.SpectralOperator1D): input_coord_type = Coordinate input_basis_type = Sine subaxis_dependence = [True] - subaxis_coupling = [False] + subaxis_coupling = [True] @staticmethod def _output_basis(input_basis): - return Cosine(input_basis.coord, input_basis.size, input_basis.bounds, input_basis.dealias, input_basis.library) + return None @staticmethod - def _group_matrix(group, input_basis, output_basis): - # Rescale group (native wavenumber) to get physical wavenumber - k = group / input_basis.COV.stretch - #integral(sin(n*x), 0, pi) = (2 / n) * (n % 2) - # integ sin(k*x) = 0 - if (group%2): - return np.array([[2/k]]) - else: - return np.array([[0]]) + def _full_matrix(input_basis, output_basis): + # Build native integration vector + k = input_basis.native_wavenumbers + integ_vector = np.zeros(k.size) + integ_vector[1::2] = 2/k[1::2]/np.pi + # Return with shape (1, N) + return integ_vector[None, :] class IntegrateCosine(operators.Integrate, operators.SpectralOperator1D): @@ -1544,7 +1542,7 @@ def _group_matrix(group, input_basis, output_basis): L = input_basis.COV.problem_length return np.array([[L]]) else: - return np.array([[0]]) + raise ValueError("This should never happen.") class DifferentiateSine(operators.Differentiate, operators.SpectralOperator1D): From 2ce34663497dcc5059fc18d101b76e77a004273c Mon Sep 17 00:00:00 2001 From: "J. S. Oishi" Date: Fri, 26 Sep 2025 15:39:41 -0400 Subject: [PATCH 04/22] Updated sin/cos with running 2D RBC example script - stress-free boundaries --- dedalus/core/basis.py | 19 +++ .../plot_snapshots.py | 79 ++++++++++++ .../stress_free_rbc.py | 112 ++++++++++++++++++ 3 files changed, 210 insertions(+) create mode 100644 examples/ivp_2d_stress_free_rayleigh_benard/plot_snapshots.py create mode 100644 examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index c39a1040..10f81efb 100644 --- a/dedalus/core/basis.py +++ b/dedalus/core/basis.py @@ -1564,6 +1564,25 @@ def _group_matrix(group, input_basis, output_basis): return np.array([[k]]) +class ConvertConstantCosine(operators.ConvertConstant, operators.SpectralOperator1D): + """Upcast constants to Cosine.""" + + output_basis_type = Cosine + subaxis_dependence = [True] + subaxis_coupling = [False] + + @staticmethod + def _group_matrix(group, input_basis, output_basis): + # Rescale group (native wavenumber) to get physical wavenumber + k = group / output_basis.COV.stretch + # 1 = cos(0*x) + if k == 0: + unit_amplitude = 1 / output_basis.constant_mode_value + return np.array([[unit_amplitude]]) + else: + return np.zeros(shape=(1, 0)) + + class DifferentiateCosine(operators.Differentiate, operators.SpectralOperator1D): """Cosine series differentiation.""" diff --git a/examples/ivp_2d_stress_free_rayleigh_benard/plot_snapshots.py b/examples/ivp_2d_stress_free_rayleigh_benard/plot_snapshots.py new file mode 100644 index 00000000..c46b8d77 --- /dev/null +++ b/examples/ivp_2d_stress_free_rayleigh_benard/plot_snapshots.py @@ -0,0 +1,79 @@ +""" +Plot 2D cartesian snapshots. + +Usage: + plot_snapshots.py ... [--output=] + +Options: + --output= Output directory [default: ./frames] + +""" + +import h5py +import numpy as np +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt +from dedalus.extras import plot_tools + + +def main(filename, start, count, output): + """Save plot of specified tasks for given range of analysis writes.""" + + # Plot settings + tasks = ['buoyancy', 'vorticity'] + scale = 1.5 + dpi = 200 + title_func = lambda sim_time: 't = {:.3f}'.format(sim_time) + savename_func = lambda write: 'write_{:06}.png'.format(write) + + # Layout + nrows, ncols = 2, 1 + image = plot_tools.Box(4, 1) + pad = plot_tools.Frame(0.3, 0, 0, 0) + margin = plot_tools.Frame(0.2, 0.1, 0, 0) + + # Create multifigure + mfig = plot_tools.MultiFigure(nrows, ncols, image, pad, margin, scale) + fig = mfig.figure + + # Plot writes + with h5py.File(filename, mode='r') as file: + for index in range(start, start+count): + for n, task in enumerate(tasks): + # Build subfigure axes + i, j = divmod(n, ncols) + axes = mfig.add_axes(i, j, [0, 0, 1, 1]) + # Call 3D plotting helper, slicing in time + dset = file['tasks'][task] + plot_tools.plot_bot_3d(dset, 0, index, axes=axes, title=task, even_scale=True, visible_axes=False) + # Add time title + title = title_func(file['scales/sim_time'][index]) + title_height = 1 - 0.5 * mfig.margin.top / mfig.fig.y + fig.suptitle(title, x=0.44, y=title_height, ha='left') + # Save figure + savename = savename_func(file['scales/write_number'][index]) + savepath = output.joinpath(savename) + fig.savefig(str(savepath), dpi=dpi) + fig.clear() + plt.close(fig) + + +if __name__ == "__main__": + + import pathlib + from docopt import docopt + from dedalus.tools import logging + from dedalus.tools import post + from dedalus.tools.parallel import Sync + + args = docopt(__doc__) + + output_path = pathlib.Path(args['--output']).absolute() + # Create output directory if needed + with Sync() as sync: + if sync.comm.rank == 0: + if not output_path.exists(): + output_path.mkdir() + post.visit_writes(args[''], main, output=output_path) + diff --git a/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py b/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py new file mode 100644 index 00000000..dbcd16c0 --- /dev/null +++ b/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py @@ -0,0 +1,112 @@ +""" +Dedalus script simulating 2D horizontally-periodic Rayleigh-Benard convection. +This script demonstrates solving a 2D Cartesian initial value problem. It can +be ran serially or in parallel, and uses the built-in analysis framework to save +data snapshots to HDF5 files. The `plot_snapshots.py` script can be used to +produce plots from the saved data. It should take about 5 cpu-minutes to run. + +The problem is non-dimensionalized using the box height and freefall time, so +the resulting thermal diffusivity and viscosity are related to the Prandtl +and Rayleigh numbers as: + + kappa = (Rayleigh * Prandtl)**(-1/2) + nu = (Rayleigh / Prandtl)**(-1/2) + +For incompressible hydro with two boundaries, we need two tau terms for each the +velocity and buoyancy. Here we choose to use a first-order formulation, putting +one tau term each on auxiliary first-order gradient variables and the others in +the PDE, and lifting them all to the first derivative basis. This formulation puts +a tau term in the divergence constraint, as required for this geometry. + +To run and plot using e.g. 4 processes: + $ mpiexec -n 4 python3 rayleigh_benard.py + $ mpiexec -n 4 python3 plot_snapshots.py snapshots/*.h5 +""" + +import numpy as np +import dedalus.public as d3 +import logging +logger = logging.getLogger(__name__) + + +# Parameters +Lx, Lz = 4, 1 +Nx, Nz = 256, 64 +Rayleigh = 2e6 +Prandtl = 1 +dealias = 3/2 +stop_sim_time = 50 +timestepper = d3.RK222 +max_timestep = 0.125 +dtype = np.float64 + +# Bases +coords = d3.CartesianCoordinates('x', 'z') +dist = d3.Distributor(coords, dtype=dtype) +xbasis = d3.RealFourier(coords['x'], size=Nx, bounds=(0, Lx), dealias=dealias) +zsbasis = d3.Sine(coords['z'], size=Nz, bounds=(0, Lz), dealias=dealias, library='matrix') +zcbasis = d3.Cosine(coords['z'], size=Nz, bounds=(0, Lz), dealias=dealias, library='matrix') + +# Fields +p = dist.Field(name='p', bases=(xbasis,zcbasis)) +b = dist.Field(name='b', bases=(xbasis,zsbasis)) +u = dist.Field(name='u', bases=(xbasis,zcbasis)) +w = dist.Field(name='w', bases=(xbasis,zsbasis)) +tau_p = dist.Field(name='tau_p') + +# Substitutions +kappa = (Rayleigh * Prandtl)**(-1/2) +nu = (Rayleigh / Prandtl)**(-1/2) +x, z = dist.local_grids(xbasis, zcbasis) + +dx = lambda A: d3.Differentiate(A,coords['x']) +dz = lambda A: d3.Differentiate(A,coords['z']) +lap = lambda A: dx(dx(A)) + dz(dz(A)) +# Problem +# First-order form: "div(f)" becomes "trace(grad_f)" +# First-order form: "lap(f)" becomes "div(grad_f)" +problem = d3.IVP([p, b, u, w, tau_p], namespace=locals()) +problem.add_equation("dx(u) + dz(w) + tau_p = 0") +problem.add_equation("dt(b) - kappa*lap(b) = - u*dx(b) - w*dz(b)") +problem.add_equation("dt(u) - nu*lap(u) + dx(p) = - u*dx(u) - w*dz(u)") +problem.add_equation("dt(w) - nu*lap(w) + dz(p) - b = - u*dx(w) - w*dz(w)") +problem.add_equation("integ(p) = 0") # Pressure gauge + +# Solver +solver = problem.build_solver(timestepper) +solver.stop_sim_time = stop_sim_time + +# Initial conditions +b.fill_random('g', seed=42, distribution='normal', scale=1e-3) # Random noise +b['g'] *= z * (Lz - z) # Damp noise at walls +b['g'] += Lz - z # Add linear background + +# Analysis +snapshots = solver.evaluator.add_file_handler('snapshots', sim_dt=0.25, max_writes=50) +snapshots.add_task(b, name='buoyancy') +snapshots.add_task(dz(u) - dx(w), name='vorticity') + +# CFL +#CFL = d3.CFL(solver, initial_dt=max_timestep, cadence=10, safety=0.5, threshold=0.05, +# max_change=1.5, min_change=0.5, max_dt=max_timestep) +#CFL.add_velocity(u) + +# Flow properties +flow = d3.GlobalFlowProperty(solver, cadence=10) +flow.add_property((u*u + w*w)/nu, name='Re') + +# Main loop +try: + logger.info('Starting main loop') + while solver.proceed: + #timestep = CFL.compute_timestep() + timestep = 1e-3 + solver.step(timestep) + if (solver.iteration-1) % 10 == 0: + max_Re = flow.max('Re') + logger.info('Iteration=%i, Time=%e, dt=%e, max(Re)=%f' %(solver.iteration, solver.sim_time, timestep, max_Re)) +except: + logger.error('Exception raised, triggering end of main loop.') + raise +finally: + solver.log_stats() From 0f3f38d61e2170abba86a3d56cdf17fbaa57cb73 Mon Sep 17 00:00:00 2001 From: "J. S. Oishi" Date: Sat, 27 Sep 2025 18:10:01 -0400 Subject: [PATCH 05/22] Added fast sine/cosine transforms - updated stress free RBC script to no longer include background... --- dedalus/core/transforms.py | 99 ++++++++++++++++++- .../stress_free_rbc.py | 33 +++---- 2 files changed, 113 insertions(+), 19 deletions(-) diff --git a/dedalus/core/transforms.py b/dedalus/core/transforms.py index 02e91f65..1fc2566d 100644 --- a/dedalus/core/transforms.py +++ b/dedalus/core/transforms.py @@ -711,6 +711,101 @@ def backward_matrix(self): return np.asarray(functions, order='C') +class FastSineTransform(SineTransform): + """Abstract base class for fast sine transforms.""" + + def __init__(self, *args, **kw): + super().__init__(*args, **kw) + # Standard scaling factors for unit-amplitude normalization from DST-II + self.forward_rescale_zero = 0 + self.forward_rescale_pos = 1 / self.N + self.backward_rescale_zero = 0 + self.backward_rescale_pos = 1 / 2 + + def resize_rescale_forward(self, data_in, data_out, axis, Kmax): + """Resize by padding/trunction and rescale to unit amplitude.""" + + # there is no zero frequency in the DST... + zerofreq = axslice(axis, 0, 1) + Nyfreq = axslice(axis,Kmax,Kmax+1) + data_out[zerofreq] *= self.forward_rescale_zero + #data_out[Nyfreq] *= 0 + if Kmax > 0: + in_posfreq = axslice(axis, 0, Kmax) + out_posfreq = axslice(axis, 1, Kmax+1) + np.multiply(data_in[in_posfreq], self.forward_rescale_pos, data_out[out_posfreq]) + if self.KM > Kmax: + badfreq = axslice(axis, Kmax+1, None) + data_out[badfreq] = 0 + + def resize_rescale_backward(self, data_in, data_out, axis, Kmax): + """Resize by padding/trunction and rescale to unit amplitude.""" + # zerofreq = axslice(axis, 0, 1) + # np.multiply(data_in[zerofreq], self.backward_rescale_zero, data_out[zerofreq]) + in_posfreq = axslice(axis, 1, Kmax+1) + out_posfreq = axslice(axis, 0, Kmax) + Nyfreq = axslice(axis,Kmax,Kmax+1) + data_out[Nyfreq] = 0 + if Kmax > 0: + np.multiply(data_in[in_posfreq], self.backward_rescale_pos, data_out[out_posfreq]) + + if self.KN > Kmax: + badfreq = axslice(axis, Kmax+1, None) + data_out[badfreq] = 0 + + +@register_transform(basis.Sine, 'scipy') +class ScipyDST(FastSineTransform): + """Fast sine transform using scipy.fft.""" + + def forward(self, gdata, cdata, axis): + """Apply forward transform along specified axis.""" + # Call DST + temp = scipy.fft.dst(gdata, type=2, axis=axis) # Creates temporary + # Resize and rescale for unit-ampltidue normalization + self.resize_rescale_forward(temp, cdata, axis, self.Kmax) + + def backward(self, cdata, gdata, axis): + """Apply backward transform along specified axis.""" + # Resize and rescale for unit-amplitude normalization + # Need temporary to avoid overwriting problems + temp = np.empty_like(gdata) # Creates temporary + self.resize_rescale_backward(cdata, temp, axis, self.Kmax) + # Call IDST + temp = scipy.fft.dst(temp, type=3, axis=axis, overwrite_x=True) # Creates temporary + np.copyto(gdata, temp) + + +@register_transform(basis.Sine, 'fftw') +class FFTWDST(FFTWBase, FastSineTransform): + """Fast sine transform using FFTW.""" + + @CachedMethod + def _build_fftw_plan(self, dtype, gshape, axis): + """Build FFTW plans and temporary arrays.""" + logger.debug("Building FFTW DST plan for (dtype, gshape, axis) = (%s, %s, %s)" %(dtype, gshape, axis)) + flags = ['FFTW_'+self.rigor.upper()] + plan = fftw.DiscreteSineTransform(dtype, gshape, axis, flags=flags) + temp = fftw.create_array(gshape, dtype) + return plan, temp + + def forward(self, gdata, cdata, axis): + """Apply forward transform along specified axis.""" + plan, temp = self._build_fftw_plan(gdata.dtype, gdata.shape, axis) + # Execute FFTW plan + plan.forward(gdata, temp) + # Resize and rescale for unit-ampltidue normalization + self.resize_rescale_forward(temp, cdata, axis, self.Kmax) + + def backward(self, cdata, gdata, axis): + """Apply backward transform along specified axis.""" + plan, temp = self._build_fftw_plan(gdata.dtype, gdata.shape, axis) + # Resize and rescale for unit-amplitude normalization + self.resize_rescale_backward(cdata, temp, axis, self.Kmax) + # Execute FFTW plan + plan.backward(temp, gdata) + + class CosineTransform(SeparableTransform): """ Abstract base class for cosine transforms. @@ -828,7 +923,7 @@ def resize_rescale_backward(self, data_in, data_out, axis, Kmax): data_out[badfreq] = 0 -#@register_transform(basis.Cosine, 'scipy') +@register_transform(basis.Cosine, 'scipy') class ScipyDCT(FastCosineTransform): """Fast cosine transform using scipy.fft.""" @@ -850,7 +945,7 @@ def backward(self, cdata, gdata, axis): np.copyto(gdata, temp) -#@register_transform(basis.Cosine, 'fftw') +@register_transform(basis.Cosine, 'fftw') class FFTWDCT(FFTWBase, FastCosineTransform): """Fast cosine transform using FFTW.""" diff --git a/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py b/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py index dbcd16c0..f8e1de61 100644 --- a/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py +++ b/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py @@ -1,9 +1,11 @@ """ -Dedalus script simulating 2D horizontally-periodic Rayleigh-Benard convection. -This script demonstrates solving a 2D Cartesian initial value problem. It can -be ran serially or in parallel, and uses the built-in analysis framework to save -data snapshots to HDF5 files. The `plot_snapshots.py` script can be used to -produce plots from the saved data. It should take about 5 cpu-minutes to run. +Dedalus script simulating 2D horizontally-periodic Rayleigh-Benard +convection with stress-free boundary conditions using Sine/Cosine +bases in z. This script demonstrates solving a 2D Cartesian initial +value problem. It can be ran serially or in parallel, and uses the +built-in analysis framework to save data snapshots to HDF5 files. The +`plot_snapshots.py` script can be used to produce plots from the saved +data. It should take about 5 cpu-minutes to run. The problem is non-dimensionalized using the box height and freefall time, so the resulting thermal diffusivity and viscosity are related to the Prandtl @@ -12,15 +14,13 @@ kappa = (Rayleigh * Prandtl)**(-1/2) nu = (Rayleigh / Prandtl)**(-1/2) -For incompressible hydro with two boundaries, we need two tau terms for each the -velocity and buoyancy. Here we choose to use a first-order formulation, putting -one tau term each on auxiliary first-order gradient variables and the others in -the PDE, and lifting them all to the first derivative basis. This formulation puts -a tau term in the divergence constraint, as required for this geometry. +Note that unlike the no-slip Cheybshev example, here we solve for the +buoyancy *perturbation* instead of the total buoyancy field. To run and plot using e.g. 4 processes: $ mpiexec -n 4 python3 rayleigh_benard.py $ mpiexec -n 4 python3 plot_snapshots.py snapshots/*.h5 + """ import numpy as np @@ -39,13 +39,14 @@ timestepper = d3.RK222 max_timestep = 0.125 dtype = np.float64 - +library = 'fftw' # 'scipy', 'matrix' +logger.info(f"Running with {library} DCT/DST library") # Bases coords = d3.CartesianCoordinates('x', 'z') dist = d3.Distributor(coords, dtype=dtype) xbasis = d3.RealFourier(coords['x'], size=Nx, bounds=(0, Lx), dealias=dealias) -zsbasis = d3.Sine(coords['z'], size=Nz, bounds=(0, Lz), dealias=dealias, library='matrix') -zcbasis = d3.Cosine(coords['z'], size=Nz, bounds=(0, Lz), dealias=dealias, library='matrix') +zsbasis = d3.Sine(coords['z'], size=Nz, bounds=(0, Lz), dealias=dealias, library=library) +zcbasis = d3.Cosine(coords['z'], size=Nz, bounds=(0, Lz), dealias=dealias, library=library) # Fields p = dist.Field(name='p', bases=(xbasis,zcbasis)) @@ -62,12 +63,11 @@ dx = lambda A: d3.Differentiate(A,coords['x']) dz = lambda A: d3.Differentiate(A,coords['z']) lap = lambda A: dx(dx(A)) + dz(dz(A)) + # Problem -# First-order form: "div(f)" becomes "trace(grad_f)" -# First-order form: "lap(f)" becomes "div(grad_f)" problem = d3.IVP([p, b, u, w, tau_p], namespace=locals()) problem.add_equation("dx(u) + dz(w) + tau_p = 0") -problem.add_equation("dt(b) - kappa*lap(b) = - u*dx(b) - w*dz(b)") +problem.add_equation("dt(b) - kappa*lap(b) - w = - u*dx(b) - w*dz(b)") problem.add_equation("dt(u) - nu*lap(u) + dx(p) = - u*dx(u) - w*dz(u)") problem.add_equation("dt(w) - nu*lap(w) + dz(p) - b = - u*dx(w) - w*dz(w)") problem.add_equation("integ(p) = 0") # Pressure gauge @@ -79,7 +79,6 @@ # Initial conditions b.fill_random('g', seed=42, distribution='normal', scale=1e-3) # Random noise b['g'] *= z * (Lz - z) # Damp noise at walls -b['g'] += Lz - z # Add linear background # Analysis snapshots = solver.evaluator.add_file_handler('snapshots', sim_dt=0.25, max_writes=50) From cd6b7f7856c00f3490f10e5640f473a294403fc5 Mon Sep 17 00:00:00 2001 From: "J. S. Oishi" Date: Mon, 6 Oct 2025 10:30:53 -0400 Subject: [PATCH 06/22] Fixed typo in skew matrix --- dedalus/core/operators.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dedalus/core/operators.py b/dedalus/core/operators.py index 9c5875a9..f700e0eb 100644 --- a/dedalus/core/operators.py +++ b/dedalus/core/operators.py @@ -2143,7 +2143,7 @@ def subproblem_matrix(self, subproblem): factors[self.index] = 1j * factors[self.index] else: azimuth_index = len(self.tensorsig) + self.azimuth_axis - id_m = sparse.identity(shape[self.azimuth_axis]//2, format='csr') + id_m = sparse.identity(shape[azimuth_index]//2, format='csr') mul_1j = np.array([[0, -1], [1, 0]]) factors[azimuth_index] = sparse.kron(id_m, mul_1j) return reduce(sparse.kron, factors, 1).tocsr() From 481b496e3319ec689655461207f1e979a8bfb895 Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Mon, 10 Nov 2025 15:57:53 -0500 Subject: [PATCH 07/22] Change sine/cosine to even/odd parity and extend to work for tensor fields --- dedalus/core/basis.py | 442 ++++++++++++------ dedalus/core/operators.py | 9 +- dedalus/core/transforms.py | 12 +- .../stress_free_rbc.py | 6 +- 4 files changed, 308 insertions(+), 161 deletions(-) diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index 10f81efb..3773799e 100644 --- a/dedalus/core/basis.py +++ b/dedalus/core/basis.py @@ -9,10 +9,10 @@ from . import operators from ..libraries import spin_recombination from ..tools.array import kron, axslice, apply_matrix, permute_axis -from ..tools.cache import CachedAttribute, CachedMethod, CachedClass +from ..tools.cache import CachedAttribute, CachedMethod, CachedClass, CachedFunction from ..tools import jacobi from ..tools import clenshaw -from ..tools.array import reshape_vector, axindex, axslice, interleave_matrices +from ..tools.array import reshape_vector, axindex, axslice, interleave_matrices, sparse_block_diag from ..tools.dispatch import MultiClass, SkipDispatchException from ..tools.general import unify, DeferredTuple from .coords import Coordinate, CartesianCoordinates, S2Coordinates, SphericalCoordinates, PolarCoordinates, AzimuthalCoordinate, DirectProduct @@ -36,8 +36,9 @@ 'RealFourier', 'ComplexFourier', 'Fourier', - 'Sine', - 'Cosine', + 'EvenParity', + 'OddParity', + 'Parity', 'DiskBasis', 'AnnulusBasis', 'SphereBasis', @@ -696,8 +697,8 @@ def _group_matrix(group, input_basis, output_basis): unit_amplitude = 1 / output_basis.constant_mode_value return np.array([[unit_amplitude]]) else: - # Constructor should only loop over group 0. - raise ValueError("This should never happen.") + # Constructor should only loop over group 0 + raise ValueError(f"This should never happen: group = {group}") class DifferentiateJacobi(operators.Differentiate, operators.SpectralOperator1D): @@ -1008,13 +1009,12 @@ class ConvertConstantComplexFourier(operators.ConvertConstant, operators.Spectra @staticmethod def _group_matrix(group, input_basis, output_basis): - # Rescale group (native wavenumber) to get physical wavenumber - k = group / output_basis.COV.stretch # 1 = exp(1j*0*x) - if k == 0: + if group == 0: unit_amplitude = 1 / output_basis.constant_mode_value return np.array([[unit_amplitude]]) else: + # Return zero-width column for subproblem construction return np.zeros(shape=(1, 0)) @@ -1080,8 +1080,8 @@ def _group_matrix(group, input_basis, output_basis): L = input_basis.COV.problem_length return np.array([[L]]) else: - # Constructor should only loop over group 0. - raise ValueError("This should never happen.") + # Constructor should only loop over group 0 + raise ValueError(f"This should never happen: group = {group}") class AverageComplexFourier(operators.Average, operators.SpectralOperator1D): @@ -1104,8 +1104,8 @@ def _group_matrix(group, input_basis, output_basis): if k == 0: return np.array([[1]]) else: - # Constructor should only loop over group 0. - raise ValueError("This should never happen.") + # Constructor should only loop over group 0 + raise ValueError(f"This should never happen: group = {group}") class RealFourier(FourierBase, metaclass=CachedClass): @@ -1195,14 +1195,12 @@ class ConvertConstantRealFourier(operators.ConvertConstant, operators.SpectralOp @staticmethod def _group_matrix(group, input_basis, output_basis): - # Rescale group (native wavenumber) to get physical wavenumber - k = group / output_basis.COV.stretch # 1 = cos(0*x) - if k == 0: + if group == 0: unit_amplitude = 1 / output_basis.constant_mode_value - return np.array([[unit_amplitude], - [0]]) + return np.array([[unit_amplitude], [0]]) else: + # Return zero-width column for subproblem construction return np.zeros(shape=(2, 0)) @@ -1274,8 +1272,8 @@ def _group_matrix(group, input_basis, output_basis): L = input_basis.COV.problem_length return np.array([[L, 0]]) else: - # Constructor should only loop over group 0. - raise ValueError("This should never happen.") + # Constructor should only loop over group 0 + raise ValueError(f"This should never happen: group = {group}") class AverageRealFourier(operators.Average, operators.SpectralOperator1D): @@ -1299,8 +1297,8 @@ def _group_matrix(group, input_basis, output_basis): if k == 0: return np.array([[1, 0]]) else: - # Constructor should only loop over group 0. - raise ValueError("This should never happen.") + # Constructor should only loop over group 0 + raise ValueError(f"This should never happen: group = {group}") # class HilbertTransformFourier(operators.HilbertTransform): @@ -1335,8 +1333,8 @@ def _group_matrix(group, input_basis, output_basis): # return 0 -class Sine(FourierBase, metaclass=CachedClass): - """Sine series basis.""" +class ParityBase(FourierBase): + """Base class for parity-based bases.""" native_bounds = (0, np.pi) default_library = "fftw" @@ -1349,23 +1347,79 @@ def _native_wavenumbers(self): kmax = self.size - 1 return np.arange(0, kmax+1) + def _native_grid(self, scale): + """Native flat global grid.""" + N, = self.grid_shape((scale,)) + return (np.pi / N) * (1/2 + np.arange(N)) + + @staticmethod + @CachedFunction + def _transform_plan(grid_size, coeff_size, parity, library): + """Caching layer to share plans across even/odd parity bases.""" + # Use matrix transforms for trivial cases + if (grid_size == 1 or coeff_size == 1) and (library != "matrix"): + return ParityBase._transform_plan(grid_size, coeff_size, parity, "matrix") + parity_name = {1: "cos", -1: "sin"}[parity] + return ParityBase.transforms[f"{library}-{parity_name}"](grid_size, coeff_size) + + def transform_plan(self, grid_size, parity): + """Build transform plan.""" + return self._transform_plan(grid_size, self.size, parity, self.library) + + def forward_transform(self, field, axis, gdata, cdata): + # Get plans + data_axis = len(field.tensorsig) + axis + grid_size = gdata.shape[data_axis] + P = self.component_parity(field.tensorsig) + parity_plans = {p: self.transform_plan(grid_size, p) for p in np.unique(P)} + # Transform component-by-component + for i, p in np.ndenumerate(P): + parity_plans[p].forward(gdata[i], cdata[i], axis) + # Permute coefficients + if self.forward_coeff_permutation is not None: + permute_axis(cdata, axis+len(field.tensorsig), self.forward_coeff_permutation, out=cdata) + + def backward_transform(self, field, axis, cdata, gdata): + # Get plans + data_axis = len(field.tensorsig) + axis + grid_size = gdata.shape[data_axis] + P = self.component_parity(field.tensorsig) + parity_plans = {p: self.transform_plan(grid_size, p) for p in np.unique(P)} + # Permute coefficients + if self.backward_coeff_permutation is not None: + permute_axis(cdata, axis+len(field.tensorsig), self.backward_coeff_permutation, out=cdata) + # Transform component-by-component + P = self.component_parity(field.tensorsig) + for i, p in np.ndenumerate(P): + parity_plans[p].backward(cdata[i], gdata[i], axis) + + +def Parity(*args, parity=None, **kw): + """Factory function dispatching to EvenParity and OddParity based on provided parity.""" + if parity is None: + raise ValueError("parity must be specified") + elif parity == 1: + return EvenParity(*args, **kw) + elif parity == -1: + return OddParity(*args, **kw) + else: + raise ValueError(f"Unrecognized parity: {parity}") + + +class EvenParity(ParityBase, metaclass=CachedClass): + """Even parity basis (cosine series for scalars).""" + def valid_elements(self, tensorsig, grid_space, elements): vshape = tuple(cs.dim for cs in tensorsig) + elements[0].shape valid = np.ones(shape=vshape, dtype=bool) - if not grid_space[0]: - # Drop sine part of k=0 for all cartesian components and spin scalars - if not isinstance(self.coord, AzimuthalCoordinate) or not tensorsig: - # Drop sine part of k=0 - groups = self.elements_to_groups(grid_space, elements) - allcomps = tuple(slice(None) for cs in tensorsig) - selection = (groups[0] == 0) - valid[allcomps + (selection,)] = False return valid def __add__(self, other): + if other is None: + return self if other is self: return self - if isinstance(other, Sine): + if isinstance(other, EvenParity): if self.grid_params == other.grid_params: size = max(self.size, other.size) return self.clone_with(size=size) @@ -1375,60 +1429,49 @@ def __mul__(self, other): if other is None: return self if other is self: - return Cosine(self.coord, self.size, self.bounds, self.dealias, self.library) - if isinstance(other, Sine): + return self + if isinstance(other, OddParity): if self.grid_params == other.grid_params: size = max(self.size, other.size) - return Cosine(self.coord, size, self.bounds, self.dealias, self.library) - if isinstance(other, Cosine): + return OddParity(self.coord, size, self.bounds, self.dealias, self.library) + if isinstance(other, EvenParity): if self.grid_params == other.grid_params: size = max(self.size, other.size) - return Sine(self.coord, size, self.bounds, self.dealias, self.library) + return EvenParity(self.coord, size, self.bounds, self.dealias, self.library) return NotImplemented - def _native_grid(self, scale): - """Native flat global grid.""" - N, = self.grid_shape((scale,)) - return (np.pi / N) * (1/2 + np.arange(N)) - def __pow__(self, other): return NotImplemented -# def __pow__(self, other): -# space = self.space -# if (other % 2) == 0: -# return space.Cosine -# elif (other % 2) == 1: -# return space.Sine -# else: -# return NotImplemented - + @CachedMethod + def component_parity(self, tensorsig): + P = np.ones([cs.dim for cs in tensorsig], dtype=int) + coord = self.coord + for i, cs in enumerate(tensorsig): + if coord in cs.coords: + start = cs.coords.index(coord) + P[axslice(i, start, start+1)] *= -1 + return P -class Cosine(FourierBase, metaclass=CachedClass): - """Cosine series basis.""" - native_bounds = (0, np.pi) - default_library = "fftw" - group_shape = (1,) - transforms = {} - - @CachedAttribute - def _native_wavenumbers(self): - # Excludes Nyquist mode - kmax = self.size - 1 - return np.arange(0, kmax+1) +class OddParity(ParityBase, metaclass=CachedClass): + """Odd parity basis (sine series for scalars).""" def valid_elements(self, tensorsig, grid_space, elements): vshape = tuple(cs.dim for cs in tensorsig) + elements[0].shape valid = np.ones(shape=vshape, dtype=bool) + if not grid_space[0]: + # Drop sine part of k=0 + groups = self.elements_to_groups(grid_space, elements) + allcomps = tuple(slice(None) for cs in tensorsig) + selection = (groups[0] == 0) + valid[allcomps + (selection,)] = False return valid def __add__(self, other): - if other is None: - return self if other is self: return self - if isinstance(other, Cosine): + if isinstance(other, OddParity): if self.grid_params == other.grid_params: size = max(self.size, other.size) return self.clone_with(size=size) @@ -1438,52 +1481,122 @@ def __mul__(self, other): if other is None: return self if other is self: - return self - if isinstance(other, Sine): + return EvenParity(self.coord, self.size, self.bounds, self.dealias, self.library) + if isinstance(other, OddParity): if self.grid_params == other.grid_params: size = max(self.size, other.size) - return Sine(self.coord, size, self.bounds, self.dealias, self.library) - if isinstance(other, Cosine): + return EvenParity(self.coord, size, self.bounds, self.dealias, self.library) + if isinstance(other, EvenParity): if self.grid_params == other.grid_params: size = max(self.size, other.size) - return Cosine(self.coord, size, self.bounds, self.dealias, self.library) + return OddParity(self.coord, size, self.bounds, self.dealias, self.library) return NotImplemented - def _native_grid(self, scale): - """Native flat global grid.""" - N, = self.grid_shape((scale,)) - return (np.pi / N) * (1/2 + np.arange(N)) - def __pow__(self, other): return NotImplemented + @CachedMethod + def component_parity(self, tensorsig): + P = np.ones([cs.dim for cs in tensorsig], dtype=int) + coord = self.coord + for i, cs in enumerate(tensorsig): + if coord in cs.coords: + start = cs.coords.index(coord) + P[axslice(i, start, start+1)] *= -1 + return -P + -class InterpolateSine(operators.Interpolate, operators.SpectralOperator1D): - """RealFourier interpolation.""" +class SpectralOperatorParity(operators.SpectralOperator1D): + """Base class for spectral operators that operate on parity-based bases.""" - input_basis_type = Sine - basis_subaxis = 0 + @CachedAttribute + def operand_component_parity(self): + return self.input_basis.component_parity(self.operand.tensorsig) + + def subproblem_matrix(self, subproblem): + """Build operator matrix for a specific subproblem.""" + axis = self.last_axis + group = subproblem.group[axis] + # Build matrices for each parity + P = self.operand_component_parity + if group is None: + parity_matrices = {p: self.subspace_matrix(self.dist.coeff_layout, p) for p in np.unique(P)} + else: + parity_matrices = {p: self.group_matrix(group, p) for p in np.unique(P)} + # Kronecker up to proper size + shape = subproblem.coeff_shape(self.domain) + N_before = prod(shape[:axis]) + if N_before > 1: + I_before = sparse.identity(N_before, format='coo') # COO faster for kron + for p in parity_matrices: + parity_matrices[p] = sparse.kron(I_before, parity_matrices[p]) + N_after = prod(shape[axis+1:]) + if N_after > 1: + I_after = sparse.identity(N_after, format='coo') # COO faster for kron + for p in parity_matrices: + parity_matrices[p] = sparse.kron(parity_matrices[p], I_after) + # Block diagonalize over components + blocks = [parity_matrices[p] for p in P.ravel()] + return sparse_block_diag(blocks) + + def subspace_matrix(self, layout, parity): + """Build matrix operating on local subspace data.""" + # Caching layer to allow insertion of other arguments + return self._subspace_matrix(layout, self.input_basis, self.output_basis, self.first_axis, parity) + + def operate(self, out): + operand = self.args[0] + input_basis = self.input_basis + data_axis = len(operand.tensorsig) + self.last_axis + # Set output layout + out.preset_layout(operand.layout) + # Apply operator to each component + P = self.operand_component_parity + parity_matrices = {p: self.subspace_matrix(operand.layout, p) for p in np.unique(P)} + for i, p in np.ndenumerate(P): + apply_matrix(parity_matrices[p], operand.data[i], axis=data_axis, out=out.data[i]) + + +class ConvertConstantEvenParity(operators.ConvertConstant, SpectralOperatorParity): + """ + Upcast constants to EvenParity. + Note: this is a single implementation because tensor fields have both even and odd components. + """ + + output_basis_type = EvenParity subaxis_dependence = [True] - subaxis_coupling = [True] + subaxis_coupling = [False] - @staticmethod - def _output_basis(input_basis, position): - return None + def __init__(self, operand, output_basis, out=None): + super().__init__(operand, output_basis, out=out) + for cs in self.operand.tensorsig: + if self.coords in cs.coords: + raise NotImplementedError("Converting constant tensor fields to EvenParity is not supported.") + + @CachedAttribute + def operand_component_parity(self): + return self.output_basis.component_parity(self.operand.tensorsig) @staticmethod - def _full_matrix(input_basis, output_basis, position): - # Build native interpolation vector - x = input_basis.COV.native_coord(position) - k = input_basis.native_wavenumbers - interp_vector = np.zeros(k.size) - interp_vector[1:] = np.sin(k[1:] * x) - # Return with shape (1, N) - return interp_vector[None, :] + def _group_matrix(group, input_basis, output_basis, parity): + if parity != 1: + raise ValueError(f"This should never happen: parity = {parity}") + # 1 = cos(0*x) + if group == 0: + unit_amplitude = 1 / output_basis.constant_mode_value + return np.array([[unit_amplitude]]) + else: + # Return zero-width column for subproblem construction + return np.zeros(shape=(1, 0)) -class InterpolateCosine(operators.Interpolate, operators.SpectralOperator1D): - """RealFourier interpolation.""" - input_basis_type = Cosine +class InterpolateParity(operators.Interpolate, SpectralOperatorParity): + """ + Parity basis interpolation. + Note: this is a single implementation because tensor fields have both even and odd components. + """ + + input_basis_type = ParityBase basis_subaxis = 0 subaxis_dependence = [True] subaxis_coupling = [True] @@ -1492,42 +1605,64 @@ class InterpolateCosine(operators.Interpolate, operators.SpectralOperator1D): def _output_basis(input_basis, position): return None + def subspace_matrix(self, layout, parity): + """Build matrix operating on global subspace data.""" + return self._subspace_matrix(layout, self.input_basis, self.output_basis, self.first_axis, self.position, parity) + @staticmethod - def _full_matrix(input_basis, output_basis, position): + def _full_matrix(input_basis, output_basis, position, parity): # Build native interpolation vector x = input_basis.COV.native_coord(position) k = input_basis.native_wavenumbers - interp_vector = np.cos(k * x) + if parity == 1: + interp_vector = np.cos(k * x) + elif parity == -1: + interp_vector = np.sin(k * x) + else: + raise ValueError(f"This should never happen: parity = {parity}") # Return with shape (1, N) return interp_vector[None, :] -class IntegrateSine(operators.Integrate, operators.SpectralOperator1D): - """RealFourier integration.""" - input_coord_type = Coordinate - input_basis_type = Sine +class DifferentiateParity(operators.Differentiate, SpectralOperatorParity): + """ + Parity basis differentiation. + Note: this is a single implementation because tensor fields have both even and odd components. + """ + + input_basis_type = ParityBase subaxis_dependence = [True] - subaxis_coupling = [True] + subaxis_coupling = [False] @staticmethod def _output_basis(input_basis): - return None + # Swap parity + if isinstance(input_basis, EvenParity): + return OddParity(input_basis.coord, input_basis.size, input_basis.bounds, input_basis.dealias, input_basis.library) + elif isinstance(input_basis, OddParity): + return EvenParity(input_basis.coord, input_basis.size, input_basis.bounds, input_basis.dealias, input_basis.library) + else: + raise ValueError(f"This should never happen: input_basis = {input_basis}") @staticmethod - def _full_matrix(input_basis, output_basis): - # Build native integration vector - k = input_basis.native_wavenumbers - integ_vector = np.zeros(k.size) - integ_vector[1::2] = 2/k[1::2]/np.pi - # Return with shape (1, N) - return integ_vector[None, :] + def _group_matrix(group, input_basis, output_basis, parity): + # Rescale group (native wavenumber) to get physical wavenumber + k = group / input_basis.COV.stretch + if parity == 1: + # dx cos(k*x) = -k * sin(k*x) + return np.array([[-k]]) + elif parity == -1: + # dx sin(k*x) = k * cos(k*x) + return np.array([[k]]) + else: + raise ValueError(f"This should never happen: parity = {parity}") -class IntegrateCosine(operators.Integrate, operators.SpectralOperator1D): - """RealFourier integration.""" +class IntegrateEvenParity(operators.Integrate, operators.SpectralOperator1D): + """EvenParity basis integration.""" input_coord_type = Coordinate - input_basis_type = Cosine + input_basis_type = EvenParity subaxis_dependence = [True] subaxis_coupling = [False] @@ -1537,69 +1672,84 @@ def _output_basis(input_basis): @staticmethod def _group_matrix(group, input_basis, output_basis): - # integ cos(k*x) = L * δ(k, 0) + # \int_{0}^{\pi} cos(kx) dx = \pi (k=0) if group == 0: L = input_basis.COV.problem_length return np.array([[L]]) else: - raise ValueError("This should never happen.") + # Constructor should only loop over group 0 + raise ValueError(f"This should never happen: group = {group}") -class DifferentiateSine(operators.Differentiate, operators.SpectralOperator1D): - """Sine series differentiation.""" +class IntegrateOddParity(operators.Integrate, operators.SpectralOperator1D): + """OddParity basis integration.""" - input_basis_type = Sine + input_coord_type = Coordinate + input_basis_type = OddParity subaxis_dependence = [True] - subaxis_coupling = [False] + subaxis_coupling = [True] @staticmethod def _output_basis(input_basis): - return Cosine(input_basis.coord, input_basis.size, input_basis.bounds, input_basis.dealias, input_basis.library) + return None @staticmethod - def _group_matrix(group, input_basis, output_basis): - # Rescale group (native wavenumber) to get physical wavenumber - k = group / input_basis.COV.stretch - # dx sin(k*x) = k * cos(k*x) - return np.array([[k]]) + def _full_matrix(input_basis, output_basis): + # Build native integration vector + # \int_{0}^{\pi} sin(kx) dx = (1 - (-1)^k) / k + k = input_basis.native_wavenumbers + k_odd = (k % 2 == 1) + integ_vector = np.zeros(k.size) + integ_vector[k_odd] = 2 / k[k_odd] * input_basis.COV.stretch + # Return with shape (1, N) + return integ_vector[None, :] -class ConvertConstantCosine(operators.ConvertConstant, operators.SpectralOperator1D): - """Upcast constants to Cosine.""" +class AverageEvenParity(operators.Average, operators.SpectralOperator1D): + """EvenParity basis averaging.""" - output_basis_type = Cosine + input_coord_type = Coordinate + input_basis_type = EvenParity subaxis_dependence = [True] subaxis_coupling = [False] + @staticmethod + def _output_basis(input_basis): + return None + @staticmethod def _group_matrix(group, input_basis, output_basis): - # Rescale group (native wavenumber) to get physical wavenumber - k = group / output_basis.COV.stretch - # 1 = cos(0*x) - if k == 0: - unit_amplitude = 1 / output_basis.constant_mode_value - return np.array([[unit_amplitude]]) + # \int_{0}^{\pi} cos(kx) dx = \pi (k=0) + if group == 0: + return np.array([[1]]) else: - return np.zeros(shape=(1, 0)) + # Constructor should only loop over group 0 + raise ValueError(f"This should never happen: group = {group}") -class DifferentiateCosine(operators.Differentiate, operators.SpectralOperator1D): - """Cosine series differentiation.""" +class AverageOddParity(operators.Average, operators.SpectralOperator1D): + """OddParity basis averaging.""" - input_basis_type = Cosine + input_coord_type = Coordinate + input_basis_type = OddParity subaxis_dependence = [True] - subaxis_coupling = [False] + subaxis_coupling = [True] @staticmethod def _output_basis(input_basis): - return Sine(input_basis.coord, input_basis.size, input_basis.bounds, input_basis.dealias, input_basis.library) + return None @staticmethod - def _group_matrix(group, input_basis, output_basis): - # Rescale group (native wavenumber) to get physical wavenumber - k = group / input_basis.COV.stretch - # dx cos(k*x) = -k * sin(k*x) - return np.array([[-k]]) + def _full_matrix(input_basis, output_basis): + # Build native integration vector + # \int_{0}^{\pi} sin(kx) dx = (1 - (-1)^k) / k + k = input_basis.native_wavenumbers + k_odd = (k % 2 == 1) + integ_vector = np.zeros(k.size) + integ_vector[k_odd] = 2 / k[k_odd] + ave_vector = integ_vector / np.pi + # Return with shape (1, N) + return ave_vector[None, :] # class HilbertTransformSine(operators.HilbertTransform): diff --git a/dedalus/core/operators.py b/dedalus/core/operators.py index f700e0eb..20c2d0ed 100644 --- a/dedalus/core/operators.py +++ b/dedalus/core/operators.py @@ -990,6 +990,7 @@ def operate(self, out): # Apply matrix if arg.data.size and out.data.size: data_axis = self.last_axis + len(arg.tensorsig) + # TODO: special case to vector multiply when group size is 1 apply_matrix(self.subspace_matrix(layout), arg.data, data_axis, out=out.data) else: out.data.fill(0) @@ -1196,7 +1197,7 @@ def __init__(self, operand, coord): SpectralOperator.__init__(self, operand) # Require integrand is a scalar if coord in operand.tensorsig: - raise ValueError("Can only integrate scalars.") + raise ValueError("Can only integrate scalar fields.") # SpectralOperator requirements self.coord = coord self.input_basis = operand.domain.get_basis(coord) @@ -1269,7 +1270,7 @@ def __init__(self, operand, coord): SpectralOperator.__init__(self, operand) # Require integrand is a scalar if coord in operand.tensorsig: - raise ValueError("Can only average scalars.") + raise ValueError("Can only average scalar fields.") # SpectralOperator requirements self.coord = coord self.input_basis = operand.domain.get_basis(coord) @@ -1621,10 +1622,6 @@ def replace(self, old, new): # # Simplify operand, skipping conversion # return self.operand.simplify(*vars) - def subspace_matrix(self, layout): - """Build matrix operating on global subspace data.""" - return self._subspace_matrix(layout, self.input_basis, self.output_basis, self.first_axis) - def operate(self, out): """Perform operation.""" arg = self.args[0] diff --git a/dedalus/core/transforms.py b/dedalus/core/transforms.py index 1fc2566d..bf5a12a5 100644 --- a/dedalus/core/transforms.py +++ b/dedalus/core/transforms.py @@ -676,7 +676,7 @@ def wavenumbers(self): return np.arange(self.KM + 1) -@register_transform(basis.Sine, 'matrix') +@register_transform(basis.ParityBase, 'matrix-sin') class SineMMT(SineTransform, SeparableMatrixTransform): """Sine MMT.""" @@ -754,7 +754,7 @@ def resize_rescale_backward(self, data_in, data_out, axis, Kmax): data_out[badfreq] = 0 -@register_transform(basis.Sine, 'scipy') +@register_transform(basis.ParityBase, 'scipy-sin') class ScipyDST(FastSineTransform): """Fast sine transform using scipy.fft.""" @@ -776,7 +776,7 @@ def backward(self, cdata, gdata, axis): np.copyto(gdata, temp) -@register_transform(basis.Sine, 'fftw') +@register_transform(basis.ParityBase, 'fftw-sin') class FFTWDST(FFTWBase, FastSineTransform): """Fast sine transform using FFTW.""" @@ -853,7 +853,7 @@ def wavenumbers(self): return np.arange(self.KM + 1) -@register_transform(basis.Cosine, 'matrix') +@register_transform(basis.ParityBase, 'matrix-cos') class CosineMMT(CosineTransform, SeparableMatrixTransform): """Cosine MMT.""" @@ -923,7 +923,7 @@ def resize_rescale_backward(self, data_in, data_out, axis, Kmax): data_out[badfreq] = 0 -@register_transform(basis.Cosine, 'scipy') +@register_transform(basis.ParityBase, 'scipy-cos') class ScipyDCT(FastCosineTransform): """Fast cosine transform using scipy.fft.""" @@ -945,7 +945,7 @@ def backward(self, cdata, gdata, axis): np.copyto(gdata, temp) -@register_transform(basis.Cosine, 'fftw') +@register_transform(basis.ParityBase, 'fftw-cos') class FFTWDCT(FFTWBase, FastCosineTransform): """Fast cosine transform using FFTW.""" diff --git a/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py b/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py index f8e1de61..91389554 100644 --- a/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py +++ b/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py @@ -31,7 +31,7 @@ # Parameters Lx, Lz = 4, 1 -Nx, Nz = 256, 64 +Nx, Nz = 8, 8 Rayleigh = 2e6 Prandtl = 1 dealias = 3/2 @@ -45,8 +45,8 @@ coords = d3.CartesianCoordinates('x', 'z') dist = d3.Distributor(coords, dtype=dtype) xbasis = d3.RealFourier(coords['x'], size=Nx, bounds=(0, Lx), dealias=dealias) -zsbasis = d3.Sine(coords['z'], size=Nz, bounds=(0, Lz), dealias=dealias, library=library) -zcbasis = d3.Cosine(coords['z'], size=Nz, bounds=(0, Lz), dealias=dealias, library=library) +zsbasis = d3.OddParity(coords['z'], size=Nz, bounds=(0, Lz), dealias=dealias, library=library) +zcbasis = d3.EvenParity(coords['z'], size=Nz, bounds=(0, Lz), dealias=dealias, library=library) # Fields p = dist.Field(name='p', bases=(xbasis,zcbasis)) From a15acad4d47bce3cecb75bfe36a246c5148ab2e9 Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Mon, 10 Nov 2025 18:22:19 -0500 Subject: [PATCH 08/22] Add parity operator tests --- dedalus/tests/test_fourier_operators.py | 132 ++++++++++++++++++++++-- 1 file changed, 125 insertions(+), 7 deletions(-) diff --git a/dedalus/tests/test_fourier_operators.py b/dedalus/tests/test_fourier_operators.py index 018861e3..bb50f196 100644 --- a/dedalus/tests/test_fourier_operators.py +++ b/dedalus/tests/test_fourier_operators.py @@ -21,6 +21,15 @@ def build_fourier(N, bounds, dealias, dtype): return c, d, b, x +@CachedMethod +def build_parity(N, bounds, dealias, parity, dtype): + c = d3.Coordinate('x') + d = d3.Distributor(c, dtype=dtype) + b = d3.Parity(c, size=N, bounds=bounds, dealias=dealias, parity=parity) + x = d.local_grid(b, scale=1) + return c, d, b, x + + @pytest.mark.parametrize('N', N_range) @pytest.mark.parametrize('bounds', bounds_range) @pytest.mark.parametrize('dealias', dealias_range) @@ -36,6 +45,21 @@ def test_fourier_convert_constant(N, bounds, dealias, dtype, layout): assert np.allclose(g['g'], f['g']) +@pytest.mark.parametrize('N', N_range) +@pytest.mark.parametrize('bounds', bounds_range) +@pytest.mark.parametrize('dealias', dealias_range) +@pytest.mark.parametrize('dtype', dtype_range) +@pytest.mark.parametrize('layout', ['g', 'c']) +def test_parity_convert_constant(N, bounds, dealias, dtype, layout): + """Test conversion from constant to EvenParity basis.""" + c, d, b, x = build_parity(N, bounds, dealias, 1, dtype) + f = d.Field() + f['g'] = 1 + f.change_layout(layout) + g = d3.Convert(f, b).evaluate() + assert np.allclose(g['g'], f['g']) + + @pytest.mark.parametrize('N', N_range) @pytest.mark.parametrize('bounds', bounds_range) @pytest.mark.parametrize('dealias', dealias_range) @@ -44,12 +68,35 @@ def test_fourier_differentiate(N, bounds, dealias, dtype): """Test differentiation in Fourier basis.""" c, d, b, x = build_fourier(N, bounds, dealias, dtype) f = d.Field(bases=b) - k = 4 * np.pi / (bounds[1] - bounds[0]) + L = bounds[1] - bounds[0] + k = 4 * np.pi / L f['g'] = 1 + np.sin(k*x+0.1) g = d3.Differentiate(f, c).evaluate() assert np.allclose(g['g'], k*np.cos(k*x+0.1)) +@pytest.mark.parametrize('N', N_range) +@pytest.mark.parametrize('bounds', bounds_range) +@pytest.mark.parametrize('dealias', dealias_range) +@pytest.mark.parametrize('dtype', dtype_range) +@pytest.mark.parametrize('parity', [1, -1]) +def test_parity_differentiate(N, bounds, dealias, dtype, parity): + """Test differentiation in Parity bases.""" + c, d, b, x = build_parity(N, bounds, dealias, parity, dtype) + f = d.Field(bases=b) + x0 = bounds[0] + L = bounds[1] - bounds[0] + k = 3 * np.pi / L + if parity == 1: + f['g'] = 1 + np.cos(k*(x-x0)) + g = d3.Differentiate(f, c).evaluate() + assert np.allclose(g['g'], -k*np.sin(k*(x-x0))) + elif parity == -1: + f['g'] = np.sin(k*(x-x0)) + g = d3.Differentiate(f, c).evaluate() + assert np.allclose(g['g'], k*np.cos(k*(x-x0))) + + @pytest.mark.parametrize('N', N_range) @pytest.mark.parametrize('bounds', bounds_range) @pytest.mark.parametrize('dealias', dealias_range) @@ -58,12 +105,38 @@ def test_fourier_interpolate(N, bounds, dealias, dtype): """Test interpolation in Fourier basis.""" c, d, b, x = build_fourier(N, bounds, dealias, dtype) f = d.Field(bases=b) - k = 4 * np.pi / (bounds[1] - bounds[0]) - f['g'] = 1 + np.sin(k*x+0.1) + L = bounds[1] - bounds[0] + k = 4 * np.pi / L + f0 = lambda x: 1 + np.sin(k*x+0.1) + f['g'] = f0(x) + results = [] + for p in [bounds[0], bounds[1], bounds[0] + L*np.random.rand()]: + g = d3.Interpolate(f, c, p).evaluate() + results.append(np.allclose(g['g'], f0(p))) + assert all(results) + + +@pytest.mark.parametrize('N', N_range) +@pytest.mark.parametrize('bounds', bounds_range) +@pytest.mark.parametrize('dealias', dealias_range) +@pytest.mark.parametrize('dtype', dtype_range) +@pytest.mark.parametrize('parity', [1, -1]) +def test_parity_interpolate(N, bounds, dealias, dtype, parity): + """Test interpolation in Parity bases.""" + c, d, b, x = build_parity(N, bounds, dealias, parity, dtype) + f = d.Field(bases=b) + x0 = bounds[0] + L = bounds[1] - bounds[0] + k = 3 * np.pi / L + if parity == 1: + f0 = lambda x: 1 + np.cos(k*(x-x0)) + elif parity == -1: + f0 = lambda x: np.sin(k*(x-x0)) + f['g'] = f0(x) results = [] for p in [bounds[0], bounds[1], bounds[0] + (bounds[1] - bounds[0]) * np.random.rand()]: g = d3.Interpolate(f, c, p).evaluate() - results.append(np.allclose(g['g'], 1 + np.sin(k*p+0.1))) + results.append(np.allclose(g['g'], f0(p))) assert all(results) @@ -75,10 +148,33 @@ def test_fourier_integrate(N, bounds, dealias, dtype): """Test integration in Fourier basis.""" c, d, b, x = build_fourier(N, bounds, dealias, dtype) f = d.Field(bases=b) - k = 4 * np.pi / (bounds[1] - bounds[0]) + L = bounds[1] - bounds[0] + k = 4 * np.pi / L f['g'] = 1 + np.sin(k*x+0.1) g = d3.Integrate(f, c).evaluate() - assert np.allclose(g['g'], bounds[1] - bounds[0]) + assert np.allclose(g['g'], L) + + +@pytest.mark.parametrize('N', N_range) +@pytest.mark.parametrize('bounds', bounds_range) +@pytest.mark.parametrize('dealias', dealias_range) +@pytest.mark.parametrize('dtype', dtype_range) +@pytest.mark.parametrize('parity', [1, -1]) +def test_parity_integrate(N, bounds, dealias, dtype, parity): + """Test integration in Parity bases.""" + c, d, b, x = build_parity(N, bounds, dealias, parity, dtype) + f = d.Field(bases=b) + x0 = bounds[0] + L = bounds[1] - bounds[0] + k = 3 * np.pi / L + if parity == 1: + f['g'] = 1 + np.cos(k*(x-x0)) + g = d3.Integrate(f, c).evaluate() + assert np.allclose(g['g'], L) + elif parity == -1: + f['g'] = np.sin(k*(x-x0)) + g = d3.Integrate(f, c).evaluate() + assert np.allclose(g['g'], L*2/3/np.pi) @pytest.mark.parametrize('N', N_range) @@ -89,8 +185,30 @@ def test_fourier_average(N, bounds, dealias, dtype): """Test averaging in Fourier basis.""" c, d, b, x = build_fourier(N, bounds, dealias, dtype) f = d.Field(bases=b) - k = 4 * np.pi / (bounds[1] - bounds[0]) + L = bounds[1] - bounds[0] + k = 4 * np.pi / L f['g'] = 1 + np.sin(k*x+0.1) g = d3.Average(f, c).evaluate() assert np.allclose(g['g'], 1) + +@pytest.mark.parametrize('N', N_range) +@pytest.mark.parametrize('bounds', bounds_range) +@pytest.mark.parametrize('dealias', dealias_range) +@pytest.mark.parametrize('dtype', dtype_range) +@pytest.mark.parametrize('parity', [1, -1]) +def test_parity_average(N, bounds, dealias, dtype, parity): + """Test averaging in Parity bases.""" + c, d, b, x = build_parity(N, bounds, dealias, parity, dtype) + f = d.Field(bases=b) + x0 = bounds[0] + L = bounds[1] - bounds[0] + k = 3 * np.pi / L + if parity == 1: + f['g'] = 1 + np.cos(k*(x-x0)) + g = d3.Average(f, c).evaluate() + assert np.allclose(g['g'], 1) + elif parity == -1: + f['g'] = np.sin(k*(x-x0)) + g = d3.Average(f, c).evaluate() + assert np.allclose(g['g'], 2/3/np.pi) From 3c855f6599e7d6d9591f34d533b12171e0784898 Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Mon, 10 Nov 2025 18:22:36 -0500 Subject: [PATCH 09/22] Clean up parity transforms and add transform tests --- dedalus/core/transforms.py | 100 +++++++++++++++---------------- dedalus/tests/test_transforms.py | 44 ++++++++++++++ 2 files changed, 91 insertions(+), 53 deletions(-) diff --git a/dedalus/core/transforms.py b/dedalus/core/transforms.py index bf5a12a5..8c1376b3 100644 --- a/dedalus/core/transforms.py +++ b/dedalus/core/transforms.py @@ -642,21 +642,24 @@ class SineTransform(SeparableTransform): Notes ----- - Let KN = (N - 1) be the maximum (Nyquist) mode on the grid. + Let KN = (N - 1) be the maximum fully resolved (non-Nyquist) mode on the grid. Let KM = (M - 1) be the maximum retained mode in coeff space. Then K = min(KN, KM) is the maximum wavenumber used in the transforms. A unit-amplitude normalization is used. + Grid: + x(n) = \pi (1/2 + n) / N, n = 0 .. (N-1) + Forward transform: if k == 0: a(k) = 0 elif k <= K: - a(k) = (2/N) \sum_{x=0}^{N-1} f(x) \sin(\pi k x / N) + a(k) = (2/N) \sum_{n=0}^{N-1} f(n) \sin(k x(n)) else: a(k) = 0 Backward transform: - f(x) = \sum_{k=1}^{K} a(k) \sin(\pi k x / N) + f(n) = \sum_{k=1}^{K} a(k) \sin(k x(n)) Coefficient ordering: The sine coefficients are ordered simply as @@ -673,7 +676,7 @@ def __init__(self, grid_size, coeff_size): @property def wavenumbers(self): """One-dimensional global wavenumber array.""" - return np.arange(self.KM + 1) + return np.arange(self.M) @register_transform(basis.ParityBase, 'matrix-sin') @@ -684,8 +687,6 @@ class SineMMT(SineTransform, SeparableMatrixTransform): def forward_matrix(self): """Build forward transform matrix.""" N = self.N - M = self.M - Kmax = self.Kmax K = self.wavenumbers[:, None] X = np.arange(N)[None, :] + 1/2 dX = N / np.pi @@ -699,8 +700,6 @@ def forward_matrix(self): def backward_matrix(self): """Build backward transform matrix.""" N = self.N - M = self.M - Kmax = self.Kmax K = self.wavenumbers[None, :] X = np.arange(N)[:, None] + 1/2 dX = N / np.pi @@ -717,41 +716,35 @@ class FastSineTransform(SineTransform): def __init__(self, *args, **kw): super().__init__(*args, **kw) # Standard scaling factors for unit-amplitude normalization from DST-II - self.forward_rescale_zero = 0 - self.forward_rescale_pos = 1 / self.N - self.backward_rescale_zero = 0 - self.backward_rescale_pos = 1 / 2 + self.forward_rescale = 1 / self.N + self.backward_rescale = 1 / 2 def resize_rescale_forward(self, data_in, data_out, axis, Kmax): """Resize by padding/trunction and rescale to unit amplitude.""" - - # there is no zero frequency in the DST... zerofreq = axslice(axis, 0, 1) - Nyfreq = axslice(axis,Kmax,Kmax+1) - data_out[zerofreq] *= self.forward_rescale_zero - #data_out[Nyfreq] *= 0 + data_out[zerofreq] = 0 if Kmax > 0: - in_posfreq = axslice(axis, 0, Kmax) - out_posfreq = axslice(axis, 1, Kmax+1) - np.multiply(data_in[in_posfreq], self.forward_rescale_pos, data_out[out_posfreq]) - if self.KM > Kmax: - badfreq = axslice(axis, Kmax+1, None) - data_out[badfreq] = 0 + # Shift to account for zero frequency in output + posfreq_in = axslice(axis, 0, Kmax) + posfreq_out = axslice(axis, 1, Kmax+1) + np.multiply(data_in[posfreq_in], self.forward_rescale, data_out[posfreq_out]) + if self.KM > Kmax: + badfreq = axslice(axis, Kmax+1, None) + data_out[badfreq] = 0 def resize_rescale_backward(self, data_in, data_out, axis, Kmax): """Resize by padding/trunction and rescale to unit amplitude.""" - # zerofreq = axslice(axis, 0, 1) - # np.multiply(data_in[zerofreq], self.backward_rescale_zero, data_out[zerofreq]) - in_posfreq = axslice(axis, 1, Kmax+1) - out_posfreq = axslice(axis, 0, Kmax) - Nyfreq = axslice(axis,Kmax,Kmax+1) - data_out[Nyfreq] = 0 - if Kmax > 0: - np.multiply(data_in[in_posfreq], self.backward_rescale_pos, data_out[out_posfreq]) - - if self.KN > Kmax: - badfreq = axslice(axis, Kmax+1, None) - data_out[badfreq] = 0 + if Kmax == 0: + zerofreq = axslice(axis, 0, 1) + data_out[zerofreq] = 0 + else: + # Shift to account for zero frequency in input + posfreq_in = axslice(axis, 1, Kmax+1) + posfreq_out = axslice(axis, 0, Kmax) + np.multiply(data_in[posfreq_in], self.backward_rescale, data_out[posfreq_out]) + if self.KN >= Kmax: + badfreq = axslice(axis, Kmax, None) + data_out[badfreq] = 0 @register_transform(basis.ParityBase, 'scipy-sin') @@ -761,6 +754,7 @@ class ScipyDST(FastSineTransform): def forward(self, gdata, cdata, axis): """Apply forward transform along specified axis.""" # Call DST + # Avoid overwrite_x to prevent overwriting problems temp = scipy.fft.dst(gdata, type=2, axis=axis) # Creates temporary # Resize and rescale for unit-ampltidue normalization self.resize_rescale_forward(temp, cdata, axis, self.Kmax) @@ -772,7 +766,7 @@ def backward(self, cdata, gdata, axis): temp = np.empty_like(gdata) # Creates temporary self.resize_rescale_backward(cdata, temp, axis, self.Kmax) # Call IDST - temp = scipy.fft.dst(temp, type=3, axis=axis, overwrite_x=True) # Creates temporary + temp = scipy.fft.dst(temp, type=3, axis=axis, overwrite_x=True) np.copyto(gdata, temp) @@ -819,21 +813,24 @@ class CosineTransform(SeparableTransform): Notes ----- - Let KN = (N - 1) be the maximum (Nyquist) mode on the grid. + Let KN = (N - 1) be the maximum fully resolved (non-Nyquist) mode on the grid. Let KM = (M - 1) be the maximum retained mode in coeff space. Then K = min(KN, KM) is the maximum wavenumber used in the transforms. A unit-amplitude normalization is used. + Grid: + x(n) = \pi (1/2 + n) / N, n = 0 .. (N-1) + Forward transform: if k == 0: - a(k) = (1/N) \sum_{x=0}^{N-1} f(x) + a(k) = (1/N) \sum_{n=0}^{N-1} f(n) elif k <= K: - a(k) = (2/N) \sum_{x=0}^{N-1} f(x) \cos(\pi k x / N) + a(k) = (2/N) \sum_{n=0}^{N-1} f(n) \cos(k x(n)) else: a(k) = 0 Backward transform: - f(x) = \sum_{k=0}^{K} a(k) \cos(\pi k x / N) + f(n) = \sum_{k=0}^{K} a(k) \cos(k x(n)) Coefficient ordering: The cosine coefficients are ordered simply as @@ -850,7 +847,7 @@ def __init__(self, grid_size, coeff_size): @property def wavenumbers(self): """One-dimensional global wavenumber array.""" - return np.arange(self.KM + 1) + return np.arange(self.M) @register_transform(basis.ParityBase, 'matrix-cos') @@ -861,8 +858,6 @@ class CosineMMT(CosineTransform, SeparableMatrixTransform): def forward_matrix(self): """Build forward transform matrix.""" N = self.N - M = self.M - Kmax = self.Kmax K = self.wavenumbers[:, None] X = np.arange(N)[None, :] + 1/2 dX = N / np.pi @@ -877,8 +872,6 @@ def forward_matrix(self): def backward_matrix(self): """Build backward transform matrix.""" N = self.N - M = self.M - Kmax = self.Kmax K = self.wavenumbers[None, :] X = np.arange(N)[:, None] + 1/2 dX = N / np.pi @@ -907,9 +900,9 @@ def resize_rescale_forward(self, data_in, data_out, axis, Kmax): if Kmax > 0: posfreq = axslice(axis, 1, Kmax+1) np.multiply(data_in[posfreq], self.forward_rescale_pos, data_out[posfreq]) - if self.KM > Kmax: - badfreq = axslice(axis, Kmax+1, None) - data_out[badfreq] = 0 + if self.KM > Kmax: + badfreq = axslice(axis, Kmax+1, None) + data_out[badfreq] = 0 def resize_rescale_backward(self, data_in, data_out, axis, Kmax): """Resize by padding/trunction and rescale to unit amplitude.""" @@ -918,9 +911,9 @@ def resize_rescale_backward(self, data_in, data_out, axis, Kmax): if Kmax > 0: posfreq = axslice(axis, 1, Kmax+1) np.multiply(data_in[posfreq], self.backward_rescale_pos, data_out[posfreq]) - if self.KN > Kmax: - badfreq = axslice(axis, Kmax+1, None) - data_out[badfreq] = 0 + if self.KN > Kmax: + badfreq = axslice(axis, Kmax+1, None) + data_out[badfreq] = 0 @register_transform(basis.ParityBase, 'scipy-cos') @@ -930,8 +923,9 @@ class ScipyDCT(FastCosineTransform): def forward(self, gdata, cdata, axis): """Apply forward transform along specified axis.""" # Call DCT + # Avoid overwrite_x to prevent overwriting problems temp = scipy.fft.dct(gdata, type=2, axis=axis) # Creates temporary - # Resize and rescale for unit-ampltidue normalization + # Resize and rescale for unit-amplitude normalization self.resize_rescale_forward(temp, cdata, axis, self.Kmax) def backward(self, cdata, gdata, axis): @@ -941,7 +935,7 @@ def backward(self, cdata, gdata, axis): temp = np.empty_like(gdata) # Creates temporary self.resize_rescale_backward(cdata, temp, axis, self.Kmax) # Call IDCT - temp = scipy.fft.dct(temp, type=3, axis=axis, overwrite_x=True) # Creates temporary + temp = scipy.fft.dct(temp, type=3, axis=axis, overwrite_x=True) np.copyto(gdata, temp) diff --git a/dedalus/tests/test_transforms.py b/dedalus/tests/test_transforms.py index c576ae95..df478109 100644 --- a/dedalus/tests/test_transforms.py +++ b/dedalus/tests/test_transforms.py @@ -57,6 +57,50 @@ def test_real_fourier_libraries_forward(N, dealias, dtype, library): assert np.allclose(u_mat['c'], u_lib['c']) +@pytest.mark.parametrize('N', [16]) +@pytest.mark.parametrize('parity', [1, -1]) +@pytest.mark.parametrize('dealias', [0.5, 1, 1.5]) +@pytest.mark.parametrize('dtype', [np.float64, np.complex128]) +@pytest.mark.parametrize('library', ['scipy', 'fftw']) +def test_parity_libraries_backward(N, parity, dealias, dtype, library): + """Tests that fast real Fourier transforms match matrix transforms.""" + c = coords.Coordinate('x') + d = distributor.Distributor([c]) + # Matrix + b_mat = basis.Parity(c, size=N, bounds=(0, 2*np.pi), parity=parity, dealias=dealias, library='matrix') + u_mat = field.Field(dist=d, bases=(b_mat,), dtype=dtype) + u_mat.preset_scales(dealias) + u_mat['c'] = np.random.randn(N) + # Library + b_lib = basis.Parity(c, size=N, bounds=(0, 2*np.pi), parity=parity, dealias=dealias, library=library) + u_lib = field.Field(dist=d, bases=(b_lib,), dtype=dtype) + u_lib.preset_scales(dealias) + u_lib['c'] = u_mat['c'] + assert np.allclose(u_mat['g'], u_lib['g']) + + +@pytest.mark.parametrize('N', [16]) +@pytest.mark.parametrize('parity', [1, -1]) +@pytest.mark.parametrize('dealias', [0.5, 1, 1.5]) +@pytest.mark.parametrize('dtype', [np.float64, np.complex128]) +@pytest.mark.parametrize('library', ['scipy', 'fftw']) +def test_parity_libraries_forward(N, parity, dealias, dtype, library): + """Tests that fast real Fourier transforms match matrix transforms.""" + c = coords.Coordinate('x') + d = distributor.Distributor([c]) + # Matrix + b_mat = basis.Parity(c, size=N, bounds=(0, 2*np.pi), parity=parity, dealias=dealias, library='matrix') + u_mat = field.Field(dist=d, bases=(b_mat,), dtype=dtype) + u_mat.preset_scales(dealias) + u_mat['g'] = np.random.randn(int(np.ceil(dealias * N))) + # Library + b_lib = basis.Parity(c, size=N, bounds=(0, 2*np.pi), parity=parity, dealias=dealias, library=library) + u_lib = field.Field(dist=d, bases=(b_lib,), dtype=dtype) + u_lib.preset_scales(dealias) + u_lib['g'] = u_mat['g'] + assert np.allclose(u_mat['c'], u_lib['c']) + + @pytest.mark.parametrize('N', N_range) @pytest.mark.parametrize('dealias', dealias_range) def test_CF_scalar_roundtrip(N, dealias): From 19dde86a33af8cee8a2c66c5c5dcdfd80fe5a25a Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Wed, 17 Apr 2024 22:18:09 -0400 Subject: [PATCH 10/22] Implement Hilbert transforms for fourier bases --- dedalus/core/basis.py | 40 +++++++++++++++++++++++++++++++++++++ dedalus/core/operators.py | 42 +++++++++++++++++++++++++++++++-------- 2 files changed, 74 insertions(+), 8 deletions(-) diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index 3773799e..752c0da4 100644 --- a/dedalus/core/basis.py +++ b/dedalus/core/basis.py @@ -1037,6 +1037,25 @@ def _group_matrix(group, input_basis, output_basis): return np.array([[1j*k]]) +class HilbertTransformComplexFourier(operators.HilbertTransform, operators.SpectralOperator1D): + """ComplexFourier Hilbert transform.""" + + input_basis_type = ComplexFourier + subaxis_dependence = [True] + subaxis_coupling = [False] + + @staticmethod + def _output_basis(input_basis): + return input_basis + + @staticmethod + def _group_matrix(group, input_basis, output_basis): + # Rescale group (native wavenumber) to get physical wavenumber + k = group / input_basis.COV.stretch + # Hx exp(1j*k*x) = -1j * sgn(k) * exp(1j*k*x) + return np.array([[-1j*np.sign(k)]]) + + class InterpolateComplexFourier(operators.Interpolate, operators.SpectralOperator1D): """ComplexFourier interpolation.""" @@ -1225,6 +1244,27 @@ def _group_matrix(group, input_basis, output_basis): [k, 0]]) +class HilbertTransformRealFourier(operators.HilbertTransform, operators.SpectralOperator1D): + """RealFourier Hilbert transform.""" + + input_basis_type = RealFourier + subaxis_dependence = [True] + subaxis_coupling = [False] + + @staticmethod + def _output_basis(input_basis): + return input_basis + + @staticmethod + def _group_matrix(group, input_basis, output_basis): + # Rescale group (native wavenumber) to get physical wavenumber + k = group / input_basis.COV.stretch + # Hx cos(n*x) = sin(n*x) + # Hx -sin(n*x) = cos(n*x) + return np.array([[0, 1], + [1, 0]]) + + class InterpolateRealFourier(operators.Interpolate, operators.SpectralOperator1D): """RealFourier interpolation.""" diff --git a/dedalus/core/operators.py b/dedalus/core/operators.py index 20c2d0ed..5dbd0674 100644 --- a/dedalus/core/operators.py +++ b/dedalus/core/operators.py @@ -1442,33 +1442,59 @@ class HilbertTransform(SpectralOperator1D, metaclass=MultiClass): """ + name = "Hilbert" + + def __init__(self, operand, coord, out=None): + super().__init__(operand, out=out) + # SpectralOperator requirements + self.coord = coord + self.input_basis = operand.domain.get_basis(coord) + self.output_basis = self._output_basis(self.input_basis) + self.first_axis = self.dist.get_axis(coord) + self.last_axis = self.first_axis + self.axis = self.first_axis + # LinearOperator requirements + self.operand = operand + # FutureField requirements + self.domain = operand.domain.substitute_basis(self.input_basis, self.output_basis) + self.tensorsig = operand.tensorsig + self.dtype = operand.dtype + @classmethod - def _check_args(cls, operand, space, out=None): + def _check_args(cls, operand, coord, out=None): # Dispatch by operand basis if isinstance(operand, Operand): - if isinstance(operand.get_basis(space), cls.input_basis_type): + basis = operand.domain.get_basis(coord) + if isinstance(basis, cls.input_basis_type): return True return False - @property - def base(self): - return HilbertTransform + def new_operand(self, operand, **kw): + return HilbertTransform(operand, self.coord, **kw) + + @staticmethod + def _output_basis(input_basis): + # Subclasses must implement + raise NotImplementedError() + + def __str__(self): + return 'H{!s}({!s})'.format(self.coord.name, self.operand) class HilbertTransformConstant(HilbertTransform): """Constant Hilbert transform.""" @classmethod - def _check_args(cls, operand, space, out=None): + def _check_args(cls, operand, coord, out=None): # Dispatch for numbers of constant bases if isinstance(operand, Number): return True if isinstance(operand, Operand): - if operand.get_basis(space) is None: + if operand.domain.get_basis(coord) is None: return True return False - def __new__(cls, operand, space, out=None): + def __new__(cls, operand, coord, out=None): return 0 From 5559fd8f2098d487848cd5597cd9b670059a7016 Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Wed, 17 Apr 2024 23:00:50 -0400 Subject: [PATCH 11/22] Add order parameter to Differentiate operator --- dedalus/core/basis.py | 42 ++++++++++++++++++++++++++------------- dedalus/core/operators.py | 38 ++++++++++++++++++++++------------- 2 files changed, 52 insertions(+), 28 deletions(-) diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index 752c0da4..01d8742a 100644 --- a/dedalus/core/basis.py +++ b/dedalus/core/basis.py @@ -708,16 +708,26 @@ class DifferentiateJacobi(operators.Differentiate, operators.SpectralOperator1D) subaxis_dependence = [True] subaxis_coupling = [True] + @classmethod + def _check_args(cls, operand, coord, order=1, out=None): + # Only integer derivatives implemented + if float(order).is_integer() and order > 0: + return super()._check_args(operand, coord, order=order, out=out) + return False + @staticmethod - def _output_basis(input_basis): - return input_basis.derivative_basis(order=1) + def _output_basis(input_basis, order): + return input_basis.derivative_basis(order=int(order)) @staticmethod @CachedMethod - def _full_matrix(input_basis, output_basis): + def _full_matrix(input_basis, output_basis, order): N = input_basis.size a, b = input_basis.a, input_basis.b - matrix = jacobi.differentiation_matrix(N, a, b) / input_basis.COV.stretch + matrix = jacobi.differentiation_matrix(N, a, b) + for i in range(1, int(order)): + matrix = jacobi.differentiation_matrix(N, a+i, b+i) @ matrix + matrix *= input_basis.COV.stretch ** (-order) return matrix.tocsr() @@ -1026,15 +1036,16 @@ class DifferentiateComplexFourier(operators.Differentiate, operators.SpectralOpe subaxis_coupling = [False] @staticmethod - def _output_basis(input_basis): + def _output_basis(input_basis, order): return input_basis @staticmethod - def _group_matrix(group, input_basis, output_basis): + def _group_matrix(group, input_basis, output_basis, order): # Rescale group (native wavenumber) to get physical wavenumber k = group / input_basis.COV.stretch - # dx exp(1j*k*x) = 1j * k * exp(1j*k*x) - return np.array([[1j*k]]) + # dx**n exp(1j*k*x) = (1j*k)**n * exp(1j*k*x) + S = (1j * k) ** order + return np.array([[S]]) class HilbertTransformComplexFourier(operators.HilbertTransform, operators.SpectralOperator1D): @@ -1231,17 +1242,20 @@ class DifferentiateRealFourier(operators.Differentiate, operators.SpectralOperat subaxis_coupling = [False] @staticmethod - def _output_basis(input_basis): + def _output_basis(input_basis, order): return input_basis @staticmethod - def _group_matrix(group, input_basis, output_basis): + def _group_matrix(group, input_basis, output_basis, order): # Rescale group (native wavenumber) to get physical wavenumber k = group / input_basis.COV.stretch - # dx cos(k*x) = k * -sin(k*x) - # dx -sin(k*x) = -k * cos(k*x) - return np.array([[0, -k], - [k, 0]]) + # dx**n exp(1j*k*x) = (1j*k)**n * exp(1j*k*x) = S * exp(1j*k*x) + # dx**n [cos(k*x) + 1j*sin(k*x)] = S * [cos(k*x) + 1j*sin(k*x)] + # dx**n cos(k*x) = R(S) * cos(k*x) + I(S) * -sin(k*x) + # dx**n -sin(k*x) = R(S) * -sin(k*x) - I(S) * cos(k*x) + S = (1j * k) ** order + return np.array([[S.real, -S.imag], + [S.imag, S.real]]) class HilbertTransformRealFourier(operators.HilbertTransform, operators.SpectralOperator1D): diff --git a/dedalus/core/operators.py b/dedalus/core/operators.py index 5dbd0674..e5758943 100644 --- a/dedalus/core/operators.py +++ b/dedalus/core/operators.py @@ -1355,12 +1355,13 @@ class Differentiate(SpectralOperator1D, metaclass=MultiClass): name = "Diff" - def __init__(self, operand, coord, out=None): + def __init__(self, operand, coord, order=1, out=None): super().__init__(operand, out=out) + self.order = order # SpectralOperator requirements self.coord = coord self.input_basis = operand.domain.get_basis(coord) - self.output_basis = self._output_basis(self.input_basis) + self.output_basis = self._output_basis(self.input_basis, self.order) self.first_axis = self.dist.get_axis(coord) self.last_axis = self.first_axis self.axis = self.first_axis @@ -1372,7 +1373,7 @@ def __init__(self, operand, coord, out=None): self.dtype = operand.dtype @classmethod - def _check_args(cls, operand, coord, out=None): + def _check_args(cls, operand, coord, order=1, out=None): # Dispatch by operand basis if isinstance(operand, Operand): basis = operand.domain.get_basis(coord) @@ -1381,29 +1382,38 @@ def _check_args(cls, operand, coord, out=None): return False def new_operand(self, operand, **kw): - return Differentiate(operand, self.coord, **kw) + return Differentiate(operand, self.coord, self.order, **kw) + + def subspace_matrix(self, layout): + return self._subspace_matrix(layout, self.input_basis, self.output_basis, self.first_axis, self.order) + + def group_matrix(self, group): + return self._group_matrix(group, self.input_basis, self.output_basis, self.order) @staticmethod - def _output_basis(input_basis): + def _output_basis(input_basis, order): # Subclasses must implement raise NotImplementedError() def __str__(self): - return 'd{!s}({!s})'.format(self.coord.name, self.operand) + if self.order == 1: + return 'd{!s}({!s})'.format(self.coord.name, self.operand) + else: + return 'd{!s}({!s},{!s})'.format(self.coord.name, self.operand, self.order) - def _expand_multiply(self, operand, vars): - """Expand over multiplication.""" - args = operand.args - # Apply product rule to factors - partial_diff = lambda i: prod([self.new_operand(arg) if i==j else arg for j,arg in enumerate(args)]) - return sum((partial_diff(i) for i in range(len(args)))) + # def _expand_multiply(self, operand, vars): + # """Expand over multiplication.""" + # args = operand.args + # # Apply product rule to factors + # partial_diff = lambda i: prod([self.new_operand(arg) if i==j else arg for j,arg in enumerate(args)]) + # return sum((partial_diff(i) for i in range(len(args)))) class DifferentiateConstant(Differentiate): """Constant differentiation.""" @classmethod - def _check_args(cls, operand, coord, out=None): + def _check_args(cls, operand, coord, order=1, out=None): # Dispatch for numbers of constant bases if isinstance(operand, Number): return True @@ -1412,7 +1422,7 @@ def _check_args(cls, operand, coord, out=None): return True return False - def __new__(cls, operand, coord, out=None): + def __new__(cls, operand, coord, order=1, out=None): return 0 From 44579a73eeaee76bcc7cd354b00cedaf800c20ac Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Thu, 18 Apr 2024 12:25:26 -0400 Subject: [PATCH 12/22] Change Cartesian Laplacian to use 2nd order derivatives --- dedalus/core/operators.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dedalus/core/operators.py b/dedalus/core/operators.py index e5758943..ee3adf2e 100644 --- a/dedalus/core/operators.py +++ b/dedalus/core/operators.py @@ -4054,7 +4054,7 @@ def __init__(self, operand, coordsys, out=None): # Wrap to handle gradient wrt single coordinate if isinstance(coordsys, coords.Coordinate): coordsys = coords.CartesianCoordinates(coordsys.name) - parts = [Differentiate(Differentiate(operand, c), c) for c in coordsys.coords] + parts = [Differentiate(operand, c, order=2) for c in coordsys.coords] arg = sum(parts) LinearOperator.__init__(self, arg, out=out) self.coordsys = coordsys From 2823db454d8e79ac98238438b830f1e5bd4650de Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Thu, 18 Apr 2024 12:32:58 -0400 Subject: [PATCH 13/22] Add Riesz derivative --- dedalus/core/basis.py | 44 +++++++++++++++++++++++++++++ dedalus/core/operators.py | 59 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 103 insertions(+) diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index 01d8742a..e2ae353f 100644 --- a/dedalus/core/basis.py +++ b/dedalus/core/basis.py @@ -1048,6 +1048,26 @@ def _group_matrix(group, input_basis, output_basis, order): return np.array([[S]]) +class RieszDerivativeComplexFourier(operators.RieszDerivative, operators.SpectralOperator1D): + """ComplexFourier Riesz derivative.""" + + input_basis_type = ComplexFourier + subaxis_dependence = [True] + subaxis_coupling = [False] + + @staticmethod + def _output_basis(input_basis, order): + return input_basis + + @staticmethod + def _group_matrix(group, input_basis, output_basis, order): + # Rescale group (native wavenumber) to get physical wavenumber + k = group / input_basis.COV.stretch + # R_a exp(1j*k*x) = - |k|**a * exp(1j*k*x) + S = - abs(k) ** order + return np.array([[S]]) + + class HilbertTransformComplexFourier(operators.HilbertTransform, operators.SpectralOperator1D): """ComplexFourier Hilbert transform.""" @@ -1258,6 +1278,30 @@ def _group_matrix(group, input_basis, output_basis, order): [S.imag, S.real]]) +class RieszDerivativeRealFourier(operators.RieszDerivative, operators.SpectralOperator1D): + """RealFourier Riesz derivative.""" + + input_basis_type = RealFourier + subaxis_dependence = [True] + subaxis_coupling = [False] + + @staticmethod + def _output_basis(input_basis, order): + return input_basis + + @staticmethod + def _group_matrix(group, input_basis, output_basis, order): + # Rescale group (native wavenumber) to get physical wavenumber + k = group / input_basis.COV.stretch + # R_a exp(1j*k*x) = - |k|**n * exp(1j*k*x) = S * exp(1j*k*x) + # R_a [cos(k*x) + 1j*sin(k*x)] = S * [cos(k*x) + 1j*sin(k*x)] + # R_a cos(k*x) = S * cos(k*x) + # R_a -sin(k*x) = S * -sin(k*x) + S = - abs(k) ** order + return np.array([[S, 0], + [0, S]]) + + class HilbertTransformRealFourier(operators.HilbertTransform, operators.SpectralOperator1D): """RealFourier Hilbert transform.""" diff --git a/dedalus/core/operators.py b/dedalus/core/operators.py index ee3adf2e..67fee47d 100644 --- a/dedalus/core/operators.py +++ b/dedalus/core/operators.py @@ -1508,6 +1508,65 @@ def __new__(cls, operand, coord, out=None): return 0 +class RieszDerivative(SpectralOperator1D, metaclass=MultiClass): + """ + Riesz derivative along one dimension. + + Parameters + ---------- + operand : number or Operand object + space : Space object + + Notes + ----- + R_a exp(1j*k*x) = - |k|**a * exp(1j*k*x). + + """ + + name = "Riesz" + + def __init__(self, operand, coord, order, out=None): + super().__init__(operand, out=out) + self.order = order + # SpectralOperator requirements + self.coord = coord + self.input_basis = operand.domain.get_basis(coord) + self.output_basis = self._output_basis(self.input_basis, order) + self.first_axis = self.dist.get_axis(coord) + self.last_axis = self.first_axis + self.axis = self.first_axis + # LinearOperator requirements + self.operand = operand + # FutureField requirements + self.domain = operand.domain.substitute_basis(self.input_basis, self.output_basis) + self.tensorsig = operand.tensorsig + self.dtype = operand.dtype + + @classmethod + def _check_args(cls, operand, coord, out=None): + # Dispatch by operand basis + if isinstance(operand, Operand): + basis = operand.domain.get_basis(coord) + if isinstance(basis, cls.input_basis_type): + return True + return False + + def new_operand(self, operand, **kw): + return RieszDerivative(operand, self.coord, self.order, **kw) + + def subspace_matrix(self, layout): + return self._subspace_matrix(layout, self.input_basis, self.output_basis, self.first_axis, self.order) + + def group_matrix(self, group): + return self._group_matrix(group, self.input_basis, self.output_basis, self.order) + + def __str__(self): + if self.order == 1: + return 'R{!s}({!s})'.format(self.coord.name, self.operand) + else: + return 'R{!s}({!s},{!s})'.format(self.coord.name, self.operand, self.order) + + def convert(arg, bases): # Skip for numbers if isinstance(arg, Number): From f7d30fd57bdfcd0ae0a8be24fa56173fc57459d1 Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Fri, 19 Apr 2024 15:50:46 -0400 Subject: [PATCH 14/22] Fix real hilbert symbol --- dedalus/core/basis.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index e2ae353f..b0a4f578 100644 --- a/dedalus/core/basis.py +++ b/dedalus/core/basis.py @@ -1319,8 +1319,8 @@ def _group_matrix(group, input_basis, output_basis): k = group / input_basis.COV.stretch # Hx cos(n*x) = sin(n*x) # Hx -sin(n*x) = cos(n*x) - return np.array([[0, 1], - [1, 0]]) + return np.array([[0, -1], + [1, 0]]) class InterpolateRealFourier(operators.Interpolate, operators.SpectralOperator1D): From 67b22bb4fe601257a0178a3012056a6aa0d74784 Mon Sep 17 00:00:00 2001 From: Daniel Lecoanet Date: Sun, 21 Apr 2024 18:42:57 -0500 Subject: [PATCH 15/22] Fix real hilbert symbol --- dedalus/core/basis.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index b0a4f578..51a9d609 100644 --- a/dedalus/core/basis.py +++ b/dedalus/core/basis.py @@ -1319,8 +1319,8 @@ def _group_matrix(group, input_basis, output_basis): k = group / input_basis.COV.stretch # Hx cos(n*x) = sin(n*x) # Hx -sin(n*x) = cos(n*x) - return np.array([[0, -1], - [1, 0]]) + return np.array([[ 0, 1], + [-1, 0]]) class InterpolateRealFourier(operators.Interpolate, operators.SpectralOperator1D): From daad73d03058dfaf5cf2c2690ef7696be3a3cd64 Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Wed, 12 Nov 2025 08:32:36 -0500 Subject: [PATCH 16/22] Add high order parity differentiate and fourier/parity high order differentiation tests --- dedalus/core/basis.py | 44 +++++++++++++++++-------- dedalus/tests/test_fourier_operators.py | 33 ++++++++++++++----- 2 files changed, 55 insertions(+), 22 deletions(-) diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index 51a9d609..a8fd968d 100644 --- a/dedalus/core/basis.py +++ b/dedalus/core/basis.py @@ -1084,7 +1084,8 @@ def _group_matrix(group, input_basis, output_basis): # Rescale group (native wavenumber) to get physical wavenumber k = group / input_basis.COV.stretch # Hx exp(1j*k*x) = -1j * sgn(k) * exp(1j*k*x) - return np.array([[-1j*np.sign(k)]]) + S = -1j * np.sign(k) + return np.array([[S]]) class InterpolateComplexFourier(operators.Interpolate, operators.SpectralOperator1D): @@ -1315,8 +1316,6 @@ def _output_basis(input_basis): @staticmethod def _group_matrix(group, input_basis, output_basis): - # Rescale group (native wavenumber) to get physical wavenumber - k = group / input_basis.COV.stretch # Hx cos(n*x) = sin(n*x) # Hx -sin(n*x) = cos(n*x) return np.array([[ 0, 1], @@ -1507,6 +1506,14 @@ def Parity(*args, parity=None, **kw): class EvenParity(ParityBase, metaclass=CachedClass): """Even parity basis (cosine series for scalars).""" + def derivative_basis(self, order=1): + if order % 2 == 0: + return self + elif order % 2 == 1: + return OddParity(self.coord, self.size, self.bounds, self.dealias, self.library) + else: + raise ValueError(f"Invalid derivative order: {order}") + def valid_elements(self, tensorsig, grid_space, elements): vshape = tuple(cs.dim for cs in tensorsig) + elements[0].shape valid = np.ones(shape=vshape, dtype=bool) @@ -1555,6 +1562,14 @@ def component_parity(self, tensorsig): class OddParity(ParityBase, metaclass=CachedClass): """Odd parity basis (sine series for scalars).""" + def derivative_basis(self, order=1): + if order % 2 == 0: + return self + elif order % 2 == 1: + return EvenParity(self.coord, self.size, self.bounds, self.dealias, self.library) + else: + raise ValueError(f"Invalid derivative order: {order}") + def valid_elements(self, tensorsig, grid_space, elements): vshape = tuple(cs.dim for cs in tensorsig) + elements[0].shape valid = np.ones(shape=vshape, dtype=bool) @@ -1733,25 +1748,26 @@ class DifferentiateParity(operators.Differentiate, SpectralOperatorParity): subaxis_coupling = [False] @staticmethod - def _output_basis(input_basis): - # Swap parity - if isinstance(input_basis, EvenParity): - return OddParity(input_basis.coord, input_basis.size, input_basis.bounds, input_basis.dealias, input_basis.library) - elif isinstance(input_basis, OddParity): - return EvenParity(input_basis.coord, input_basis.size, input_basis.bounds, input_basis.dealias, input_basis.library) - else: - raise ValueError(f"This should never happen: input_basis = {input_basis}") + def _output_basis(input_basis, order): + return input_basis.derivative_basis(order) + + def subspace_matrix(self, layout, parity): + """Build matrix operating on local subspace data.""" + # Caching layer to allow insertion of other arguments + return self._subspace_matrix(layout, self.input_basis, self.output_basis, self.first_axis, self.order, parity) @staticmethod - def _group_matrix(group, input_basis, output_basis, parity): + def _group_matrix(group, input_basis, output_basis, order, parity): # Rescale group (native wavenumber) to get physical wavenumber k = group / input_basis.COV.stretch if parity == 1: # dx cos(k*x) = -k * sin(k*x) - return np.array([[-k]]) + S = k**order * (-1)**((order+1)//2) + return np.array([[S]]) elif parity == -1: # dx sin(k*x) = k * cos(k*x) - return np.array([[k]]) + S = k**order * (-1)**(order//2) + return np.array([[S]]) else: raise ValueError(f"This should never happen: parity = {parity}") diff --git a/dedalus/tests/test_fourier_operators.py b/dedalus/tests/test_fourier_operators.py index bb50f196..31d194b3 100644 --- a/dedalus/tests/test_fourier_operators.py +++ b/dedalus/tests/test_fourier_operators.py @@ -64,15 +64,21 @@ def test_parity_convert_constant(N, bounds, dealias, dtype, layout): @pytest.mark.parametrize('bounds', bounds_range) @pytest.mark.parametrize('dealias', dealias_range) @pytest.mark.parametrize('dtype', dtype_range) -def test_fourier_differentiate(N, bounds, dealias, dtype): +@pytest.mark.parametrize('order', [1, 2, 3]) +def test_fourier_differentiate(N, bounds, dealias, dtype, order): """Test differentiation in Fourier basis.""" c, d, b, x = build_fourier(N, bounds, dealias, dtype) f = d.Field(bases=b) L = bounds[1] - bounds[0] k = 4 * np.pi / L f['g'] = 1 + np.sin(k*x+0.1) - g = d3.Differentiate(f, c).evaluate() - assert np.allclose(g['g'], k*np.cos(k*x+0.1)) + g = d3.Differentiate(f, c, order=order).evaluate() + if order == 1: + assert np.allclose(g['g'], k*np.cos(k*x+0.1)) + elif order == 2: + assert np.allclose(g['g'], -k**2*np.sin(k*x+0.1)) + elif order == 3: + assert np.allclose(g['g'], -k**3*np.cos(k*x+0.1)) @pytest.mark.parametrize('N', N_range) @@ -80,7 +86,8 @@ def test_fourier_differentiate(N, bounds, dealias, dtype): @pytest.mark.parametrize('dealias', dealias_range) @pytest.mark.parametrize('dtype', dtype_range) @pytest.mark.parametrize('parity', [1, -1]) -def test_parity_differentiate(N, bounds, dealias, dtype, parity): +@pytest.mark.parametrize('order', [1, 2, 3]) +def test_parity_differentiate(N, bounds, dealias, dtype, parity, order): """Test differentiation in Parity bases.""" c, d, b, x = build_parity(N, bounds, dealias, parity, dtype) f = d.Field(bases=b) @@ -89,12 +96,22 @@ def test_parity_differentiate(N, bounds, dealias, dtype, parity): k = 3 * np.pi / L if parity == 1: f['g'] = 1 + np.cos(k*(x-x0)) - g = d3.Differentiate(f, c).evaluate() - assert np.allclose(g['g'], -k*np.sin(k*(x-x0))) + g = d3.Differentiate(f, c, order=order).evaluate() + if order == 1: + assert np.allclose(g['g'], -k*np.sin(k*(x-x0))) + elif order == 2: + assert np.allclose(g['g'], -k**2*np.cos(k*(x-x0))) + elif order == 3: + assert np.allclose(g['g'], k**3*np.sin(k*(x-x0))) elif parity == -1: f['g'] = np.sin(k*(x-x0)) - g = d3.Differentiate(f, c).evaluate() - assert np.allclose(g['g'], k*np.cos(k*(x-x0))) + g = d3.Differentiate(f, c, order=order).evaluate() + if order == 1: + assert np.allclose(g['g'], k*np.cos(k*(x-x0))) + elif order == 2: + assert np.allclose(g['g'], -k**2*np.sin(k*(x-x0))) + elif order == 3: + assert np.allclose(g['g'], -k**3*np.cos(k*(x-x0))) @pytest.mark.parametrize('N', N_range) From 2b9855cc1640c066c45b7404dd58ea23b30cc039 Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Wed, 12 Nov 2025 19:50:12 -0500 Subject: [PATCH 17/22] Add riesz derivative and hilbert transform for parity bases --- dedalus/core/basis.py | 170 ++++++++++++++++++-------------------- dedalus/core/operators.py | 18 +--- 2 files changed, 84 insertions(+), 104 deletions(-) diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index a8fd968d..ba64960f 100644 --- a/dedalus/core/basis.py +++ b/dedalus/core/basis.py @@ -1270,10 +1270,9 @@ def _output_basis(input_basis, order): def _group_matrix(group, input_basis, output_basis, order): # Rescale group (native wavenumber) to get physical wavenumber k = group / input_basis.COV.stretch - # dx**n exp(1j*k*x) = (1j*k)**n * exp(1j*k*x) = S * exp(1j*k*x) - # dx**n [cos(k*x) + 1j*sin(k*x)] = S * [cos(k*x) + 1j*sin(k*x)] - # dx**n cos(k*x) = R(S) * cos(k*x) + I(S) * -sin(k*x) - # dx**n -sin(k*x) = R(S) * -sin(k*x) - I(S) * cos(k*x) + # dx^n exp(ikx) = (ik)^n exp(ikx) = S exp(ikx) + # dx^n cos(kx) = R(S) cos(kx) - I(S) sin(kx) + # dx^n sin(kx) = R(S) sin(kx) + I(S) cos(kx) S = (1j * k) ** order return np.array([[S.real, -S.imag], [S.imag, S.real]]) @@ -1294,10 +1293,9 @@ def _output_basis(input_basis, order): def _group_matrix(group, input_basis, output_basis, order): # Rescale group (native wavenumber) to get physical wavenumber k = group / input_basis.COV.stretch - # R_a exp(1j*k*x) = - |k|**n * exp(1j*k*x) = S * exp(1j*k*x) - # R_a [cos(k*x) + 1j*sin(k*x)] = S * [cos(k*x) + 1j*sin(k*x)] - # R_a cos(k*x) = S * cos(k*x) - # R_a -sin(k*x) = S * -sin(k*x) + # R_a exp(ikx) = - |k|^n exp(ikx) = S exp(ikx) + # R_a cos(kx) = S cos(kx) + # R_a sin(kx) = S sin(kx) S = - abs(k) ** order return np.array([[S, 0], [0, S]]) @@ -1316,8 +1314,8 @@ def _output_basis(input_basis): @staticmethod def _group_matrix(group, input_basis, output_basis): - # Hx cos(n*x) = sin(n*x) - # Hx -sin(n*x) = cos(n*x) + # Hx cos(kx) = sin(kx) + # Hx sin(kx) = -cos(kx) return np.array([[ 0, 1], [-1, 0]]) @@ -1337,7 +1335,6 @@ def _output_basis(input_basis, position): @staticmethod def _full_matrix(input_basis, output_basis, position): # Build native interpolation vector - # Interleaved cos(k*x), -sin(k*x) x = input_basis.COV.native_coord(position) k = input_basis.native_wavenumbers interp_vector = np.zeros(k.size) @@ -1363,8 +1360,8 @@ def _output_basis(input_basis): def _group_matrix(group, input_basis, output_basis): # Rescale group (native wavenumber) to get physical wavenumber k = group / input_basis.COV.stretch - # integ cos(k*x) = L * δ(k, 0) - # integ -sin(k*x) = 0 + # integ cos(kx) dx = L (k==0) + # integ sin(kx) dx = 0 if k == 0: L = input_basis.COV.problem_length return np.array([[L, 0]]) @@ -1389,8 +1386,8 @@ def _output_basis(input_basis): def _group_matrix(group, input_basis, output_basis): # Rescale group (native wavenumber) to get physical wavenumber k = group / input_basis.COV.stretch - # integ cos(k*x) / L = δ(k, 0) - # integ -sin(k*x) / L = 0 + # integ cos(kx) dx / L = (k==0) + # integ sin(kx) dx / L = 0 if k == 0: return np.array([[1, 0]]) else: @@ -1398,38 +1395,6 @@ def _group_matrix(group, input_basis, output_basis): raise ValueError(f"This should never happen: group = {group}") -# class HilbertTransformFourier(operators.HilbertTransform): -# """Fourier series Hilbert transform.""" - -# input_basis_type = Fourier -# bands = [-1, 1] -# separable = True - -# @staticmethod -# def output_basis(space, input_basis): -# return space.Fourier - -# @staticmethod -# def _build_subspace_entry(i, j, space, input_basis): -# # Hx(cos(n*x)) = sin(n*x) -# # Hx(sin(n*x)) = -cos(n*x) -# n = j // 2 -# if n == 0: -# return 0 -# elif (j % 2) == 0: -# # Hx(cos(n*x)) = sin(n*x) -# if i == (j + 1): -# return 1 -# else: -# return 0 -# else: -# # Hx(sin(n*x)) = -cos(n*x) -# if i == (j - 1): -# return (-1) -# else: -# return 0 - - class ParityBase(FourierBase): """Base class for parity-based bases.""" @@ -1761,17 +1726,84 @@ def _group_matrix(group, input_basis, output_basis, order, parity): # Rescale group (native wavenumber) to get physical wavenumber k = group / input_basis.COV.stretch if parity == 1: - # dx cos(k*x) = -k * sin(k*x) + # dx^1 cos(kx) = -k^1 sin(kx) + # dx^2 cos(kx) = -k^2 cos(kx) + # dx^3 cos(kx) = k^3 sin(kx) + # dx^4 cos(kx) = k^4 cos(kx) S = k**order * (-1)**((order+1)//2) return np.array([[S]]) elif parity == -1: - # dx sin(k*x) = k * cos(k*x) + # dx^1 sin(kx) = k^1 cos(kx) + # dx^2 sin(kx) = -k^2 sin(kx) + # dx^3 sin(kx) = -k^3 cos(kx) + # dx^4 sin(kx) = k^4 sin(kx) S = k**order * (-1)**(order//2) return np.array([[S]]) else: raise ValueError(f"This should never happen: parity = {parity}") +class RieszDerivativeParity(operators.RieszDerivative, SpectralOperatorParity): + """ + Parity basis Riesz derivative. + Note: this is a single implementation because tensor fields have both even and odd components. + """ + + input_basis_type = ParityBase + subaxis_dependence = [True] + subaxis_coupling = [False] + + @staticmethod + def _output_basis(input_basis, order): + return input_basis + + def subspace_matrix(self, layout, parity): + """Build matrix operating on local subspace data.""" + # Caching layer to allow insertion of other arguments + return self._subspace_matrix(layout, self.input_basis, self.output_basis, self.first_axis, self.order, parity) + + @staticmethod + def _group_matrix(group, input_basis, output_basis, order, parity): + # Rescale group (native wavenumber) to get physical wavenumber + k = group / input_basis.COV.stretch + # R_a exp(ikx) = - |k|^n exp(ikx) = S exp(ikx) + S = - abs(k) ** order + if parity == 1: + # R_a cos(kx) = S cos(kx) + return np.array([[S]]) + elif parity == -1: + # R_a sin(kx) = S sin(kx) + return np.array([[S]]) + else: + raise ValueError(f"This should never happen: parity = {parity}") + + +class HilbertTransformParity(operators.HilbertTransform, SpectralOperatorParity): + """ + Parity basis Hilbert transform. + Note: this is a single implementation because tensor fields have both even and odd components. + """ + + input_basis_type = ParityBase + subaxis_dependence = [True] + subaxis_coupling = [False] + + @staticmethod + def _output_basis(input_basis): + return input_basis.derivative_basis() + + @staticmethod + def _group_matrix(group, input_basis, output_basis, parity): + if parity == 1: + # Hx cos(kx) = sin(kx) + return np.array([[1]]) + elif parity == -1: + # Hx sin(kx) = -cos(kx) + return np.array([[-1]]) + else: + raise ValueError(f"This should never happen: parity = {parity}") + + class IntegrateEvenParity(operators.Integrate, operators.SpectralOperator1D): """EvenParity basis integration.""" @@ -1866,46 +1898,6 @@ def _full_matrix(input_basis, output_basis): return ave_vector[None, :] -# class HilbertTransformSine(operators.HilbertTransform): -# """Sine series Hilbert transform.""" - -# input_basis_type = Sine -# bands = [0] -# separable = True - -# @staticmethod -# def output_basis(space, input_basis): -# return space.Cosine - -# @staticmethod -# def _build_subspace_entry(i, j, space, input_basis): -# # Hx(sin(n*x)) = -cos(n*x) -# if i == j: -# return (-1) -# else: -# return 0 - - -# class HilbertTransformCosine(operators.HilbertTransform): -# """Cosine series Hilbert transform.""" - -# input_basis_type = Cosine -# bands = [0] -# separable = True - -# @staticmethod -# def output_basis(space, input_basis): -# return space.Sine - -# @staticmethod -# def _build_subspace_entry(i, j, space, input_basis): -# # Hx(cos(n*x)) = sin(n*x) -# if i == j: -# return 1 -# else: -# return 0 - - class MultidimensionalBasis(Basis): def forward_transform(self, field, axis, gdata, cdata): diff --git a/dedalus/core/operators.py b/dedalus/core/operators.py index 67fee47d..ed77ad23 100644 --- a/dedalus/core/operators.py +++ b/dedalus/core/operators.py @@ -1342,6 +1342,7 @@ def __init__(self, operand, coord): # return arg +@alias("diff", "D") class Differentiate(SpectralOperator1D, metaclass=MultiClass): """ Differentiation along one dimension. @@ -1426,21 +1427,7 @@ def __new__(cls, operand, coord, order=1, out=None): return 0 -# @prefix('H') -# @parseable('hilbert_transform', 'hilbert', 'H') -# def hilbert_transform(arg, *spaces, **space_kw): -# # Parse space/order keywords into space list -# for space, order in space_kw.items(): -# spaces += (space,) * order -# # Identify domain -# domain = unify_attributes((arg,)+spaces, 'domain', require=False) -# # Apply iteratively -# for space in spaces: -# space = domain.get_space_object(space) -# arg = HilbertTransform(arg, space) -# return arg - - +@alias("hilbert", "H") class HilbertTransform(SpectralOperator1D, metaclass=MultiClass): """ Hilbert transform along one dimension. @@ -1508,6 +1495,7 @@ def __new__(cls, operand, coord, out=None): return 0 +@alias("riesz", "R") class RieszDerivative(SpectralOperator1D, metaclass=MultiClass): """ Riesz derivative along one dimension. From 6c0e17fba55ba40d9f952442b6d5fed65c917848 Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Thu, 13 Nov 2025 14:15:48 -0500 Subject: [PATCH 18/22] Add Riesz and Hilbert tests --- dedalus/core/operators.py | 6 +- dedalus/tests/test_fourier_operators.py | 76 +++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 2 deletions(-) diff --git a/dedalus/core/operators.py b/dedalus/core/operators.py index ed77ad23..ccd56931 100644 --- a/dedalus/core/operators.py +++ b/dedalus/core/operators.py @@ -39,6 +39,8 @@ 'Integrate', 'Average', 'Differentiate', + 'RieszDerivative', + 'HilbertTransform', 'Convert', 'TransposeComponents', 'RadialComponent', @@ -1513,7 +1515,7 @@ class RieszDerivative(SpectralOperator1D, metaclass=MultiClass): name = "Riesz" - def __init__(self, operand, coord, order, out=None): + def __init__(self, operand, coord, order=1, out=None): super().__init__(operand, out=out) self.order = order # SpectralOperator requirements @@ -1531,7 +1533,7 @@ def __init__(self, operand, coord, order, out=None): self.dtype = operand.dtype @classmethod - def _check_args(cls, operand, coord, out=None): + def _check_args(cls, operand, coord, order=1, out=None): # Dispatch by operand basis if isinstance(operand, Operand): basis = operand.domain.get_basis(coord) diff --git a/dedalus/tests/test_fourier_operators.py b/dedalus/tests/test_fourier_operators.py index 31d194b3..dd596276 100644 --- a/dedalus/tests/test_fourier_operators.py +++ b/dedalus/tests/test_fourier_operators.py @@ -114,6 +114,82 @@ def test_parity_differentiate(N, bounds, dealias, dtype, parity, order): assert np.allclose(g['g'], -k**3*np.cos(k*(x-x0))) +@pytest.mark.parametrize('N', N_range) +@pytest.mark.parametrize('bounds', bounds_range) +@pytest.mark.parametrize('dealias', dealias_range) +@pytest.mark.parametrize('dtype', dtype_range) +@pytest.mark.parametrize('order', [1, 1.5, 2]) +def test_fourier_riesz(N, bounds, dealias, dtype, order): + """Test Riesz derivative in Fourier basis.""" + c, d, b, x = build_fourier(N, bounds, dealias, dtype) + f = d.Field(bases=b) + L = bounds[1] - bounds[0] + k = 4 * np.pi / L + f['g'] = 1 + np.sin(k*x+0.1) + g = d3.RieszDerivative(f, c, order=order).evaluate() + assert np.allclose(g['g'], -abs(k)**order*np.sin(k*x+0.1)) + + +@pytest.mark.parametrize('N', N_range) +@pytest.mark.parametrize('bounds', bounds_range) +@pytest.mark.parametrize('dealias', dealias_range) +@pytest.mark.parametrize('dtype', dtype_range) +@pytest.mark.parametrize('parity', [1, -1]) +@pytest.mark.parametrize('order', [1, 1.5, 2]) +def test_parity_riesz(N, bounds, dealias, dtype, parity, order): + """Test Riesz derivative in Fourier basis.""" + c, d, b, x = build_parity(N, bounds, dealias, parity, dtype) + f = d.Field(bases=b) + x0 = bounds[0] + L = bounds[1] - bounds[0] + k = 4 * np.pi / L + if parity == 1: + f['g'] = 1 + np.cos(k*(x-x0)) + g = d3.RieszDerivative(f, c, order=order).evaluate() + assert np.allclose(g['g'], -abs(k)**order*np.cos(k*(x-x0))) + elif parity == -1: + f['g'] = np.sin(k*(x-x0)) + g = d3.RieszDerivative(f, c, order=order).evaluate() + assert np.allclose(g['g'], -abs(k)**order*np.sin(k*(x-x0))) + + +@pytest.mark.parametrize('N', N_range) +@pytest.mark.parametrize('bounds', bounds_range) +@pytest.mark.parametrize('dealias', dealias_range) +@pytest.mark.parametrize('dtype', dtype_range) +def test_fourier_hilbert(N, bounds, dealias, dtype): + """Test Hilbert transform in Fourier basis.""" + c, d, b, x = build_fourier(N, bounds, dealias, dtype) + f = d.Field(bases=b) + L = bounds[1] - bounds[0] + k = 4 * np.pi / L + f['g'] = 1 + np.sin(k*x+0.1) + g = d3.HilbertTransform(f, c).evaluate() + assert np.allclose(g['g'], -np.cos(k*x+0.1)) + + +@pytest.mark.parametrize('N', N_range) +@pytest.mark.parametrize('bounds', bounds_range) +@pytest.mark.parametrize('dealias', dealias_range) +@pytest.mark.parametrize('dtype', dtype_range) +@pytest.mark.parametrize('parity', [1, -1]) +def test_parity_hilbert(N, bounds, dealias, dtype, parity): + """Test Hilbert transform in Parity bases.""" + c, d, b, x = build_parity(N, bounds, dealias, parity, dtype) + f = d.Field(bases=b) + x0 = bounds[0] + L = bounds[1] - bounds[0] + k = 4 * np.pi / L + if parity == 1: + f['g'] = 1 + np.cos(k*(x-x0)) + g = d3.HilbertTransform(f, c).evaluate() + assert np.allclose(g['g'], np.sin(k*(x-x0))) + elif parity == -1: + f['g'] = np.sin(k*(x-x0)) + g = d3.HilbertTransform(f, c).evaluate() + assert np.allclose(g['g'], -np.cos(k*(x-x0))) + + @pytest.mark.parametrize('N', N_range) @pytest.mark.parametrize('bounds', bounds_range) @pytest.mark.parametrize('dealias', dealias_range) From da71d22759bd8d0224033d9c1f8d718867cae1ff Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Fri, 14 Nov 2025 09:47:31 -0500 Subject: [PATCH 19/22] Pass order for differentiate parity matrices --- dedalus/core/basis.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index ba64960f..b4de951d 100644 --- a/dedalus/core/basis.py +++ b/dedalus/core/basis.py @@ -1617,6 +1617,9 @@ def subproblem_matrix(self, subproblem): blocks = [parity_matrices[p] for p in P.ravel()] return sparse_block_diag(blocks) + def group_matrix(self, group, parity): + return self._group_matrix(group, self.input_basis, self.output_basis, parity) + def subspace_matrix(self, layout, parity): """Build matrix operating on local subspace data.""" # Caching layer to allow insertion of other arguments @@ -1721,6 +1724,9 @@ def subspace_matrix(self, layout, parity): # Caching layer to allow insertion of other arguments return self._subspace_matrix(layout, self.input_basis, self.output_basis, self.first_axis, self.order, parity) + def group_matrix(self, group, parity): + return self._group_matrix(group, self.input_basis, self.output_basis, self.order, parity) + @staticmethod def _group_matrix(group, input_basis, output_basis, order, parity): # Rescale group (native wavenumber) to get physical wavenumber @@ -1762,6 +1768,9 @@ def subspace_matrix(self, layout, parity): # Caching layer to allow insertion of other arguments return self._subspace_matrix(layout, self.input_basis, self.output_basis, self.first_axis, self.order, parity) + def group_matrix(self, group, parity): + return self._group_matrix(group, self.input_basis, self.output_basis, self.order, parity) + @staticmethod def _group_matrix(group, input_basis, output_basis, order, parity): # Rescale group (native wavenumber) to get physical wavenumber From d0153fb724073a9fa8da3120890933c0357b2594 Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Wed, 17 Dec 2025 17:45:02 -0500 Subject: [PATCH 20/22] More work on vector parity operators and RBC example --- dedalus/core/basis.py | 94 +++++++++++-------- dedalus/core/operators.py | 54 +++++++++-- .../plot_snapshots.py | 4 +- .../stress_free_rbc.py | 53 ++++++----- 4 files changed, 132 insertions(+), 73 deletions(-) diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index b4de951d..c7fdf335 100644 --- a/dedalus/core/basis.py +++ b/dedalus/core/basis.py @@ -895,6 +895,18 @@ def __rmatmul__(self, other): def __pow__(self, other): return self + def derivative_basis(self, order=1): + return self + + def gradient_basis(self): + return self + + def component_basis(self): + return self + + def skew_basis(self): + return self + def elements_to_groups(self, grid_space, elements): if grid_space[0]: groups = elements @@ -1414,9 +1426,8 @@ def _native_grid(self, scale): N, = self.grid_shape((scale,)) return (np.pi / N) * (1/2 + np.arange(N)) - @staticmethod - @CachedFunction - def _transform_plan(grid_size, coeff_size, parity, library): + @CachedMethod + def _transform_plan(self, grid_size, coeff_size, parity, library): """Caching layer to share plans across even/odd parity bases.""" # Use matrix transforms for trivial cases if (grid_size == 1 or coeff_size == 1) and (library != "matrix"): @@ -1435,6 +1446,7 @@ def forward_transform(self, field, axis, gdata, cdata): P = self.component_parity(field.tensorsig) parity_plans = {p: self.transform_plan(grid_size, p) for p in np.unique(P)} # Transform component-by-component + gdata = gdata.copy() # Copy to avoid overwrite errors with dealiasing for i, p in np.ndenumerate(P): parity_plans[p].forward(gdata[i], cdata[i], axis) # Permute coefficients @@ -1451,10 +1463,27 @@ def backward_transform(self, field, axis, cdata, gdata): if self.backward_coeff_permutation is not None: permute_axis(cdata, axis+len(field.tensorsig), self.backward_coeff_permutation, out=cdata) # Transform component-by-component - P = self.component_parity(field.tensorsig) + cdata = cdata.copy() # Copy to avoid overwrite errors with dealiasing for i, p in np.ndenumerate(P): parity_plans[p].backward(cdata[i], gdata[i], axis) + def derivative_basis(self, order=1): + if order % 2 == 0: + return self + elif order % 2 == 1: + return self.opposite_parity() + else: + raise ValueError(f"Invalid derivative order: {order}") + + def gradient_basis(self): + return self + + def component_basis(self): + return self.opposite_parity() + + def skew_basis(self): + return self.opposite_parity() + def Parity(*args, parity=None, **kw): """Factory function dispatching to EvenParity and OddParity based on provided parity.""" @@ -1471,19 +1500,6 @@ def Parity(*args, parity=None, **kw): class EvenParity(ParityBase, metaclass=CachedClass): """Even parity basis (cosine series for scalars).""" - def derivative_basis(self, order=1): - if order % 2 == 0: - return self - elif order % 2 == 1: - return OddParity(self.coord, self.size, self.bounds, self.dealias, self.library) - else: - raise ValueError(f"Invalid derivative order: {order}") - - def valid_elements(self, tensorsig, grid_space, elements): - vshape = tuple(cs.dim for cs in tensorsig) + elements[0].shape - valid = np.ones(shape=vshape, dtype=bool) - return valid - def __add__(self, other): if other is None: return self @@ -1513,6 +1529,14 @@ def __mul__(self, other): def __pow__(self, other): return NotImplemented + def valid_elements(self, tensorsig, grid_space, elements): + vshape = tuple(cs.dim for cs in tensorsig) + elements[0].shape + valid = np.ones(shape=vshape, dtype=bool) + return valid + + def opposite_parity(self): + return OddParity(self.coord, self.size, self.bounds, self.dealias, self.library) + @CachedMethod def component_parity(self, tensorsig): P = np.ones([cs.dim for cs in tensorsig], dtype=int) @@ -1527,25 +1551,6 @@ def component_parity(self, tensorsig): class OddParity(ParityBase, metaclass=CachedClass): """Odd parity basis (sine series for scalars).""" - def derivative_basis(self, order=1): - if order % 2 == 0: - return self - elif order % 2 == 1: - return EvenParity(self.coord, self.size, self.bounds, self.dealias, self.library) - else: - raise ValueError(f"Invalid derivative order: {order}") - - def valid_elements(self, tensorsig, grid_space, elements): - vshape = tuple(cs.dim for cs in tensorsig) + elements[0].shape - valid = np.ones(shape=vshape, dtype=bool) - if not grid_space[0]: - # Drop sine part of k=0 - groups = self.elements_to_groups(grid_space, elements) - allcomps = tuple(slice(None) for cs in tensorsig) - selection = (groups[0] == 0) - valid[allcomps + (selection,)] = False - return valid - def __add__(self, other): if other is self: return self @@ -1573,6 +1578,20 @@ def __mul__(self, other): def __pow__(self, other): return NotImplemented + def valid_elements(self, tensorsig, grid_space, elements): + vshape = tuple(cs.dim for cs in tensorsig) + elements[0].shape + valid = np.ones(shape=vshape, dtype=bool) + if not grid_space[0]: + # Drop sine part of k=0 + groups = self.elements_to_groups(grid_space, elements) + allcomps = tuple(slice(None) for cs in tensorsig) + selection = (groups[0] == 0) + valid[allcomps + (selection,)] = False + return valid + + def opposite_parity(self): + return EvenParity(self.coord, self.size, self.bounds, self.dealias, self.library) + @CachedMethod def component_parity(self, tensorsig): P = np.ones([cs.dim for cs in tensorsig], dtype=int) @@ -1628,13 +1647,14 @@ def subspace_matrix(self, layout, parity): def operate(self, out): operand = self.args[0] input_basis = self.input_basis - data_axis = len(operand.tensorsig) + self.last_axis + data_axis = self.last_axis # Set output layout out.preset_layout(operand.layout) # Apply operator to each component P = self.operand_component_parity parity_matrices = {p: self.subspace_matrix(operand.layout, p) for p in np.unique(P)} for i, p in np.ndenumerate(P): + # TODO: check shapes on first try apply_matrix(parity_matrices[p], operand.data[i], axis=data_axis, out=out.data[i]) diff --git a/dedalus/core/operators.py b/dedalus/core/operators.py index ccd56931..0a53fe77 100644 --- a/dedalus/core/operators.py +++ b/dedalus/core/operators.py @@ -2155,10 +2155,22 @@ def __init__(self, operand, index=0, out=None): # LinearOperator requirements self.operand = operand # FutureField requirements - self.domain = operand.domain + bases = self._build_bases(operand, index) + self.domain = Domain(operand.dist, bases) self.tensorsig = operand.tensorsig self.dtype = operand.dtype + def _build_bases(self, operand, index): + input_bases = operand.domain.bases + output_bases = [] + coordsys = operand.tensorsig[index] + for basis in input_bases: + if basis.coord in coordsys.coords: + output_bases.append(basis.skew_basis()) + else: + output_bases.append(basis) + return tuple(output_bases) + def new_operand(self, operand, **kw): return Skew(operand, index=self.index, **kw) @@ -2434,11 +2446,10 @@ def __init__(self, operand, coordsys, out=None): if args[i] == 0: args[i] = 2*operand args[i].args[0] = 0 - original_args = list(args[i].original_args) - original_args[0] = 0 - args[i].original_args = tuple(original_args) - bases = self._build_bases(*args) - args = [convert(arg, bases) for arg in args] + args[i].original_args = (0, args[i].original_args[1]) + # Convert along orthogonal bases + bases = self._build_bases(operand, coordsys) + args = self._convert_args(coordsys, args, bases) LinearOperator.__init__(self, *args, out=out) self.coordsys = coordsys # LinearOperator requirements @@ -2448,8 +2459,22 @@ def __init__(self, operand, coordsys, out=None): self.tensorsig = (coordsys,) + operand.tensorsig self.dtype = operand.dtype - def _build_bases(self, *args): - return sum(args).domain.bases + def _build_bases(self, operand, coordsys): + input_bases = operand.domain.bases + output_bases = [] + for basis in input_bases: + if basis.coord is coordsys or basis.coord in coordsys.coords: + output_bases.append(basis.gradient_basis()) + else: + output_bases.append(basis) + return tuple(output_bases) + + def _convert_args(self, coordsys, args, bases): + converted_args = [] + for coord, arg in zip(coordsys.coords, args): + arg_bases = [b for b in bases if b.coord != coord] + converted_args.append(convert(arg, arg_bases)) + return converted_args def matrix_dependence(self, *vars): arg_vals = [arg.matrix_dependence(self, *vars) for arg in self.args] @@ -3405,10 +3430,21 @@ def __init__(self, operand, index, comp, out=None): # LinearOperator requirements self.operand = operand # FutureField requirements - self.domain = operand.domain + bases = self._build_bases(operand, index, comp) + self.domain = Domain(operand.dist, bases) self.tensorsig = operand.tensorsig[:index] + operand.tensorsig[index+1:] self.dtype = operand.dtype + def _build_bases(self, operand, index, comp): + input_bases = operand.domain.bases + output_bases = [] + for basis in input_bases: + if basis.coord is comp: + output_bases.append(basis.component_basis()) + else: + output_bases.append(basis) + return tuple(output_bases) + def check_conditions(self): """Check that operands are in a proper layout.""" # Any layout diff --git a/examples/ivp_2d_stress_free_rayleigh_benard/plot_snapshots.py b/examples/ivp_2d_stress_free_rayleigh_benard/plot_snapshots.py index c46b8d77..8fdca5da 100644 --- a/examples/ivp_2d_stress_free_rayleigh_benard/plot_snapshots.py +++ b/examples/ivp_2d_stress_free_rayleigh_benard/plot_snapshots.py @@ -21,14 +21,14 @@ def main(filename, start, count, output): """Save plot of specified tasks for given range of analysis writes.""" # Plot settings - tasks = ['buoyancy', 'vorticity'] + tasks = ['buoyancy', 'vorticity', 'ux', 'uz'] scale = 1.5 dpi = 200 title_func = lambda sim_time: 't = {:.3f}'.format(sim_time) savename_func = lambda write: 'write_{:06}.png'.format(write) # Layout - nrows, ncols = 2, 1 + nrows, ncols = 4, 1 image = plot_tools.Box(4, 1) pad = plot_tools.Frame(0.3, 0, 0, 0) margin = plot_tools.Frame(0.2, 0.1, 0, 0) diff --git a/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py b/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py index 91389554..18c57884 100644 --- a/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py +++ b/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py @@ -31,45 +31,47 @@ # Parameters Lx, Lz = 4, 1 -Nx, Nz = 8, 8 -Rayleigh = 2e6 +Nx, Nz = 256, 64 +Rayleigh = 1e5 Prandtl = 1 dealias = 3/2 -stop_sim_time = 50 +stop_sim_time = 40 timestepper = d3.RK222 +timestep = 1e-2 max_timestep = 0.125 dtype = np.float64 -library = 'fftw' # 'scipy', 'matrix' -logger.info(f"Running with {library} DCT/DST library") + # Bases +# cx = d3.Coordinate('x') +# cz = d3.Coordinate('z') +# coords = d3.DirectProduct(cx, cz) coords = d3.CartesianCoordinates('x', 'z') +cx, cz = coords.coords dist = d3.Distributor(coords, dtype=dtype) -xbasis = d3.RealFourier(coords['x'], size=Nx, bounds=(0, Lx), dealias=dealias) -zsbasis = d3.OddParity(coords['z'], size=Nz, bounds=(0, Lz), dealias=dealias, library=library) -zcbasis = d3.EvenParity(coords['z'], size=Nz, bounds=(0, Lz), dealias=dealias, library=library) +xbasis = d3.RealFourier(cx, size=Nx, bounds=(0, Lx), dealias=dealias) +zobasis = d3.OddParity(cz, size=Nz, bounds=(0, Lz), dealias=dealias) +zebasis = d3.EvenParity(cz, size=Nz, bounds=(0, Lz), dealias=dealias) # Fields -p = dist.Field(name='p', bases=(xbasis,zcbasis)) -b = dist.Field(name='b', bases=(xbasis,zsbasis)) -u = dist.Field(name='u', bases=(xbasis,zcbasis)) -w = dist.Field(name='w', bases=(xbasis,zsbasis)) +p = dist.Field(name='p', bases=(xbasis,zebasis)) +b = dist.Field(name='b', bases=(xbasis,zobasis)) +u = dist.VectorField(coords, name='u', bases=(xbasis,zebasis)) tau_p = dist.Field(name='tau_p') # Substitutions kappa = (Rayleigh * Prandtl)**(-1/2) nu = (Rayleigh / Prandtl)**(-1/2) -x, z = dist.local_grids(xbasis, zcbasis) - -dx = lambda A: d3.Differentiate(A,coords['x']) -dz = lambda A: d3.Differentiate(A,coords['z']) -lap = lambda A: dx(dx(A)) + dz(dz(A)) +x, z = dist.local_grids(xbasis, zebasis) +ex = dist.VectorField(coords, bases=(zebasis)) +ex['g'][0] = 1 +ez = dist.VectorField(coords, bases=(zobasis)) +ez['g'][1] = 1 # Problem -problem = d3.IVP([p, b, u, w, tau_p], namespace=locals()) -problem.add_equation("dx(u) + dz(w) + tau_p = 0") -problem.add_equation("dt(b) - kappa*lap(b) - w = - u*dx(b) - w*dz(b)") -problem.add_equation("dt(u) - nu*lap(u) + dx(p) = - u*dx(u) - w*dz(u)") -problem.add_equation("dt(w) - nu*lap(w) + dz(p) - b = - u*dx(w) - w*dz(w)") +problem = d3.IVP([p, b, u, tau_p], namespace=locals()) +problem.add_equation("div(u) + tau_p = 0") +problem.add_equation("dt(b) - kappa*lap(b) = - u@grad(b) + ez@u") +problem.add_equation("dt(u) - nu*lap(u) + grad(p) = - u@grad(u) + b*ez") problem.add_equation("integ(p) = 0") # Pressure gauge # Solver @@ -83,7 +85,9 @@ # Analysis snapshots = solver.evaluator.add_file_handler('snapshots', sim_dt=0.25, max_writes=50) snapshots.add_task(b, name='buoyancy') -snapshots.add_task(dz(u) - dx(w), name='vorticity') +snapshots.add_task(ex@u, name='ux') +snapshots.add_task(ez@u, name='uz') +snapshots.add_task(-d3.div(d3.skew(u)), name='vorticity') # CFL #CFL = d3.CFL(solver, initial_dt=max_timestep, cadence=10, safety=0.5, threshold=0.05, @@ -92,14 +96,13 @@ # Flow properties flow = d3.GlobalFlowProperty(solver, cadence=10) -flow.add_property((u*u + w*w)/nu, name='Re') +flow.add_property(np.sqrt(u@u)/nu, name='Re') # Main loop try: logger.info('Starting main loop') while solver.proceed: #timestep = CFL.compute_timestep() - timestep = 1e-3 solver.step(timestep) if (solver.iteration-1) % 10 == 0: max_Re = flow.max('Re') From 8eb7006b1da91119a4f705bf53c4454b00fd9895 Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Thu, 18 Dec 2025 10:03:43 -0500 Subject: [PATCH 21/22] Get explicit multiplication working with size=1 parity vectors --- dedalus/core/arithmetic.py | 3 ++- dedalus/core/basis.py | 2 +- .../ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py | 4 ++-- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/dedalus/core/arithmetic.py b/dedalus/core/arithmetic.py index 64daa530..01717354 100644 --- a/dedalus/core/arithmetic.py +++ b/dedalus/core/arithmetic.py @@ -874,7 +874,8 @@ def __init__(self, domain, layout, broadcast_dims): self.layout = layout self.broadcast_dims = broadcast_dims # Determine deployment dimensions - deploy_dims_ext = np.array(broadcast_dims) & np.array(domain.constant) + constant_input_dims = np.array(domain.global_shape(layout, domain.dealias)) == 1 # includes None and size=1 bases + deploy_dims_ext = np.array(broadcast_dims) & constant_input_dims deploy_dims = deploy_dims_ext[~layout.local] # Build subcomm or skip casting if any(deploy_dims): diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index c7fdf335..ea6c33d9 100644 --- a/dedalus/core/basis.py +++ b/dedalus/core/basis.py @@ -1431,7 +1431,7 @@ def _transform_plan(self, grid_size, coeff_size, parity, library): """Caching layer to share plans across even/odd parity bases.""" # Use matrix transforms for trivial cases if (grid_size == 1 or coeff_size == 1) and (library != "matrix"): - return ParityBase._transform_plan(grid_size, coeff_size, parity, "matrix") + return self._transform_plan(grid_size, coeff_size, parity, "matrix") parity_name = {1: "cos", -1: "sin"}[parity] return ParityBase.transforms[f"{library}-{parity_name}"](grid_size, coeff_size) diff --git a/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py b/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py index 18c57884..786254ed 100644 --- a/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py +++ b/examples/ivp_2d_stress_free_rayleigh_benard/stress_free_rbc.py @@ -62,9 +62,9 @@ kappa = (Rayleigh * Prandtl)**(-1/2) nu = (Rayleigh / Prandtl)**(-1/2) x, z = dist.local_grids(xbasis, zebasis) -ex = dist.VectorField(coords, bases=(zebasis)) +ex = dist.VectorField(coords, bases=zebasis.clone_with(size=1)) ex['g'][0] = 1 -ez = dist.VectorField(coords, bases=(zobasis)) +ez = dist.VectorField(coords, bases=zobasis.clone_with(size=1)) ez['g'][1] = 1 # Problem From 5496f971ab503f98c695505b8282e90d0323ffa5 Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Thu, 18 Dec 2025 17:56:51 -0500 Subject: [PATCH 22/22] Fix skew for disk --- dedalus/core/basis.py | 3 +++ dedalus/core/operators.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index ea6c33d9..bd157583 100644 --- a/dedalus/core/basis.py +++ b/dedalus/core/basis.py @@ -2792,6 +2792,9 @@ def __rmatmul__(self, other): return self.clone_with(shape=shape, k=k) return NotImplemented + def skew_basis(self): + return self + def global_grid_radius(self, dist, scale): r = self.radial_COV.problem_coord(self._native_radius_grid(scale)) return reshape_vector(r, dim=dist.dim, axis=dist.get_basis_axis(self)+1) diff --git a/dedalus/core/operators.py b/dedalus/core/operators.py index 0a53fe77..ef1205b8 100644 --- a/dedalus/core/operators.py +++ b/dedalus/core/operators.py @@ -2165,7 +2165,7 @@ def _build_bases(self, operand, index): output_bases = [] coordsys = operand.tensorsig[index] for basis in input_bases: - if basis.coord in coordsys.coords: + if (basis.coordsys is coordsys) or (basis.coord in coordsys.coords): output_bases.append(basis.skew_basis()) else: output_bases.append(basis)