Skip to content

SETDRK Integrator

The implementation of SETDRK method is modified from exponax

StableETDRK¤

torchfsm.integrator._stable_etdrk.SETDRKIntegrator ¤

Bases: Enum

Stable ETDRK Integrator Provides a unified interface for all ETDRK methods.

Source code in torchfsm/integrator/_stable_etdrk/__init__.py 
 96
 97
 98
 99
100
101
102
103
104
105
class SETDRKIntegrator(Enum):
    """
    Stable ETDRK Integrator
    Provides a unified interface for all ETDRK methods.
    """

    SETDRK1 = SETDRK1Wrapper()
    SETDRK2 = SETDRK2Wrapper()
    SETDRK3 = SETDRK3Wrapper()
    SETDRK4 = SETDRK4Wrapper()
SETDRK1 class-attribute instance-attribute ¤
SETDRK1 = SETDRK1Wrapper()
SETDRK2 class-attribute instance-attribute ¤
SETDRK2 = SETDRK2Wrapper()
SETDRK3 class-attribute instance-attribute ¤
SETDRK3 = SETDRK3Wrapper()
SETDRK4 class-attribute instance-attribute ¤
SETDRK4 = SETDRK4Wrapper()

Wrappers¤

torchfsm.integrator._stable_etdrk.SETDRK1Wrapper ¤

Source code in torchfsm/integrator/_stable_etdrk/__init__.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
class SETDRK1Wrapper:

    def __call__(self,
                 dt: float,
                linear_coef: torch.Tensor,
                nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
                n_integration_points: int = 16,
                integration_radius: float = 1.0,
                cpu_cached: bool = False,):
        if cpu_cached:
            return CachedSETDRK1(
                dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
            )
        else:
            return UnCachedSETDRK1(
                dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
            )
__call__ ¤
__call__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
    cpu_cached: bool = False,
)
Source code in torchfsm/integrator/_stable_etdrk/__init__.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def __call__(self,
             dt: float,
            linear_coef: torch.Tensor,
            nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
            n_integration_points: int = 16,
            integration_radius: float = 1.0,
            cpu_cached: bool = False,):
    if cpu_cached:
        return CachedSETDRK1(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )
    else:
        return UnCachedSETDRK1(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )

torchfsm.integrator._stable_etdrk.SETDRK2Wrapper ¤

Source code in torchfsm/integrator/_stable_etdrk/__init__.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
class SETDRK2Wrapper:

    def __call__(self,
                 dt: float,
                linear_coef: torch.Tensor,
                nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
                n_integration_points: int = 16,
                integration_radius: float = 1.0,
                cpu_cached: bool = False,):
        if cpu_cached:
            return CachedSETDRK2(
                dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
            )
        else:
            return UnCachedSETDRK2(
                dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
            )
__call__ ¤
__call__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
    cpu_cached: bool = False,
)
Source code in torchfsm/integrator/_stable_etdrk/__init__.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def __call__(self,
             dt: float,
            linear_coef: torch.Tensor,
            nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
            n_integration_points: int = 16,
            integration_radius: float = 1.0,
            cpu_cached: bool = False,):
    if cpu_cached:
        return CachedSETDRK2(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )
    else:
        return UnCachedSETDRK2(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )

torchfsm.integrator._stable_etdrk.SETDRK3Wrapper ¤

Source code in torchfsm/integrator/_stable_etdrk/__init__.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
class SETDRK3Wrapper:

    def __call__(self,
                 dt: float,
                linear_coef: torch.Tensor,
                nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
                n_integration_points: int = 16,
                integration_radius: float = 1.0,
                cpu_cached: bool = False,):
        if cpu_cached:
            return CachedSETDRK3(
                dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
            )
        else:
            return UnCachedSETDRK3(
                dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
            )
__call__ ¤
__call__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
    cpu_cached: bool = False,
)
Source code in torchfsm/integrator/_stable_etdrk/__init__.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def __call__(self,
             dt: float,
            linear_coef: torch.Tensor,
            nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
            n_integration_points: int = 16,
            integration_radius: float = 1.0,
            cpu_cached: bool = False,):
    if cpu_cached:
        return CachedSETDRK3(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )
    else:
        return UnCachedSETDRK3(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )

torchfsm.integrator._stable_etdrk.SETDRK4Wrapper ¤

Source code in torchfsm/integrator/_stable_etdrk/__init__.py 
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
class SETDRK4Wrapper:

    def __call__(self,
                 dt: float,
                linear_coef: torch.Tensor,
                nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
                n_integration_points: int = 16,
                integration_radius: float = 1.0,
                cpu_cached: bool = False,):
        if cpu_cached:
            return CachedSETDRK4(
                dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
            )
        else:
            return UnCachedSETDRK4(
                dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
            )
__call__ ¤
__call__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
    cpu_cached: bool = False,
)
Source code in torchfsm/integrator/_stable_etdrk/__init__.py 
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def __call__(self,
             dt: float,
            linear_coef: torch.Tensor,
            nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
            n_integration_points: int = 16,
            integration_radius: float = 1.0,
            cpu_cached: bool = False,):
    if cpu_cached:
        return CachedSETDRK4(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )
    else:
        return UnCachedSETDRK4(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )

Integration steps¤

torchfsm.integrator._stable_etdrk._setdrk_step.setdrk1_step ¤

setdrk1_step(
    u_hat: Tensor,
    exp_term: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    coef_1: Tensor,
)
Source code in torchfsm/integrator/_stable_etdrk/_setdrk_step.py 
 5
 6
 7
 8
 9
10
11
def setdrk1_step(
    u_hat: torch.Tensor,
    exp_term: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    coef_1: torch.Tensor,
):
    return exp_term * u_hat + coef_1 * nonlinear_func(u_hat)

torchfsm.integrator._stable_etdrk._setdrk_step.setdrk2_step ¤

setdrk2_step(
    u_hat: Tensor,
    exp_term: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    coef_1: Tensor,
    coef_2: Tensor,
)
Source code in torchfsm/integrator/_stable_etdrk/_setdrk_step.py 
14
15
16
17
18
19
20
21
22
23
24
25
def setdrk2_step(
    u_hat: torch.Tensor,
    exp_term: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    coef_1: torch.Tensor,
    coef_2: torch.Tensor,
):
    u_nonlin_hat = nonlinear_func(u_hat)
    u_stage_1_hat = exp_term * u_hat + coef_1 * u_nonlin_hat
    u_stage_1_nonlin_hat = nonlinear_func(u_stage_1_hat)
    u_next_hat = u_stage_1_hat + coef_2 * (u_stage_1_nonlin_hat - u_nonlin_hat)
    return u_next_hat

torchfsm.integrator._stable_etdrk._setdrk_step.setdrk3_step ¤

setdrk3_step(
    u_hat: Tensor,
    exp_term: Tensor,
    half_exp_term: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    coef_1: Tensor,
    coef_2: Tensor,
    coef_3: Tensor,
    coef_4: Tensor,
    coef_5: Tensor,
)
Source code in torchfsm/integrator/_stable_etdrk/_setdrk_step.py 
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
def setdrk3_step(
    u_hat: torch.Tensor,
    exp_term: torch.Tensor,
    half_exp_term: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    coef_1: torch.Tensor,
    coef_2: torch.Tensor,
    coef_3: torch.Tensor,
    coef_4: torch.Tensor,
    coef_5: torch.Tensor,
):
    u_nonlin_hat = nonlinear_func(u_hat)
    u_stage_1_hat = half_exp_term * u_hat + coef_1 * u_nonlin_hat
    u_stage_1_nonlin_hat = nonlinear_func(u_stage_1_hat)
    u_stage_2_hat = exp_term * u_hat + coef_2 * (
        2 * u_stage_1_nonlin_hat - u_nonlin_hat
    )
    u_stage_2_nonlin_hat = nonlinear_func(u_stage_2_hat)
    u_next_hat = (
        exp_term * u_hat
        + coef_3 * u_nonlin_hat
        + coef_4 * u_stage_1_nonlin_hat
        + coef_5 * u_stage_2_nonlin_hat
    )
    return u_next_hat

torchfsm.integrator._stable_etdrk._setdrk_step.setdrk4_step ¤

setdrk4_step(
    u_hat: Tensor,
    exp_term: Tensor,
    half_exp_term: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    coef_1: Tensor,
    coef_2: Tensor,
    coef_3: Tensor,
    coef_4: Tensor,
    coef_5: Tensor,
    coef_6: Tensor,
)
Source code in torchfsm/integrator/_stable_etdrk/_setdrk_step.py 
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
def setdrk4_step(
    u_hat: torch.Tensor,
    exp_term: torch.Tensor,
    half_exp_term: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    coef_1: torch.Tensor,
    coef_2: torch.Tensor,
    coef_3: torch.Tensor,
    coef_4: torch.Tensor,
    coef_5: torch.Tensor,
    coef_6: torch.Tensor,
):
    u_nonlin_hat = nonlinear_func(u_hat)
    u_stage_1_hat = half_exp_term * u_hat + coef_1 * u_nonlin_hat
    u_stage_1_nonlin_hat = nonlinear_func(u_stage_1_hat)
    u_stage_2_hat = half_exp_term * u_hat + coef_2 * u_stage_1_nonlin_hat
    u_stage_2_nonlin_hat = nonlinear_func(u_stage_2_hat)
    u_stage_3_hat = half_exp_term * u_stage_1_hat + coef_3 * (
        2 * u_stage_2_nonlin_hat - u_nonlin_hat
    )
    u_stage_3_nonlin_hat = nonlinear_func(u_stage_3_hat)
    u_next_hat = (
        exp_term * u_hat
        + coef_4 * u_nonlin_hat
        + coef_5 * 2 * (u_stage_1_nonlin_hat + u_stage_2_nonlin_hat)
        + coef_6 * u_stage_3_nonlin_hat
    )
    return u_next_hat

