Skip to content

Chebyshev Filters

ChebyshevFilter1D

Bases: Module

1D Chebyshev spectral filter for smoothing and numerical stabilization.

Mathematical Formulation:

Filters are applied as a multiplicative mask in Chebyshev coefficient space:

a_k_filtered = F(k) * a_k

where a_k are Chebyshev expansion coefficients and F(k) is the filter kernel.

Exponential filter

F(k) = exp(-alpha * (k/N)^power)

Hyperviscosity filter

F(k) = exp(-nu_h * k^power * dt)

grid : ChebyshevGrid1D
    The 1D Chebyshev grid for forward/inverse transforms.
Source code in spectraldiffx/_src/chebyshev/filters.py 
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
class ChebyshevFilter1D(eqx.Module):
    """
    1D Chebyshev spectral filter for smoothing and numerical stabilization.

    Mathematical Formulation:
    -------------------------
    Filters are applied as a multiplicative mask in Chebyshev coefficient space:

        a_k_filtered = F(k) * a_k

    where a_k are Chebyshev expansion coefficients and F(k) is the filter kernel.

    Exponential filter:
        F(k) = exp(-alpha * (k/N)^power)

    Hyperviscosity filter:
        F(k) = exp(-nu_h * k^power * dt)

    Attributes:
    -----------
        grid : ChebyshevGrid1D
            The 1D Chebyshev grid for forward/inverse transforms.
    """

    grid: ChebyshevGrid1D

    def exponential_filter(
        self,
        u: Array,
        alpha: float = 36.0,
        power: int = 16,
        spectral: bool = False,
    ) -> Array:
        """
        Apply 1D exponential filter in Chebyshev mode space.

            F(k) = exp(-alpha * (k/N)^power)

        This filter is near unity for low modes and falls off sharply near
        the highest Chebyshev mode k=N, removing poorly-resolved content.

        Parameters:
        -----------
        u : Array [N1]
            Physical-space field or Chebyshev coefficients.
        alpha : float
            Damping strength. Default 36.0 (≈ machine-epsilon damping at k=N).
        power : int
            Sharpening exponent (even integer). Default 16.
        spectral : bool
            If True, u is treated as Chebyshev coefficients. Default False.

        Returns:
        --------
        Array [N1]
            Filtered field (physical if spectral=False, else coefficients).
        """
        a = u if spectral else self.grid.transform(u)
        N = self.grid.N
        n_modes = N + 1 if self.grid.node_type == "gauss-lobatto" else N
        # Normalize by the highest mode index so F(k_max) = exp(-alpha).
        k_max = max(1, n_modes - 1)
        real_dtype = _coeff_dtype(a)
        k = jnp.arange(n_modes, dtype=real_dtype)
        filter_mask = jnp.exp(-alpha * (k / k_max) ** power)
        a_f = a * filter_mask
        return a_f if spectral else self.grid.transform(a_f, inverse=True)

    def hyperviscosity(
        self,
        u: Array,
        nu_hyper: float,
        dt: float,
        power: int = 4,
        spectral: bool = False,
    ) -> Array:
        """
        Apply 1D hyperviscous damping in Chebyshev mode space.

            F(k) = exp(-nu_h * k^power * dt)

        Simulates high-order diffusion: du/dt = (-1)^{p/2} nu_h d^p u/dx^p.

        Parameters:
        -----------
        u : Array [N1]
            Physical-space field or Chebyshev coefficients.
        nu_hyper : float
            Hyperviscosity coefficient.
        dt : float
            Time step for the damping.
        power : int
            Diffusion order. Default 4 (biharmonic).
        spectral : bool
            If True, u is treated as Chebyshev coefficients. Default False.

        Returns:
        --------
        Array [N1]
        """
        a = u if spectral else self.grid.transform(u)
        N = self.grid.N
        n_modes = N + 1 if self.grid.node_type == "gauss-lobatto" else N
        real_dtype = _coeff_dtype(a)
        k = jnp.arange(n_modes, dtype=real_dtype)
        filter_mask = jnp.exp(-nu_hyper * k**power * dt)
        a_f = a * filter_mask
        return a_f if spectral else self.grid.transform(a_f, inverse=True)

