Skip to content

Chebyshev Grids

ChebyshevGrid1D

Bases: Module

1D Chebyshev grid on the domain [-L, L].

Mathematical Framework:

Maps the reference domain [-1, 1] to the physical domain [-L, L] via x = L·ξ.

Gauss-Lobatto nodes (includes endpoints): ξⱼ = cos(πj/N), j = 0, ..., N (N+1 points; ξ₀=1, ξₙ=-1) xⱼ = L·ξⱼ

Gauss nodes (excludes endpoints): ξⱼ = cos(π(2j+1)/(2N)), j = 0, ..., N-1 (N points) xⱼ = L·ξⱼ

Differentiation (chain rule): d/dx = (1/L) · d/dξ

So the physical differentiation matrix is: D_phys = D_ref / L

Attributes:

Name Type Description
N int 
Polynomial degree (Gauss-Lobatto: N+1 points; Gauss: N points).

L : float Physical domain half-length. Domain is [-L, L]. node_type : str Node distribution: 'gauss-lobatto' (default) or 'gauss'. dealias : str or None Dealiasing strategy: '2/3' truncates top 1/3 of modes; None keeps all.
Source code in spectraldiffx/_src/chebyshev/grid.py 
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
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
class ChebyshevGrid1D(eqx.Module):
    """
    1D Chebyshev grid on the domain [-L, L].

    Mathematical Framework:
    -----------------------
    Maps the reference domain [-1, 1] to the physical domain [-L, L] via x = L·ξ.

    Gauss-Lobatto nodes (includes endpoints):
        ξⱼ = cos(πj/N),  j = 0, ..., N     (N+1 points; ξ₀=1, ξₙ=-1)
        xⱼ = L·ξⱼ

    Gauss nodes (excludes endpoints):
        ξⱼ = cos(π(2j+1)/(2N)),  j = 0, ..., N-1    (N points)
        xⱼ = L·ξⱼ

    Differentiation (chain rule):
        d/dx = (1/L) · d/dξ

    So the physical differentiation matrix is:
        D_phys = D_ref / L

    Attributes
    ----------
        N : int
            Polynomial degree (Gauss-Lobatto: N+1 points; Gauss: N points).
        L : float
            Physical domain half-length. Domain is [-L, L].
        node_type : str
            Node distribution: 'gauss-lobatto' (default) or 'gauss'.
        dealias : str or None
            Dealiasing strategy: '2/3' truncates top 1/3 of modes; None keeps all.
    """

    N: int
    L: float
    node_type: str
    dealias: Literal["2/3", None] | None
    _D: Array  # differentiation matrix on [-L, L], shape (n_pts, n_pts)

    def __init__(
        self,
        N: int,
        L: float = 1.0,
        node_type: Literal["gauss-lobatto", "gauss"] = "gauss-lobatto",
        dealias: Literal["2/3", None] | None = "2/3",
    ):
        """
        Parameters
        ----------
        N : int
            Polynomial degree (≥ 1).
        L : float
            Domain half-length (domain is [-L, L]). Default 1.0.
        node_type : str
            'gauss-lobatto' (default) or 'gauss'.
        dealias : str or None
            '2/3' to zero the top third of modes; None for no dealiasing.
        """
        self.N = N
        self.L = L
        self.node_type = node_type
        self.dealias = dealias

        if node_type == "gauss-lobatto":
            D_np = _cheb_diff_matrix_gl(N) / L
        else:
            D_np = _cheb_diff_matrix_gauss(N) / L
        self._D = jnp.array(D_np)

    # ------------------------------------------------------------------
    # Factory class methods
    # ------------------------------------------------------------------

    @classmethod
    def from_N_L(
        cls,
        N: int,
        L: float,
        node_type: Literal["gauss-lobatto", "gauss"] = "gauss-lobatto",
        dealias: Literal["2/3", None] | None = "2/3",
    ) -> ChebyshevGrid1D:
        """
        Initialize from polynomial degree N and domain half-length L.

        Parameters
        ----------
        N : int
            Polynomial degree.
        L : float
            Domain half-length (domain is [-L, L]).

        Example:
        --------
        >>> grid = ChebyshevGrid1D.from_N_L(N=16, L=1.0)
        """
        return cls(N=N, L=L, node_type=node_type, dealias=dealias)

    @classmethod
    def from_N_dx(
        cls,
        N: int,
        dx: float,
        node_type: Literal["gauss-lobatto", "gauss"] = "gauss-lobatto",
        dealias: Literal["2/3", None] | None = "2/3",
    ) -> ChebyshevGrid1D:
        """
        Initialize from polynomial degree N and average grid spacing dx.

        Computes: L = N * dx / 2  (since dx ≈ 2L/N for the full domain [-L, L])

        Parameters
        ----------
        N : int
            Polynomial degree.
        dx : float
            Average grid spacing (dx = 2L/N).

        Example:
        --------
        >>> grid = ChebyshevGrid1D.from_N_dx(N=16, dx=0.125)  # L = 1.0
        """
        L = N * dx / 2.0
        return cls(N=N, L=L, node_type=node_type, dealias=dealias)

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def x(self) -> Float[Array, "N1"]:
        """
        Physical grid nodes on [-L, L].

        Gauss-Lobatto: xⱼ = L·cos(πj/N), j = 0,...,N  (N+1 points, decreasing)
        Gauss:         xⱼ = L·cos(π(2j+1)/(2N)), j=0,...,N-1 (N points, decreasing)

        Returns
        -------
        x : Array [N+1] for GL, [N] for Gauss
        """
        if self.node_type == "gauss-lobatto":
            return self.L * jnp.cos(jnp.pi * jnp.arange(self.N + 1) / self.N)
        else:
            j = jnp.arange(self.N)
            return self.L * jnp.cos(jnp.pi * (2 * j + 1) / (2 * self.N))

    @property
    def D(self) -> Float[Array, "N1 N1"]:
        """
        Chebyshev differentiation matrix on [-L, L].

        Satisfies: (D @ u)ᵢ ≈ du/dx at xᵢ (exact for polynomials of degree ≤ N).

        Rows sum to zero: D @ ones = 0 (derivative of constant = 0).

        Returns
        -------
        D : Array [N+1, N+1] for GL, [N, N] for Gauss
        """
        return self._D

    # ------------------------------------------------------------------
    # Transform
    # ------------------------------------------------------------------

    def transform(self, u: Array, inverse: bool = False) -> Array:
        """
        Chebyshev spectral transform (forward or inverse).

        Forward (physical → spectral):
            Computes Chebyshev expansion coefficients aₖ from nodal values uⱼ.
            Uses DCT-I (FFT-based) for Gauss-Lobatto nodes.
            Uses DCT-II (FFT-based) for Gauss nodes.

        Inverse (spectral → physical):
            Recovers nodal values uⱼ from coefficients aₖ.

        Convention (both node types):
            uⱼ = Σₖ aₖ Tₖ(xⱼ/L)

        i.e. ``a`` holds the true Chebyshev coefficients, so a constant
        field u ≡ c maps to a = [c, 0, …, 0] and Tₙ(x/L) maps to the unit
        vector eₙ.

        Parameters
        ----------
        u : Array
            Physical-space values (forward) or spectral coefficients (inverse).
        inverse : bool
            If True, compute physical values from spectral coefficients.

        Returns
        -------
        Array
            Spectral coefficients (forward) or physical values (inverse).
        """
        N = self.N
        if self.node_type == "gauss-lobatto":
            return self._transform_gl(u, N, inverse)
        else:
            return self._transform_gauss(u, N, inverse)

    @staticmethod
    def _transform_gl(u: Array, N: int, inverse: bool) -> Array:
        """FFT-based DCT-I Chebyshev transform for Gauss-Lobatto nodes.

        Forward:
            aₖ = (2 / (N cₖ)) Σⱼ'' uⱼ cos(πjk/N),   c₀ = c_N = 2, cₖ = 1 otherwise

        (the double prime halves the j = 0 and j = N terms).  This is a
        DCT-I, evaluated as a length-2N real FFT of the even extension.
        Inverse:
            uⱼ = Σₖ aₖ cos(πjk/N)

        Acts on the last axis of ``u`` (shape ``(..., N+1)``), so batched
        and tensor-product transforms need no ``vmap``.
        """
        if not inverse:
            # Even extension y = [u_0, …, u_N, u_{N-1}, …, u_1] (length 2N)
            y = jnp.concatenate([u, u[..., -2:0:-1]], axis=-1)
            a = jnp.fft.rfft(y, axis=-1).real / N
            # rfft/N gives 2/cₖ-weighted values: halve the two end modes.
            return a.at[..., 0].multiply(0.5).at[..., N].multiply(0.5)
        else:
            # irfft(X, n=2N)[j] = (1/2N)[X₀ + X_N(−1)ʲ + 2 Σ_{0<k<N} Xₖ cos(πjk/N)],
            # so X = N·a with the end modes doubled reproduces Σₖ aₖ cos(πjk/N).
            X = N * u.at[..., 0].multiply(2.0).at[..., N].multiply(2.0)
            y = jnp.fft.irfft(X + 0j, n=2 * N, axis=-1)
            return y[..., : N + 1]

    @staticmethod
    def _transform_gauss(u: Array, N: int, inverse: bool) -> Array:
        """FFT-based DCT-II / synthesis transform for Gauss nodes.

        Forward convention:
            a[0]  = (1/N) Σⱼ uⱼ              (k=0: halved)
            a[k]  = (2/N) Σⱼ uⱼ cos(πk(2j+1)/(2N))  (k > 0)

        Inverse:
            u[n] = Σₖ a[k] Tₖ(xₙ)   (direct synthesis, no adjustment to a[0])

        Dtypes are derived from the input to preserve precision in x64 mode.
        Acts on the last axis of ``u`` (shape ``(..., N)``).
        """
        # Derive real and complex dtypes from the input to preserve precision.
        # NB: ``jnp.result_type(u, jnp.float32)`` would downcast float64 to
        # float32 (the scalar type is treated as a strong dtype), so promote
        # the dtypes explicitly.
        real_dtype = jnp.promote_types(u.dtype, jnp.float32)
        complex_dtype = jnp.promote_types(real_dtype, jnp.complex64)

        if not inverse:
            # Forward DCT-II: a[k] = (2/N) Σⱼ uⱼ cos(πk(2j+1)/(2N))
            # Via FFT: zero-pad, apply half-sample phase shift
            pad = jnp.zeros((*u.shape[:-1], N), dtype=real_dtype)
            u_pad = jnp.concatenate([u.astype(real_dtype), pad], axis=-1)  # length 2N
            Y = jnp.fft.rfft(u_pad, axis=-1)  # complex, length N+1
            k_idx = jnp.arange(N + 1, dtype=real_dtype)
            phase = jnp.exp((-1j * jnp.pi * k_idx / (2 * N)).astype(complex_dtype))
            Z = Y * phase  # Z[k] = Σⱼ uⱼ exp(-iπ(2j+1)k/(2N))
            a = Z[..., :N].real * (2.0 / N)
            # Halve k=0 so that a[0] = (1/N) Σ uⱼ
            return a.at[..., 0].multiply(0.5)
        else:
            # Inverse (synthesis): u[n] = Σₖ a[k] Tₖ(xₙ)
            # = Re[Σₖ a[k] exp(iπk(2n+1)/(2N))]
            # = Re[2N · IFFT_{2N}([h₀,...,h_{N-1},0,...,0])[n]]
            # where h[k] = a[k] · exp(iπk/(2N))
            # NOTE: a[0] is NOT doubled here; the halving in the forward
            # is already accounted for since the synthesis uses a[k] directly.
            k_idx = jnp.arange(N, dtype=real_dtype)
            phase = jnp.exp((1j * jnp.pi * k_idx / (2 * N)).astype(complex_dtype))
            h = u.astype(complex_dtype) * phase  # complex, length N
            pad = jnp.zeros((*u.shape[:-1], N), dtype=complex_dtype)
            H_full = jnp.concatenate([h, pad], axis=-1)
            y = jnp.fft.ifft(H_full, axis=-1)
            return (2 * N * y.real)[..., :N]

    # ------------------------------------------------------------------
    # Dealiasing
    # ------------------------------------------------------------------

    def dealias_filter(self) -> Float[Array, "N1"]:
        """
        Dealiasing mask in Chebyshev mode space.

        2/3 rule: keep modes 0,...,floor(2N/3), zero out higher modes.
        This prevents aliasing in quadratic nonlinearities.

        Returns
        -------
        mask : Array [N+1] for GL, [N] for Gauss
            1.0 for kept modes, 0.0 for removed modes.
        """
        n_modes = self.N + 1 if self.node_type == "gauss-lobatto" else self.N
        if self.dealias == "2/3":
            cutoff = int(2 * self.N / 3)
            mask = jnp.arange(n_modes) <= cutoff
            return mask.astype(self._D.dtype)
        else:
            return jnp.ones(n_modes)

    # ------------------------------------------------------------------
    # Consistency check
    # ------------------------------------------------------------------

    def check_consistency(self, rtol: float = 1e-5) -> bool:
        """
        Verify that N ≥ 1 and L > 0.

        Parameters
        ----------
        rtol : float
            Unused (kept for API consistency with FourierGrid).

        Returns
        -------
        bool
            True if consistent, raises ValueError otherwise.
        """
        if self.N < 1:
            raise ValueError(f"N must be ≥ 1, got N={self.N}")
        if self.L <= 0:
            raise ValueError(f"L must be > 0, got L={self.L}")
        return True

