gistfile1.txt

import itertools
from collections import namedtuple
import numpy as np
import scipy.sparse.linalg

def step():
    W = 0
    '''
    div E = rho / eta0
    div B = 0
    rot E = - dB/dt
    rot B = mu0 ( J + eta0 dE/dt )
    '''

    '''
    div E = rho / eta0
    dEx/dx + dEy/dy + dEz/dz = rho / eta0
    '''
    # A[divE] += 1
    # b += rho / eta0
    for w, (i, j, k) in enumerate(all_points()):
        A[W + w, idxE(0, i - 1, j    , k    )] += -0.5 / dx
        A[W + w, idxE(0, i + 1, j    , k    )] += +0.5 / dx
        A[W + w, idxE(1, i    , j - 1, k    )] += -0.5 / dy
        A[W + w, idxE(1, i    , j + 1, k    )] += +0.5 / dy
        A[W + w, idxE(2, i    , j    , k - 1)] += -0.5 / dz
        A[W + w, idxE(2, i    , j    , k + 1)] += +0.5 / dz
        b[W + w] += rho[idxS(i, j, k)] / eta0
    W += w + 1
    assert W == Nx * Ny * Nz

    '''
    div B = 0
    dBx/dx + dBy/dy + dBz/dz = 0
    '''
    # A[divE] += 1
    # b += 0
    for w, (i, j, k) in enumerate(all_points()):
        A[W + w, idxB(0, i - 1, j    , k    )] -= 0.5 / dx
        A[W + w, idxB(0, i + 1, j    , k    )] += 0.5 / dx
        A[W + w, idxB(1, i    , j - 1, k    )] -= 0.5 / dy
        A[W + w, idxB(1, i    , j + 1, k    )] += 0.5 / dy
        A[W + w, idxB(2, i    , j    , k - 1)] -= 0.5 / dz
        A[W + w, idxB(2, i    , j    , k + 1)] += 0.5 / dz
        b[W + w] += 0
    W += w + 1
    assert W == 2 * Nx * Ny * Nz

    '''
    rot E = - dB/dt
    dEz/dy - dEy/dz = - dBx/dt
    dEx/dz - dEz/dx = - dBy/dt
    dEy/dx - dEx/dy = - dBz/dt
    '''
    # A[rotE] += 1
    # A[B] += 1/dt
    # b += B0/dt
    for w, (c, i, j, k) in enumerate(all_components()):
        if c == 0:
            A[W + w, idxE(2, i    , j - 1, k    )] += - 0.5 / dy
            A[W + w, idxE(2, i    , j + 1, k    )] += + 0.5 / dy
            A[W + w, idxE(1, i    , j    , k - 1)] += + 0.5 / dz
            A[W + w, idxE(1, i    , j    , k + 1)] += - 0.5 / dz
        elif c == 1:
            A[W + w, idxE(0, i    , j    , k - 1)] += - 0.5 / dz
            A[W + w, idxE(0, i    , j    , k + 1)] += + 0.5 / dz
            A[W + w, idxE(2, i - 1, j    , k    )] += + 0.5 / dx
            A[W + w, idxE(2, i + 1, j    , k    )] += - 0.5 / dx
        else:
            A[W + w, idxE(1, i - 1, j    , k    )] += - 0.5 / dx
            A[W + w, idxE(1, i + 1, j    , k    )] += + 0.5 / dx
            A[W + w, idxE(0, i    , j - 1, k    )] += + 0.5 / dy
            A[W + w, idxE(0, i    , j + 1, k    )] += - 0.5 / dy

        A[W + w, idxB(c, i, j, k)] += 1 / dt
        b[W + w] += B0[idxV(c, i, j, k)] / dt
    W += w + 1
    assert W == 2 * Nx * Ny * Nz + 3 * Nx * Ny * Nz

    '''
    rot B = mu0 ( J + eta0 dE/dt )
    dBz/dy - dBy/dz = mu0 ( Jx + eta0 dEx/dt )
    dBx/dz - dBz/dx = mu0 ( Jy + eta0 dEy/dt )
    dBy/dx - dBx/dy = mu0 ( Jz + eta0 dEz/dt )
    '''
    # A[rotB] += 1
    # A[E] -= mu0 * eta0 / dt
    # b += mu0 ( J + eta0 * E0 / dt)
    for w, (c, i, j, k) in enumerate(all_components()):
        if c == 0:
            A[W + w, idxB(2, i    , j - 1, k    )] += - 0.5 / dy
            A[W + w, idxB(2, i    , j + 1, k    )] += + 0.5 / dy
            A[W + w, idxB(1, i    , j    , k - 1)] += + 0.5 / dz
            A[W + w, idxB(1, i    , j    , k + 1)] += - 0.5 / dz
        elif c == 1:
            A[W + w, idxB(0, i    , j    , k - 1)] += - 0.5 / dz
            A[W + w, idxB(0, i    , j    , k + 1)] += + 0.5 / dz
            A[W + w, idxB(2, i - 1, j    , k    )] += + 0.5 / dx
            A[W + w, idxB(2, i + 1, j    , k    )] += - 0.5 / dx
        else:
            A[W + w, idxB(1, i - 1, j    , k    )] += - 0.5 / dx
            A[W + w, idxB(1, i + 1, j    , k    )] += + 0.5 / dx
            A[W + w, idxB(0, i    , j - 1, k    )] += + 0.5 / dy
            A[W + w, idxB(0, i    , j + 1, k    )] += - 0.5 / dy

        A[W + w, idxE(c, i, j, k)] += - mu0 * eta0 / dt
        b[W + w] += mu0 * ( J[idxV(c, i, j, k)] - eta0 * E0[idxV(c, i, j, k)] / dt)
    W += w + 1
    assert W == 2 * Nx * Ny * Nz + 2 * 3 * Nx * Ny * Nz