Functions

exponential_filter(u, alpha=36.0, power=16, spectral=False)

Apply 1D exponential filter in Chebyshev mode space.

F(k) = exp(-alpha * (k/N)^power)

This filter is near unity for low modes and falls off sharply near the highest Chebyshev mode k=N, removing poorly-resolved content.

Parameters:

u : Array [N1] Physical-space field or Chebyshev coefficients. alpha : float Damping strength. Default 36.0 (≈ machine-epsilon damping at k=N). power : int Sharpening exponent (even integer). Default 16. spectral : bool If True, u is treated as Chebyshev coefficients. Default False.

Returns:

Array [N1] Filtered field (physical if spectral=False, else coefficients).

Source code in spectraldiffx/_src/chebyshev/filters.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def exponential_filter(
    self,
    u: Array,
    alpha: float = 36.0,
    power: int = 16,
    spectral: bool = False,
) -> Array:
    """
    Apply 1D exponential filter in Chebyshev mode space.

        F(k) = exp(-alpha * (k/N)^power)

    This filter is near unity for low modes and falls off sharply near
    the highest Chebyshev mode k=N, removing poorly-resolved content.

    Parameters:
    -----------
    u : Array [N1]
        Physical-space field or Chebyshev coefficients.
    alpha : float
        Damping strength. Default 36.0 (≈ machine-epsilon damping at k=N).
    power : int
        Sharpening exponent (even integer). Default 16.
    spectral : bool
        If True, u is treated as Chebyshev coefficients. Default False.

    Returns:
    --------
    Array [N1]
        Filtered field (physical if spectral=False, else coefficients).
    """
    a = u if spectral else self.grid.transform(u)
    N = self.grid.N
    n_modes = N + 1 if self.grid.node_type == "gauss-lobatto" else N
    # Normalize by the highest mode index so F(k_max) = exp(-alpha).
    k_max = max(1, n_modes - 1)
    real_dtype = _coeff_dtype(a)
    k = jnp.arange(n_modes, dtype=real_dtype)
    filter_mask = jnp.exp(-alpha * (k / k_max) ** power)
    a_f = a * filter_mask
    return a_f if spectral else self.grid.transform(a_f, inverse=True)

hyperviscosity(u, nu_hyper, dt, power=4, spectral=False)

Apply 1D hyperviscous damping in Chebyshev mode space.

F(k) = exp(-nu_h * k^power * dt)

Simulates high-order diffusion: du/dt = (-1)^{p/2} nu_h d^p u/dx^p.

Parameters:

u : Array [N1] Physical-space field or Chebyshev coefficients. nu_hyper : float Hyperviscosity coefficient. dt : float Time step for the damping. power : int Diffusion order. Default 4 (biharmonic). spectral : bool If True, u is treated as Chebyshev coefficients. Default False.

Returns:

Array [N1]

Source code in spectraldiffx/_src/chebyshev/filters.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
def hyperviscosity(
    self,
    u: Array,
    nu_hyper: float,
    dt: float,
    power: int = 4,
    spectral: bool = False,
) -> Array:
    """
    Apply 1D hyperviscous damping in Chebyshev mode space.

        F(k) = exp(-nu_h * k^power * dt)

    Simulates high-order diffusion: du/dt = (-1)^{p/2} nu_h d^p u/dx^p.

    Parameters:
    -----------
    u : Array [N1]
        Physical-space field or Chebyshev coefficients.
    nu_hyper : float
        Hyperviscosity coefficient.
    dt : float
        Time step for the damping.
    power : int
        Diffusion order. Default 4 (biharmonic).
    spectral : bool
        If True, u is treated as Chebyshev coefficients. Default False.

    Returns:
    --------
    Array [N1]
    """
    a = u if spectral else self.grid.transform(u)
    N = self.grid.N
    n_modes = N + 1 if self.grid.node_type == "gauss-lobatto" else N
    real_dtype = _coeff_dtype(a)
    k = jnp.arange(n_modes, dtype=real_dtype)
    filter_mask = jnp.exp(-nu_hyper * k**power * dt)
    a_f = a * filter_mask
    return a_f if spectral else self.grid.transform(a_f, inverse=True)

