[Feature] Add HamEvo for time-dependent Hamiltonians

[jpmoutinho]

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)

[jpmoutinho]

This was only open to save the code snippet above for potential future use. It can be closed once the Dormand-Prince method is implemented in the new solver repo.