def plot(V):
    from mpl_toolkits.mplot3d import axes3d
    import matplotlib.pyplot as plt
    import numpy as np

    fig = plt.figure()
    ax = fig.gca(projection='3d')

    z, y, x = np.meshgrid(np.arange(Nz),
                          np.arange(Ny),
                          np.arange(Nx))
    pos = np.array([xyz for xyz in all_points()])
    pos = np.reshape(pos, [Nz, Ny, Nx, 3])
    z = pos[..., 0]
    y = pos[..., 1]
    x = pos[..., 2]
    
    V = np.reshape(V, [Nz, Ny, Nx, 3])
    u = V[..., 0]
    v = V[..., 1]
    w = V[..., 2]

    if True:
        from mayavi import mlab
        obj = mlab.quiver3d(x, y, z, u, v, w, line_width=3, scale_factor=1)
        mlab.show()
    else:
        ax.quiver(x, y, z, u, v, w, length=0.1)

        plt.show()

Nx, Ny, Nz = (12, 15, 18) # number of grid points
NE = 3 * Nx * Ny * Nz
NB = 3 * Nx * Ny * Nz
N = NE + NB # total number of variables
M = 2 * Nx * Ny * Nz + (NE + NB)

all_points = lambda: itertools.product(range(Nz), range(Ny), range(Nx))
all_components = lambda: itertools.product(range(Nz), range(Ny), range(Nx), range(3))

# modulo operation (%) means periodic boundary conditions
idxS = lambda i, j, k: ((k % Nz) * Ny + (j % Ny)) * Nx + (i % Nx)
idxV = lambda c, i, j, k: idxS(i, j, k) * 3 + c
idxE = lambda c, i, j, k: idxV(c, i, j, k)
idxB = lambda c, i, j, k: NE + idxV(c, i, j, k)

dt = 0.1
dx = 0.05
dy = dx
dz = dx
eta0 = 8.854187817e-12
mu0 = 4e-7 * np.pi
dtype = 'float32'

rho = np.zeros([Nx * Ny * Nz], dtype=dtype)
E0 = np.zeros([3 * Nx * Ny * Nz], dtype=dtype)
B0 = np.zeros([3 * Nx * Ny * Nz], dtype=dtype)
J = np.zeros([3 * Nx * Ny * Nz], dtype=dtype)
#for k in range(Nz):
#    J[idxV(2, Nx // 2, Ny // 2, k)] = 1
for i in (0, 1):
    for j in (0, 1):
        for k in (0, 1):
            rho[idxS(Nx // 2 + i, Ny // 2 + j, Nz // 2 + k)] = 1e2 * eta0

A = np.zeros([M, N], dtype=dtype)
b = np.zeros([M], dtype=dtype)

Solution = namedtuple('Solution', ['x', 'istop', 'itn', 'r1norm', 'r2norm', 'anorm', 'acond', 'arnorm', 'xnorm', 'var'])

for t in range(1):
    print(t)
    step()
    A = scipy.sparse.csc_matrix(A)
    y = scipy.sparse.linalg.lsqr(A, b)
    res = Solution(*y)
    print(res.r1norm)
    E0, B0 = np.split(res.x, 2)
    print(E0[idxV(0, Nx // 2, Ny // 2, Nz // 2)])
    print(E0[idxV(1, Nx // 2, Ny // 2, Nz // 2)])
    print(E0[idxV(2, Nx // 2, Ny // 2, Nz // 2)])
    print(np.amax(E0), np.mean(E0), np.amin(E0))
    print(E0.shape, NE)

plot(E0)