Cached SETDRK¤

torchfsm.integrator._stable_etdrk._cached._CachedSETDRKBase ¤

Bases: ABC

Source code in torchfsm/integrator/_stable_etdrk/_cached.py 
 9
10
11
12
13
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
class _CachedSETDRKBase(ABC):

    def __init__(
        self,
        dt: float,
        linear_coef: torch.Tensor,
        nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
        n_integration_points: int = 16,
        integration_radius: float = 1.0,
    ):
        self.dt = dt
        self._nonlinear_func = nonlinear_func
        self._exp_term = torch.exp(self.dt * linear_coef)
        self.n_integration_points = n_integration_points
        self.integration_radius = integration_radius

    def _get_lr(self, linear_coef):
        return (
            self.integration_radius
            * roots_of_unity(
                self.n_integration_points, device="cpu", dtype=linear_coef.real.dtype
            )
            + linear_coef.unsqueeze(-1).cpu() * self.dt
        )

    @abstractmethod
    def step(
        self,
        u_hat,
    ):
        """
        Advance the state in Fourier space.
        """
        raise NotImplementedError("This method should be implemented by subclasses.")
dt instance-attribute ¤
dt = dt
n_integration_points instance-attribute ¤
n_integration_points = n_integration_points
integration_radius instance-attribute ¤
integration_radius = integration_radius
__init__ ¤
__init__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
)
Source code in torchfsm/integrator/_stable_etdrk/_cached.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
def __init__(
    self,
    dt: float,
    linear_coef: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
):
    self.dt = dt
    self._nonlinear_func = nonlinear_func
    self._exp_term = torch.exp(self.dt * linear_coef)
    self.n_integration_points = n_integration_points
    self.integration_radius = integration_radius
step abstractmethod ¤
step(u_hat)

Advance the state in Fourier space.

Source code in torchfsm/integrator/_stable_etdrk/_cached.py 
34
35
36
37
38
39
40
41
42
@abstractmethod
def step(
    self,
    u_hat,
):
    """
    Advance the state in Fourier space.
    """
    raise NotImplementedError("This method should be implemented by subclasses.")

torchfsm.integrator._stable_etdrk._cached.CachedSETDRK1 ¤

Bases: _CachedSETDRKBase

First-order CachedSETDRK method.

Source code in torchfsm/integrator/_stable_etdrk/_cached.py 
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
class CachedSETDRK1(_CachedSETDRKBase):
    """
    First-order CachedSETDRK method.
    """

    def __init__(
        self,
        dt: float,
        linear_coef: torch.Tensor,
        nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
        n_integration_points: int = 16,
        integration_radius: float = 1.0,
    ):
        super().__init__(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )
        lr = self._get_lr(linear_coef)
        self._coef_1 = dt * torch.mean((torch.exp(lr) - 1) / lr, axis=-1).real.to(
            linear_coef.device
        )

    def step(
        self,
        u_hat,
    ):
        return setdrk1_step(
            u_hat,
            self._exp_term,
            self._nonlinear_func,
            self._coef_1,
        )
dt instance-attribute ¤
dt = dt
n_integration_points instance-attribute ¤
n_integration_points = n_integration_points
integration_radius instance-attribute ¤
integration_radius = integration_radius
__init__ ¤
__init__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
)
Source code in torchfsm/integrator/_stable_etdrk/_cached.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def __init__(
    self,
    dt: float,
    linear_coef: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
):
    super().__init__(
        dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
    )
    lr = self._get_lr(linear_coef)
    self._coef_1 = dt * torch.mean((torch.exp(lr) - 1) / lr, axis=-1).real.to(
        linear_coef.device
    )
step ¤
step(u_hat)
Source code in torchfsm/integrator/_stable_etdrk/_cached.py 
66
67
68
69
70
71
72
73
74
75
def step(
    self,
    u_hat,
):
    return setdrk1_step(
        u_hat,
        self._exp_term,
        self._nonlinear_func,
        self._coef_1,
    )

torchfsm.integrator._stable_etdrk._cached.CachedSETDRK2 ¤

Bases: _CachedSETDRKBase

Second-order CachedSETDRK method.

Source code in torchfsm/integrator/_stable_etdrk/_cached.py 
 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
class CachedSETDRK2(_CachedSETDRKBase):
    """
    Second-order CachedSETDRK method.
    """

    def __init__(
        self,
        dt: float,
        linear_coef: torch.Tensor,
        nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
        n_integration_points: int = 16,
        integration_radius: float = 1.0,
    ):
        super().__init__(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )
        lr = self._get_lr(linear_coef)
        self._coef_1 = dt * torch.mean((torch.exp(lr) - 1) / lr, axis=-1).real.to(
            linear_coef.device
        )
        self._coef_2 = dt * torch.mean(
            (torch.exp(lr) - 1 - lr) / lr**2, axis=-1
        ).real.to(linear_coef.device)

    def step(
        self,
        u_hat,
    ):
        return setdrk2_step(
            u_hat,
            self._exp_term,
            self._nonlinear_func,
            self._coef_1,
            self._coef_2,
        )
