qsim.py

from abc import ABC, abstractmethod
from argparse import ArgumentParser, BooleanOptionalAction

import numpy as np
from matplotlib import pyplot as plt
from scipy import linalg

pauli_z = np.array([[1, 0], [0, -1]], dtype=complex)
pauli_x = np.array([[1, 0], [0, -1]], dtype=complex)


class System(ABC):

    time_step: float

    @abstractmethod
    def evolve(self, state):
        pass

    @abstractmethod
    def expval(self, state):
        pass


class Qubit(System):

    def __init__(self, time_step) -> None:
        self.energy = - pauli_z / 2
        self.time_step = time_step
        self.time_evolution = linalg.expm(-1j * self.energy * self.time_step)

    def evolve(self, state):
        return self.time_evolution @ state

    def expval(self, state):
        return state.conj().T @ pauli_x @ state


class ParticleSimulator:

    def __init__(self, system: System):
        self.system = system
        self.times = []
        self.states = []
        self.expvals = []

    def run(self, state, duration):
        time = 0

        while time < duration:
            self.times.append(time)
            self.states.append(state)
            self.expvals.append(self.system.expval(state))

            time += self.system.time_step
            state = self.system.evolve(state)


class EnsembleSimulator:

    def __init__(self, system: System, size: int):
        self.system = system
        self.particles = [ParticleSimulator(system) for _ in range(size)]
        self.times = []
        self.expvals = []

    def run(self, state, duration):
        expvals = []

        for particle in self.particles:
            particle.run(state, duration)
            expvals.append(particle.expvals)

        self.times = self.particles[0].times
        self.expvals = np.array(expvals)


def main(args):
    sim = EnsembleSimulator(Qubit(np.pi / 2), 1000)
    sim.run(np.array([1, 1]) / np.sqrt(2), 20 * np.pi)

    if args.plot_particle:
        plt.plot(sim.times, sim.expvals[0], '--')
        plt.show()
    if args.plot_ensemble:
        plt.plot(sim.times, sim.expvals.mean(axis=0), '--')
        plt.show()


if __name__ == '__main__':
    parser = ArgumentParser('qsim', description='Quantum simulator')
    parser.add_argument('--plot-particle', action=BooleanOptionalAction)
    parser.add_argument('--plot-ensemble', action=BooleanOptionalAction)

    main(parser.parse_args())