Attributes

x property

Physical grid nodes on [-L, L].

Gauss-Lobatto: xⱼ = L·cos(πj/N), j = 0,...,N (N+1 points, decreasing) Gauss: xⱼ = L·cos(π(2j+1)/(2N)), j=0,...,N-1 (N points, decreasing)

Returns:

Name Type Description
x Array [N+1] for GL, [N] for Gauss

D property

Chebyshev differentiation matrix on [-L, L].

Satisfies: (D @ u)ᵢ ≈ du/dx at xᵢ (exact for polynomials of degree ≤ N).

Rows sum to zero: D @ ones = 0 (derivative of constant = 0).

Returns:

Name Type Description
D Array [N+1, N+1] for GL, [N, N] for Gauss

Functions

__init__(N, L=1.0, node_type='gauss-lobatto', dealias='2/3')

Parameters:

Name Type Description Default
N int

Polynomial degree (≥ 1).

 required
L float

Domain half-length (domain is [-L, L]). Default 1.0.

 1.0
node_type str

'gauss-lobatto' (default) or 'gauss'.

 'gauss-lobatto'
dealias str or None

'2/3' to zero the top third of modes; None for no dealiasing.

 '2/3'
Source code in spectraldiffx/_src/chebyshev/grid.py 
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
def __init__(
    self,
    N: int,
    L: float = 1.0,
    node_type: Literal["gauss-lobatto", "gauss"] = "gauss-lobatto",
    dealias: Literal["2/3", None] | None = "2/3",
):
    """
    Parameters
    ----------
    N : int
        Polynomial degree (≥ 1).
    L : float
        Domain half-length (domain is [-L, L]). Default 1.0.
    node_type : str
        'gauss-lobatto' (default) or 'gauss'.
    dealias : str or None
        '2/3' to zero the top third of modes; None for no dealiasing.
    """
    self.N = N
    self.L = L
    self.node_type = node_type
    self.dealias = dealias

    if node_type == "gauss-lobatto":
        D_np = _cheb_diff_matrix_gl(N) / L
    else:
        D_np = _cheb_diff_matrix_gauss(N) / L
    self._D = jnp.array(D_np)