dt instance-attribute ¤
dt = dt
n_integration_points instance-attribute ¤
n_integration_points = n_integration_points
integration_radius instance-attribute ¤
integration_radius = integration_radius
__init__ ¤
__init__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
)
Source code in torchfsm/integrator/_stable_etdrk/_cached.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def __init__(
    self,
    dt: float,
    linear_coef: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
):
    super().__init__(
        dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
    )
    lr = self._get_lr(linear_coef)
    self._coef_1 = dt * torch.mean((torch.exp(lr) - 1) / lr, axis=-1).real.to(
        linear_coef.device
    )
    self._coef_2 = dt * torch.mean(
        (torch.exp(lr) - 1 - lr) / lr**2, axis=-1
    ).real.to(linear_coef.device)
step ¤
step(u_hat)
Source code in torchfsm/integrator/_stable_etdrk/_cached.py 
102
103
104
105
106
107
108
109
110
111
112
def step(
    self,
    u_hat,
):
    return setdrk2_step(
        u_hat,
        self._exp_term,
        self._nonlinear_func,
        self._coef_1,
        self._coef_2,
    )

torchfsm.integrator._stable_etdrk._cached.CachedSETDRK3 ¤

Bases: _CachedSETDRKBase

Third-order CachedSETDRK method.

Source code in torchfsm/integrator/_stable_etdrk/_cached.py 
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
class CachedSETDRK3(_CachedSETDRKBase):
    """
    Third-order CachedSETDRK method.
    """

    def __init__(
        self,
        dt: float,
        linear_coef: torch.Tensor,
        nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
        n_integration_points: int = 16,
        integration_radius: float = 1.0,
    ):
        super().__init__(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )
        lr = self._get_lr(linear_coef)
        self._half_exp_term = torch.exp(0.5 * dt * linear_coef)
        self._coef_1 = dt * torch.mean((torch.exp(lr / 2) - 1) / lr, axis=-1).real.to(
            linear_coef.device
        )
        self._coef_2 = dt * torch.mean((torch.exp(lr) - 1) / lr, axis=-1).real.to(
            linear_coef.device
        )
        self._coef_3 = dt * torch.mean(
            (-4 - lr + torch.exp(lr) * (4 - 3 * lr + lr**2)) / (lr**3), axis=-1
        ).real.to(linear_coef.device)
        self._coef_4 = dt * torch.mean(
            (4.0 * (2.0 + lr + torch.exp(lr) * (-2 + lr))) / (lr**3), axis=-1
        ).real.to(linear_coef.device)
        self._coef_5 = dt * torch.mean(
            (-4 - 3 * lr - lr**2 + torch.exp(lr) * (4 - lr)) / (lr**3), axis=-1
        ).real.to(linear_coef.device)

    def step(
        self,
        u_hat,
    ):
        return setdrk3_step(
            u_hat,
            self._exp_term,
            self._half_exp_term,
            self._nonlinear_func,
            self._coef_1,
            self._coef_2,
            self._coef_3,
            self._coef_4,
            self._coef_5,
        )
dt instance-attribute ¤
dt = dt
n_integration_points instance-attribute ¤
n_integration_points = n_integration_points
integration_radius instance-attribute ¤
integration_radius = integration_radius
__init__ ¤
__init__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
)
Source code in torchfsm/integrator/_stable_etdrk/_cached.py 
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
def __init__(
    self,
    dt: float,
    linear_coef: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
):
    super().__init__(
        dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
    )
    lr = self._get_lr(linear_coef)
    self._half_exp_term = torch.exp(0.5 * dt * linear_coef)
    self._coef_1 = dt * torch.mean((torch.exp(lr / 2) - 1) / lr, axis=-1).real.to(
        linear_coef.device
    )
    self._coef_2 = dt * torch.mean((torch.exp(lr) - 1) / lr, axis=-1).real.to(
        linear_coef.device
    )
    self._coef_3 = dt * torch.mean(
        (-4 - lr + torch.exp(lr) * (4 - 3 * lr + lr**2)) / (lr**3), axis=-1
    ).real.to(linear_coef.device)
    self._coef_4 = dt * torch.mean(
        (4.0 * (2.0 + lr + torch.exp(lr) * (-2 + lr))) / (lr**3), axis=-1
    ).real.to(linear_coef.device)
    self._coef_5 = dt * torch.mean(
        (-4 - 3 * lr - lr**2 + torch.exp(lr) * (4 - lr)) / (lr**3), axis=-1
    ).real.to(linear_coef.device)
step ¤
step(u_hat)
Source code in torchfsm/integrator/_stable_etdrk/_cached.py 
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
def step(
    self,
    u_hat,
):
    return setdrk3_step(
        u_hat,
        self._exp_term,
        self._half_exp_term,
        self._nonlinear_func,
        self._coef_1,
        self._coef_2,
        self._coef_3,
        self._coef_4,
        self._coef_5,
    )

torchfsm.integrator._stable_etdrk._cached.CachedSETDRK4 ¤

