utils/EBU_Quantum/with_body/fdtd_matrix_zeroed.py

#
# Generate FDTD equation matrices for a simple 2D rectangular grid
# problem with PMC boundaries at the edges.  The material is homogeneous.
#
# Two matrices are produced in different output files:  A_E and A_H.
# These can be combined to form a full system matrix:
#
# [   0   A_E  ]
# [  A_H   0   ]
#
# Usage:
#
# generate_fdtd_matrix.py --nx <num> --ny <num> -g <num> --e_out <file> --h_out <file> [--eps_r <value>]
#

import os
import sys
import math
import numpy as np
import argparse
from square_hole_grid import *

def off_diagonal_stack(A, B):
    M, N = A.shape[0], B.shape[0]
    print(f"Shape of A {A.shape}")
    print(f"Shape of B {B.shape}")

    zero_MM = np.zeros((M, M), dtype=A.dtype)
    zero_NN = np.zeros((N, N), dtype=B.dtype)

    top = np.hstack((zero_MM, A))
    bottom = np.hstack((B, zero_NN))
    stacked = np.vstack((top, bottom))
    return stacked


def next_power_of_two(n):
    return 1 if n == 0 else 2 ** (n - 1).bit_length()


def keepmaskToFlatInd(i, j, keep_mask):
    if ~keep_mask[j][i]:
        raise Exception("Grid point at index (i, j) lies strictly inside a holed portion of the geometry.")

    nJ, nI = keep_mask.shape
    flatind = -1

    for jj in range(j + 1):
        if jj == j:
            for ii in range(i + 1):
                flatind += keep_mask[jj][ii]
        else:
            for ii in range(nI):
                flatind += keep_mask[jj][ii]

    return flatind


def GetEzDofNum(i, j, nx, ny, mask):
    if i < 0 or j < 0 or i >= nx or j >= ny:
        return -1
    if ~mask[j, i]:
        return -1
    return j * nx + i


def GetHxDofNum(i, j, nx, ny, mask):
    if i < 0 or j < 0 or i >= nx or j >= ny - 1:
        return -1
    if ~mask[j, i]:
        return -1
    return j * nx + i


def GetHyDofNum(i, j, nx, ny, mask):
    if i < 0 or j < 0 or i >= nx - 1 or j >= ny:
        return -1
    if ~mask[j, i]:
        return -1
    return j * (nx - 1) + i + nx * (ny - 1)


def MathematicaPrint(ofname, mat):
    rows, cols = np.shape(mat)
    with open(ofname, "w") as of:
        of.write("{")
        for j in range(rows):
            if j > 0:
                of.write(",\n")
            of.write("{")
            for i in range(cols):
                num = mat[j][i]
                numstr = f"{num:17.10e}".replace("e+", "*10^").replace("e-", "*10^-")
                if i > 0:
                    of.write(", ")
                of.write(numstr)
            of.write("}")
        of.write("}\n")