from_N_L(N, L, node_type='gauss-lobatto', dealias='2/3') classmethod

Initialize from polynomial degree N and domain half-length L.

Parameters:

Name Type Description Default
N int

Polynomial degree.

 required
L float

Domain half-length (domain is [-L, L]).

 required
Example:

grid = ChebyshevGrid1D.from_N_L(N=16, L=1.0)

Source code in spectraldiffx/_src/chebyshev/grid.py 
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
@classmethod
def from_N_L(
    cls,
    N: int,
    L: float,
    node_type: Literal["gauss-lobatto", "gauss"] = "gauss-lobatto",
    dealias: Literal["2/3", None] | None = "2/3",
) -> ChebyshevGrid1D:
    """
    Initialize from polynomial degree N and domain half-length L.

    Parameters
    ----------
    N : int
        Polynomial degree.
    L : float
        Domain half-length (domain is [-L, L]).

    Example:
    --------
    >>> grid = ChebyshevGrid1D.from_N_L(N=16, L=1.0)
    """
    return cls(N=N, L=L, node_type=node_type, dealias=dealias)

from_N_dx(N, dx, node_type='gauss-lobatto', dealias='2/3') classmethod

Initialize from polynomial degree N and average grid spacing dx.

Computes: L = N * dx / 2 (since dx ≈ 2L/N for the full domain [-L, L])

Parameters:

Name Type Description Default
N int

Polynomial degree.

 required
dx float

Average grid spacing (dx = 2L/N).

 required
Example:

grid = ChebyshevGrid1D.from_N_dx(N=16, dx=0.125) # L = 1.0

Source code in spectraldiffx/_src/chebyshev/grid.py 
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
@classmethod
def from_N_dx(
    cls,
    N: int,
    dx: float,
    node_type: Literal["gauss-lobatto", "gauss"] = "gauss-lobatto",
    dealias: Literal["2/3", None] | None = "2/3",
) -> ChebyshevGrid1D:
    """
    Initialize from polynomial degree N and average grid spacing dx.

    Computes: L = N * dx / 2  (since dx ≈ 2L/N for the full domain [-L, L])

    Parameters
    ----------
    N : int
        Polynomial degree.
    dx : float
        Average grid spacing (dx = 2L/N).

    Example:
    --------
    >>> grid = ChebyshevGrid1D.from_N_dx(N=16, dx=0.125)  # L = 1.0
    """
    L = N * dx / 2.0
    return cls(N=N, L=L, node_type=node_type, dealias=dealias)

transform(u, inverse=False)

Chebyshev spectral transform (forward or inverse).

Forward (physical → spectral): Computes Chebyshev expansion coefficients aₖ from nodal values uⱼ. Uses DCT-I (FFT-based) for Gauss-Lobatto nodes. Uses DCT-II (FFT-based) for Gauss nodes.

Inverse (spectral → physical): Recovers nodal values uⱼ from coefficients aₖ.

Convention (both node types): uⱼ = Σₖ aₖ Tₖ(xⱼ/L)

i.e. a holds the true Chebyshev coefficients, so a constant field u ≡ c maps to a = [c, 0, …, 0] and Tₙ(x/L) maps to the unit vector eₙ.

Parameters:

Name Type Description Default
u Array

Physical-space values (forward) or spectral coefficients (inverse).

 required
inverse bool

If True, compute physical values from spectral coefficients.

 False

Returns:

Type Description
Array

Spectral coefficients (forward) or physical values (inverse).
Source code in spectraldiffx/_src/chebyshev/grid.py 
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
def transform(self, u: Array, inverse: bool = False) -> Array:
    """
    Chebyshev spectral transform (forward or inverse).

    Forward (physical → spectral):
        Computes Chebyshev expansion coefficients aₖ from nodal values uⱼ.
        Uses DCT-I (FFT-based) for Gauss-Lobatto nodes.
        Uses DCT-II (FFT-based) for Gauss nodes.

    Inverse (spectral → physical):
        Recovers nodal values uⱼ from coefficients aₖ.

    Convention (both node types):
        uⱼ = Σₖ aₖ Tₖ(xⱼ/L)

    i.e. ``a`` holds the true Chebyshev coefficients, so a constant
    field u ≡ c maps to a = [c, 0, …, 0] and Tₙ(x/L) maps to the unit
    vector eₙ.

    Parameters
    ----------
    u : Array
        Physical-space values (forward) or spectral coefficients (inverse).
    inverse : bool
        If True, compute physical values from spectral coefficients.

    Returns
    -------
    Array
        Spectral coefficients (forward) or physical values (inverse).
    """
    N = self.N
    if self.node_type == "gauss-lobatto":
        return self._transform_gl(u, N, inverse)
    else:
        return self._transform_gauss(u, N, inverse)

dealias_filter()

Dealiasing mask in Chebyshev mode space.