Bases: _CachedSETDRKBase

Fourth-order CachedSETDRK method.

Source code in torchfsm/integrator/_stable_etdrk/_cached.py 
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
class CachedSETDRK4(_CachedSETDRKBase):
    """
    Fourth-order CachedSETDRK method.
    """

    def __init__(
        self,
        dt: float,
        linear_coef: torch.Tensor,
        nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
        n_integration_points: int = 16,
        integration_radius: float = 1.0,
    ):
        super().__init__(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )
        lr = self._get_lr(linear_coef)
        self._half_exp_term = torch.exp(0.5 * dt * linear_coef)
        self._coef_1 = dt * torch.mean((torch.exp(lr / 2) - 1) / lr, axis=-1).real.to(
            linear_coef.device
        )
        self._coef_2 = self._coef_1
        self._coef_3 = self._coef_1
        self._coef_4 = dt * torch.mean(
            (-4 - lr + torch.exp(lr) * (4 - 3 * lr + lr**2)) / (lr**3), axis=-1
        ).real.to(linear_coef.device)
        self._coef_5 = dt * torch.mean(
            (2 + lr + torch.exp(lr) * (-2 + lr)) / (lr**3), axis=-1
        ).real.to(linear_coef.device)
        self._coef_6 = dt * torch.mean(
            (-4 - 3 * lr - lr**2 + torch.exp(lr) * (4 - lr)) / (lr**3), axis=-1
        ).real.to(linear_coef.device)

    def step(
        self,
        u_hat,
    ):
        return setdrk4_step(
            u_hat,
            self._exp_term,
            self._half_exp_term,
            self._nonlinear_func,
            self._coef_1,
            self._coef_2,
            self._coef_3,
            self._coef_4,
            self._coef_5,
            self._coef_6,
        )
dt instance-attribute ¤
dt = dt
n_integration_points instance-attribute ¤
n_integration_points = n_integration_points
integration_radius instance-attribute ¤
integration_radius = integration_radius
__init__ ¤
__init__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
)
Source code in torchfsm/integrator/_stable_etdrk/_cached.py 
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
def __init__(
    self,
    dt: float,
    linear_coef: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
):
    super().__init__(
        dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
    )
    lr = self._get_lr(linear_coef)
    self._half_exp_term = torch.exp(0.5 * dt * linear_coef)
    self._coef_1 = dt * torch.mean((torch.exp(lr / 2) - 1) / lr, axis=-1).real.to(
        linear_coef.device
    )
    self._coef_2 = self._coef_1
    self._coef_3 = self._coef_1
    self._coef_4 = dt * torch.mean(
        (-4 - lr + torch.exp(lr) * (4 - 3 * lr + lr**2)) / (lr**3), axis=-1
    ).real.to(linear_coef.device)
    self._coef_5 = dt * torch.mean(
        (2 + lr + torch.exp(lr) * (-2 + lr)) / (lr**3), axis=-1
    ).real.to(linear_coef.device)
    self._coef_6 = dt * torch.mean(
        (-4 - 3 * lr - lr**2 + torch.exp(lr) * (4 - lr)) / (lr**3), axis=-1
    ).real.to(linear_coef.device)
step ¤
step(u_hat)
Source code in torchfsm/integrator/_stable_etdrk/_cached.py 
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
def step(
    self,
    u_hat,
):
    return setdrk4_step(
        u_hat,
        self._exp_term,
        self._half_exp_term,
        self._nonlinear_func,
        self._coef_1,
        self._coef_2,
        self._coef_3,
        self._coef_4,
        self._coef_5,
        self._coef_6,
    )

Uncached SETDRK¤

torchfsm.integrator._stable_etdrk._uncached._UnCachedSETDRKBase ¤

Bases: ABC

Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
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
class _UnCachedSETDRKBase(ABC):

    def __init__(
        self,
        dt: float,
        linear_coef: torch.Tensor,
        nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
        n_integration_points: int = 16,
        integration_radius: float = 1.0,
    ):
        self.dt = dt
        self._nonlinear_func = nonlinear_func
        self._exp_term = torch.exp(self.dt * linear_coef)
        self.n_integration_points = n_integration_points
        self.integration_radius = integration_radius

    def _get_lr(self, linear_coef):
        return (
            self.integration_radius
            * roots_of_unity(
                self.n_integration_points,
                device=linear_coef.device,
                dtype=linear_coef.real.dtype,
            )
            + linear_coef.unsqueeze(-1) * self.dt
        )

    @abstractmethod
    def step(
        self,
        u_hat,
    ):
        """
        Advance the state in Fourier space.
        """
dt instance-attribute ¤
dt = dt
n_integration_points instance-attribute ¤
n_integration_points = n_integration_points
integration_radius instance-attribute ¤
integration_radius = integration_radius
__init__ ¤
__init__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
)
Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
def __init__(
    self,
    dt: float,
    linear_coef: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
):
    self.dt = dt
    self._nonlinear_func = nonlinear_func
    self._exp_term = torch.exp(self.dt * linear_coef)
    self.n_integration_points = n_integration_points
    self.integration_radius = integration_radius
