Skip to content

Chebyshev Derivative Operators

ChebyshevDerivative1D

Bases: Module

1D Chebyshev derivative operator using the precomputed differentiation matrix.

Mathematical Formulation

For a function u(x) sampled at Chebyshev nodes xⱼ on [−L, L]:

(du/dx)ⱼ = Σₖ D_{jk} uₖ

where D is the (N+1)×(N+1) (Gauss–Lobatto) or N×N (Gauss) differentiation matrix precomputed in :class:ChebyshevGrid1D. Higher-order derivatives are matrix powers:

d²u/dx² = D · D · u = D² · u

Attributes

grid : ChebyshevGrid1D 1D Chebyshev grid carrying the differentiation matrix D.

Examples

Derivative of sin(πx) on [−1, 1]:

import jax.numpy as jnp grid = ChebyshevGrid1D.from_N_L(N=32, L=1.0) deriv = ChebyshevDerivative1D(grid=grid) u = jnp.sin(jnp.pi * grid.x) du_dx = deriv(u) # ≈ π cos(πx) d2u_dx2 = deriv(u, order=2) # ≈ −π² sin(πx)

Source code in spectraldiffx/_src/chebyshev/operators.py 
17
18
19
20
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
class ChebyshevDerivative1D(eqx.Module):
    """1D Chebyshev derivative operator using the precomputed differentiation matrix.

    Mathematical Formulation
    ------------------------
    For a function u(x) sampled at Chebyshev nodes xⱼ on [−L, L]:

        (du/dx)ⱼ = Σₖ D_{jk} uₖ

    where D is the (N+1)×(N+1) (Gauss–Lobatto) or N×N (Gauss) differentiation
    matrix precomputed in :class:`ChebyshevGrid1D`.  Higher-order derivatives
    are matrix powers:

        d²u/dx² = D · D · u = D² · u

    Attributes
    ----------
    grid : ChebyshevGrid1D
        1D Chebyshev grid carrying the differentiation matrix D.

    Examples
    --------
    Derivative of sin(πx) on [−1, 1]:

    >>> import jax.numpy as jnp
    >>> grid = ChebyshevGrid1D.from_N_L(N=32, L=1.0)
    >>> deriv = ChebyshevDerivative1D(grid=grid)
    >>> u = jnp.sin(jnp.pi * grid.x)
    >>> du_dx = deriv(u)  # ≈ π cos(πx)
    >>> d2u_dx2 = deriv(u, order=2)  # ≈ −π² sin(πx)
    """

    grid: ChebyshevGrid1D

    def __call__(self, u: Num[Array, "Npts"], order: int = 1) -> Float[Array, "Npts"]:
        """Apply the n-th derivative Dⁿ to a nodal field.

        Parameters
        ----------
        u : Float[Array, "Npts"]
            Nodal values at Chebyshev nodes (Npts = N+1 for GL, N for Gauss).
        order : int
            Derivative order (≥ 0).  ``order=0`` returns a copy of ``u``.

        Returns
        -------
        Float[Array, "Npts"]
            n-th derivative at the Chebyshev nodes.
        """
        if order < 0:
            raise ValueError(f"order must be >= 0, got {order}")
        D = self.grid.D
        result = u
        for _ in range(order):
            result = D @ result
        return result

    def gradient(self, u: Num[Array, "Npts"]) -> Float[Array, "Npts"]:
        """First derivative ``du/dx`` at Chebyshev nodes."""
        return self(u, order=1)

    def laplacian(self, u: Num[Array, "Npts"]) -> Float[Array, "Npts"]:
        """Second derivative ``d²u/dx²`` at Chebyshev nodes."""
        return self(u, order=2)

Functions

__call__(u, order=1)

Apply the n-th derivative Dⁿ to a nodal field.

Parameters

u : Float[Array, "Npts"] Nodal values at Chebyshev nodes (Npts = N+1 for GL, N for Gauss). order : int Derivative order (≥ 0). order=0 returns a copy of u.