2/3 rule: keep modes 0,...,floor(2N/3), zero out higher modes. This prevents aliasing in quadratic nonlinearities.

Returns:

Name Type Description
mask Array [N+1] for GL, [N] for Gauss

1.0 for kept modes, 0.0 for removed modes.
Source code in spectraldiffx/_src/chebyshev/grid.py 
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
def dealias_filter(self) -> Float[Array, "N1"]:
    """
    Dealiasing mask in Chebyshev mode space.

    2/3 rule: keep modes 0,...,floor(2N/3), zero out higher modes.
    This prevents aliasing in quadratic nonlinearities.

    Returns
    -------
    mask : Array [N+1] for GL, [N] for Gauss
        1.0 for kept modes, 0.0 for removed modes.
    """
    n_modes = self.N + 1 if self.node_type == "gauss-lobatto" else self.N
    if self.dealias == "2/3":
        cutoff = int(2 * self.N / 3)
        mask = jnp.arange(n_modes) <= cutoff
        return mask.astype(self._D.dtype)
    else:
        return jnp.ones(n_modes)

check_consistency(rtol=1e-05)

Verify that N ≥ 1 and L > 0.

Parameters:

Name Type Description Default
rtol float

Unused (kept for API consistency with FourierGrid).

 1e-05

Returns:

Type Description
bool

True if consistent, raises ValueError otherwise.
Source code in spectraldiffx/_src/chebyshev/grid.py 
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
def check_consistency(self, rtol: float = 1e-5) -> bool:
    """
    Verify that N ≥ 1 and L > 0.

    Parameters
    ----------
    rtol : float
        Unused (kept for API consistency with FourierGrid).

    Returns
    -------
    bool
        True if consistent, raises ValueError otherwise.
    """
    if self.N < 1:
        raise ValueError(f"N must be ≥ 1, got N={self.N}")
    if self.L <= 0:
        raise ValueError(f"L must be > 0, got L={self.L}")
    return True

ChebyshevGrid2D

Bases: Module

2D Chebyshev grid on [-Lx, Lx] × [-Ly, Ly].

Tensor product of two 1D Chebyshev grids. Physical arrays have shape (Ny_pts, Nx_pts) where Ny_pts = Ny+1 (GL) or Ny (Gauss), similarly for x.

Mathematical Framework:

For u(x, y) on [-Lx, Lx] × [-Ly, Ly]:

Partial derivatives via differentiation matrices:
    ∂u/∂x [j,i] = (u @ Dxᵀ)[j,i]     (Dx applied along axis 1)
    ∂u/∂y [j,i] = (Dy @ u)[j,i]        (Dy applied along axis 0)

Attributes:

Name Type Description
Nx, Ny int 
Polynomial degrees in x and y directions.