step abstractmethod ¤
step(u_hat)

Advance the state in Fourier space.

Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
45
46
47
48
49
50
51
52
@abstractmethod
def step(
    self,
    u_hat,
):
    """
    Advance the state in Fourier space.
    """

torchfsm.integrator._stable_etdrk._uncached.UnCachedSETDRK1 ¤

Bases: _UnCachedSETDRKBase

First-order UnCachedSETDRK method.

Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
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
class UnCachedSETDRK1(_UnCachedSETDRKBase):
    """
    First-order UnCachedSETDRK method.
    """

    def __init__(
        self,
        dt: float,
        linear_coef: torch.Tensor,
        nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
        n_integration_points: int = 16,
        integration_radius: float = 1.0,
    ):
        super().__init__(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )
        lr = self._get_lr(linear_coef)
        self._coef_1 = dt * torch.mean((torch.exp(lr) - 1) / lr, axis=-1).real

    def step(
        self,
        u_hat,
    ):
        return setdrk1_step(
            u_hat,
            self._exp_term,
            self._nonlinear_func,
            self._coef_1,
        )
dt instance-attribute ¤
dt = dt
n_integration_points instance-attribute ¤
n_integration_points = n_integration_points
integration_radius instance-attribute ¤
integration_radius = integration_radius
__init__ ¤
__init__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
)
Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
def __init__(
    self,
    dt: float,
    linear_coef: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
):
    super().__init__(
        dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
    )
    lr = self._get_lr(linear_coef)
    self._coef_1 = dt * torch.mean((torch.exp(lr) - 1) / lr, axis=-1).real
step ¤
step(u_hat)
Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
74
75
76
77
78
79
80
81
82
83
def step(
    self,
    u_hat,
):
    return setdrk1_step(
        u_hat,
        self._exp_term,
        self._nonlinear_func,
        self._coef_1,
    )

torchfsm.integrator._stable_etdrk._uncached.UnCachedSETDRK2 ¤

Bases: _UnCachedSETDRKBase

Second-order UnCachedSETDRK method.

Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
 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
class UnCachedSETDRK2(_UnCachedSETDRKBase):
    """
    Second-order UnCachedSETDRK method.
    """

    def __init__(
        self,
        dt: float,
        linear_coef: torch.Tensor,
        nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
        n_integration_points: int = 16,
        integration_radius: float = 1.0,
    ):
        super().__init__(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )
        lr = self._get_lr(linear_coef)
        self._coef_1 = dt * torch.mean((torch.exp(lr) - 1) / lr, axis=-1).real
        self._coef_2 = dt * torch.mean((torch.exp(lr) - 1 - lr) / lr**2, axis=-1).real

    def step(
        self,
        u_hat,
    ):
        return setdrk2_step(
            u_hat,
            self._exp_term,
            self._nonlinear_func,
            self._coef_1,
            self._coef_2,
        )
dt instance-attribute ¤
dt = dt
n_integration_points instance-attribute ¤
n_integration_points = n_integration_points
integration_radius instance-attribute ¤
integration_radius = integration_radius
__init__ ¤
__init__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
)
Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def __init__(
    self,
    dt: float,
    linear_coef: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
):
    super().__init__(
        dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
    )
    lr = self._get_lr(linear_coef)
    self._coef_1 = dt * torch.mean((torch.exp(lr) - 1) / lr, axis=-1).real
    self._coef_2 = dt * torch.mean((torch.exp(lr) - 1 - lr) / lr**2, axis=-1).real
step ¤
step(u_hat)
Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
106
107
108
109
110
111
112
113
114
115
116
def step(
    self,
    u_hat,
):
    return setdrk2_step(
        u_hat,
        self._exp_term,
        self._nonlinear_func,
        self._coef_1,
        self._coef_2,
    )

torchfsm.integrator._stable_etdrk._uncached.UnCachedSETDRK3 ¤

Bases: _UnCachedSETDRKBase

Third-order UnCachedSETDRK method.

Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
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
class UnCachedSETDRK3(_UnCachedSETDRKBase):
    """
    Third-order UnCachedSETDRK method.
    """

    def __init__(
        self,
        dt: float,
        linear_coef: torch.Tensor,
        nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
        n_integration_points: int = 16,
        integration_radius: float = 1.0,
    ):
        super().__init__(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )
        lr = self._get_lr(linear_coef)
        self._half_exp_term = torch.exp(0.5 * dt * linear_coef)
        self._coef_1 = dt * torch.mean((torch.exp(lr / 2) - 1) / lr, axis=-1).real
        self._coef_2 = dt * torch.mean((torch.exp(lr) - 1) / lr, axis=-1).real
        self._coef_3 = (
            dt
            * torch.mean(
                (-4 - lr + torch.exp(lr) * (4 - 3 * lr + lr**2)) / (lr**3), axis=-1
            ).real
        )
        self._coef_4 = (
            dt
            * torch.mean(
                (4.0 * (2.0 + lr + torch.exp(lr) * (-2 + lr))) / (lr**3), axis=-1
            ).real
        )
        self._coef_5 = (
            dt
            * torch.mean(
                (-4 - 3 * lr - lr**2 + torch.exp(lr) * (4 - lr)) / (lr**3), axis=-1
            ).real
        )

    def step(
        self,
        u_hat,
    ):
        return setdrk3_step(
            u_hat,
            self._exp_term,
            self._half_exp_term,
            self._nonlinear_func,
            self._coef_1,
            self._coef_2,
            self._coef_3,
            self._coef_4,
            self._coef_5,
        )