ChebyshevFilter2D

Bases: Module

2D Chebyshev spectral filter for smoothing on [-Lx, Lx] × [-Ly, Ly].

Applies separable 1D exponential or hyperviscosity filters along each axis.

grid : ChebyshevGrid2D
    The 2D Chebyshev grid.
Source code in spectraldiffx/_src/chebyshev/filters.py 
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
class ChebyshevFilter2D(eqx.Module):
    """
    2D Chebyshev spectral filter for smoothing on [-Lx, Lx] × [-Ly, Ly].

    Applies separable 1D exponential or hyperviscosity filters along each axis.

    Attributes:
    -----------
        grid : ChebyshevGrid2D
            The 2D Chebyshev grid.
    """

    grid: ChebyshevGrid2D

    def exponential_filter(
        self,
        u: Array,
        alpha: float = 36.0,
        power: int = 16,
        spectral: bool = False,
    ) -> Array:
        """
        Apply 2D separable exponential filter.

            F(kx, ky) = exp(-alpha * (kx/Nx)^power) * exp(-alpha * (ky/Ny)^power)

        Parameters:
        -----------
        u : Array [Ny_pts, Nx_pts]
            Physical-space field or Chebyshev coefficients.
        alpha : float
            Damping strength. Default 36.0.
        power : int
            Sharpening exponent. Default 16.
        spectral : bool
            If True, u is treated as spectral coefficients. Default False.

        Returns:
        --------
        Array [Ny_pts, Nx_pts]
        """
        a = u if spectral else self.grid.transform(u)
        Nx, Ny = self.grid.Nx, self.grid.Ny
        nx_modes = Nx + 1 if self.grid.node_type == "gauss-lobatto" else Nx
        ny_modes = Ny + 1 if self.grid.node_type == "gauss-lobatto" else Ny
        # Normalize by the highest mode index; derive dtype from coefficients.
        kx_max = max(1, nx_modes - 1)
        ky_max = max(1, ny_modes - 1)
        real_dtype = _coeff_dtype(a)
        kx = jnp.arange(nx_modes, dtype=real_dtype)
        ky = jnp.arange(ny_modes, dtype=real_dtype)
        Fx = jnp.exp(-alpha * (kx / kx_max) ** power)
        Fy = jnp.exp(-alpha * (ky / ky_max) ** power)
        # Apply separable filter: F = Fy[:, None] * Fx[None, :]
        a_f = a * Fy[:, None] * Fx[None, :]
        return a_f if spectral else self.grid.transform(a_f, inverse=True)

    def hyperviscosity(
        self,
        u: Array,
        nu_hyper: float,
        dt: float,
        power: int = 4,
        spectral: bool = False,
    ) -> Array:
        """
        Apply 2D separable hyperviscosity filter.

            F(kx, ky) = exp(-nu_h * (kx^power + ky^power) * dt)

        Parameters:
        -----------
        u : Array [Ny_pts, Nx_pts]
        nu_hyper : float
        dt : float
        power : int
        spectral : bool

        Returns:
        --------
        Array [Ny_pts, Nx_pts]
        """
        a = u if spectral else self.grid.transform(u)
        Nx, Ny = self.grid.Nx, self.grid.Ny
        nx_modes = Nx + 1 if self.grid.node_type == "gauss-lobatto" else Nx
        ny_modes = Ny + 1 if self.grid.node_type == "gauss-lobatto" else Ny
        real_dtype = _coeff_dtype(a)
        kx = jnp.arange(nx_modes, dtype=real_dtype)
        ky = jnp.arange(ny_modes, dtype=real_dtype)
        # Separable: exp(-nu * kx^p * dt) * exp(-nu * ky^p * dt)
        Fx = jnp.exp(-nu_hyper * kx**power * dt)
        Fy = jnp.exp(-nu_hyper * ky**power * dt)
        a_f = a * Fy[:, None] * Fx[None, :]
        return a_f if spectral else self.grid.transform(a_f, inverse=True)

Functions

exponential_filter(u, alpha=36.0, power=16, spectral=False)

Apply 2D separable exponential filter.