Lx, Ly : float Physical domain half-lengths. node_type : str Node type for both directions. dealias : str or None Dealiasing strategy.
Source code in spectraldiffx/_src/chebyshev/grid.py 
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
class ChebyshevGrid2D(eqx.Module):
    """
    2D Chebyshev grid on [-Lx, Lx] × [-Ly, Ly].

    Tensor product of two 1D Chebyshev grids. Physical arrays have shape
    (Ny_pts, Nx_pts) where Ny_pts = Ny+1 (GL) or Ny (Gauss), similarly for x.

    Mathematical Framework:
    -----------------------
    For u(x, y) on [-Lx, Lx] × [-Ly, Ly]:

        Partial derivatives via differentiation matrices:
            ∂u/∂x [j,i] = (u @ Dxᵀ)[j,i]     (Dx applied along axis 1)
            ∂u/∂y [j,i] = (Dy @ u)[j,i]        (Dy applied along axis 0)

    Attributes
    ----------
        Nx, Ny : int
            Polynomial degrees in x and y directions.
        Lx, Ly : float
            Physical domain half-lengths.
        node_type : str
            Node type for both directions.
        dealias : str or None
            Dealiasing strategy.
    """

    Nx: int
    Ny: int
    Lx: float
    Ly: float
    node_type: str
    dealias: Literal["2/3", None] | None
    _Dx: Array  # x-direction differentiation matrix
    _Dy: Array  # y-direction differentiation matrix
    _Dx2: Array  # x-direction second-derivative matrix (Dx @ Dx)
    _Dy2: Array  # y-direction second-derivative matrix (Dy @ Dy)

    def __init__(
        self,
        Nx: int,
        Ny: int,
        Lx: float = 1.0,
        Ly: float = 1.0,
        node_type: Literal["gauss-lobatto", "gauss"] = "gauss-lobatto",
        dealias: Literal["2/3", None] | None = "2/3",
    ):
        self.Nx = Nx
        self.Ny = Ny
        self.Lx = Lx
        self.Ly = Ly
        self.node_type = node_type
        self.dealias = dealias

        if node_type == "gauss-lobatto":
            Dx_np = _cheb_diff_matrix_gl(Nx) / Lx
            Dy_np = _cheb_diff_matrix_gl(Ny) / Ly
        else:
            Dx_np = _cheb_diff_matrix_gauss(Nx) / Lx
            Dy_np = _cheb_diff_matrix_gauss(Ny) / Ly

        self._Dx = jnp.array(Dx_np)
        self._Dy = jnp.array(Dy_np)
        self._Dx2 = jnp.array(Dx_np @ Dx_np)
        self._Dy2 = jnp.array(Dy_np @ Dy_np)

    @classmethod
    def from_N_L(
        cls,
        Nx: int,
        Ny: int,
        Lx: float,
        Ly: float,
        node_type: Literal["gauss-lobatto", "gauss"] = "gauss-lobatto",
        dealias: Literal["2/3", None] | None = "2/3",
    ) -> ChebyshevGrid2D:
        """Initialize from polynomial degrees and domain half-lengths."""
        return cls(Nx=Nx, Ny=Ny, Lx=Lx, Ly=Ly, node_type=node_type, dealias=dealias)

    @classmethod
    def from_N_dx(
        cls,
        Nx: int,
        Ny: int,
        dx: float,
        dy: float,
        node_type: Literal["gauss-lobatto", "gauss"] = "gauss-lobatto",
        dealias: Literal["2/3", None] | None = "2/3",
    ) -> ChebyshevGrid2D:
        """Initialize from polynomial degrees and average grid spacings."""
        Lx = Nx * dx / 2.0
        Ly = Ny * dy / 2.0
        return cls(Nx=Nx, Ny=Ny, Lx=Lx, Ly=Ly, node_type=node_type, dealias=dealias)

    @property
    def x(self) -> Float[Array, "Nx1"]:
        """Physical x-nodes on [-Lx, Lx]."""
        if self.node_type == "gauss-lobatto":
            return self.Lx * jnp.cos(jnp.pi * jnp.arange(self.Nx + 1) / self.Nx)
        else:
            j = jnp.arange(self.Nx)
            return self.Lx * jnp.cos(jnp.pi * (2 * j + 1) / (2 * self.Nx))

    @property
    def y(self) -> Float[Array, "Ny1"]:
        """Physical y-nodes on [-Ly, Ly]."""
        if self.node_type == "gauss-lobatto":
            return self.Ly * jnp.cos(jnp.pi * jnp.arange(self.Ny + 1) / self.Ny)
        else:
            j = jnp.arange(self.Ny)
            return self.Ly * jnp.cos(jnp.pi * (2 * j + 1) / (2 * self.Ny))

    @property
    def X(self) -> tuple[Float[Array, "Ny1 Nx1"], Float[Array, "Ny1 Nx1"]]:
        """
        2D meshgrid (X, Y) with shapes (Ny_pts, Nx_pts).

        Indexing: X[j, i] = x[i], Y[j, i] = y[j].
        """
        result = jnp.meshgrid(self.x, self.y, indexing="xy")
        return result[0], result[1]

    @property
    def Dx(self) -> Float[Array, "Nx1 Nx1"]:
        """x-direction differentiation matrix on [-Lx, Lx]."""
        return self._Dx

    @property
    def Dy(self) -> Float[Array, "Ny1 Ny1"]:
        """y-direction differentiation matrix on [-Ly, Ly]."""
        return self._Dy

    @property
    def Dx2(self) -> Float[Array, "Nx1 Nx1"]:
        """Precomputed x-direction second-derivative matrix: Dx @ Dx."""
        return self._Dx2

    @property
    def Dy2(self) -> Float[Array, "Ny1 Ny1"]:
        """Precomputed y-direction second-derivative matrix: Dy @ Dy."""
        return self._Dy2

    def dealias_filter(
        self,
    ) -> tuple[Float[Array, "Nx1"], Float[Array, "Ny1"]]:
        """
        1D dealiasing masks for x and y mode spaces.

        Returns
        -------
        (mask_x, mask_y) : tuple of Arrays
            Masks of shape (Nx_pts,) and (Ny_pts,).
        """
        nx_modes = self.Nx + 1 if self.node_type == "gauss-lobatto" else self.Nx
        ny_modes = self.Ny + 1 if self.node_type == "gauss-lobatto" else self.Ny
        if self.dealias == "2/3":
            cutoff_x = int(2 * self.Nx / 3)
            cutoff_y = int(2 * self.Ny / 3)
            mask_x = (jnp.arange(nx_modes) <= cutoff_x).astype(self._Dx.dtype)
            mask_y = (jnp.arange(ny_modes) <= cutoff_y).astype(self._Dy.dtype)
        else:
            mask_x = jnp.ones(nx_modes)
            mask_y = jnp.ones(ny_modes)
        return mask_x, mask_y

    def transform(self, u: Array, inverse: bool = False) -> Array:
        """
        2D Chebyshev transform (tensor product of 1D transforms).

        Forward: applies 1D transform along x-axis (axis 1), then y-axis (axis 0).
        Inverse: applies 1D inverse along y-axis (axis 0), then x-axis (axis 1).

        Parameters
        ----------
        u : Array [Ny_pts, Nx_pts]
            Physical-space field (forward) or spectral coefficients (inverse).
        inverse : bool
            If True, inverse transform.

        Returns
        -------
        Array [Ny_pts, Nx_pts]
        """
        Nx, Ny = self.Nx, self.Ny
        is_gl = self.node_type == "gauss-lobatto"
        _gl = ChebyshevGrid1D._transform_gl
        _gauss = ChebyshevGrid1D._transform_gauss

        def tx(row):
            return _gl(row, Nx, inverse) if is_gl else _gauss(row, Nx, inverse)

        def ty(col):
            return _gl(col, Ny, inverse) if is_gl else _gauss(col, Ny, inverse)

        if not inverse:
            # Transform along x (rows), then y (cols)
            u_x = jax.vmap(tx, in_axes=0, out_axes=0)(u)
            return jax.vmap(ty, in_axes=1, out_axes=1)(u_x)
        else:
            # Untransform along y (cols), then x (rows)
            u_y = jax.vmap(ty, in_axes=1, out_axes=1)(u)
            return jax.vmap(tx, in_axes=0, out_axes=0)(u_y)

    def check_consistency(self, rtol: float = 1e-5) -> bool:
        """Verify that Nx, Ny ≥ 1 and Lx, Ly > 0."""
        errors = []
        if self.Nx < 1:
            errors.append(f"Nx must be ≥ 1, got Nx={self.Nx}")
        if self.Ny < 1:
            errors.append(f"Ny must be ≥ 1, got Ny={self.Ny}")
        if self.Lx <= 0:
            errors.append(f"Lx must be > 0, got Lx={self.Lx}")
        if self.Ly <= 0:
            errors.append(f"Ly must be > 0, got Ly={self.Ly}")
        if errors:
            raise ValueError("\n".join(errors))
        return True

Attributes

x property

Physical x-nodes on [-Lx, Lx].

y property

Physical y-nodes on [-Ly, Ly].

X property

2D meshgrid (X, Y) with shapes (Ny_pts, Nx_pts).

Indexing: X[j, i] = x[i], Y[j, i] = y[j].

Dx property

x-direction differentiation matrix on [-Lx, Lx].

Dy property

y-direction differentiation matrix on [-Ly, Ly].

Dx2 property

Precomputed x-direction second-derivative matrix: Dx @ Dx.

Dy2 property

Precomputed y-direction second-derivative matrix: Dy @ Dy.

Functions

from_N_L(Nx, Ny, Lx, Ly, node_type='gauss-lobatto', dealias='2/3') classmethod

Initialize from polynomial degrees and domain half-lengths.

Source code in spectraldiffx/_src/chebyshev/grid.py 
536
537
538
539
540
541
542
543
544
545
546
547
@classmethod
def from_N_L(
    cls,
    Nx: int,
    Ny: int,
    Lx: float,
    Ly: float,
    node_type: Literal["gauss-lobatto", "gauss"] = "gauss-lobatto",
    dealias: Literal["2/3", None] | None = "2/3",
) -> ChebyshevGrid2D:
    """Initialize from polynomial degrees and domain half-lengths."""
    return cls(Nx=Nx, Ny=Ny, Lx=Lx, Ly=Ly, node_type=node_type, dealias=dealias)

from_N_dx(Nx, Ny, dx, dy, node_type='gauss-lobatto', dealias='2/3') classmethod

Initialize from polynomial degrees and average grid spacings.

Source code in spectraldiffx/_src/chebyshev/grid.py 
549
550
551
552
553
554
555
556
557
558
559
560
561
562
@classmethod
def from_N_dx(
    cls,
    Nx: int,
    Ny: int,
    dx: float,
    dy: float,
    node_type: Literal["gauss-lobatto", "gauss"] = "gauss-lobatto",
    dealias: Literal["2/3", None] | None = "2/3",
) -> ChebyshevGrid2D:
    """Initialize from polynomial degrees and average grid spacings."""
    Lx = Nx * dx / 2.0
    Ly = Ny * dy / 2.0
    return cls(Nx=Nx, Ny=Ny, Lx=Lx, Ly=Ly, node_type=node_type, dealias=dealias)