dt instance-attribute ¤
dt = dt
n_integration_points instance-attribute ¤
n_integration_points = n_integration_points
integration_radius instance-attribute ¤
integration_radius = integration_radius
__init__ ¤
__init__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
)
Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
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
def __init__(
    self,
    dt: float,
    linear_coef: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
):
    super().__init__(
        dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
    )
    lr = self._get_lr(linear_coef)
    self._half_exp_term = torch.exp(0.5 * dt * linear_coef)
    self._coef_1 = dt * torch.mean((torch.exp(lr / 2) - 1) / lr, axis=-1).real
    self._coef_2 = dt * torch.mean((torch.exp(lr) - 1) / lr, axis=-1).real
    self._coef_3 = (
        dt
        * torch.mean(
            (-4 - lr + torch.exp(lr) * (4 - 3 * lr + lr**2)) / (lr**3), axis=-1
        ).real
    )
    self._coef_4 = (
        dt
        * torch.mean(
            (4.0 * (2.0 + lr + torch.exp(lr) * (-2 + lr))) / (lr**3), axis=-1
        ).real
    )
    self._coef_5 = (
        dt
        * torch.mean(
            (-4 - 3 * lr - lr**2 + torch.exp(lr) * (4 - lr)) / (lr**3), axis=-1
        ).real
    )
step ¤
step(u_hat)
Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
def step(
    self,
    u_hat,
):
    return setdrk3_step(
        u_hat,
        self._exp_term,
        self._half_exp_term,
        self._nonlinear_func,
        self._coef_1,
        self._coef_2,
        self._coef_3,
        self._coef_4,
        self._coef_5,
    )

torchfsm.integrator._stable_etdrk._uncached.UnCachedSETDRK4 ¤

Bases: _UnCachedSETDRKBase

Fourth-order UnCachedSETDRK method.

Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
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
class UnCachedSETDRK4(_UnCachedSETDRKBase):
    """
    Fourth-order UnCachedSETDRK method.
    """

    def __init__(
        self,
        dt: float,
        linear_coef: torch.Tensor,
        nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
        n_integration_points: int = 16,
        integration_radius: float = 1.0,
    ):
        super().__init__(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )
        lr = self._get_lr(linear_coef)
        self._half_exp_term = torch.exp(0.5 * dt * linear_coef)
        self._coef_1 = dt * torch.mean((torch.exp(lr / 2) - 1) / lr, axis=-1).real
        self._coef_2 = self._coef_1
        self._coef_3 = self._coef_1
        self._coef_4 = (
            dt
            * torch.mean(
                (-4 - lr + torch.exp(lr) * (4 - 3 * lr + lr**2)) / (lr**3), axis=-1
            ).real
        )
        self._coef_5 = (
            dt
            * torch.mean((2 + lr + torch.exp(lr) * (-2 + lr)) / (lr**3), axis=-1).real
        )
        self._coef_6 = (
            dt
            * torch.mean(
                (-4 - 3 * lr - lr**2 + torch.exp(lr) * (4 - lr)) / (lr**3), axis=-1
            ).real
        )

    def step(
        self,
        u_hat,
    ):
        return setdrk4_step(
            u_hat,
            self._exp_term,
            self._half_exp_term,
            self._nonlinear_func,
            self._coef_1,
            self._coef_2,
            self._coef_3,
            self._coef_4,
            self._coef_5,
            self._coef_6,
        )
dt instance-attribute ¤
dt = dt
n_integration_points instance-attribute ¤
n_integration_points = n_integration_points
integration_radius instance-attribute ¤
integration_radius = integration_radius
__init__ ¤
__init__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
)
Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
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
def __init__(
    self,
    dt: float,
    linear_coef: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
):
    super().__init__(
        dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
    )
    lr = self._get_lr(linear_coef)
    self._half_exp_term = torch.exp(0.5 * dt * linear_coef)
    self._coef_1 = dt * torch.mean((torch.exp(lr / 2) - 1) / lr, axis=-1).real
    self._coef_2 = self._coef_1
    self._coef_3 = self._coef_1
    self._coef_4 = (
        dt
        * torch.mean(
            (-4 - lr + torch.exp(lr) * (4 - 3 * lr + lr**2)) / (lr**3), axis=-1
        ).real
    )
    self._coef_5 = (
        dt
        * torch.mean((2 + lr + torch.exp(lr) * (-2 + lr)) / (lr**3), axis=-1).real
    )
    self._coef_6 = (
        dt
        * torch.mean(
            (-4 - 3 * lr - lr**2 + torch.exp(lr) * (4 - lr)) / (lr**3), axis=-1
        ).real
    )
