Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Spectral updates #35

Merged
merged 26 commits into from
Sep 5, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
26 commits
Select commit Hold shift + click to select a range
835a1d4
Preliminary function to compute the spectrum
Ceyron Sep 4, 2024
43f2c35
Potential improvement to spectrum computation
Ceyron Sep 4, 2024
ea7238b
Add hint on enstrophy
Ceyron Sep 4, 2024
d48f504
Fix on wavenumber norm computation
Ceyron Sep 5, 2024
c504695
Rename to better reflect its effect on oddball mode
Ceyron Sep 5, 2024
82476f5
Add tests on filter mask
Ceyron Sep 5, 2024
b6f7ecd
Improve first half of spectral documentation
Ceyron Sep 5, 2024
22f19bd
Merge common subexpressions
Ceyron Sep 5, 2024
84a0023
Extend and fix documentation
Ceyron Sep 5, 2024
fb609ad
Add experimental convinience function to extract the Fourier coeffici…
Ceyron Sep 5, 2024
55a1e17
Unify the creation of scaling arrays
Ceyron Sep 5, 2024
d52a480
Test scaling array for norm_compensation
Ceyron Sep 5, 2024
e96be3f
Add tests for coefficient extraction in 1D
Ceyron Sep 5, 2024
6fd79f4
Test for coefficient extration in 2D
Ceyron Sep 5, 2024
30e5163
Enhance docstring
Ceyron Sep 5, 2024
bae8fe5
Remove experimental function
Ceyron Sep 5, 2024
7be9b69
Fix spectrum function
Ceyron Sep 5, 2024
92ff949
Export get_spectrum instead of make_incompressible
Ceyron Sep 5, 2024
8e0a8a8
Adapt docs
Ceyron Sep 5, 2024
ebd9674
Use proper links
Ceyron Sep 5, 2024
6777cb8
Enhance docstring
Ceyron Sep 5, 2024
abb0a47
Enhance docstring
Ceyron Sep 5, 2024
da40239
Test the spectrum creation
Ceyron Sep 5, 2024
bd1eb67
Restructure docstring
Ceyron Sep 5, 2024
375623e
Tests helper utilities for fft setup
Ceyron Sep 5, 2024
f639e14
Merge branch 'dev' into spectral-updates
Ceyron Sep 5, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/api/utilities/derivatives.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,4 +4,4 @@

---

::: exponax.make_incompressible
::: exponax.spectral.make_incompressible
10 changes: 9 additions & 1 deletion docs/api/utilities/spectral.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,4 +4,12 @@

---

::: exponax.ifft
::: exponax.ifft

---

::: exponax.get_spectrum

---

::: exponax.spectral.build_scaling_array
4 changes: 2 additions & 2 deletions exponax/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from ._forced_stepper import ForcedStepper
from ._interpolation import FourierInterpolator, map_between_resolutions
from ._repeated_stepper import RepeatedStepper
from ._spectral import derivative, fft, ifft, make_incompressible
from ._spectral import derivative, fft, get_spectrum, ifft
from ._utils import (
build_ic_set,
make_grid,
Expand All @@ -26,7 +26,7 @@
"derivative",
"fft",
"ifft",
"make_incompressible",
"get_spectrum",
"make_grid",
"rollout",
"repeat",
Expand Down
16 changes: 10 additions & 6 deletions exponax/_interpolation.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,12 @@
from jaxtyping import Array, Complex, Float

from ._spectral import (
build_reconstructional_scaling_array,
build_scaled_wavenumbers,
build_scaling_array,
fft,
get_modes_slices,
ifft,
nyquist_filter_mask,
oddball_filter_mask,
space_indices,
wavenumber_shape,
)
Expand Down Expand Up @@ -84,8 +83,11 @@ def __init__(
self.num_points = state.shape[-1]

self.state_hat_scaled = fft(state, num_spatial_dims=self.num_spatial_dims) / (
build_reconstructional_scaling_array(
self.num_spatial_dims, self.num_points, indexing=indexing
build_scaling_array(
self.num_spatial_dims,
self.num_points,
mode="reconstruction",
indexing=indexing,
)
)
self.wavenumbers = build_scaled_wavenumbers(
Expand Down Expand Up @@ -242,12 +244,13 @@ def map_between_resolutions(
) / build_scaling_array(
num_spatial_dims,
old_num_points,
mode="norm_compensation",
)

if new_num_points > old_num_points:
# Upscaling
if old_num_points % 2 == 0 and oddball_zero:
old_state_hat_scaled *= nyquist_filter_mask(
old_state_hat_scaled *= oddball_filter_mask(
num_spatial_dims, old_num_points
)

Expand All @@ -269,11 +272,12 @@ def map_between_resolutions(
new_state_hat = new_state_hat_scaled * build_scaling_array(
num_spatial_dims,
new_num_points,
mode="norm_compensation",
)
if old_num_points > new_num_points:
# Downscaling
if new_num_points % 2 == 0 and oddball_zero:
new_state_hat *= nyquist_filter_mask(num_spatial_dims, new_num_points)
new_state_hat *= oddball_filter_mask(num_spatial_dims, new_num_points)

new_state = ifft(
new_state_hat,
Expand Down
Loading
Loading