dealias_filter()

1D dealiasing masks for x and y mode spaces.

Returns:

Type Description
(mask_x, mask_y) : tuple of Arrays

Masks of shape (Nx_pts,) and (Ny_pts,).
Source code in spectraldiffx/_src/chebyshev/grid.py 
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def dealias_filter(
    self,
) -> tuple[Float[Array, "Nx1"], Float[Array, "Ny1"]]:
    """
    1D dealiasing masks for x and y mode spaces.

    Returns
    -------
    (mask_x, mask_y) : tuple of Arrays
        Masks of shape (Nx_pts,) and (Ny_pts,).
    """
    nx_modes = self.Nx + 1 if self.node_type == "gauss-lobatto" else self.Nx
    ny_modes = self.Ny + 1 if self.node_type == "gauss-lobatto" else self.Ny
    if self.dealias == "2/3":
        cutoff_x = int(2 * self.Nx / 3)
        cutoff_y = int(2 * self.Ny / 3)
        mask_x = (jnp.arange(nx_modes) <= cutoff_x).astype(self._Dx.dtype)
        mask_y = (jnp.arange(ny_modes) <= cutoff_y).astype(self._Dy.dtype)
    else:
        mask_x = jnp.ones(nx_modes)
        mask_y = jnp.ones(ny_modes)
    return mask_x, mask_y

transform(u, inverse=False)

2D Chebyshev transform (tensor product of 1D transforms).

Forward: applies 1D transform along x-axis (axis 1), then y-axis (axis 0). Inverse: applies 1D inverse along y-axis (axis 0), then x-axis (axis 1).

Parameters:

Name Type Description Default
u Array[Ny_pts, Nx_pts]

Physical-space field (forward) or spectral coefficients (inverse).

 required
inverse bool

If True, inverse transform.

 False

Returns:

Type Description
Array[Ny_pts, Nx_pts]
Source code in spectraldiffx/_src/chebyshev/grid.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
def transform(self, u: Array, inverse: bool = False) -> Array:
    """
    2D Chebyshev transform (tensor product of 1D transforms).

    Forward: applies 1D transform along x-axis (axis 1), then y-axis (axis 0).
    Inverse: applies 1D inverse along y-axis (axis 0), then x-axis (axis 1).

    Parameters
    ----------
    u : Array [Ny_pts, Nx_pts]
        Physical-space field (forward) or spectral coefficients (inverse).
    inverse : bool
        If True, inverse transform.

    Returns
    -------
    Array [Ny_pts, Nx_pts]
    """
    Nx, Ny = self.Nx, self.Ny
    is_gl = self.node_type == "gauss-lobatto"
    _gl = ChebyshevGrid1D._transform_gl
    _gauss = ChebyshevGrid1D._transform_gauss

    def tx(row):
        return _gl(row, Nx, inverse) if is_gl else _gauss(row, Nx, inverse)

    def ty(col):
        return _gl(col, Ny, inverse) if is_gl else _gauss(col, Ny, inverse)

    if not inverse:
        # Transform along x (rows), then y (cols)
        u_x = jax.vmap(tx, in_axes=0, out_axes=0)(u)
        return jax.vmap(ty, in_axes=1, out_axes=1)(u_x)
    else:
        # Untransform along y (cols), then x (rows)
        u_y = jax.vmap(ty, in_axes=1, out_axes=1)(u)
        return jax.vmap(tx, in_axes=0, out_axes=0)(u_y)

check_consistency(rtol=1e-05)

Verify that Nx, Ny ≥ 1 and Lx, Ly > 0.

Source code in spectraldiffx/_src/chebyshev/grid.py 
673
674
675
676
677
678
679
680
681
682
683
684
685
686
def check_consistency(self, rtol: float = 1e-5) -> bool:
    """Verify that Nx, Ny ≥ 1 and Lx, Ly > 0."""
    errors = []
    if self.Nx < 1:
        errors.append(f"Nx must be ≥ 1, got Nx={self.Nx}")
    if self.Ny < 1:
        errors.append(f"Ny must be ≥ 1, got Ny={self.Ny}")
    if self.Lx <= 0:
        errors.append(f"Lx must be > 0, got Lx={self.Lx}")
    if self.Ly <= 0:
        errors.append(f"Ly must be > 0, got Ly={self.Ly}")
    if errors:
        raise ValueError("\n".join(errors))
    return True

ChebyshevGrid3D

Bases: Module

3D Chebyshev grid on [-Lz, Lz] × [-Ly, Ly] × [-Lx, Lx].

Tensor product of three 1D Chebyshev grids. Axis ordering mirrors :class:FourierGrid3D: physical arrays have shape (Nz_pts, Ny_pts, Nx_pts) with

axis 0 : z,   axis 1 : y,   axis 2 : x

where N_pts = N+1 (Gauss–Lobatto) or N* (Gauss).

Mathematical Framework:

For u(x, y, z) sampled at the tensor-product nodes u[k, j, i] = u(xᵢ, yⱼ, z_k):

∂u/∂x = u ×₂ Dx     (Dx contracted with axis 2)
∂u/∂y = u ×₁ Dy     (Dy contracted with axis 1)
∂u/∂z = u ×₀ Dz     (Dz contracted with axis 0)

Each 1D matrix is (n+1)×(n+1) for GL nodes, so a single partial derivative costs O(Nx·Ny·Nz·N) rather than the O((Nx·Ny·Nz)²) of a dense Kronecker operator.

Attributes:

Name Type Description
Nx, Ny, Nz int 
Polynomial degrees in x, y and z.

Lx, Ly, Lz : float Physical domain half-lengths. node_type : str Node type for all three directions. dealias : str or None Dealiasing strategy ('2/3' or None).
Example:

grid = ChebyshevGrid3D.from_N_L(Nx=8, Ny=8, Nz=8, Lx=1.0, Ly=1.0, Lz=1.0) Z, Y, X = grid.X # each of shape (9, 9, 9)