step ¤
step(u_hat)
Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
def step(
    self,
    u_hat,
):
    return setdrk4_step(
        u_hat,
        self._exp_term,
        self._half_exp_term,
        self._nonlinear_func,
        self._coef_1,
        self._coef_2,
        self._coef_3,
        self._coef_4,
        self._coef_5,
        self._coef_6,
    )

torchfsm.integrator._stable_etdrk._uncached.UnCachedSETDRK4 ¤

Bases: _UnCachedSETDRKBase

Fourth-order UnCachedSETDRK method.

Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
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
class UnCachedSETDRK4(_UnCachedSETDRKBase):
    """
    Fourth-order UnCachedSETDRK method.
    """

    def __init__(
        self,
        dt: float,
        linear_coef: torch.Tensor,
        nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
        n_integration_points: int = 16,
        integration_radius: float = 1.0,
    ):
        super().__init__(
            dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
        )
        lr = self._get_lr(linear_coef)
        self._half_exp_term = torch.exp(0.5 * dt * linear_coef)
        self._coef_1 = dt * torch.mean((torch.exp(lr / 2) - 1) / lr, axis=-1).real
        self._coef_2 = self._coef_1
        self._coef_3 = self._coef_1
        self._coef_4 = (
            dt
            * torch.mean(
                (-4 - lr + torch.exp(lr) * (4 - 3 * lr + lr**2)) / (lr**3), axis=-1
            ).real
        )
        self._coef_5 = (
            dt
            * torch.mean((2 + lr + torch.exp(lr) * (-2 + lr)) / (lr**3), axis=-1).real
        )
        self._coef_6 = (
            dt
            * torch.mean(
                (-4 - 3 * lr - lr**2 + torch.exp(lr) * (4 - lr)) / (lr**3), axis=-1
            ).real
        )

    def step(
        self,
        u_hat,
    ):
        return setdrk4_step(
            u_hat,
            self._exp_term,
            self._half_exp_term,
            self._nonlinear_func,
            self._coef_1,
            self._coef_2,
            self._coef_3,
            self._coef_4,
            self._coef_5,
            self._coef_6,
        )
dt instance-attribute ¤
dt = dt
n_integration_points instance-attribute ¤
n_integration_points = n_integration_points
integration_radius instance-attribute ¤
integration_radius = integration_radius
__init__ ¤
__init__(
    dt: float,
    linear_coef: Tensor,
    nonlinear_func: Callable[[Tensor], Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
)
Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
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
def __init__(
    self,
    dt: float,
    linear_coef: torch.Tensor,
    nonlinear_func: Callable[[torch.Tensor], torch.Tensor],
    n_integration_points: int = 16,
    integration_radius: float = 1.0,
):
    super().__init__(
        dt, linear_coef, nonlinear_func, n_integration_points, integration_radius
    )
    lr = self._get_lr(linear_coef)
    self._half_exp_term = torch.exp(0.5 * dt * linear_coef)
    self._coef_1 = dt * torch.mean((torch.exp(lr / 2) - 1) / lr, axis=-1).real
    self._coef_2 = self._coef_1
    self._coef_3 = self._coef_1
    self._coef_4 = (
        dt
        * torch.mean(
            (-4 - lr + torch.exp(lr) * (4 - 3 * lr + lr**2)) / (lr**3), axis=-1
        ).real
    )
    self._coef_5 = (
        dt
        * torch.mean((2 + lr + torch.exp(lr) * (-2 + lr)) / (lr**3), axis=-1).real
    )
    self._coef_6 = (
        dt
        * torch.mean(
            (-4 - 3 * lr - lr**2 + torch.exp(lr) * (4 - lr)) / (lr**3), axis=-1
        ).real
    )
step ¤
step(u_hat)
Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
def step(
    self,
    u_hat,
):
    return setdrk4_step(
        u_hat,
        self._exp_term,
        self._half_exp_term,
        self._nonlinear_func,
        self._coef_1,
        self._coef_2,
        self._coef_3,
        self._coef_4,
        self._coef_5,
        self._coef_6,
    )

torchfsm.integrator._stable_etdrk._uncached.roots_of_unity ¤

roots_of_unity(
    M: int, device=None, dtype=None
) -> torch.Tensor

Return (complex-valued) array with M roots of unity.

Source code in torchfsm/integrator/_stable_etdrk/_uncached.py 
 8
 9
10
11
12
13
14
15
def roots_of_unity(M: int, device=None, dtype=None) -> torch.Tensor:
    """
    Return (complex-valued) array with M roots of unity.
    """
    # return torch.exp(1j * torch.pi * (torch.arange(1, M+1) - 0.5) / M)
    return torch.exp(
        2j * torch.pi * (torch.arange(1, M + 1, device=device, dtype=dtype) - 0.5) / M
    )