Merge pull request #178 from gautierronan/master

Allow complex numbers in odeint

Author: Zymrael

test/models/test_ode.py

@@ -258,4 +258,40 @@ def forward(self, t, x, u, v, z, args={}):
     t_eval, sol2 = odeprob(x0, t_span=torch.linspace(0, 5, 10))
 
     assert (sol1==sol2).all()
-    grad(sol2.sum(), x0)
+    grad(sol2.sum(), x0)
+
+
+def test_complex_ode():
+    """Test odeint for complex numbers with a simple complex-valued ODE, corresponding
+    to Rabi oscillations of quantum two-level system."""
+    class Rabi(nn.Module):
+        def __init__(self, omega):
+            super().__init__()
+            self.sx = torch.tensor([[0, 1], [1, 0]], dtype=torch.complex128)
+            self.omega = omega
+            return
+        def forward(self, t, x):
+            dx = -1.0j * self.omega * self.sx @ x
+            dx += dx.adjoint()
+            return dx
+
+    # Odeint parameters
+    omega = torch.randn(1)
+    rabi = Rabi(omega)
+    tspan = torch.linspace(0., 2., 10)
+    
+    # Random initial state
+    x0 = torch.rand(2, 2, dtype=torch.complex128)
+    x0 = 0.5 * (x0 + x0.adjoint()) / torch.real(x0.trace())
+    # Solve the ODE problem
+    t_eval, sol = odeint(f=rabi, x=x0, t_span=tspan, solver="dopri5", atol=1e-8, rtol=1e-6)
+    
+    # Expected solution
+    sx = torch.tensor([[0, 1], [1, 0]], dtype=torch.complex128)
+    si = torch.tensor([[1, 0], [0, 1]], dtype=torch.complex128)
+    U_t = torch.cos(omega * t_eval)[:, None, None] * si 
+    U_t += -1j * torch.sin(omega * t_eval)[:, None, None] * sx
+    sol_exp = U_t @ x0 @ U_t.adjoint()
+    
+    # Check result
+    assert torch.allclose(sol, sol_exp, rtol=1e-5, atol=1e-5)

torchdyn/numerics/odeint.py

@@ -65,11 +65,6 @@ def odeint(f:Callable, x:Tensor, t_span:Union[List, Tensor], solver:Union[str, n
 	x, t_span = solver.sync_device_dtype(x, t_span)
 	stepping_class = solver.stepping_class
 
-	# instantiate save_at tensor
-	if len(save_at) == 0: save_at = t_span
-	if not isinstance(save_at, torch.Tensor):
-		save_at = torch.tensor(save_at)
-
 	# instantiate the interpolator similar to the solver steps above
 	if isinstance(solver, Tsitouras45):
 		if verbose: warn("Running interpolation not yet implemented for `tsit5`")
@@ -87,6 +82,10 @@ def odeint(f:Callable, x:Tensor, t_span:Union[List, Tensor], solver:Union[str, n
 		if stepping_class == 'fixed':
 			if atol != odeint.__defaults__[0] or rtol != odeint.__defaults__[1]:
 				warn("Setting tolerances has no effect on fixed-step methods")
+			# instantiate save_at tensor
+			if len(save_at) == 0: save_at = t_span
+			if not isinstance(save_at, torch.Tensor):
+				save_at = torch.tensor(save_at)
 			return _fixed_odeint(f_, x, t_span, solver, save_at=save_at, args=args)
 		elif stepping_class == 'adaptive':
 			t = t_span[0]

torchdyn/numerics/utils.py

@@ -31,23 +31,23 @@ def norm_(aug_state):
 
 
 def hairer_norm(tensor):
-    return tensor.pow(2).mean().sqrt()
+    return tensor.abs().pow(2).mean().sqrt()
 
 
 def init_step(f, f0, x0, t0, order, atol, rtol):
     scale = atol + torch.abs(x0) * rtol
     d0, d1 = hairer_norm(x0 / scale), hairer_norm(f0 / scale)
 
     if d0 < 1e-5 or d1 < 1e-5:
-        h0 = torch.tensor(1e-6, dtype=x0.dtype, device=x0.device)
+        h0 = torch.tensor(1e-6, dtype=t0.dtype, device=t0.device)
     else:
         h0 = 0.01 * d0 / d1
 
     x_new = x0 + h0 * f0
     f_new = f(t0 + h0, x_new)
     d2 = hairer_norm((f_new - f0) / scale) / h0
     if d1 <= 1e-15 and d2 <= 1e-15:
-        h1 = torch.max(torch.tensor(1e-6, dtype=x0.dtype, device=x0.device), h0 * 1e-3)
+        h1 = torch.max(torch.tensor(1e-6, dtype=t0.dtype, device=t0.device), h0 * 1e-3)
     else:
         h1 = (0.01 / max(d1, d2)) ** (1. / float(order + 1))
     dt = torch.min(100 * h0, h1).to(t0)