def build_amat(L, nx, ny, hole_size, cx, cy, eps=1, mu0=1):

    (
        X_holed,
        Y_holed,
        keep_mask,
        XY_kept,
        hole_bounds,
        inner_nodes,
        outer_nodes,
    ) = square_grid_with_square_hole_points(
        N_outer=nx,
        L=L,
        hole_center=(cx, cy),
        hole_size=hole_size,
        align="snap",
        include_outer_boundary=True,
    )

    h = X_holed[0][1] - X_holed[0][0]
    ezMask = keep_mask.copy()
    hxMask = ezMask[:-1, :] * ezMask[1:, :]
    hyMask = ezMask[:, :-1] * ezMask[:, 1:]

    n_e = nx * ny

    e_h = np.array([])
    for k, it in inner_nodes.items():
        e_h = np.concatenate((e_h, it["flat"]))
    e_h = np.unique(e_h)
    n_e_h = e_h.shape[0]

    e_o = np.array([])
    for k, it in outer_nodes.items():
        e_o = np.concatenate((e_o, it["flat"]))
    e_o = np.unique(e_o)
    n_e_o = e_o.shape[0]

    n_hx = nx * (ny - 1)
    n_hy = (nx - 1) * ny

    d = int(n_e + n_hx + n_hy)
    N = next_power_of_two(d)
    nq = math.ceil(math.log2(N))

    H_coef = 1.0 / (h * mu0)
    E_coef = 1.0 / (h * eps)

    AE_mat = np.zeros((n_e, n_hx + n_hy))

    for j in range(ny):
        for i in range(nx):
            if ~ezMask[j, i]:
                ez_ix = GetEzDofNum(i, j, nx, ny, ezMask)
                AE_mat[ez_ix][:] = 0
            else:
                ez_ix = GetEzDofNum(i, j, nx, ny, ezMask)

                # dHy/dx
                hy_ix_plus = GetHyDofNum(i, j, nx, ny, hyMask)
                hy_ix_minus = GetHyDofNum(i - 1, j, nx, ny, hyMask)

                if hy_ix_plus == -1:
                    AE_mat[ez_ix][hy_ix_minus] = -2.0 * E_coef
                elif hy_ix_minus == -1:
                    AE_mat[ez_ix][hy_ix_plus] = 2.0 * E_coef
                else:
                    AE_mat[ez_ix][hy_ix_plus] = E_coef
                    AE_mat[ez_ix][hy_ix_minus] = -E_coef

                # -dHx/dy
                hx_ix_plus = GetHxDofNum(i, j, nx, ny, hxMask)
                hx_ix_minus = GetHxDofNum(i, j - 1, nx, ny, hxMask)

                if hx_ix_plus == -1:
                    AE_mat[ez_ix][hx_ix_minus] = 2.0 * E_coef
                elif hx_ix_minus == -1:
                    AE_mat[ez_ix][hx_ix_plus] = -2.0 * E_coef
                else:
                    AE_mat[ez_ix][hx_ix_plus] = -E_coef
                    AE_mat[ez_ix][hx_ix_minus] = E_coef

    AH_mat = np.zeros((n_hx + n_hy, n_e))

    for j in range(ny - 1):
        for i in range(nx):
            if ~hxMask[j, i]:
                hx_ix = GetHxDofNum(i, j, nx, ny, hxMask)
                AH_mat[hx_ix][:] = 0
            else:
                hx_ix = GetHxDofNum(i, j, nx, ny, hxMask)
                ez_ix_plus = GetEzDofNum(i, j + 1, nx, ny, ezMask)
                ez_ix_minus = GetEzDofNum(i, j, nx, ny, ezMask)
                AH_mat[hx_ix][ez_ix_plus] = -H_coef
                AH_mat[hx_ix][ez_ix_minus] = H_coef

    for j in range(ny):
        for i in range(nx - 1):
            if ~hyMask[j, i]:
                hy_ix = GetHyDofNum(i, j, nx, ny, hyMask)
                AH_mat[hy_ix][:] = 0
            else:
                hy_ix = GetHyDofNum(i, j, nx, ny, hyMask)
                ez_ix_plus = GetEzDofNum(i + 1, j, nx, ny, ezMask)
                ez_ix_minus = GetEzDofNum(i, j, nx, ny, ezMask)
                AH_mat[hy_ix][ez_ix_plus] = H_coef
                AH_mat[hy_ix][ez_ix_minus] = -H_coef

    small_a = off_diagonal_stack(AE_mat, AH_mat)

    print("amat constructed.")
    return N, nq