Returns

Float[Array, "Npts"] n-th derivative at the Chebyshev nodes.

Source code in spectraldiffx/_src/chebyshev/operators.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
def __call__(self, u: Num[Array, "Npts"], order: int = 1) -> Float[Array, "Npts"]:
    """Apply the n-th derivative Dⁿ to a nodal field.

    Parameters
    ----------
    u : Float[Array, "Npts"]
        Nodal values at Chebyshev nodes (Npts = N+1 for GL, N for Gauss).
    order : int
        Derivative order (≥ 0).  ``order=0`` returns a copy of ``u``.

    Returns
    -------
    Float[Array, "Npts"]
        n-th derivative at the Chebyshev nodes.
    """
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    D = self.grid.D
    result = u
    for _ in range(order):
        result = D @ result
    return result

gradient(u)

First derivative du/dx at Chebyshev nodes.

Source code in spectraldiffx/_src/chebyshev/operators.py 
74
75
76
def gradient(self, u: Num[Array, "Npts"]) -> Float[Array, "Npts"]:
    """First derivative ``du/dx`` at Chebyshev nodes."""
    return self(u, order=1)

laplacian(u)

Second derivative d²u/dx² at Chebyshev nodes.

Source code in spectraldiffx/_src/chebyshev/operators.py 
78
79
80
def laplacian(self, u: Num[Array, "Npts"]) -> Float[Array, "Npts"]:
    """Second derivative ``d²u/dx²`` at Chebyshev nodes."""
    return self(u, order=2)

ChebyshevDerivative2D

Bases: Module

2D Chebyshev derivative operators on [−Lx, Lx] × [−Ly, Ly].

Mathematical Formulation

For u(x, y) on a (Nypts, Nxpts) grid with differentiation matrices Dx, Dy stored on the grid:

(∂u/∂x)[j, i] = (u · Dxᵀ)[j, i]   # applied along axis 1 (x)
(∂u/∂y)[j, i] = (Dy · u)[j, i]    # applied along axis 0 (y)

The scalar Laplacian and 2D divergence/curl follow directly:

∇²u   = ∂²u/∂x² + ∂²u/∂y²
∇·V  = ∂vₓ/∂x + ∂vᵧ/∂y
(∇×V)_z = ∂vᵧ/∂x − ∂vₓ/∂y

Attributes

grid : ChebyshevGrid2D 2D Chebyshev grid carrying Dx, Dy and the precomputed Dx², Dy².

Examples

Laplacian of u(x, y) = sin(πx)·sin(πy) on [−1, 1]²:

import jax.numpy as jnp grid = ChebyshevGrid2D.from_N_L(Nx=24, Ny=24, Lx=1.0, Ly=1.0) deriv = ChebyshevDerivative2D(grid=grid) X, Y = grid.X u = jnp.sin(jnp.pi * X) * jnp.sin(jnp.pi * Y) lap_u = deriv.laplacian(u) # ≈ −2π² u

Divergence of V = (y, −x) (should be ~0):

vx, vy = Y, -X div = deriv.divergence(vx, vy) # ≈ 0

Source code in spectraldiffx/_src/chebyshev/operators.py 
 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
