Skip to content

Poisson Solver¤

(for completion - not a time-dependent stepper)

exponax.poisson.Poisson ¤

Bases: Module

Source code in exponax/_poisson.py 
 14
 15
 16
 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
 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 Poisson(eqx.Module):
    num_spatial_dims: int
    domain_extent: float
    num_points: int
    dx: float

    _inv_operator: Complex[Array, "1 ... N"]

    def __init__(
        self,
        num_spatial_dims: int,
        domain_extent: float,
        num_points: int,
        *,
        order=2,
    ):
        """
        Exactly solve the Poisson equation with periodic boundary conditions.

        This "stepper" is different from all other steppers in this package in
        that it does not solve a time-dependent PDE. Instead, it solves the
        Poisson equation

        $$ u_{xx} = - f $$

        for a given right hand side $f$.

        It is included for completion.

        **Arguments:**

        - `num_spatial_dims`: The number of spatial dimensions.
        - `domain_extent`: The extent of the domain.
        - `num_points`: The number of points in each spatial dimension.
        - `order`: The order of the Poisson equation. Defaults to 2. You can
            also set `order=4` for the biharmonic equation.
        """
        self.num_spatial_dims = num_spatial_dims
        self.domain_extent = domain_extent
        self.num_points = num_points

        # Uses the convention that N does **not** include the right boundary
        # point
        self.dx = domain_extent / num_points

        derivative_operator = build_derivative_operator(
            num_spatial_dims, domain_extent, num_points
        )
        operator = build_laplace_operator(derivative_operator, order=order)

        # Uses mean zero solution
        self._inv_operator = jnp.where(operator == 0, 0.0, 1 / operator)

    def step_fourier(
        self,
        f_hat: Complex[Array, "C ... (N//2)+1"],
    ) -> Complex[Array, "C ... (N//2)+1"]:
        """
        Solve the Poisson equation in Fourier space.

        **Arguments:**

        - `f_hat`: The Fourier transform of the right hand side.

        **Returns:**

        - `u_hat`: The Fourier transform of the solution.
        """
        return -self._inv_operator * f_hat

    def step(
        self,
        f: Float[Array, "C ... N"],
    ) -> Float[Array, "C ... N"]:
        """
        Solve the Poisson equation in state space.

        **Arguments:**

        - `f`: The right hand side.

        **Returns:**

        - `u`: The solution.
        """
        f_hat = fft(f, num_spatial_dims=self.num_spatial_dims)
        u_hat = self.step_fourier(f_hat)
        u = ifft(
            u_hat,
            num_spatial_dims=self.num_spatial_dims,
            num_points=self.num_points,
        )
        return u

    def __call__(
        self,
        f: Float[Array, "C ... N"],
    ) -> Float[Array, "C ... N"]:
        """
        Solve the Poisson equation in state space.

        **Arguments:**

        - `f`: The right hand side.

        **Returns:**

        - `u`: The solution.
        """
        if f.shape[1:] != spatial_shape(self.num_spatial_dims, self.num_points):
            raise ValueError(
                f"""Shape of f[1:] is {f.shape[1:]} but should be
                 {spatial_shape(self.num_spatial_dims, self.num_points)}"""
            )
        return self.step(f)
__init__ ¤
__init__(
    num_spatial_dims: int,
    domain_extent: float,
    num_points: int,
    *,
    order=2
)

Exactly solve the Poisson equation with periodic boundary conditions.

This "stepper" is different from all other steppers in this package in that it does not solve a time-dependent PDE. Instead, it solves the Poisson equation

\[ u_{xx} = - f \]

for a given right hand side \(f\).

It is included for completion.

Arguments:

  • num_spatial_dims: The number of spatial dimensions.
  • domain_extent: The extent of the domain.
  • num_points: The number of points in each spatial dimension.
  • order: The order of the Poisson equation. Defaults to 2. You can also set order=4 for the biharmonic equation.
Source code in exponax/_poisson.py 
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
def __init__(
    self,
    num_spatial_dims: int,
    domain_extent: float,
    num_points: int,
    *,
    order=2,
):
    """
    Exactly solve the Poisson equation with periodic boundary conditions.

    This "stepper" is different from all other steppers in this package in
    that it does not solve a time-dependent PDE. Instead, it solves the
    Poisson equation

    $$ u_{xx} = - f $$

    for a given right hand side $f$.

    It is included for completion.

    **Arguments:**

    - `num_spatial_dims`: The number of spatial dimensions.
    - `domain_extent`: The extent of the domain.
    - `num_points`: The number of points in each spatial dimension.
    - `order`: The order of the Poisson equation. Defaults to 2. You can
        also set `order=4` for the biharmonic equation.
    """
    self.num_spatial_dims = num_spatial_dims
    self.domain_extent = domain_extent
    self.num_points = num_points

    # Uses the convention that N does **not** include the right boundary
    # point
    self.dx = domain_extent / num_points

    derivative_operator = build_derivative_operator(
        num_spatial_dims, domain_extent, num_points
    )
    operator = build_laplace_operator(derivative_operator, order=order)

    # Uses mean zero solution
    self._inv_operator = jnp.where(operator == 0, 0.0, 1 / operator)
__call__ ¤
__call__(
    f: Float[Array, "C ... N"],
) -> Float[Array, "C ... N"]

Solve the Poisson equation in state space.

Arguments:

  • f: The right hand side.

Returns:

  • u: The solution.
Source code in exponax/_poisson.py 
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
def __call__(
    self,
    f: Float[Array, "C ... N"],
) -> Float[Array, "C ... N"]:
    """
    Solve the Poisson equation in state space.

    **Arguments:**

    - `f`: The right hand side.

    **Returns:**

    - `u`: The solution.
    """
    if f.shape[1:] != spatial_shape(self.num_spatial_dims, self.num_points):
        raise ValueError(
            f"""Shape of f[1:] is {f.shape[1:]} but should be
             {spatial_shape(self.num_spatial_dims, self.num_points)}"""
        )
    return self.step(f)