-
Notifications
You must be signed in to change notification settings - Fork 242
/
test_integrator.py
51 lines (40 loc) · 1.71 KB
/
test_integrator.py
1
2
3
4
5
6
7
8
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
43
44
45
46
47
48
49
50
51
"""Test for the integrators."""
import geomstats.backend as gs
import geomstats.integrator as integrator
import geomstats.tests
from geomstats.geometry.euclidean import Euclidean
from geomstats.geometry.matrices import Matrices
class TestIntegrator(geomstats.tests.TestCase):
def setUp(self):
self.dimension = 4
self.dt = 0.1
self.euclidean = Euclidean(self.dimension)
self.matrices = Matrices(self.dimension, self.dimension)
self.intercept = self.euclidean.random_uniform(1)
self.slope = Matrices.to_symmetric(self.matrices.random_uniform(1))
def function_linear(self, _, vector):
return - gs.dot(self.slope, vector)
@geomstats.tests.np_and_pytorch_only
def test_symplectic_euler_step(self):
state = (self.intercept, self.slope)
result = len(integrator._symplectic_euler_step(
state, self.function_linear, self.dt))
expected = len(state)
self.assertAllClose(result, expected)
@geomstats.tests.np_and_pytorch_only
def test_rk4_step(self):
state = (self.intercept, self.slope)
result = len(integrator.rk4_step(
state, self.function_linear, self.dt))
expected = len(state)
self.assertAllClose(result, expected)
@geomstats.tests.np_and_pytorch_only
def test_integrator(self):
initial_state = self.euclidean.random_uniform(2)
def function(_, velocity):
return gs.zeros_like(velocity)
for step in ['euler', 'rk4']:
flow, _ = integrator.integrate(function, initial_state, step=step)
result = flow[-1]
expected = initial_state[0] + initial_state[1]
self.assertAllClose(result, expected)