bsspot.py

def crank(V, L1, R1x, L2, R2x, dt, n, crumbs=[], callback=None):
    V = V.copy()
    dt *= 0.5

    L1e = flatten_tensor(L1)
    L1i = L1e.copy()
    R1 = np.array(R1x).T

    L2e = flatten_tensor(L2)
    L2i = L2e.copy()
    R2 = np.array(R2x)

    m = 2

    # L  = (As + Ass - r*np.eye(nspots))*-dt + np.eye(nspots)
    L1e.data *= dt
    L1e.data[m, :] += 1
    L1i.data *= -dt
    L1i.data[m, :] += 1
    R1 *= dt

    L2e.data *= dt
    L2e.data[m, :] += 1
    L2i.data *= -dt
    L2i.data[m, :] += 1
    R2 *= dt

    offsets1 = (abs(min(L1i.offsets)), abs(max(L1i.offsets)))
    offsets2 = (abs(min(L2i.offsets)), abs(max(L2i.offsets)))

    print_step = max(1, int(n / 10))
    to_percent = 100.0 / n
    utils.tic("Crank:")
    for k in xrange(n):
        if not k % print_step:
            if isnan(V).any():
                print "Crank fail @ t = %f (%i steps)" % (dt * k, k)
                return crumbs
            print int(k * to_percent),
        if callback is not None:
            callback(V, ((n - k) * dt))
        V = (L1e.dot(V.T.flat).reshape(V.shape[::-1]).T) + R1
        V = spl.solve_banded(offsets2, L2i.data,
                             (V + R2).flat, overwrite_b=True).reshape(V.shape)
        V = L2e.dot(V.flat).reshape(V.shape) + R2
        V = spl.solve_banded(offsets1, L1i.data,
                             (V + R1).T.flat, overwrite_b=True).reshape(V.shape[::-1]).T
        crumbs.append(V.copy())
    utils.toc()
    return crumbs