F(kx, ky) = exp(-alpha * (kx/Nx)^power) * exp(-alpha * (ky/Ny)^power)
Parameters:

u : Array [Ny_pts, Nx_pts] Physical-space field or Chebyshev coefficients. alpha : float Damping strength. Default 36.0. power : int Sharpening exponent. Default 16. spectral : bool If True, u is treated as spectral coefficients. Default False.

Returns:

Array [Ny_pts, Nx_pts]

Source code in spectraldiffx/_src/chebyshev/filters.py 
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
def exponential_filter(
    self,
    u: Array,
    alpha: float = 36.0,
    power: int = 16,
    spectral: bool = False,
) -> Array:
    """
    Apply 2D separable exponential filter.

        F(kx, ky) = exp(-alpha * (kx/Nx)^power) * exp(-alpha * (ky/Ny)^power)

    Parameters:
    -----------
    u : Array [Ny_pts, Nx_pts]
        Physical-space field or Chebyshev coefficients.
    alpha : float
        Damping strength. Default 36.0.
    power : int
        Sharpening exponent. Default 16.
    spectral : bool
        If True, u is treated as spectral coefficients. Default False.

    Returns:
    --------
    Array [Ny_pts, Nx_pts]
    """
    a = u if spectral else self.grid.transform(u)
    Nx, Ny = self.grid.Nx, self.grid.Ny
    nx_modes = Nx + 1 if self.grid.node_type == "gauss-lobatto" else Nx
    ny_modes = Ny + 1 if self.grid.node_type == "gauss-lobatto" else Ny
    # Normalize by the highest mode index; derive dtype from coefficients.
    kx_max = max(1, nx_modes - 1)
    ky_max = max(1, ny_modes - 1)
    real_dtype = _coeff_dtype(a)
    kx = jnp.arange(nx_modes, dtype=real_dtype)
    ky = jnp.arange(ny_modes, dtype=real_dtype)
    Fx = jnp.exp(-alpha * (kx / kx_max) ** power)
    Fy = jnp.exp(-alpha * (ky / ky_max) ** power)
    # Apply separable filter: F = Fy[:, None] * Fx[None, :]
    a_f = a * Fy[:, None] * Fx[None, :]
    return a_f if spectral else self.grid.transform(a_f, inverse=True)

hyperviscosity(u, nu_hyper, dt, power=4, spectral=False)

Apply 2D separable hyperviscosity filter.

F(kx, ky) = exp(-nu_h * (kx^power + ky^power) * dt)
Parameters:

u : Array [Ny_pts, Nx_pts] nu_hyper : float dt : float power : int spectral : bool

Returns:

Array [Ny_pts, Nx_pts]

Source code in spectraldiffx/_src/chebyshev/filters.py 
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
def hyperviscosity(
    self,
    u: Array,
    nu_hyper: float,
    dt: float,
    power: int = 4,
    spectral: bool = False,
) -> Array:
    """
    Apply 2D separable hyperviscosity filter.

        F(kx, ky) = exp(-nu_h * (kx^power + ky^power) * dt)

    Parameters:
    -----------
    u : Array [Ny_pts, Nx_pts]
    nu_hyper : float
    dt : float
    power : int
    spectral : bool

    Returns:
    --------
    Array [Ny_pts, Nx_pts]
    """
    a = u if spectral else self.grid.transform(u)
    Nx, Ny = self.grid.Nx, self.grid.Ny
    nx_modes = Nx + 1 if self.grid.node_type == "gauss-lobatto" else Nx
    ny_modes = Ny + 1 if self.grid.node_type == "gauss-lobatto" else Ny
    real_dtype = _coeff_dtype(a)
    kx = jnp.arange(nx_modes, dtype=real_dtype)
    ky = jnp.arange(ny_modes, dtype=real_dtype)
    # Separable: exp(-nu * kx^p * dt) * exp(-nu * ky^p * dt)
    Fx = jnp.exp(-nu_hyper * kx**power * dt)
    Fy = jnp.exp(-nu_hyper * ky**power * dt)
    a_f = a * Fy[:, None] * Fx[None, :]
    return a_f if spectral else self.grid.transform(a_f, inverse=True)