pyqtorch
pyqtorch copied to clipboard
[Feature] Add HamEvo for time-dependent Hamiltonians
Below is the previous code calculating HamEvo
based on a 4th order Runge Kutta method. It could be useful in the future for a simple implementation for time-dependent Hamiltonians. To note that Runge Kutta methods do not conserve wavefunction normalization and may quickly diverge, leading to wrong results.
class HamEvo(torch.nn.Module):
"""
Base class for Hamiltonian evolution classes, performing the evolution using RK4 method.
Args:
H (tensor): Hamiltonian tensor.
t (tensor): Time tensor.
qubits (Any): Qubits for operation.
n_qubits (int): Number of qubits.
n_steps (int): Number of steps to be performed in RK4-based evolution. Defaults to 100.
"""
def __init__(
self,
H: torch.Tensor,
t: torch.Tensor,
qubit_support: list[int],
n_qubits: int,
n_steps: int = 100,
):
super().__init__()
self.qubit_support = qubit_support
self.H: torch.Tensor
self.t: torch.Tensor
self.n_qubits = n_qubits
self.n_steps = n_steps
if H.ndim == 2:
H = H.unsqueeze(2)
if H.size(-1) == t.size(0) or t.size(0) == 1:
self.register_buffer("H", H)
self.register_buffer("t", t)
elif H.size(-1) == 1:
(x, y, _) = H.size()
self.register_buffer("H", H.expand(x, y, t.size(0)))
self.register_buffer("t", t)
else:
msg = "H and t batchsizes either have to match or (one of them has to) be equal to one."
raise ValueError(msg)
def apply(self, state: torch.Tensor) -> torch.Tensor:
"""
Applies the Hamiltonian evolution operation on the given state using RK4 method.
Args:
state (tensor): Input quantum state.
Returns:
tensor: Output state after Hamiltonian evolution.
"""
batch_size = max(state.size(-1), self.H.size(-1))
if state.size(-1) == 1:
state = state.repeat(*[1 for _ in range(len(state.size()) - 1)], batch_size)
h = self.t.reshape((1, -1)) / self.n_steps
for _ in range(self.n_qubits - 1):
h = h.unsqueeze(0)
h = h.expand_as(state)
_state = state.clone()
for _ in range(self.n_steps):
k1 = -1j * _apply_einsum(_state, self.H, self.qubit_support, self.n_qubits, batch_size)
k2 = -1j * _apply_einsum(
_state + h / 2 * k1, self.H, self.qubit_support, self.n_qubits, batch_size
)
k3 = -1j * _apply_einsum(
_state + h / 2 * k2, self.H, self.qubit_support, self.n_qubits, batch_size
)
k4 = -1j * _apply_einsum(
_state + h * k3, self.H, self.qubit_support, self.n_qubits, batch_size
)
_state += h / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
return _state
def forward(self, state: torch.Tensor) -> torch.Tensor:
return self.apply(state)