Skip to content

ETDRK Integrator

ETDRK¤

torchfsm.integrator.ETDRKIntegrator ¤

Bases: Enum

ETDRK Integrator Provides a unified interface for all ETDRK methods.

Source code in torchfsm/integrator/_etdrk.py 
85
86
87
88
89
90
91
92
93
class ETDRKIntegrator(Enum):
    """
    ETDRK Integrator
    Provides a unified interface for all ETDRK methods.
    """

    ETDRK0 = ETDRK0
    ETDRK1 = ETDRK1
    ETDRK2 = ETDRK2
ETDRK0 class-attribute instance-attribute ¤
ETDRK0 = ETDRK0
ETDRK1 class-attribute instance-attribute ¤
ETDRK1 = ETDRK1
ETDRK2 class-attribute instance-attribute ¤
ETDRK2 = ETDRK2

torchfsm.integrator._etdrk.ETDRK0 ¤

Exactly solve a linear PDE in Fourier space

Source code in torchfsm/integrator/_etdrk.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class ETDRK0:
    """
    Exactly solve a linear PDE in Fourier space
    """

    def __init__(
        self,
        dt: float,
        linear_coef: torch.Tensor,
    ):
        self.dt = dt
        self._exp_term = torch.exp(self.dt * linear_coef)

    def step(
        self,
        u_hat,
    ):
        return self._exp_term * u_hat
dt instance-attribute ¤
dt = dt
__init__ ¤
__init__(dt: float, linear_coef: Tensor)
Source code in torchfsm/integrator/_etdrk.py 
15
16
17
18
19
20
21
def __init__(
    self,
    dt: float,
    linear_coef: torch.Tensor,
):
    self.dt = dt
    self._exp_term = torch.exp(self.dt * linear_coef)
step ¤
step(u_hat)
Source code in torchfsm/integrator/_etdrk.py 
23
24
25
26
27
def step(
    self,
    u_hat,
):
    return self._exp_term * u_hat

torchfsm.integrator._etdrk.ETDRK1 ¤

Bases: ETDRK0

First-order ETDRK method.

Source code in torchfsm/integrator/_etdrk.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
class ETDRK1(ETDRK0):
    """
    First-order ETDRK method.
    """

    def __init__(
        self,
        dt: float,
        linear_coef: torch.Tensor,
        nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    ):
        super().__init__(dt, linear_coef)
        self._nonlinear_func = nonlinear_func
        self._coef_1 = torch.where(
            linear_coef == 0, self.dt, (self._exp_term - 1) / linear_coef
        )

    def step(
        self,
        u_hat,
    ):
        return self._exp_term * u_hat + self._coef_1 * self._nonlinear_func(u_hat)
dt instance-attribute ¤
dt = dt
__init__ ¤
__init__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
)
Source code in torchfsm/integrator/_etdrk.py 
35
36
37
38
39
40
41
42
43
44
45
def __init__(
    self,
    dt: float,
    linear_coef: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
):
    super().__init__(dt, linear_coef)
    self._nonlinear_func = nonlinear_func
    self._coef_1 = torch.where(
        linear_coef == 0, self.dt, (self._exp_term - 1) / linear_coef
    )
step ¤
step(u_hat)
Source code in torchfsm/integrator/_etdrk.py 
47
48
49
50
51
def step(
    self,
    u_hat,
):
    return self._exp_term * u_hat + self._coef_1 * self._nonlinear_func(u_hat)

torchfsm.integrator._etdrk.ETDRK2 ¤

Bases: ETDRK1

Second-order ETDRK method.

Source code in torchfsm/integrator/_etdrk.py 
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
class ETDRK2(ETDRK1):
    """
    Second-order ETDRK method.
    """

    def __init__(
        self,
        dt: float,
        linear_coef: torch.Tensor,
        nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    ):
        super().__init__(dt, linear_coef, nonlinear_func)
        self._coef_2 = torch.where(
            linear_coef == 0,
            self.dt / 2,
            (self._exp_term - 1 - linear_coef * self.dt) / (linear_coef**2 * self.dt),
        )

    def step(
        self,
        u_hat,
    ):
        u_nonlin_hat = self._nonlinear_func(u_hat)
        u_stage_1_hat = self._exp_term * u_hat + self._coef_1 * u_nonlin_hat
        u_stage_1_nonlin_hat = self._nonlinear_func(u_stage_1_hat)
        u_next_hat = u_stage_1_hat + self._coef_2 * (
            u_stage_1_nonlin_hat - u_nonlin_hat
        )
        return u_next_hat
dt instance-attribute ¤
dt = dt
__init__ ¤
__init__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
)
Source code in torchfsm/integrator/_etdrk.py 
59
60
61
62
63
64
65
66
67
68
69
70
def __init__(
    self,
    dt: float,
    linear_coef: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
):
    super().__init__(dt, linear_coef, nonlinear_func)
    self._coef_2 = torch.where(
        linear_coef == 0,
        self.dt / 2,
        (self._exp_term - 1 - linear_coef * self.dt) / (linear_coef**2 * self.dt),
    )
step ¤
step(u_hat)
Source code in torchfsm/integrator/_etdrk.py 
72
73
74
75
76
77
78
79
80
81
82
def step(
    self,
    u_hat,
):
    u_nonlin_hat = self._nonlinear_func(u_hat)
    u_stage_1_hat = self._exp_term * u_hat + self._coef_1 * u_nonlin_hat
    u_stage_1_nonlin_hat = self._nonlinear_func(u_stage_1_hat)
    u_next_hat = u_stage_1_hat + self._coef_2 * (
        u_stage_1_nonlin_hat - u_nonlin_hat
    )
    return u_next_hat