def build_amat_dict(L, nx, ny, hole_size, cx, cy, eps=1, mu0=1):

    (
        X_holed,
        Y_holed,
        keep_mask,
        XY_kept,
        hole_bounds,
        inner_nodes,
        outer_nodes,
    ) = square_grid_with_square_hole_points(
        N_outer=nx,
        L=L,
        hole_center=(cx, cy),
        hole_size=hole_size,
        align="snap",
        include_outer_boundary=True,
    )

    h = X_holed[0][1] - X_holed[0][0]
    ezMask = keep_mask.copy()
    hxMask = ezMask[:-1, :] * ezMask[1:, :]
    hyMask = ezMask[:, :-1] * ezMask[:, 1:]

    n_e = nx * ny

    e_h = np.array([])
    for k, it in inner_nodes.items():
        e_h = np.concatenate((e_h, it["flat"]))
    e_h = np.unique(e_h)
    n_e_h = e_h.shape[0]

    e_o = np.array([])
    for k, it in outer_nodes.items():
        e_o = np.concatenate((e_o, it["flat"]))
    e_o = np.unique(e_o)
    n_e_o = e_o.shape[0]

    n_hx = nx * (ny - 1)
    n_hy = (nx - 1) * ny

    d = int(n_e + n_hx + n_hy)
    N = next_power_of_two(d)
    nq = math.ceil(math.log2(N))

    H_coef = 1.0 / (h * mu0)
    E_coef = 1.0 / (h * eps)

    small_a_dict = {}

    # A_E block
    for j in range(ny):
        for i in range(nx):
            if ~ezMask[j, i]:
                continue
            ez_ix = GetEzDofNum(i, j, nx, ny, ezMask)
            row = ez_ix

            hy_ix_plus = GetHyDofNum(i, j, nx, ny, hyMask)
            hy_ix_minus = GetHyDofNum(i - 1, j, nx, ny, hyMask)

            if hy_ix_plus == -1:
                col = n_e + hy_ix_minus
                small_a_dict[(row, col)] = -2.0 * E_coef
            elif hy_ix_minus == -1:
                col = n_e + hy_ix_plus
                small_a_dict[(row, col)] = 2.0 * E_coef
            else:
                small_a_dict[(row, n_e + hy_ix_plus)] = E_coef
                small_a_dict[(row, n_e + hy_ix_minus)] = -E_coef

            hx_ix_plus = GetHxDofNum(i, j, nx, ny, hxMask)
            hx_ix_minus = GetHxDofNum(i, j - 1, nx, ny, hxMask)

            if hx_ix_plus == -1:
                col = n_e + hx_ix_minus
                small_a_dict[(row, col)] = 2.0 * E_coef
            elif hx_ix_minus == -1:
                col = n_e + hx_ix_plus
                small_a_dict[(row, col)] = -2.0 * E_coef
            else:
                small_a_dict[(row, n_e + hx_ix_plus)] = -E_coef
                small_a_dict[(row, n_e + hx_ix_minus)] = E_coef

    # A_H block
    for j in range(ny - 1):
        for i in range(nx):
            if ~hxMask[j, i]:
                continue
            hx_ix = GetHxDofNum(i, j, nx, ny, hxMask)
            row = n_e + hx_ix

            ez_ix_plus = GetEzDofNum(i, j + 1, nx, ny, ezMask)
            ez_ix_minus = GetEzDofNum(i, j, nx, ny, ezMask)

            if ez_ix_plus != -1:
                small_a_dict[(row, ez_ix_plus)] = -H_coef
            if ez_ix_minus != -1:
                small_a_dict[(row, ez_ix_minus)] = H_coef

    for j in range(ny):
        for i in range(nx - 1):
            if ~hyMask[j, i]:
                continue
            hy_ix = GetHyDofNum(i, j, nx, ny, hyMask)
            row = n_e + hy_ix

            ez_ix_plus = GetEzDofNum(i + 1, j, nx, ny, ezMask)
            ez_ix_minus = GetEzDofNum(i, j, nx, ny, ezMask)

            if ez_ix_plus != -1:
                small_a_dict[(row, ez_ix_plus)] = H_coef
            if ez_ix_minus != -1:
                small_a_dict[(row, ez_ix_minus)] = -H_coef

    print("amat_dict constructed.")
    return N, nq, small_a_dict