from multiprocessing import Pool

import numpy as np
import scipy.linalg as LA

k = 1.
g1 = 1
dt = 0.01
Htilde = k * np.ones((2, 2), dtype=np.complex64)
A1 = np.zeros((2, 2), dtype=np.complex64)
A1[0][0] = np.sqrt(2 * g1)
AdA = np.matmul(A1.T, A1)
H = Htilde - 1.j * AdA / 2
U = LA.expm(-1.j * H * dt)


def simulate(psi, ln=100):
    r = np.random.rand(ln)
    for st in range(ln):
        phi = U.dot(psi)
        n = LA.norm(phi)
        if n > r[st]:
            psi = phi / n
        else:
            psi = A1.dot(psi)
            psi /= LA.norm(psi)
    return psi


psi_init = np.array([0, 1])
with Pool() as pool:
    pool.map(simulate, [psi_init for k in range(10000)])