129
130
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
class ChebyshevDerivative2D(eqx.Module):
    """2D Chebyshev derivative operators on [−Lx, Lx] × [−Ly, Ly].

    Mathematical Formulation
    ------------------------
    For u(x, y) on a (Nypts, Nxpts) grid with differentiation matrices Dx, Dy
    stored on the grid:

        (∂u/∂x)[j, i] = (u · Dxᵀ)[j, i]   # applied along axis 1 (x)
        (∂u/∂y)[j, i] = (Dy · u)[j, i]    # applied along axis 0 (y)

    The scalar Laplacian and 2D divergence/curl follow directly:

        ∇²u   = ∂²u/∂x² + ∂²u/∂y²
        ∇·V  = ∂vₓ/∂x + ∂vᵧ/∂y
        (∇×V)_z = ∂vᵧ/∂x − ∂vₓ/∂y

    Attributes
    ----------
    grid : ChebyshevGrid2D
        2D Chebyshev grid carrying Dx, Dy and the precomputed Dx², Dy².

    Examples
    --------
    Laplacian of u(x, y) = sin(πx)·sin(πy) on [−1, 1]²:

    >>> import jax.numpy as jnp
    >>> grid = ChebyshevGrid2D.from_N_L(Nx=24, Ny=24, Lx=1.0, Ly=1.0)
    >>> deriv = ChebyshevDerivative2D(grid=grid)
    >>> X, Y = grid.X
    >>> u = jnp.sin(jnp.pi * X) * jnp.sin(jnp.pi * Y)
    >>> lap_u = deriv.laplacian(u)  # ≈ −2π² u

    Divergence of V = (y, −x) (should be ~0):

    >>> vx, vy = Y, -X
    >>> div = deriv.divergence(vx, vy)  # ≈ 0
    """

    grid: ChebyshevGrid2D

    def gradient(
        self, u: Num[Array, "Nypts Nxpts"]
    ) -> tuple[Float[Array, "Nypts Nxpts"], Float[Array, "Nypts Nxpts"]]:
        """Partial derivatives (∂u/∂x, ∂u/∂y) of a 2D nodal field."""
        Dx = self.grid.Dx
        Dy = self.grid.Dy
        du_dx = u @ Dx.T  # apply Dx along x-axis (axis 1)
        du_dy = Dy @ u  # apply Dy along y-axis (axis 0)
        return du_dx, du_dy

    def laplacian(self, u: Num[Array, "Nypts Nxpts"]) -> Float[Array, "Nypts Nxpts"]:
        """2D Laplacian ∇²u = ∂²u/∂x² + ∂²u/∂y².

        Uses precomputed Dx² and Dy² from the grid so the per-call cost is
        two matrix–matrix multiplies and an add (no O(N³) recomputation).
        """
        d2u_dx2 = u @ self.grid.Dx2.T
        d2u_dy2 = self.grid.Dy2 @ u
        return d2u_dx2 + d2u_dy2

    def divergence(
        self,
        vx: Num[Array, "Nypts Nxpts"],
        vy: Num[Array, "Nypts Nxpts"],
    ) -> Float[Array, "Nypts Nxpts"]:
        """Cartesian divergence ∇·V = ∂vₓ/∂x + ∂vᵧ/∂y."""
        Dx = self.grid.Dx
        Dy = self.grid.Dy
        dvx_dx = vx @ Dx.T
        dvy_dy = Dy @ vy
        return dvx_dx + dvy_dy

    def curl(
        self,
        vx: Num[Array, "Nypts Nxpts"],
        vy: Num[Array, "Nypts Nxpts"],
    ) -> Float[Array, "Nypts Nxpts"]:
        """Scalar curl ζ = ∂vᵧ/∂x − ∂vₓ/∂y (z-component of ∇×V)."""
        Dx = self.grid.Dx
        Dy = self.grid.Dy
        dvy_dx = vy @ Dx.T
        dvx_dy = Dy @ vx
        return dvy_dx - dvx_dy

    def advection_scalar(
        self,
        vx: Num[Array, "Nypts Nxpts"],
        vy: Num[Array, "Nypts Nxpts"],
        q: Num[Array, "Nypts Nxpts"],
    ) -> Float[Array, "Nypts Nxpts"]:
        """Scalar advection (V·∇)q = vₓ·∂q/∂x + vᵧ·∂q/∂y."""
        dq_dx, dq_dy = self.gradient(q)
        return vx * dq_dx + vy * dq_dy

Functions

gradient(u)

Partial derivatives (∂u/∂x, ∂u/∂y) of a 2D nodal field.