Source code in spectraldiffx/_src/chebyshev/grid.py 
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
class ChebyshevGrid3D(eqx.Module):
    """
    3D Chebyshev grid on [-Lz, Lz] × [-Ly, Ly] × [-Lx, Lx].

    Tensor product of three 1D Chebyshev grids.  Axis ordering mirrors
    :class:`FourierGrid3D`: physical arrays have shape
    ``(Nz_pts, Ny_pts, Nx_pts)`` with

        axis 0 : z,   axis 1 : y,   axis 2 : x

    where N*_pts = N*+1 (Gauss–Lobatto) or N* (Gauss).

    Mathematical Framework:
    -----------------------
    For u(x, y, z) sampled at the tensor-product nodes u[k, j, i] = u(xᵢ, yⱼ, z_k):

        ∂u/∂x = u ×₂ Dx     (Dx contracted with axis 2)
        ∂u/∂y = u ×₁ Dy     (Dy contracted with axis 1)
        ∂u/∂z = u ×₀ Dz     (Dz contracted with axis 0)

    Each 1D matrix is (n+1)×(n+1) for GL nodes, so a single partial
    derivative costs O(Nx·Ny·Nz·N) rather than the O((Nx·Ny·Nz)²) of a
    dense Kronecker operator.

    Attributes
    ----------
        Nx, Ny, Nz : int
            Polynomial degrees in x, y and z.
        Lx, Ly, Lz : float
            Physical domain half-lengths.
        node_type : str
            Node type for all three directions.
        dealias : str or None
            Dealiasing strategy ('2/3' or None).

    Example:
    --------
    >>> grid = ChebyshevGrid3D.from_N_L(Nx=8, Ny=8, Nz=8, Lx=1.0, Ly=1.0, Lz=1.0)
    >>> Z, Y, X = grid.X  # each of shape (9, 9, 9)
    """

    Nx: int
    Ny: int
    Nz: int
    Lx: float
    Ly: float
    Lz: float
    node_type: str
    dealias: Literal["2/3", None] | None
    _Dx: Array
    _Dy: Array
    _Dz: Array
    _Dx2: Array
    _Dy2: Array
    _Dz2: Array

    def __init__(
        self,
        Nx: int,
        Ny: int,
        Nz: int,
        Lx: float = 1.0,
        Ly: float = 1.0,
        Lz: float = 1.0,
        node_type: Literal["gauss-lobatto", "gauss"] = "gauss-lobatto",
        dealias: Literal["2/3", None] | None = "2/3",
    ):
        self.Nx, self.Ny, self.Nz = Nx, Ny, Nz
        self.Lx, self.Ly, self.Lz = Lx, Ly, Lz
        self.node_type = node_type
        self.dealias = dealias

        diff = (
            _cheb_diff_matrix_gl
            if node_type == "gauss-lobatto"
            else _cheb_diff_matrix_gauss
        )
        Dx_np = diff(Nx) / Lx
        Dy_np = diff(Ny) / Ly
        Dz_np = diff(Nz) / Lz
        self._Dx = jnp.array(Dx_np)
        self._Dy = jnp.array(Dy_np)
        self._Dz = jnp.array(Dz_np)
        self._Dx2 = jnp.array(Dx_np @ Dx_np)
        self._Dy2 = jnp.array(Dy_np @ Dy_np)
        self._Dz2 = jnp.array(Dz_np @ Dz_np)

    @classmethod
    def from_N_L(
        cls,
        Nx: int,
        Ny: int,
        Nz: int,
        Lx: float,
        Ly: float,
        Lz: float,
        node_type: Literal["gauss-lobatto", "gauss"] = "gauss-lobatto",
        dealias: Literal["2/3", None] | None = "2/3",
    ) -> ChebyshevGrid3D:
        """Initialize from polynomial degrees and domain half-lengths."""
        return cls(
            Nx=Nx,
            Ny=Ny,
            Nz=Nz,
            Lx=Lx,
            Ly=Ly,
            Lz=Lz,
            node_type=node_type,
            dealias=dealias,
        )

    @property
    def x(self) -> Float[Array, "Nx1"]:
        """Physical x-nodes on [-Lx, Lx]."""
        return _cheb_nodes(self.Nx, self.Lx, self.node_type)

    @property
    def y(self) -> Float[Array, "Ny1"]:
        """Physical y-nodes on [-Ly, Ly]."""
        return _cheb_nodes(self.Ny, self.Ly, self.node_type)

    @property
    def z(self) -> Float[Array, "Nz1"]:
        """Physical z-nodes on [-Lz, Lz]."""
        return _cheb_nodes(self.Nz, self.Lz, self.node_type)

    @property
    def X(
        self,
    ) -> tuple[
        Float[Array, "Nz1 Ny1 Nx1"],
        Float[Array, "Nz1 Ny1 Nx1"],
        Float[Array, "Nz1 Ny1 Nx1"],
    ]:
        """
        3D meshgrid (Z, Y, X), each of shape (Nz_pts, Ny_pts, Nx_pts).

        Indexing: Z[k, j, i] = z[k], Y[k, j, i] = y[j], X[k, j, i] = x[i]
        (same ordering as :attr:`FourierGrid3D.X`).
        """
        Z, Y, X = jnp.meshgrid(self.z, self.y, self.x, indexing="ij")
        return Z, Y, X

    @property
    def Dx(self) -> Float[Array, "Nx1 Nx1"]:
        """x-direction differentiation matrix on [-Lx, Lx]."""
        return self._Dx

    @property
    def Dy(self) -> Float[Array, "Ny1 Ny1"]:
        """y-direction differentiation matrix on [-Ly, Ly]."""
        return self._Dy

    @property
    def Dz(self) -> Float[Array, "Nz1 Nz1"]:
        """z-direction differentiation matrix on [-Lz, Lz]."""
        return self._Dz

    @property
    def Dx2(self) -> Float[Array, "Nx1 Nx1"]:
        """Precomputed x-direction second-derivative matrix: Dx @ Dx."""
        return self._Dx2

    @property
    def Dy2(self) -> Float[Array, "Ny1 Ny1"]:
        """Precomputed y-direction second-derivative matrix: Dy @ Dy."""
        return self._Dy2

    @property
    def Dz2(self) -> Float[Array, "Nz1 Nz1"]:
        """Precomputed z-direction second-derivative matrix: Dz @ Dz."""
        return self._Dz2

    def dealias_filter(self) -> Float[Array, "Nz1 Ny1 Nx1"]:
        """
        3D dealiasing mask in Chebyshev mode space, shape (Nz_pts, Ny_pts, Nx_pts).

        2/3 rule per axis: keep modes kₐ ≤ ⌊2Nₐ/3⌋ in every direction;
        the 3D mask is the outer product of the three 1D masks.
        """
        offset = 1 if self.node_type == "gauss-lobatto" else 0
        dtype = self._Dx.dtype

        def mask_1d(n: int) -> Array:
            k = jnp.arange(n + offset)
            if self.dealias == "2/3":
                return (k <= int(2 * n / 3)).astype(dtype)
            return jnp.ones(n + offset, dtype=dtype)

        mz, my, mx = mask_1d(self.Nz), mask_1d(self.Ny), mask_1d(self.Nx)
        return mz[:, None, None] * my[None, :, None] * mx[None, None, :]

    def transform(self, u: Array, inverse: bool = False) -> Array:
        """
        3D Chebyshev transform (tensor product of 1D transforms).

        Forward:  a[kz, ky, kx] such that u = Σ a[kz,ky,kx] T_kz(z/Lz) T_ky(y/Ly) T_kx(x/Lx).
        The 1D transforms along separate axes commute, so the order of
        application is immaterial.

        Parameters
        ----------
        u : Array [Nz_pts, Ny_pts, Nx_pts]
            Physical-space field (forward) or coefficients (inverse).
        inverse : bool
            If True, inverse transform.

        Returns
        -------
        Array [Nz_pts, Ny_pts, Nx_pts]
        """
        out = u
        for axis, n in ((2, self.Nx), (1, self.Ny), (0, self.Nz)):
            out = _transform_along_axis(out, n, self.node_type, inverse, axis)
        return out

    def check_consistency(self, rtol: float = 1e-5) -> bool:
        """Verify that Nx, Ny, Nz ≥ 1 and Lx, Ly, Lz > 0."""
        errors = []
        for name, n in (("Nx", self.Nx), ("Ny", self.Ny), ("Nz", self.Nz)):
            if n < 1:
                errors.append(f"{name} must be ≥ 1, got {name}={n}")
        for name, length in (("Lx", self.Lx), ("Ly", self.Ly), ("Lz", self.Lz)):
            if length <= 0:
                errors.append(f"{name} must be > 0, got {name}={length}")
        if errors:
            raise ValueError("\n".join(errors))
        return True

