"""Implement 3D/2D solvers."""
from typing import List, Union
import numpy as np
from scipy.linalg import eigh
from .brillouinzone import BrillouinZonePath1D as BZPath1D
from .brillouinzone import BrillouinZonePath2D as BZPath2D
from .brillouinzone import BrillouinZonePath3D as BZPath3D
from .geometry import Geometry1D, Geometry2D, Geometry3D
from .utils import convmat
class Solver1D:
"""Implement a 1D PWE Solver.
Parameters
----------
geometry : Geometry
geometry
path : BrillouinZonePath
A BrillouinZonePath object.
P : int
Number of terms in the direction of the reciprocal vector b1.
"""
def __init__(
self, geometry: Geometry1D, path: BZPath1D, P: int = 1, pol: str = "TM"
):
"""Initialize the PWE Solver."""
self.geo = geometry
self.path = path
self.P = P
self.pol = pol
self.eps_rc: np.ndarray
self.mu_rc: np.ndarray
self.wn: List[np.ndarray] = []
self.modes: List[np.ndarray] = []
def run(self):
"""Calculate the eigen-wavenumbers and modes."""
raise NotImplementedError
[docs]class Solver2D:
"""Implement a 2D PWE Solver.
Parameters
----------
geometry : Geometry
geometry
path : BrillouinZonePath
A BrillouinZonePath object.
P : int
Number of terms in the direction of the reciprocal vector b1.
Q : int
Number of terms in the direction of the reciprocal vector b2.
"""
def __init__(
self,
geometry: Geometry2D,
path: BZPath2D,
P: int = 1,
Q: int = 1,
pol: str = "TM",
):
"""Initialize the PWE Solver."""
self.geo = geometry
self.path = path
self.P, self.Q = P, Q
self.pol = pol
self.eps_rc: np.ndarray
self.mu_rc: np.ndarray
self.wn: List[np.ndarray] = []
self.modes: List[np.ndarray] = []
[docs] def run(self):
"""Calculate the eigen-wavenumbers and modes."""
self.eps_rc = convmat(self.geo.eps_r, self.P, self.Q)
self.mu_rc = convmat(self.geo.mu_r, self.P, self.Q)
mu_rci = np.linalg.inv(self.mu_rc)
eps_rci = np.linalg.inv(self.eps_rc)
p = np.arange(-(self.P // 2), self.P // 2 + 1)
q = np.arange(-(self.Q // 2), self.Q // 2 + 1)
P0, Q0 = np.meshgrid(p, q)
b1, b2 = self.path.b1, self.path.b2
G_k = P0.flatten() * b1[:, None] + Q0.flatten() * b2[:, None]
beta = self.path.betas.values
for col in range(beta.shape[1]):
k = beta[:, col, None] - G_k
KX = np.diag(k[0, :])
KY = np.diag(k[1, :])
if self.pol == "TM":
A = KX @ mu_rci @ KX + KY @ mu_rci @ KY
B = self.eps_rc
else:
A = KX @ eps_rci @ KX + KY @ eps_rci @ KY
B = self.mu_rc
D, V = eigh(A, B)
self.wn.append(np.sqrt(np.maximum(D, 0)))
self.modes.append(V)
[docs]class Solver3D:
"""Implement a 3D PWE Solver.
Parameters
----------
geometry : Geometry
geometry
path : BrillouinZonePath
A BrillouinZonePath object.
P : int
Number of terms in the direction of the reciprocal vector b1.
Q : int
Number of terms in the direction of the reciprocal vector b2.
R : int
Number of terms in the direction of the reciprocal vector b3.
"""
def __init__(
self, geometry: Geometry3D, path: BZPath3D, P: int = 1, Q: int = 1, R: int = 1
):
"""Initialize the PWE Solver."""
self.geo = geometry
self.path = path
self.P, self.Q, self.R = P, Q, R
self.eps_rc: np.ndarray
self.mu_rc: np.ndarray
self.wn: List[np.ndarray] = []
self.modes: List[np.ndarray] = []
[docs] def run(self):
"""Calculate the eigen-wavenumbers and modes."""
self.eps_rc = convmat(self.geo.eps_r, self.P, self.Q, self.R)
self.mu_rc = convmat(self.geo.mu_r, self.P, self.Q, self.R)
eps_rk = np.kron(np.eye(3), self.eps_rc)
mu_rki = np.kron(np.eye(3), np.linalg.inv(self.mu_rc))
p = np.arange(-(self.P // 2), self.P // 2 + 1)
q = np.arange(-(self.Q // 2), self.Q // 2 + 1)
r = np.arange(-(self.R // 2), self.R // 2 + 1)
P0, Q0, R0 = np.meshgrid(p, q, r)
b1, b2, b3 = self.path.b1, self.path.b2, self.path.b3
G_k = (
b1[:, None] * P0.flatten()
+ b2[:, None] * Q0.flatten()
+ b3[:, None] * R0.flatten()
)
beta = self.path.betas.values
for col in range(beta.shape[1]):
k = beta[:, col, None] - G_k
KX = np.diag(k[0, :])
KY = np.diag(k[1, :])
KZ = np.diag(k[2, :])
K_V = np.vstack(
(
np.hstack((0 * KX, -KZ, KY)),
np.hstack((KZ, 0 * KY, -KX)),
np.hstack((-KY, KX, 0 * KZ)),
)
)
A = K_V @ mu_rki @ K_V
B = eps_rk
D, V = eigh(A, B)
self.wn.append(np.sqrt(-np.minimum(D, 0)))
self.modes.append(V)
[docs]def Solver(geometry, path, *args, **kwargs) -> Union[Solver1D, Solver2D, Solver3D]:
"""Solver factory."""
dim_obj = {1: Solver1D, 2: Solver2D, 3: Solver3D}
return dim_obj[path.dim](geometry, path, *args, **kwargs)