Source code in spectraldiffx/_src/chebyshev/operators.py 
124
125
126
127
128
129
130
131
132
def gradient(
    self, u: Num[Array, "Nypts Nxpts"]
) -> tuple[Float[Array, "Nypts Nxpts"], Float[Array, "Nypts Nxpts"]]:
    """Partial derivatives (∂u/∂x, ∂u/∂y) of a 2D nodal field."""
    Dx = self.grid.Dx
    Dy = self.grid.Dy
    du_dx = u @ Dx.T  # apply Dx along x-axis (axis 1)
    du_dy = Dy @ u  # apply Dy along y-axis (axis 0)
    return du_dx, du_dy

laplacian(u)

2D Laplacian ∇²u = ∂²u/∂x² + ∂²u/∂y².

Uses precomputed Dx² and Dy² from the grid so the per-call cost is two matrix–matrix multiplies and an add (no O(N³) recomputation).

Source code in spectraldiffx/_src/chebyshev/operators.py 
134
135
136
137
138
139
140
141
142
def laplacian(self, u: Num[Array, "Nypts Nxpts"]) -> Float[Array, "Nypts Nxpts"]:
    """2D Laplacian ∇²u = ∂²u/∂x² + ∂²u/∂y².

    Uses precomputed Dx² and Dy² from the grid so the per-call cost is
    two matrix–matrix multiplies and an add (no O(N³) recomputation).
    """
    d2u_dx2 = u @ self.grid.Dx2.T
    d2u_dy2 = self.grid.Dy2 @ u
    return d2u_dx2 + d2u_dy2

divergence(vx, vy)

Cartesian divergence ∇·V = ∂vₓ/∂x + ∂vᵧ/∂y.

Source code in spectraldiffx/_src/chebyshev/operators.py 
144
145
146
147
148
149
150
151
152
153
154
def divergence(
    self,
    vx: Num[Array, "Nypts Nxpts"],
    vy: Num[Array, "Nypts Nxpts"],
) -> Float[Array, "Nypts Nxpts"]:
    """Cartesian divergence ∇·V = ∂vₓ/∂x + ∂vᵧ/∂y."""
    Dx = self.grid.Dx
    Dy = self.grid.Dy
    dvx_dx = vx @ Dx.T
    dvy_dy = Dy @ vy
    return dvx_dx + dvy_dy

curl(vx, vy)

Scalar curl ζ = ∂vᵧ/∂x − ∂vₓ/∂y (z-component of ∇×V).

Source code in spectraldiffx/_src/chebyshev/operators.py 
156
157
158
159
160
161
162
163
164
165
166
def curl(
    self,
    vx: Num[Array, "Nypts Nxpts"],
    vy: Num[Array, "Nypts Nxpts"],
) -> Float[Array, "Nypts Nxpts"]:
    """Scalar curl ζ = ∂vᵧ/∂x − ∂vₓ/∂y (z-component of ∇×V)."""
    Dx = self.grid.Dx
    Dy = self.grid.Dy
    dvy_dx = vy @ Dx.T
    dvx_dy = Dy @ vx
    return dvy_dx - dvx_dy

advection_scalar(vx, vy, q)

Scalar advection (V·∇)q = vₓ·∂q/∂x + vᵧ·∂q/∂y.

Source code in spectraldiffx/_src/chebyshev/operators.py 
168
169
170
171
172
173
174
175
176
def advection_scalar(
    self,
    vx: Num[Array, "Nypts Nxpts"],
    vy: Num[Array, "Nypts Nxpts"],
    q: Num[Array, "Nypts Nxpts"],
) -> Float[Array, "Nypts Nxpts"]:
    """Scalar advection (V·∇)q = vₓ·∂q/∂x + vᵧ·∂q/∂y."""
    dq_dx, dq_dy = self.gradient(q)
    return vx * dq_dx + vy * dq_dy