Attributes

x property

Physical x-nodes on [-Lx, Lx].

y property

Physical y-nodes on [-Ly, Ly].

z property

Physical z-nodes on [-Lz, Lz].

X property

3D meshgrid (Z, Y, X), each of shape (Nz_pts, Ny_pts, Nx_pts).

Indexing: Z[k, j, i] = z[k], Y[k, j, i] = y[j], X[k, j, i] = x[i] (same ordering as :attr:FourierGrid3D.X).

Dx property

x-direction differentiation matrix on [-Lx, Lx].

Dy property

y-direction differentiation matrix on [-Ly, Ly].

Dz property

z-direction differentiation matrix on [-Lz, Lz].

Dx2 property

Precomputed x-direction second-derivative matrix: Dx @ Dx.

Dy2 property

Precomputed y-direction second-derivative matrix: Dy @ Dy.

Dz2 property

Precomputed z-direction second-derivative matrix: Dz @ Dz.

Functions

from_N_L(Nx, Ny, Nz, Lx, Ly, Lz, node_type='gauss-lobatto', dealias='2/3') classmethod

Initialize from polynomial degrees and domain half-lengths.

Source code in spectraldiffx/_src/chebyshev/grid.py 
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
@classmethod
def from_N_L(
    cls,
    Nx: int,
    Ny: int,
    Nz: int,
    Lx: float,
    Ly: float,
    Lz: float,
    node_type: Literal["gauss-lobatto", "gauss"] = "gauss-lobatto",
    dealias: Literal["2/3", None] | None = "2/3",
) -> ChebyshevGrid3D:
    """Initialize from polynomial degrees and domain half-lengths."""
    return cls(
        Nx=Nx,
        Ny=Ny,
        Nz=Nz,
        Lx=Lx,
        Ly=Ly,
        Lz=Lz,
        node_type=node_type,
        dealias=dealias,
    )

dealias_filter()

3D dealiasing mask in Chebyshev mode space, shape (Nz_pts, Ny_pts, Nx_pts).

2/3 rule per axis: keep modes kₐ ≤ ⌊2Nₐ/3⌋ in every direction; the 3D mask is the outer product of the three 1D masks.

Source code in spectraldiffx/_src/chebyshev/grid.py 
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
def dealias_filter(self) -> Float[Array, "Nz1 Ny1 Nx1"]:
    """
    3D dealiasing mask in Chebyshev mode space, shape (Nz_pts, Ny_pts, Nx_pts).

    2/3 rule per axis: keep modes kₐ ≤ ⌊2Nₐ/3⌋ in every direction;
    the 3D mask is the outer product of the three 1D masks.
    """
    offset = 1 if self.node_type == "gauss-lobatto" else 0
    dtype = self._Dx.dtype

    def mask_1d(n: int) -> Array:
        k = jnp.arange(n + offset)
        if self.dealias == "2/3":
            return (k <= int(2 * n / 3)).astype(dtype)
        return jnp.ones(n + offset, dtype=dtype)

    mz, my, mx = mask_1d(self.Nz), mask_1d(self.Ny), mask_1d(self.Nx)
    return mz[:, None, None] * my[None, :, None] * mx[None, None, :]

transform(u, inverse=False)

3D Chebyshev transform (tensor product of 1D transforms).

Forward: a[kz, ky, kx] such that u = Σ a[kz,ky,kx] T_kz(z/Lz) T_ky(y/Ly) T_kx(x/Lx). The 1D transforms along separate axes commute, so the order of application is immaterial.

Parameters:

Name Type Description Default
u Array[Nz_pts, Ny_pts, Nx_pts]

Physical-space field (forward) or coefficients (inverse).

 required
inverse bool

If True, inverse transform.

 False

Returns:

Type Description
Array[Nz_pts, Ny_pts, Nx_pts]
Source code in spectraldiffx/_src/chebyshev/grid.py 
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
def transform(self, u: Array, inverse: bool = False) -> Array:
    """
    3D Chebyshev transform (tensor product of 1D transforms).

    Forward:  a[kz, ky, kx] such that u = Σ a[kz,ky,kx] T_kz(z/Lz) T_ky(y/Ly) T_kx(x/Lx).
    The 1D transforms along separate axes commute, so the order of
    application is immaterial.

    Parameters
    ----------
    u : Array [Nz_pts, Ny_pts, Nx_pts]
        Physical-space field (forward) or coefficients (inverse).
    inverse : bool
        If True, inverse transform.

    Returns
    -------
    Array [Nz_pts, Ny_pts, Nx_pts]
    """
    out = u
    for axis, n in ((2, self.Nx), (1, self.Ny), (0, self.Nz)):
        out = _transform_along_axis(out, n, self.node_type, inverse, axis)
    return out

check_consistency(rtol=1e-05)

Verify that Nx, Ny, Nz ≥ 1 and Lx, Ly, Lz > 0.

Source code in spectraldiffx/_src/chebyshev/grid.py 
935
936
937
938
939
940
941
942
943
944
945
946
def check_consistency(self, rtol: float = 1e-5) -> bool:
    """Verify that Nx, Ny, Nz ≥ 1 and Lx, Ly, Lz > 0."""
    errors = []
    for name, n in (("Nx", self.Nx), ("Ny", self.Ny), ("Nz", self.Nz)):
        if n < 1:
            errors.append(f"{name} must be ≥ 1, got {name}={n}")
    for name, length in (("Lx", self.Lx), ("Ly", self.Ly), ("Lz", self.Lz)):
        if length <= 0:
            errors.append(f"{name} must be > 0, got {name}={length}")
    if errors:
        raise ValueError("\n".join(errors))
    return True