From ddb2cbb9fe2824f4edfa875b4091ed5a4b952908 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Mon, 10 Nov 2025 11:26:29 +0100 Subject: [PATCH 01/39] WIP: logfile --- build/lib/wakis/__init__.py | 19 + build/lib/wakis/_version.py | 1 + build/lib/wakis/conductors.py | 204 +++++ build/lib/wakis/conductors3d.py | 223 ++++++ build/lib/wakis/field.py | 667 ++++++++++++++++ build/lib/wakis/geometry.py | 254 ++++++ build/lib/wakis/grid2D.py | 595 ++++++++++++++ build/lib/wakis/grid3D.py | 415 ++++++++++ build/lib/wakis/gridFIT3D.py | 841 ++++++++++++++++++++ build/lib/wakis/materials.py | 50 ++ build/lib/wakis/plotting.py | 820 ++++++++++++++++++++ build/lib/wakis/pmlBlock2D.py | 135 ++++ build/lib/wakis/pmlBlock3D.py | 265 +++++++ build/lib/wakis/routines.py | 336 ++++++++ build/lib/wakis/solver2D.py | 484 ++++++++++++ build/lib/wakis/solver3D.py | 1126 +++++++++++++++++++++++++++ build/lib/wakis/solverFIT3D.py | 1110 ++++++++++++++++++++++++++ build/lib/wakis/sources.py | 418 ++++++++++ build/lib/wakis/wakeSolver.py | 1288 +++++++++++++++++++++++++++++++ wakis/__init__.py | 2 + wakis/logger.py | 35 + wakis/solverFIT3D.py | 20 +- 22 files changed, 9306 insertions(+), 2 deletions(-) create mode 100644 build/lib/wakis/__init__.py create mode 100644 build/lib/wakis/_version.py create mode 100644 build/lib/wakis/conductors.py create mode 100644 build/lib/wakis/conductors3d.py create mode 100644 build/lib/wakis/field.py create mode 100644 build/lib/wakis/geometry.py create mode 100644 build/lib/wakis/grid2D.py create mode 100644 build/lib/wakis/grid3D.py create mode 100644 build/lib/wakis/gridFIT3D.py create mode 100644 build/lib/wakis/materials.py create mode 100644 build/lib/wakis/plotting.py create mode 100644 build/lib/wakis/pmlBlock2D.py create mode 100644 build/lib/wakis/pmlBlock3D.py create mode 100644 build/lib/wakis/routines.py create mode 100644 build/lib/wakis/solver2D.py create mode 100644 build/lib/wakis/solver3D.py create mode 100644 build/lib/wakis/solverFIT3D.py create mode 100644 build/lib/wakis/sources.py create mode 100644 build/lib/wakis/wakeSolver.py create mode 100644 wakis/logger.py diff --git a/build/lib/wakis/__init__.py b/build/lib/wakis/__init__.py new file mode 100644 index 0000000..16ebb49 --- /dev/null +++ b/build/lib/wakis/__init__.py @@ -0,0 +1,19 @@ +# copyright ################################# # +# This file is part of the wakis Package. # +# Copyright (c) CERN, 2024. # +# ########################################### # + +from . import field +from . import gridFIT3D +from . import solverFIT3D +from . import sources +from . import materials +from . import wakeSolver +from . import geometry + +from .field import Field +from .gridFIT3D import GridFIT3D +from .solverFIT3D import SolverFIT3D +from .wakeSolver import WakeSolver + +from ._version import __version__ \ No newline at end of file diff --git a/build/lib/wakis/_version.py b/build/lib/wakis/_version.py new file mode 100644 index 0000000..3966a5f --- /dev/null +++ b/build/lib/wakis/_version.py @@ -0,0 +1 @@ +__version__ = '0.6.1' \ No newline at end of file diff --git a/build/lib/wakis/conductors.py b/build/lib/wakis/conductors.py new file mode 100644 index 0000000..1f02a2d --- /dev/null +++ b/build/lib/wakis/conductors.py @@ -0,0 +1,204 @@ +import numpy as np +import scipy.optimize + +class OutRect: + def __init__(self, Lx, Ly, x_cent, y_cent): + self.Lx = Lx + self.Ly = Ly + self.x_cent = x_cent + self.y_cent = y_cent + # self.theta = theta + # self.R = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]) + # self.mR = np.array([[np.cos(-theta), -np.sin(-theta)], [np.sin(-theta), np.cos(-theta)]]) + + def out_conductor(self, x, y): + # [xx, yy] = np.dot(self.mR, np.array([x, y])) + return (-0.5 * self.Lx + self.x_cent < x < 0.5 * self.Lx + self.x_cent) and ( + -0.5 * self.Ly + self.y_cent < y < 0.5 * self.Ly + self.y_cent) + + def in_conductor(self, x, y): + return not self.out_conductor(x, y) + + def intersec_x(self, x, y): + # [xx, yy] = np.dot(self.mR, np.array([x, y])) + inters_1 = -0.5 * self.Lx + self.x_cent + inters_2 = inters_1 + self.Lx + if abs(x - inters_1) < abs(x - inters_2): + return inters_1 + else: + return inters_2 + + # [xxx, _] = np.dot(self.R, np.array([inters, yy])) + # return xxx + + def intersec_y(self, x, y): + # [xx, yy] = np.dot(self.mR, np.array([x, y])) + inters_1 = -0.5 * self.Ly + self.y_cent + inters_2 = inters_1 + self.Ly + if abs(y - inters_1) < abs(y - inters_2): + return inters_1 + else: + return inters_2 + + # [_, yyy] = np.dot(self.R, np.array([xx, inters])) + # return yyy + +class ImpFunc: + def __init__(self, func): + self.func = func + + def out_conductor(self, x, y): + return self.func(x, y) < 0 + + def in_conductor(self, x, y): + return self.func(x, y) > 0 + + def intersec_x(self, x, y): + func_x = lambda t : self.func(t, y) + + return scipy.optimize.newton_krylov(func_x, x) + + def intersec_y(self, x, y): + func_y = lambda t : self.func(x, t) + + return scipy.optimize.newton_krylov(func_y, y) + +class Plane: + def __init__(self, m_plane, q_plane, tol=0, sign=1): + self.tol = tol # 1e-16 + self.m_plane = m_plane + self.q_plane = q_plane + self.sign = sign + + def in_conductor(self, x, y): + if self.sign == 1: + return y - self.m_plane * x - self.q_plane >= self.tol + elif self.sign == -1: + return y - self.m_plane * x - self.q_plane <= self.tol + else: + print('sign must be + or - 1') + + def out_conductor(self, x, y): + return not self.in_conductor(x, y) + + def intersec_x(self, x, y): + return y / self.m_plane - self.q_plane / self.m_plane + + def intersec_y(self, x, y): + return self.m_plane * x + self.q_plane + + +class InCircle: + def __init__(self, radius, x_cent, y_cent): + self.radius = radius + self.x_cent = x_cent + self.y_cent = y_cent + + def in_conductor(self, x, y): + return np.square(x - self.x_cent) + np.square(y - self.y_cent) <= np.square(self.radius) + + def out_conductor(self, x, y): + return not self.in_conductor(x, y) + + def intersec_x(self, x, y): + if abs(y - self.y_cent) <= self.radius: + inters_1 = np.sqrt(np.square(self.radius) - np.square(y - self.y_cent)) + self.y_cent + inters_2 = -np.sqrt(np.square(self.radius) - np.square(y - self.y_cent)) + self.y_cent + if abs(x - inters_1) < abs(x - inters_2): + return inters_1 + else: + return inters_2 + else: + return np.inf + + def intersec_y(self, x, y): + if abs(x - self.x_cent) <= self.radius: + inters_1 = np.sqrt(np.square(self.radius) - np.square(x - self.x_cent)) + self.x_cent + inters_2 = -np.sqrt(np.square(self.radius) - np.square(x - self.x_cent)) + self.x_cent + if abs(y - inters_1) < abs(y - inters_2): + return inters_1 + else: + return inters_2 + else: + return np.inf + + +class OutCircle: + def __init__(self, radius, x_cent, y_cent): + self.radius = radius + self.x_cent = x_cent + self.y_cent = y_cent + + def in_conductor(self, x, y): + return np.square(x - self.x_cent) + np.square(y - self.y_cent) >= np.square(self.radius) + + def out_conductor(self, x, y): + return not self.in_conductor(x, y) + + def intersec_x(self, x, y): + # if abs(y - self.y_cent) > self.radius: + inters_1 = np.sqrt(np.square(self.radius) - np.square(y - self.y_cent)) + self.x_cent + inters_2 = -np.sqrt(np.square(self.radius) - np.square(y - self.y_cent)) + self.x_cent + if abs(x - inters_1) < abs(x - inters_2): + return inters_1 + else: + return inters_2 + # else: + # return np.inf + + def intersec_y(self, x, y): + # if abs(x - self.x_cent) > self.radius: + inters_1 = np.sqrt(np.square(self.radius) - np.square(x - self.x_cent)) + self.y_cent + inters_2 = -np.sqrt(np.square(self.radius) - np.square(x - self.x_cent)) + self.y_cent + if abs(y - inters_1) < abs(y - inters_2): + return inters_1 + else: + return inters_2 + # else: + # return np.inf + + +class ConductorsAssembly: + def __init__(self, conductors): + self.conductors = conductors + + def in_conductor(self, x, y): + for conductor in self.conductors: + if conductor.in_conductor(x, y): + return True + + return False + + def out_conductor(self, x, y): + return not self.in_conductor(x, y) + + def intersec_x(self, x, y): + list_inters = np.zeros_like(self.conductors) + for ii, conductor in enumerate(self.conductors): + list_inters[ii] = conductor.intersec_x(x, y) + + dist_inters = abs(list_inters - x) + return list_inters[np.argmin(dist_inters)] + + def intersec_y(self, x, y): + list_inters = np.zeros_like(self.conductors) + for ii, conductor in enumerate(self.conductors): + list_inters[ii] = conductor.intersec_y(x, y) + + dist_inters = abs(list_inters - y) + return list_inters[np.argmin(dist_inters)] + + +class noConductor: + def out_conductor(self, x, y): + return True + + def in_conductor(self, x, y): + return False + + def intersec_x(self, x, y): + return 1000 + + def intersec_y(self, x, y): + return 1000 + diff --git a/build/lib/wakis/conductors3d.py b/build/lib/wakis/conductors3d.py new file mode 100644 index 0000000..cbdaf24 --- /dev/null +++ b/build/lib/wakis/conductors3d.py @@ -0,0 +1,223 @@ +import numpy as np + + +class ConductorsAssembly: + def __init__(self, conductors): + self.conductors = conductors + + def in_conductor(self, x, y, z): + for conductor in self.conductors: + if conductor.in_conductor(x, y, z): + return True + + return False + + def out_conductor(self, x, y, z): + return not self.in_conductor(x, y, z) + + def intersec_x(self, x, y, z): + list_inters = np.zeros_like(self.conductors) + for ii, conductor in enumerate(self.conductors): + list_inters[ii] = conductor.intersec_x(x, y, z) + + dist_inters = abs(list_inters - x) + return list_inters[np.argmin(dist_inters)] + + def intersec_y(self, x, y, z): + list_inters = np.zeros_like(self.conductors) + for ii, conductor in enumerate(self.conductors): + list_inters[ii] = conductor.intersec_y(x, y, z) + + dist_inters = abs(list_inters - y) + return list_inters[np.argmin(dist_inters)] + + def intersec_z(self, x, y, z): + list_inters = np.zeros_like(self.conductors) + for ii, conductor in enumerate(self.conductors): + list_inters[ii] = conductor.intersec_z(x, y, z) + + dist_inters = abs(list_inters - z) + return list_inters[np.argmin(dist_inters)] + + +class InCube: + + def __init__(self, lx, ly, lz, x_cent, y_cent, z_cent): + self.lx = lx + self.ly = ly + self.lz = lz + self.x_cent = x_cent + self.y_cent = y_cent + self.z_cent = z_cent + + def out_conductor(self, x, y, z): + return (-0.5 * self.lx + self.x_cent < x < 0.5 * self.lx + self.x_cent and + -0.5 * self.ly + self.y_cent < y < 0.5 * self.ly + self.y_cent and + -0.5 * self.lz + self.z_cent < z < 0.5 * self.lz + self.z_cent) + + def in_conductor(self, x, y, z): + return not self.out_conductor(x, y, z) + + def intersec_x(self, x, y, z): + inters_1 = -0.5 * self.lx + self.x_cent + inters_2 = inters_1 + self.lx + if abs(x - inters_1) < abs(x - inters_2): + return inters_1 + else: + return inters_2 + + def intersec_y(self, x, y, z): + inters_1 = -0.5 * self.ly + self.y_cent + inters_2 = inters_1 + self.ly + if abs(y - inters_1) < abs(y - inters_2): + return inters_1 + else: + return inters_2 + + def intersec_z(self, x, y, z): + inters_1 = -0.5 * self.lz + self.z_cent + inters_2 = inters_1 + self.lz + if abs(z - inters_1) < abs(z - inters_2): + return inters_1 + else: + return inters_2 + + +class InSphere: + + def __init__(self, radius, x_cent, y_cent, z_cent): + self.radius = radius + self.x_cent = x_cent + self.y_cent = y_cent + self.z_cent = z_cent + + def in_conductor(self, x, y, z): + return (np.square(x - self.x_cent) + np.square(y - self.y_cent) + np.square(z - self.z_cent) + <= np.square(self.radius)) + + def out_conductor(self, x, y, z): + return not self.out_conductor(x, y, z) + + def intersec_x(self, x, y, z): + inters_1 = (np.sqrt(np.square(self.radius) - np.square(y - self.y_cent) + - np.square(z - self.z_cent)) + self.x_cent) + inters_2 = -(np.sqrt(np.square(self.radius) - np.square(y - self.y_cent) + - np.square(z - self.z_cent)) + self.x_cent) + + if abs(x - inters_1) < abs(x - inters_2): + return inters_1 + else: + return inters_2 + + def intersec_y(self, x, y, z): + inters_1 = (np.sqrt(np.square(self.radius) - np.square(x - self.x_cent) + - np.square(z - self.z_cent)) + self.y_cent) + inters_2 = -(np.sqrt(np.square(self.radius) - np.square(x - self.x_cent) + - np.square(z - self.z_cent)) + self.y_cent) + if abs(y - inters_1) < abs(y - inters_2): + return inters_1 + else: + return inters_2 + + def intersec_z(self, x, y, z): + inters_1 = (np.sqrt(np.square(self.radius) - np.square(x - self.x_cent) + - np.square(y - self.y_cent)) + self.z_cent) + inters_2 = -(np.sqrt(np.square(self.radius) - np.square(x - self.x_cent) + - np.square(y - self.y_cent)) + self.z_cent) + if abs(z - inters_1) < abs(z - inters_2): + return inters_1 + else: + return inters_2 + + +class noConductor: + def out_conductor(self, x, y, z): + return True + + def in_conductor(self, x, y, z): + return False + + def intersec_x(self, x, y, z): + return 1000 + + def intersec_y(self, x, y, z): + return 1000 + + def intersec_z(self, x, y, z): + return 1000 + + +class Plane: + def __init__(self, p, n): + self.p = p + self.n = n + + def in_conductor(self, x, y, z): + return self.n[0] * (x - self.p[0]) + self.n[1] * (y - self.p[1]) + self.n[2] * (z - self.p[2]) <= 0 + + def out_conductor(self, x, y, z): + return not self.in_conductor(x, y, z) + + def intersec_x(self, x, y, z): + if self.n[0] == 0: + return 1000 + else: + return -(self.n[1] * (y - self.p[1]) + self.n[2] * (z - self.p[2])) / self.n[0] + self.p[0] + + def intersec_y(self, x, y, z): + if self.n[1] == 0: + return 1000 + else: + return -(self.n[0] * (x - self.p[0]) + self.n[2] * (z - self.p[2])) / self.n[1] + self.p[1] + + def intersec_z(self, x, y, z): + if self.n[2] == 0: + return 1000 + else: + return -(self.n[0] * (x - self.p[0]) + self.n[1] * (y - self.p[1])) / self.n[2] + self.p[2] + + +class OutSphere: + def __init__(self, radius, x_cent=0, y_cent=0, z_cent=0): + self.radius = radius + self.x_cent = x_cent + self.y_cent = y_cent + self.z_cent = z_cent + + def in_conductor(self, x, y, z): + return (np.square(x - self.x_cent) + np.square(y - self.y_cent) + np.square(z - self.z_cent) + >= np.square(self.radius)) + + def out_conductor(self, x, y, z): + return not self.in_conductor(x, y, z) + + def intersec_x(self, x, y, z): + inters_1 = (np.sqrt(np.square(self.radius) - np.square(y - self.y_cent) + - np.square(z - self.z_cent)) + self.x_cent) + inters_2 = -(np.sqrt(np.square(self.radius) - np.square(y - self.y_cent) + - np.square(z - self.z_cent)) + self.x_cent) + + if abs(x - inters_1) < abs(x - inters_2): + return inters_1 + else: + return inters_2 + + def intersec_y(self, x, y, z): + inters_1 = (np.sqrt(np.square(self.radius) - np.square(x - self.x_cent) + - np.square(z - self.z_cent)) + self.y_cent) + inters_2 = -(np.sqrt(np.square(self.radius) - np.square(x - self.x_cent) + - np.square(z - self.z_cent)) + self.y_cent) + if abs(y - inters_1) < abs(y - inters_2): + return inters_1 + else: + return inters_2 + + def intersec_z(self, x, y, z): + inters_1 = (np.sqrt(np.square(self.radius) - np.square(x - self.x_cent) + - np.square(y - self.y_cent)) + self.z_cent) + inters_2 = -(np.sqrt(np.square(self.radius) - np.square(x - self.x_cent) + - np.square(y - self.y_cent)) + self.z_cent) + if abs(z - inters_1) < abs(z - inters_2): + return inters_1 + else: + return inters_2 diff --git a/build/lib/wakis/field.py b/build/lib/wakis/field.py new file mode 100644 index 0000000..3dce54b --- /dev/null +++ b/build/lib/wakis/field.py @@ -0,0 +1,667 @@ +# copyright ################################# # +# This file is part of the wakis Package. # +# Copyright (c) CERN, 2024. # +# ########################################### # + + +import numpy as xp +import copy + +try: + import cupy as xp_gpu + imported_cupy = True +except ImportError: + imported_cupy = False + +class Field: + ''' + Class to switch from 3D to collapsed notation by + defining the __getitem__ magic method + + linear numbering: + n = 1 + (i-1) + (j-1)*Nx + (k-1)*Nx*Ny + len(n) = Nx*Ny*Nz + ''' + def __init__(self, Nx, Ny, Nz, dtype=float, + use_ones=False, use_gpu=False): + + self.Nx = Nx + self.Ny = Ny + self.Nz = Nz + self.N = Nx*Ny*Nz + self.dtype = dtype + self.on_gpu = use_gpu + + if self.on_gpu: + if imported_cupy: + self.xp = xp_gpu + else: + print('*** cupy could not be imported, please CUDA check installation') + else: + self.xp = xp + + if use_ones: + self.array = self.xp.ones(self.N*3, dtype=self.dtype, order='F') + else: + self.array = self.xp.zeros(self.N*3, dtype=self.dtype, order='F') + + @property + def field_x(self): + return self.array[0:self.N] + + @property + def field_y(self): + return self.array[self.N: 2*self.N] + + @property + def field_z(self): + return self.array[2*self.N:3*self.N] + + @field_x.setter + def field_x(self, value): + if len(value.shape) > 1: + self.from_matrix(value, 'x') + else: + self.array[0:self.N] = value + + @field_y.setter + def field_y(self, value): + if len(value.shape) > 1: + self.from_matrix(value, 'y') + else: + self.array[self.N: 2*self.N] = value + + @field_z.setter + def field_z(self, value): + if len(value.shape) > 1: + self.from_matrix(value, 'z') + else: + self.array[2*self.N:3*self.N] = value + + def toarray(self): + return self.array + + def fromarray(self, array): + self.array[:] = array + + def to_matrix(self, key): + if key == 0 or key == 'x': + return self.xp.reshape(self.array[0:self.N], (self.Nx, self.Ny, self.Nz), order='F') + if key == 1 or key == 'y': + return self.xp.reshape(self.array[self.N: 2*self.N], (self.Nx, self.Ny, self.Nz), order='F') + if key == 2 or key == 'z': + return self.xp.reshape(self.array[2*self.N:3*self.N], (self.Nx, self.Ny, self.Nz), order='F') + + def from_matrix(self, mat, key): + if key == 0 or key == 'x': + self.array[0:self.N] = self.xp.reshape(mat, self.N, order='F') + elif key == 1 or key == 'y': + self.array[self.N: 2*self.N] = self.xp.reshape(mat, self.N, order='F') + elif key == 2 or key == 'z': + self.array[2*self.N:3*self.N] = self.xp.reshape(mat, self.N, order='F') + else: + raise IndexError('Component id not valid') + + def to_gpu(self): + if imported_cupy: + self.xp = xp_gpu + self.array = self.xp.asarray(self.array) # to cupy arr + self.on_gpu = True + else: + print('*** CuPy is not imported') + pass + + def from_gpu(self): + if self.on_gpu: + self.array = self.array.get() # to numpy arr + self.on_gpu = False + else: + print('*** GPU is not enabled') + pass + + def __getitem__(self, key): + + if type(key) is tuple: + if len(key) != 4: + raise IndexError('Need 3 indexes and component to access the field') + if key[3] == 0 or key[3] == 'x': + if self.on_gpu: + field = self.xp.reshape(self.array[0:self.N], (self.Nx, self.Ny, self.Nz), order='F') + return field[key[0], key[1], key[2]].get() + else: + field = self.xp.reshape(self.array[0:self.N], (self.Nx, self.Ny, self.Nz), order='F') + return field[key[0], key[1], key[2]] + elif key[3] == 1 or key[3] == 'y': + if self.on_gpu: + field = self.xp.reshape(self.array[self.N: 2*self.N], (self.Nx, self.Ny, self.Nz), order='F') + return field[key[0], key[1], key[2]].get() + else: + field = self.xp.reshape(self.array[self.N: 2*self.N], (self.Nx, self.Ny, self.Nz), order='F') + return field[key[0], key[1], key[2]] + elif key[3] == 2 or key[3] == 'z': + if self.on_gpu: + field = self.xp.reshape(self.array[2*self.N:3*self.N], (self.Nx, self.Ny, self.Nz), order='F') + return field[key[0], key[1], key[2]].get() + else: + field = self.xp.reshape(self.array[2*self.N:3*self.N], (self.Nx, self.Ny, self.Nz), order='F') + return field[key[0], key[1], key[2]] + elif type(key[3]) is str and key[3].lower() == 'abs': + field = self.get_abs() + return field[key[0], key[1], key[2]] + else: + raise IndexError('Component id not valid') + + elif type(key) is int: + if key <= self.N: + if self.on_gpu: + return self.array[key].get() + else: + return self.array[key] + else: + raise IndexError('Lexico-graphic index cannot be higher than product of dimensions') + + elif type(key) is slice: + if self.on_gpu: + return self.array[key].get() + else: + return self.array[key] + + else: + raise ValueError('key must be a 3-tuple or an integer') + + def __setitem__(self, key, value): + + if self.on_gpu: + value = self.xp.asarray(value) + + if type(key) is tuple: + if len(key) != 4: + raise IndexError('Need 3 indexes and component to access the field') + else: + field = self.to_matrix(key[3]) + field[key[0], key[1], key[2]] = value + self.from_matrix(field, key[3]) + + elif type(key) is int: + if key <= self.N: + self.array[key] = value + else: + raise IndexError('Lexico-graphic index cannot be higher than product of dimensions') + + elif type(key) is slice: + self.array[key] = value + else: + raise IndexError('key must be a 3-tuple or an integer') + + def __mul__(self, other, dtype=None): + + if dtype is None: + dtype = self.dtype + + # other is number + if type(other) is float or type(other) is int: + mulField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) + mulField.array = self.array * other + + # other is matrix + elif len(other.shape) > 1: + mulField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) + for d in ['x', 'y', 'z']: + mulField.from_matrix(self.to_matrix(d) * other, d) + + # other is 1d array + else: + mulField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) + mulField.array = self.array * other + + return mulField + + def __div__(self, other, dtype=None): + + if dtype is None: + dtype = self.dtype + + # other is number + if type(other) is float or type(other) is int: + divField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) + divField.array = self.array / other + + # other is matrix + if len(other.shape) > 1: + divField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) + for d in ['x', 'y', 'z']: + divField.from_matrix(self.to_matrix(d) / other, d) + + # other is constant or 1d array + else: + divField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) + divField.array = self.array / other + + return divField + + def __add__(self, other, dtype=None): + + if dtype is None: + dtype = self.dtype + + if type(other) is Field: + + addField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) + addField.field_x = self.field_x + other.field_x + addField.field_y = self.field_y + other.field_y + addField.field_z = self.field_z + other.field_z + + # other is matrix + elif len(other.shape) > 1: + addField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) + for d in ['x', 'y', 'z']: + addField.from_matrix(self.to_matrix(d) + other, d) + + # other is constant or 1d array + else: + addField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) + addField.array = self.array + other + + return addField + + def __repr__(self): + return 'x:\n' + self.field_x.__repr__() + '\n'+ \ + 'y:\n' + self.field_y.__repr__() + '\n'+ \ + 'z:\n' + self.field_z.__repr__() + + def __str__(self): + return 'x:\n' + self.field_x.__str__() + '\n'+ \ + 'y:\n' + self.field_y.__str__() + '\n'+ \ + 'z:\n' + self.field_z.__str__() + + def copy(self): + import copy + obj = type(self).__new__(self.__class__) # Create empty instance + + for key, value in self.__dict__.items(): + if key == "xp": + obj.xp = self.xp # Just copy reference, no need for deepcopy + elif key == "array" and self.on_gpu: + obj.array = self.xp.array(self.array) # Ensure CuPy array is copied properly + else: + obj.__dict__[key] = copy.deepcopy(value) + + return obj + + def compute_ijk(self, n): + if n > (self.N): + raise IndexError('Lexico-graphic index cannot be higher than product of dimensions') + + k = n//(self.Nx*self.Ny) + i = (n-k*self.Nx*self.Ny)%self.Nx + j = (n-k*self.Nx*self.Ny)//self.Nx + + return i, j, k + + def get_abs(self, as_matrix=True): + '''Computes the absolute or magnitude + out of the field components + ''' + if as_matrix: + if self.on_gpu: + return xp.sqrt(self.to_matrix('x')**2 + + self.to_matrix('y')**2 + + self.to_matrix('z')**2 ).get() + else: + return xp.sqrt(self.to_matrix('x')**2 + + self.to_matrix('y')**2 + + self.to_matrix('z')**2 ) + + else: # 1d array + if self.on_gpu: + return xp.sqrt(self.field_x**2 + self.field_y**2, self.field_z**2).get() + else: + return xp.sqrt(self.field_x**2 + self.field_y**2, self.field_z**2) + + def inspect(self, plane='YZ', cmap='bwr', dpi=100, figsize=[8,6], x=None, y=None, z=None, show=True, handles=False, **kwargs): + import matplotlib.pyplot as plt + from mpl_toolkits.axes_grid1 import make_axes_locatable + + if None not in (x,y,z): + pass + elif plane == 'XY': + key=[slice(0,self.Nx), slice(0,self.Ny), int(self.Nz//2)] + x, y, z = key[0], key[1], key[2] + extent = (0, self.Nx, 0, self.Ny) + xax, yax = 'nx', 'ny' + transpose = True + + elif plane == 'XZ': + key=[slice(0,self.Nx), int(self.Ny//2), slice(0,self.Nz)] + x, y, z = key[0], key[1], key[2] + extent = (0, self.Nz, 0, self.Nx) + xax, yax = 'nz', 'nx' + transpose = False + + elif plane == 'YZ': + key=[int(self.Nx//2), slice(0,self.Ny), slice(0,self.Nz)] + x, y, z = key[0], key[1], key[2] + extent = (0, self.Nz, 0, self.Ny) + xax, yax = 'nz', 'ny' + transpose = False + + fig, axs = plt.subplots(1, 3, tight_layout=True, figsize=figsize, dpi=dpi) + dims = {0:'x', 1:'y', 2:'z'} + + im = {} + + for d in [0,1,2]: + field = self.to_matrix(d) + + if self.on_gpu and hasattr(field, 'get'): + field = field.get() + + if transpose: + im[d] = axs[d].imshow(field[x,y,z].T, cmap=cmap, vmin=-field.max(), vmax=field.max(), extent=extent, origin='lower', **kwargs) + + else: + im[d] = axs[d].imshow(field[x,y,z], cmap=cmap, vmin=-field.max(), vmax=field.max(), extent=extent, origin='lower', **kwargs) + + for i, ax in enumerate(axs): + ax.set_title(f'Field {dims[i]}, plane {plane}') + fig.colorbar(im[i], cax=make_axes_locatable(ax).append_axes('right', size='5%', pad=0.1)) + ax.set_xlabel(xax) + ax.set_ylabel(yax) + + if handles: + return fig, axs + + if show: + plt.show() + + def inspect3D(self, field='all', backend='pyista', grid=None, + xmax=None, ymax=None, zmax=None, + bounding_box=True, show_grid=True, + cmap='viridis', dpi=100, show=True, handles=False): + """ + Visualize 3D field data on the structured grid using either Matplotlib + (voxel rendering) or PyVista (interactive clipping and slicing). + + This method provides two complementary visualization backends: + - **Matplotlib**: static voxel plots of the field components (x, y, z) + or all combined, useful for quick inspection, but memory intensive. + - **PyVista**: interactive 3D visualization with sliders to dynamically + clip the volume along X, Y, and Z, and optional wireframe slices. + + Parameters + ---------- + field : {'x', 'y', 'z', 'all'}, default 'all' + Which field component(s) to visualize. + - 'x', 'y', 'z': single component + - 'all': shows all three components + backend : {'matplotlib', 'pyvista'}, default 'pyvista' + Visualization backend to use: + - 'matplotlib': static voxel rendering + - 'pyvista': interactive 3D rendering with clipping sliders + grid : GridFIT3D or pyvista.StructuredGrid, optional + Structured grid object to use for visualization. If None, + a grid is constructed from the solver's internal dimensions. + xmax, ymax, zmax : int or float, optional + Maximum extents in each direction for visualization. Defaults + to the full grid dimensions if not specified. + bounding_box : bool, default True + If True, draw a wireframe bounding box of the simulation domain + (only used in PyVista backend). + show_grid : bool, default True + If True, show wireframe slice planes of the grid during + interactive visualization (PyVista backend). + cmap : str, default 'viridis' + Colormap to apply to the scalar field. + dpi : int, default 100 + Resolution of Matplotlib figures (only for Matplotlib backend). + show : bool, default True + Whether to display the figure/plot immediately. + - If False in PyVista, exports to `field.html` instead. + handles : bool, default False + If True, return figure/axes (Matplotlib) or the Plotter object + (PyVista) for further customization instead of showing directly. + + Returns + ------- + fig, axs : tuple, optional + Returned when `backend='matplotlib'` and `handles=True`. + pl : pyvista.Plotter, optional + Returned when `backend='pyvista'` and `handles=True`. + + Notes + ----- + - The PyVista backend provides interactive sliders to clip the + volume along each axis independently and inspect internal + structures of the 3D field. + - The Matplotlib backend provides a quick static voxel rendering + but is limited in interactivity and scalability. + + """ + + field = field.lower() + + # ---------- matplotlib backend --------------- + if backend.lower() == 'matplotlib': + import matplotlib.pyplot as plt + import matplotlib as mpl + from mpl_toolkits.axes_grid1 import make_axes_locatable + + fig = plt.figure(tight_layout=True, dpi=dpi, figsize=[12,6]) + + plot_x, plot_y, plot_z = False, False, False + + if field == 'all': + plot_x = True + plot_y = True + plot_z = True + + elif field.lower() == 'x': plot_x = True + elif field.lower() == 'y': plot_y = True + elif field.lower() == 'z': plot_z = True + + if xmax is None: xmax = self.Nx + if ymax is None: ymax = self.Ny + if zmax is None: zmax = self.Nz + + x,y,z = self.xp.mgrid[0:xmax+1,0:ymax+1,0:zmax+1] + axs = [] + + # field x + if plot_x: + arr = self.to_matrix('x')[0:int(xmax),0:int(ymax),0:int(zmax)] + if field == 'all': + ax = fig.add_subplot(1, 3, 1, projection='3d') + else: + ax = fig.add_subplot(1, 1, 1, projection='3d') + + vmin, vmax = -self.xp.max(self.xp.abs(arr)), +self.xp.max(self.xp.abs(arr)) + norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax) + colors = mpl.colormaps[cmap](norm(arr)) + vox = ax.voxels(x, y, z, filled=self.xp.ones_like(arr), facecolors=colors) + + m = mpl.cm.ScalarMappable(cmap=cmap, norm=norm) + m.set_array([]) + fig.colorbar(m, ax=ax, shrink=0.5, aspect=10) + ax.set_title(f'Field x') + axs.append(ax) + + # field y + if plot_y: + arr = self.to_matrix('y')[0:int(xmax),0:int(ymax),0:int(zmax)] + if field == 'all': + ax = fig.add_subplot(1, 3, 2, projection='3d') + else: + ax = fig.add_subplot(1, 1, 1, projection='3d') + + vmin, vmax = -self.xp.max(self.xp.abs(arr)), +self.xp.max(self.xp.abs(arr)) + norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax) + colors = mpl.colormaps[cmap](norm(arr)) + vox = ax.voxels(x, y, z, filled=self.xp.ones_like(arr), facecolors=colors) + + m = mpl.cm.ScalarMappable(cmap=cmap, norm=norm) + m.set_array([]) + fig.colorbar(m, ax=ax, shrink=0.5, aspect=10) + ax.set_title(f'Field y') + axs.append(ax) + + # field z + if plot_z: + arr = self.to_matrix('z')[0:int(xmax),0:int(ymax),0:int(zmax)] + if field == 'all': + ax = fig.add_subplot(1, 3, 3, projection='3d') + else: + ax = fig.add_subplot(1, 1, 1, projection='3d') + + vmin, vmax = -self.xp.max(self.xp.abs(arr)), +self.xp.max(self.xp.abs(arr)) + norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax) + colors = mpl.colormaps[cmap](norm(arr)) + vox = ax.voxels(x, y, z, filled=self.xp.ones_like(arr), facecolors=colors) + + m = mpl.cm.ScalarMappable(cmap=cmap, norm=norm) + m.set_array([]) + fig.colorbar(m, ax=ax, shrink=0.5, aspect=10) + ax.set_title(f'Field z') + axs.append(ax) + + dims = {0:'x', 1:'y', 2:'z'} + for i, ax in enumerate(axs): + ax.set_xlabel('Nx') + ax.set_ylabel('Ny') + ax.set_zlabel('Nz') + ax.xaxis.set_pane_color((1.0, 1.0, 1.0, 0.0)) + ax.yaxis.set_pane_color((1.0, 1.0, 1.0, 0.0)) + ax.zaxis.set_pane_color((1.0, 1.0, 1.0, 0.0)) + ax.set_xlim(self.Nx, 0) + ax.set_ylim(self.Ny, 0) + ax.set_zlim(self.Nz, 0) + + if handles: + return fig, axs + + if show: + plt.show() + + # ----------- pyvista backend --------------- + else: + import pyvista as pv + + if grid is not None and hasattr(grid, 'grid'): + xlo, xhi, ylo, yhi, zlo, zhi = grid.xmin, grid.xmax, grid.ymin, grid.ymax, grid.zmin, grid.zmax + grid = grid.grid + if field == 'x': + scalars = 'Field '+field + grid[scalars] = xp.reshape(self.to_matrix(field), self.N) + elif field == 'y': + scalars = 'Field '+field + grid[scalars] = xp.reshape(self.to_matrix(field), self.N) + elif field == 'z': + scalars = 'Field '+field + grid[scalars] = xp.reshape(self.to_matrix(field), self.N) + else: # for all or abs + scalars = 'Field '+'Abs' + grid[scalars] = xp.reshape(self.get_abs(), self.N) + + if xmax is None: xmax = xhi + if ymax is None: ymax = yhi + if zmax is None: zmax = zhi + + else: + print('[!] `grid` is not passed or is not a GridFIT3D object -> Using #N cells instead ') + x = xp.linspace(0, self.Nx, self.Nx+1) + y = xp.linspace(0, self.Ny, self.Ny+1) + z = xp.linspace(0, self.Nz, self.Nz+1) + xlo, xhi, ylo, yhi, zlo, zhi = x.min(), x.max(), y.min(), y.max(), z.min(), z.max() + if xmax is None: xmax = self.Nx + if ymax is None: ymax = self.Ny + if zmax is None: zmax = self.Nz + X, Y, Z = xp.meshgrid(x, y, z, indexing='ij') + grid = pv.StructuredGrid(X.transpose(), Y.transpose(), Z.transpose()) + + if field == 'x': + scalars = 'Field '+field + grid[scalars] = xp.reshape(self.to_matrix(field), self.N) + elif field == 'y': + scalars = 'Field '+field + grid[scalars] = xp.reshape(self.to_matrix(field), self.N) + elif field == 'z': + scalars = 'Field '+field + grid[scalars] = xp.reshape(self.to_matrix(field), self.N) + else: # for all or abs + scalars = 'Field '+'Abs' + grid[scalars] = xp.reshape(self.get_abs(), self.N) + + + pv.global_theme.allow_empty_mesh = True + pl = pv.Plotter() + vals = {'x':xmax, 'y':ymax, 'z':zmax} + + # --- Update function --- + def update_clip(val, axis="x"): + vals[axis] = val + # define bounds dynamically + if axis == "x": + slice_obj = grid.slice(normal="x", origin=(val, 0, 0)) + elif axis == "y": + slice_obj = grid.slice(normal="y", origin=(0, val, 0)) + else: # z + slice_obj = grid.slice(normal="z", origin=(0, 0, val)) + + # add clipped volume (scalars) + pl.add_mesh( + grid.clip_box(bounds=(xlo, vals['x'], ylo, vals['y'], zlo, vals['z']), invert=False), + scalars=scalars, + cmap=cmap, + name="clip", + ) + + # add slice wireframe (grid structure) + if show_grid: + pl.add_mesh(slice_obj, style="wireframe", color="grey", name="slice") + + # --- Sliders (placed side-by-side vertically) --- + pl.add_slider_widget( + lambda value: update_clip(value, "x"), + [xlo, xhi], + value=xmax, title="X Clip", + pointa=(0.8, 0.8), pointb=(0.95, 0.8), # top-right + style='modern', + ) + + pl.add_slider_widget( + lambda value: update_clip(value, "y"), + [ylo, yhi], + value=ymax, title="Y Clip", + pointa=(0.8, 0.6), pointb=(0.95, 0.6), # middle-right + style='modern', + ) + + pl.add_slider_widget( + lambda value: update_clip(value, "z"), + [zlo, zhi], + value=zmax, title="Z Clip", + pointa=(0.8, 0.4), pointb=(0.95, 0.4), # lower-right + style='modern', + ) + + # Camera orientation + pl.camera_position = 'zx' + pl.camera.azimuth += 30 + pl.camera.elevation += 30 + pl.set_background('mistyrose', top='white') + try: pl.add_logo_widget('../docs/img/wakis-logo-pink.png') + except: pass + pl.add_axes() + pl.enable_3_lights() + pl.enable_anti_aliasing() + + if bounding_box: + pl.add_mesh(pv.Box(bounds=(xlo, xhi, ylo, yhi, zlo, zhi)), + style="wireframe", color="black", line_width=2, name="domain_box") + + if handles: + return pl + + if not show: + pl.export_html(f'field.html') + else: + pl.show() diff --git a/build/lib/wakis/geometry.py b/build/lib/wakis/geometry.py new file mode 100644 index 0000000..3882458 --- /dev/null +++ b/build/lib/wakis/geometry.py @@ -0,0 +1,254 @@ +# copyright ################################# # +# This file is part of the wakis Package. # +# Copyright (c) CERN, 2024. # +# ########################################### # + +import numpy as np +import re + +def extract_colors_from_stp(stp_file): + """ + Extracts a mapping from solid names to RGB color values from a STEP (.stp) file. + + Args: + stp_file (str): Path to the STEP file. + + Returns: + dict[str, list[float]]: A dictionary mapping solid names to [R, G, B] colors. + """ + solids, _ = extract_names_from_stp(stp_file) + + colors = [] + stl_colors = {} + + color_pattern = re.compile(r"#\d+=COLOUR_RGB\('?',([\d.]+),([\d.]+),([\d.]+)\);") + + # Extract colors + with open(stp_file, 'r', encoding='utf-8', errors='ignore') as f: + for line in f: + color_match = color_pattern.search(line) + if color_match: + r = float(color_match.group(1)) + g = float(color_match.group(2)) + b = float(color_match.group(3)) + colors.append([r, g, b]) + + # Map solids to colors by order of appearance (colors >=solids) + for i in range(len(list(solids.keys()))): + solid = solids[list(solids.keys())[i]] + solid_re = re.sub(r'[^a-zA-Z0-9_-]', '-', solid) + stl_colors[f'{str(i).zfill(3)}_{solid_re}'] = colors[i] + + return stl_colors + +def extract_materials_from_stp(stp_file): + """ + Extracts a mapping from solid names to materials from a STEP (.stp) file. + + Args: + stp_file (str): Path to the STEP file. + + Returns: + dict[str, str]: A dictionary mapping solid names to material names. + """ + + solids, materials = extract_names_from_stp(stp_file) + stl_materials = {} + for i in range(len(list(solids.keys()))): + solid = solids[list(solids.keys())[i]] + try: + mat = materials[list(solids.keys())[i]].lower() + except: + print(f'Solid #{list(solids.keys())[i]} has no assigned material') + mat = 'None' + + # Remove problematic characters + solid_re = re.sub(r'[^a-zA-Z0-9_-]', '-', solid) + mat_re = re.sub(r'[^a-zA-Z0-9_-]', '-', mat) + stl_materials[f'{str(i).zfill(3)}_{solid_re}'] = mat_re + + return stl_materials + +def extract_solids_from_stp(stp_file, path=''): + if path and not path.endswith('/'): + path += '/' + solids, materials = extract_names_from_stp(stp_file) + stl_solids = {} + for i in range(len(list(solids.keys()))): + solid = solids[list(solids.keys())[i]] + try: + mat = materials[list(solids.keys())[i]] + except: + print(f'Solid #{list(solids.keys())[i]} has no assigned material') + mat = 'None' + + # Remove problematic characters + solid_re = re.sub(r'[^a-zA-Z0-9_-]', '-', solid) + mat_re = re.sub(r'[^a-zA-Z0-9_-]', '-', mat) + name = f'{str(i).zfill(3)}_{solid_re}_{mat_re}' + stl_solids[f'{str(i).zfill(3)}_{solid_re}'] = path + name + '.stl' + + return stl_solids + +def extract_names_from_stp(stp_file): + """ + Extracts solid names and their corresponding materials from a STEP (.stp) file. + + This function parses a given STEP file to identify solid objects and their assigned materials. + The solid names are extracted from `MANIFOLD_SOLID_BREP` statements, while the materials are + linked via `PRESENTATION_LAYER_ASSIGNMENT` statements. + + Args: + stp_file (str): Path to the STEP (.stp) file. + + Returns: + tuple[dict[int, str], dict[int, str]]: + - A dictionary mapping solid IDs to their names. + - A dictionary mapping solid IDs to their corresponding material names. + + Example: + >>> solids, materials = extract_names_from_stp("example.stp") + >>> print(solids) + {37: "Vacuum|Half_cell_dx", 39: "Be window left"} + >>> print(materials) + {37: "Vacuum", 39: "Berillium"} + """ + solid_dict = {} + material_dict = {} + + # Compile regex patterns + #solid_pattern = re.compile(r"#(\d+)=MANIFOLD_SOLID_BREP\('([^']+)'.*;") + solid_pattern = re.compile(r"#(\d+)=ADVANCED_BREP_SHAPE_REPRESENTATION\('([^']*)',\(([^)]*)\),#\d+\);") + material_pattern = re.compile(r"#(\d+)=PRESENTATION_LAYER_ASSIGNMENT\('([^']+)','[^']+',\(#([\d#,]+)\)\);") + + # First pass: extract solids + with open(stp_file, 'r', encoding='utf-8', errors='ignore') as f: + for line in f: + solid_match = solid_pattern.search(line) + if solid_match: + #solid_number = int(solid_match.group(1)) #if MANIFOLD + solid_number = int(solid_match.group(3).split(',')[0].strip().lstrip('#')) + solid_name = solid_match.group(2) + solid_dict[solid_number] = solid_name + + # Second pass: extract materials + with open(stp_file, 'r', encoding='utf-8', errors='ignore') as f: + for line in f: + material_match = material_pattern.search(line) + if material_match: + material_name = material_match.group(2) + solid_numbers = [int(num.strip("#")) for num in material_match.group(3).split(',')] + for solid_number in solid_numbers: + if solid_number in solid_dict: + material_dict[solid_number] = material_name + + return solid_dict, material_dict + +def get_stp_unit_scale(stp_file): + """ + Reads the unit definition from a STEP (.stp or .step) file and determines the + scale factor required to convert the geometry to meters. + + This function: + - Opens and scans the header section of the STEP file. + - Detects the SI base unit definition (e.g., millimeter, centimeter, meter). + - Returns a scale factor to convert the geometry to meters. + - Handles missing or unreadable unit information gracefully. + + Args: + stp_file (str): Path to the STEP (.stp or .step) file. + + Returns: + float: Scale factor to convert STEP geometry to meters. + Defaults to 1.0 if no valid unit information is found. + """ + + unit_map = { + ".MILLI.": 1e-3, + ".CENTI.": 1e-2, + ".DECI.": 1e-1, + ".KILO.": 1e3, + "$": 1.0, # '$' indicates no prefix, i.e. plain meters + } + + try: + with open(stp_file, "r", encoding="utf-8", errors="ignore") as f: + header = f.read(10000) # read only the beginning of the file + + match = re.search( + r"SI_UNIT\s*\(\s*(\.\w+\.)?\s*,\s*\.METRE\.\s*\)", + header, + re.IGNORECASE, + ) + + if match: + prefix = match.group(1).upper() if match.group(1) else "$" + scale = unit_map.get(prefix, 1.0) + print(f"Detected STEP unit: {prefix} → scale to meters: {scale}") + return scale + else: + print("No unit found, files remain in original unit.") + return 1.0 + + except Exception as exc: + print(f"Error reading unit from STEP file: {exc}") + print("Files remain in original unit.") + + return 1.0 + +def generate_stl_solids_from_stp(stp_file, results_path=''): + """ + Extracts solid objects from a STEP (.stp) file and exports them as STL files. + + This function: + - Imports the STEP file using `cadquery`. + - Extracts solid names and their materials using `extract_names_from_stp()`. + - Sanitizes solid and material names by replacing problematic characters. + - Scales the solid to meter using `get_stp_unit_scale()`. + - Saves each solid as an STL file in the current folder (default) or the given path. + + Args: + stp_file (str): Path to the STEP (.stp) file. + results_path (str) (optional): default: '', path to save the STL (.stl) files + + Raises: + Exception: If `cadquery` is not installed, it prompts the user to install it. + + Example: + >>> extract_stl_solids_from_stp("example.stp") + 000_Vacuum-Half_cell_dx_Vacuum.stl + 001_Be_window_left_Berillium.stl + """ + + try: + import cadquery as cq + except: + raise Exception('''This function needs the open-source package `cadquery` + To install it in a conda environment do: + + `pip install cadquery` + + [!] We recommend having a dedicated conda environment to avoid version issues + ''') + + stp = cq.importers.importStep(stp_file) + + scale_factor = get_stp_unit_scale(stp_file) + if scale_factor != 1.0: + print(f"Scaling geometry to meters (factor={scale_factor}).") + scaled_solids = [solid.scale(scale_factor) for solid in stp.objects[0]] + stp.objects = [scaled_solids] + + solid_dict = extract_solids_from_stp(stp_file) + + if not results_path.endswith('/'): + results_path += '/' + + print(f"Generating stl from file: {stp_file}... ") + for i, obj in enumerate(stp.objects[0]): + name = solid_dict[list(solid_dict.keys())[i]] + print(name) + obj.exportStl(results_path+name) + + return solid_dict + diff --git a/build/lib/wakis/grid2D.py b/build/lib/wakis/grid2D.py new file mode 100644 index 0000000..6d8e8e8 --- /dev/null +++ b/build/lib/wakis/grid2D.py @@ -0,0 +1,595 @@ +import numpy as np +from numba import jit +import numba + +def seg_length(x_1, y_1, x_2, y_2): + return np.linalg.norm(np.array([x_1 - x_2, y_1 - y_2])) + +@jit('b1(f8, f8)', nopython=True) +def eq(a, b): + tol = 1e-8 + return abs(a - b) < tol + +@jit('b1(f8, f8)', nopython=True) +def neq(a, b): + tol = 1e-8 + return not eq(a, b) + + +class Grid2D: + + """ + Class holding the grid info and the routines for cell extensions. + Constructor arguments: + - xmin,xmax,ymin,ymax: extent of the domain. + - nx, ny: number of cells per direction + - conductors: conductor object + - sol_type: type of solver. 'FDTD' for staircased FDTD, 'DM' for Conformal Dey-Mittra FDTD, + 'ECT' for Extended Cell Technique conformal FDTD + """ + + def __init__(self, xmin, xmax, ymin, ymax, nx, ny, conductors, sol_type): + self.xmin = xmin + self.xmax = xmax + self.ymin = ymin + self.ymax = ymax + self.nx = nx + self.ny = ny + self.dx = (xmax - xmin) / nx + self.dy = (ymax - ymin) / ny + self.conductors = conductors + + self.l_x = np.zeros((nx, ny + 1)) + self.l_y = np.zeros((nx + 1, ny)) + self.S = np.zeros((nx, ny)) + self.S_stab = np.zeros((nx, ny)) + self.S_enl = np.zeros((nx, ny)) + self.S_red = np.zeros((nx, ny)) + self.flag_unst_cell = np.zeros((nx, ny), dtype=bool) + self.flag_int_cell = np.zeros((nx, ny), dtype=bool) + self.flag_bound_cell = np.zeros((nx, ny), dtype=bool) + self.flag_avail_cell = np.zeros((nx, ny), dtype=bool) + self.flag_ext_cell = np.zeros((nx, ny), dtype=bool) + self.flag_intr_cell = np.zeros((nx, ny), dtype=bool) + self.broken = np.zeros((self.nx, self.ny), dtype=bool) + + if (sol_type is not 'FDTD') and (sol_type is not 'DM') and (sol_type is not 'ECT'): + raise ValueError("sol_type must be:\n" + + "\t'FDTD' for standard staircased FDTD\n" + + "\t'DM' for Dey-Mittra conformal FDTD\n" + + "\t'ECT' for Enlarged Cell Technique conformal FDTD") + + if sol_type is 'DM' or sol_type is 'FDTD': + self.compute_edges(in_conductor=self.conductors.in_conductor, + intersec_x=self.conductors.intersec_x, + intersec_y=self.conductors.intersec_y, l_x=self.l_x, l_y=self.l_y, + dx=self.dx, dy=self.dy, nx=self.nx, ny=self.ny, xmin=self.xmin, + ymin=self.ymin) + compute_areas(l_x=self.l_x, l_y=self.l_y, S=self.S, S_red=self.S_red, nx=self.nx, + ny=self.ny, dx=self.dx, dy=self.dy) + mark_cells(l_x=self.l_x, l_y=self.l_y, nx=self.nx, ny=self.ny, dx=self.dx, + dy=self.dy, S=self.S, flag_int_cell=self.flag_int_cell, + S_stab=self.S_stab, flag_unst_cell=self.flag_unst_cell, + flag_bound_cell=self.flag_bound_cell, + flag_avail_cell=self.flag_avail_cell) + elif sol_type is 'ECT': + self.compute_edges(in_conductor=self.conductors.in_conductor, + intersec_x=self.conductors.intersec_x, + intersec_y=self.conductors.intersec_y, l_x=self.l_x, l_y=self.l_y, + dx=self.dx, dy=self.dy, nx=self.nx, ny=self.ny, xmin=self.xmin, + ymin=self.ymin) + compute_areas(l_x=self.l_x, l_y=self.l_y, S=self.S, S_red=self.S_red, nx=self.nx, + ny=self.ny, dx=self.dx, dy=self.dy) + mark_cells(l_x=self.l_x, l_y=self.l_y, nx=self.nx, ny=self.ny, dx=self.dx, + dy=self.dy, S=self.S, flag_int_cell=self.flag_int_cell, + S_stab=self.S_stab, flag_unst_cell=self.flag_unst_cell, + flag_bound_cell=self.flag_bound_cell, + flag_avail_cell=self.flag_avail_cell) + # info about intruded cells (i,j,[(i_borrowing,j_borrowing,area_borrowing, )]) + self.borrowing = np.empty((nx, ny), dtype=object) + + for i in range(nx): + for j in range(ny): + self.borrowing[i, j] = [] + self.flag_ext_cell = self.flag_unst_cell.copy() + self.compute_extensions(nx=nx, ny=ny, S=self.S, flag_int_cell=self.flag_int_cell, + S_stab=self.S_stab, S_enl=self.S_enl, S_red=self.S_red, + flag_unst_cell=self.flag_unst_cell, + flag_avail_cell=self.flag_avail_cell, + flag_ext_cell=self.flag_ext_cell, + flag_intr_cell=self.flag_intr_cell, + borrowing=self.borrowing) + + """ + Function to compute the length of the edges of the conformal grid. + Inputs: + - tol: an edge shorter than tol will be considered as of zero length + """ + + @staticmethod + def compute_edges(tol=1e-8, in_conductor=None, intersec_x=None, intersec_y=None, l_x=None, + l_y=None, dx=None, dy=None, nx=None, ny=None, xmin=None, ymin=None): + """ + Notation: + + (x_3, y_3)----- l_x[i, j + 1] --------- + | | + | | + l_y[i, j] (i, j) l_y[i + 1, j] + | | + | | + (x_1, y_1)------- l_x[i, j] -------(x_2, y_2) + + """ + for ii in range(nx): + for jj in range(ny + 1): + x_1 = ii * dx + xmin + y_1 = jj * dy + ymin + x_2 = (ii + 1) * dx + xmin + y_2 = jj * dy + ymin + # if point 1 is in conductor + if in_conductor(x_1, y_1): + # if point 2 is in conductor, length of the l_x[i, j] is zero + if in_conductor(x_2, y_2): + l_x[ii, jj] = 0 + # if point 2 is not in conductor, length of l_x[i, j] is the fractional length + else: + l_x[ii, jj] = seg_length(intersec_x(x_2, y_2), y_1, x_2, y_2) + # if point 1 is not in conductor + else: + # if point 2 is in conductor, length of l_x[i, j] is the fractional length + if in_conductor(x_2, y_2): + l_x[ii, jj] = seg_length(x_1, y_1, intersec_x(x_1, y_1), y_2) + # if point 2 is not in conductor, length of l_x[i, j] is dx + else: + l_x[ii, jj] = dx + + for ii in range(nx + 1): + for jj in range(ny): + x_1 = ii * dx + xmin + y_1 = jj * dy + ymin + x_3 = ii * dx + xmin + y_3 = (jj + 1) * dy + ymin + # if point 1 is in conductor + if in_conductor(x_1, y_1): + # if point 3 to the right is in conductor, length of the l_y[i, j] is zero + if in_conductor(x_3, y_3): + l_y[ii, jj] = 0 + # if point 3 is not in conductor, length of l_y[i, j] is the fractional length + else: + l_y[ii, jj] = seg_length(x_1, intersec_y(x_3, y_3), x_3, y_3) + # if point 1 is not in conductor + else: + # if point 3 is in conductor, length of the l_y[i, j] is the fractional length + if in_conductor(x_3, y_3): + l_y[ii, jj] = seg_length(x_1, y_1, x_3, intersec_y(x_1, y_1)) + # if point 3 is not in conductor, length of l_y[i, j] is dy + else: + l_y[ii, jj] = dy + + # set to zero the length of very small cells + if tol > 0.: + l_x /= dx + l_y /= dy + + low_values_flags = abs(l_x) < tol + high_values_flags = abs(l_x - 1.) < tol + l_x[low_values_flags] = 0 + l_x[high_values_flags] = 1. + + low_values_flags = abs(l_y) < tol + high_values_flags = abs(l_y - 1.) < tol + l_y[low_values_flags] = 0 + l_y[high_values_flags] = 1. + + l_x *= dx + l_y *= dy + + """ + Function to compute the extension of the unstable cells + """ + + @staticmethod + def compute_extensions(nx=None, ny=None, S=None, flag_int_cell=None, + S_stab=None, S_enl=None, S_red=None, + flag_unst_cell=None, + flag_avail_cell=None, + flag_ext_cell=None, + flag_intr_cell = None, + borrowing=None, l_verbose=True, + kk=0): + + N = np.sum(flag_ext_cell) + + if l_verbose: + print(kk) + print('ext cells: %d' % N) + # Do the simple one-cell extension + Grid2D._compute_extensions_one_cell(nx=nx, ny=ny, S=S, S_stab=S_stab, S_enl=S_enl, + S_red=S_red, flag_avail_cell=flag_avail_cell, + flag_ext_cell=flag_ext_cell, + flag_intr_cell=flag_intr_cell, + borrowing=borrowing) + + N_one_cell = (N - np.sum(flag_ext_cell)) + if l_verbose: + print('one cell exts: %d' % N_one_cell) + # If any cell could not be extended do the four-cell extension + ''' + if np.sum(flag_ext_cell) > 0: + N = np.sum(flag_ext_cell) + Grid2D._compute_extensions_four_cells(nx=nx, ny=ny, S=S, flag_int_cell=flag_int_cell, + S_stab=S_stab, S_enl=S_enl, S_red=S_red, + flag_unst_cell=flag_unst_cell, + flag_avail_cell=flag_avail_cell, + flag_ext_cell=flag_ext_cell, + flag_intr_cell = flag_intr_cell, + borrowing=borrowing) + N_four_cells = (N - np.sum(flag_ext_cell)) + if 1: #l_verbose: + print('four cell exts: %d' % N_four_cells) + ''' + # If any cell could not be extended do the eight-cell extension + if np.sum(flag_ext_cell) > 0: + N = np.sum(flag_ext_cell) + Grid2D._compute_extensions_eight_cells(nx=nx, ny=ny, S=S, flag_int_cell=flag_int_cell, + S_stab=S_stab, S_enl=S_enl, S_red=S_red, + flag_unst_cell=flag_unst_cell, + flag_avail_cell=flag_avail_cell, + flag_ext_cell=flag_ext_cell, + flag_intr_cell=flag_intr_cell, + borrowing=borrowing) + N_eight_cells = (N - np.sum(flag_ext_cell)) + if 1: #l_verbose: + print('eight cell exts: %d' % N_eight_cells) + # If any cell could not be extended the algorithm failed + if np.sum(flag_ext_cell) > 0: + N = (np.sum(flag_ext_cell)) + raise RuntimeError(str(N) + ' cells could not be extended.\n' + + 'Please refine the mesh') + + """ + Function to compute the one-cell extension of the unstable cells + """ + + @staticmethod + def _compute_extensions_one_cell(nx=None, ny=None, S=None, + S_stab=None, S_enl=None, S_red=None, flag_avail_cell=None, + flag_ext_cell=None, flag_intr_cell=None, borrowing=None): + + for ii in range(0, nx): + for jj in range(0, ny): + if flag_ext_cell[ii, jj]: + S_ext = S_stab[ii, jj] - S[ii, jj] + if (S[ii - 1, jj] > S_ext and flag_avail_cell[ii - 1, jj]): + denom = S[ii - 1, jj] + patch = S_ext * S[ii - 1, jj] / denom + if S_red[ii - 1, jj] - patch > 0: + S_red[ii - 1, jj] -= patch + borrowing[ii, jj].append([ii - 1, jj, patch, None]) + flag_intr_cell[ii-1, jj] = True + S_enl[ii, jj] = S[ii, jj] + patch + flag_ext_cell[ii, jj] = False + if (S[ii, jj - 1] > S_ext and flag_avail_cell[ii, jj - 1] and flag_ext_cell[ii, jj]): + denom = S[ii, jj - 1] + patch = S_ext * S[ii, jj - 1] / denom + if S_red[ii, jj - 1] - patch > 0: + S_red[ii, jj - 1] -= patch + borrowing[ii, jj].append([ii, jj - 1, patch, None]) + flag_intr_cell[ii, jj-1] = True + S_enl[ii, jj] = S[ii, jj] + patch + flag_ext_cell[ii, jj] = False + if (S[ii, jj + 1] > S_ext and flag_avail_cell[ii, jj + 1] + and flag_ext_cell[ii, jj]): + denom = S[ii, jj + 1] + patch = S_ext * S[ii, jj + 1] / denom + if S_red[ii, jj + 1] - patch > 0: + S_red[ii, jj + 1] -= patch + borrowing[ii, jj].append([ii, jj + 1, patch, None]) + flag_intr_cell[ii, jj+1] = True + S_enl[ii, jj] = S[ii, jj] + patch + flag_ext_cell[ii, jj] = False + if (S[ii + 1, jj] > S_ext and flag_avail_cell[ii + 1, jj] + and flag_ext_cell[ii, jj]): + denom = S[ii + 1, jj] + patch = S_ext * S[ii + 1, jj] / denom + if S_red[ii + 1, jj] - patch > 0: + S_red[ii + 1, jj] -= patch + borrowing[ii, jj].append([ii + 1, jj, patch, None]) + flag_intr_cell[ii+1, jj] = True + S_enl[ii, jj] = S[ii, jj] + patch + flag_ext_cell[ii, jj] = False + + + """ + Function to compute the four-cell extension of the unstable cells + """ + + @staticmethod + def _compute_extensions_four_cells(nx=None, ny=None, S=None, flag_int_cell=None, + S_stab=None, S_enl=None, S_red=None, flag_unst_cell=None, + flag_avail_cell=None, flag_ext_cell=None, + flag_intr_cell=None, borrowing=None): + for ii in range(0, nx): + for jj in range(0, ny): + local_avail = flag_avail_cell.copy() + if (flag_unst_cell[ii, jj] and flag_int_cell[ii, jj] + and flag_ext_cell[ii, jj]): + denom = ((flag_avail_cell[ii - 1, jj]) * S[ii - 1, jj] + ( + flag_avail_cell[ii + 1, jj]) * S[ii + 1, jj] + + (flag_avail_cell[ii, jj - 1]) * S[ii, jj - 1] + ( + flag_avail_cell[ii, jj + 1]) * S[ii, jj + 1]) + S_ext = S_stab[ii, jj] - S[ii, jj] + neg_cell = True + # idea: if any cell would reach negative area it is locally not available. + # then denom has to be recomputed from scratch + + while denom >= S_ext and neg_cell: + neg_cell = False + if local_avail[ii - 1, jj]: + patch = S_ext * S[ii - 1, jj] / denom + if S_red[ii - 1, jj] - patch <= 0: + neg_cell = True + local_avail[ii - 1, jj] = False + if local_avail[ii + 1, jj]: + patch = S_ext * S[ii + 1, jj] / denom + if S_red[ii + 1, jj] - patch <= 0: + neg_cell = True + local_avail[ii + 1, jj] = False + if local_avail[ii, jj - 1]: + patch = S_ext * S[ii, jj - 1] / denom + if S_red[ii, jj - 1] - patch <= 0: + neg_cell = True + local_avail[ii, jj - 1] = False + if local_avail[ii, jj + 1]: + patch = S_ext * S[ii, jj + 1] / denom + if S_red[ii, jj + 1] - patch <= 0: + neg_cell = True + local_avail[ii, jj + 1] = False + denom = ((local_avail[ii - 1, jj]) * S[ii - 1, jj] + + (local_avail[ii + 1, jj]) * S[ii + 1, jj] + + (local_avail[ii, jj - 1]) * S[ii, jj - 1] + + (local_avail[ii, jj + 1]) * S[ii, jj + 1]) + + # If possible, do 4-cell extension + if denom >= S_ext: + S_enl[ii, jj] = S[ii, jj] + if local_avail[ii - 1, jj]: + patch = S_ext * S[ii - 1, jj] / denom + borrowing[ii, jj].append([ii - 1, jj, patch, None]) + flag_intr_cell[ii-1, jj] = True + S_enl[ii, jj] += patch + S_red[ii - 1, jj] -= patch + if local_avail[ii + 1, jj]: + patch = S_ext * S[ii + 1, jj] / denom + borrowing[ii, jj].append([ii + 1, jj, patch, None]) + flag_intr_cell[ii+1, jj] = True + S_enl[ii, jj] += patch + S_red[ii + 1, jj] -= patch + if local_avail[ii, jj - 1]: + patch = S_ext * S[ii, jj - 1] / denom + borrowing[ii, jj].append([ii, jj - 1, patch, None]) + flag_intr_cell[ii, jj-1] = True + S_enl[ii, jj] += patch + S_red[ii, jj - 1] -= patch + if local_avail[ii, jj + 1]: + patch = S_ext * S[ii, jj + 1] / denom + borrowing[ii, jj].append([ii, jj + 1, patch, None]) + flag_intr_cell[ii, jj+1] = True + S_enl[ii, jj] += patch + S_red[ii, jj + 1] -= patch + + flag_ext_cell[ii, jj] = False + + """ + Function to compute the eight-cell extension of the unstable cells + """ + + @staticmethod + def _compute_extensions_eight_cells(nx=None, ny=None, S=None, flag_int_cell=None, + S_stab=None, S_enl=None, S_red=None, flag_unst_cell=None, + flag_avail_cell=None, flag_ext_cell=None, + flag_intr_cell=None, borrowing=None): + for ii in range(0, nx): + for jj in range(0, ny): + local_avail = flag_avail_cell.copy() + if (flag_unst_cell[ii, jj] and flag_int_cell[ii, jj] + and flag_ext_cell[ii, jj]): + S_enl[ii, jj] = S[ii, jj] + S_ext = S_stab[ii, jj] - S[ii, jj] + + denom = ((flag_avail_cell[ii - 1, jj]) * S[ii - 1, jj] + + (flag_avail_cell[ii + 1, jj]) * S[ii + 1, jj] + + (flag_avail_cell[ii, jj - 1]) * S[ii, jj - 1] + + (flag_avail_cell[ii, jj + 1]) * S[ii, jj + 1] + + (flag_avail_cell[ii - 1, jj - 1]) * S[ii - 1, jj - 1] + + (flag_avail_cell[ii + 1, jj - 1]) * S[ii + 1, jj - 1] + + (flag_avail_cell[ii - 1, jj + 1]) * S[ii - 1, jj + 1] + + (flag_avail_cell[ii + 1, jj + 1]) * S[ii + 1, jj + 1]) + + neg_cell = True + while denom >= S_ext and neg_cell: + neg_cell = False + if local_avail[ii - 1, jj]: + patch = S_ext * S[ii - 1, jj] / denom + if S_red[ii - 1, jj] - patch <= 0: + neg_cell = True + local_avail[ii - 1, jj] = False + if local_avail[ii + 1, jj]: + patch = S_ext * S[ii + 1, jj] / denom + if S_red[ii + 1, jj] - patch <= 0: + neg_cell = True + local_avail[ii + 1, jj] = False + if local_avail[ii, jj - 1]: + patch = S_ext * S[ii, jj - 1] / denom + if S_red[ii, jj - 1] - patch <= 0: + neg_cell = True + local_avail[ii, jj - 1] = False + if local_avail[ii, jj + 1]: + patch = S_ext * S[ii, jj + 1] / denom + if S_red[ii, jj + 1] - patch <= 0: + neg_cell = True + local_avail[ii, jj + 1] = False + if local_avail[ii - 1, jj - 1]: + patch = S_ext * S[ii - 1, jj - 1] / denom + if S_red[ii - 1, jj - 1] - patch <= 0: + neg_cell = True + local_avail[ii - 1, jj - 1] = False + if local_avail[ii + 1, jj - 1]: + patch = S_ext * S[ii + 1, jj - 1] / denom + if S_red[ii + 1, jj - 1] - patch <= 0: + neg_cell = True + local_avail[ii + 1, jj - 1] = False + if local_avail[ii - 1, jj + 1]: + patch = S_ext * S[ii - 1, jj + 1] / denom + if S_red[ii - 1, jj + 1] - patch <= 0: + neg_cell = True + local_avail[ii - 1, jj + 1] = False + if local_avail[ii + 1, jj + 1]: + patch = S_ext * S[ii + 1, jj + 1] / denom + if S_red[ii + 1, jj + 1] - patch <= 0: + neg_cell = True + local_avail[ii + 1, jj + 1] = False + + denom = ((local_avail[ii - 1, jj]) * S[ii - 1, jj] + + (local_avail[ii + 1, jj]) * S[ii + 1, jj] + + (local_avail[ii, jj - 1]) * S[ii, jj - 1] + + (local_avail[ii, jj + 1]) * S[ii, jj + 1] + + (local_avail[ii - 1, jj - 1]) * S[ii - 1, jj - 1] + + (local_avail[ii + 1, jj - 1]) * S[ii + 1, jj - 1] + + (local_avail[ii - 1, jj + 1]) * S[ii - 1, jj + 1] + + (local_avail[ii + 1, jj + 1]) * S[ii + 1, jj + 1]) + + if denom >= S_ext: + S_enl[ii, jj] = S[ii, jj] + if local_avail[ii - 1, jj]: + patch = S_ext * S[ii - 1, jj] / denom + borrowing[ii, jj].append([ii - 1, jj, patch, None]) + flag_intr_cell[ii-1, jj] = True + S_enl[ii, jj] += patch + S_red[ii - 1, jj] -= patch + if local_avail[ii + 1, jj]: + patch = S_ext * S[ii + 1, jj] / denom + borrowing[ii, jj].append([ii + 1, jj, patch, None]) + flag_intr_cell[ii+1, jj] = True + S_enl[ii, jj] += patch + S_red[ii + 1, jj] -= patch + if local_avail[ii, jj - 1]: + patch = S_ext * S[ii, jj - 1] / denom + borrowing[ii, jj].append([ii, jj - 1, patch, None]) + flag_intr_cell[ii, jj-1] = True + S_enl[ii, jj] += patch + S_red[ii, jj - 1] -= patch + if local_avail[ii, jj + 1]: + patch = S_ext * S[ii, jj + 1] / denom + borrowing[ii, jj].append([ii, jj + 1, patch, None]) + flag_intr_cell[ii, jj+1] = True + S_enl[ii, jj] += patch + S_red[ii, jj + 1] -= patch + if local_avail[ii - 1, jj - 1]: + patch = S_ext * S[ii - 1, jj - 1] / denom + borrowing[ii, jj].append([ii - 1, jj - 1, patch, None]) + flag_intr_cell[ii-1, jj-1] = True + S_enl[ii, jj] += patch + S_red[ii - 1, jj - 1] -= patch + if local_avail[ii + 1, jj - 1]: + patch = S_ext * S[ii + 1, jj - 1] / denom + borrowing[ii, jj].append([ii + 1, jj - 1, patch, None]) + flag_intr_cell[ii+1, jj-1] = True + S_enl[ii, jj] += patch + S_red[ii + 1, jj - 1] -= patch + if local_avail[ii - 1, jj + 1]: + patch = S_ext * S[ii - 1, jj + 1] / denom + borrowing[ii, jj].append([ii - 1, jj + 1, patch, None]) + flag_intr_cell[ii-1, jj+1] = True + S_enl[ii, jj] += patch + S_red[ii - 1, jj + 1] -= patch + if local_avail[ii + 1, jj + 1]: + patch = S_ext * S[ii + 1, jj + 1] / denom + borrowing[ii, jj].append([ii + 1, jj + 1, patch, None]) + flag_intr_cell[ii+1, jj+1] = True + S_enl[ii, jj] += patch + S_red[ii + 1, jj + 1] -= patch + + flag_ext_cell[ii, jj] = False + +""" +Function to compute the area of the cells of the conformal grid. +""" +@jit('(f8[:,:], f8[:,:], f8[:,:], f8[:,:], i4, i4, f8, f8)', nopython=True) +def compute_areas(l_x, l_y, S, S_red, nx, ny, dx, dy): + # Normalize the grid lengths for robustness + l_x /= dx + l_y /= dy + + # Loop over the cells + for ii in range(nx): + for jj in range(ny): + # If at least 3 edges have full length consider the area as full + # (this also takes care of the case in which an edge lies exactly on the boundary) + #count_cane = eq(l_x[ii, jj], 1.0) + eq(l_y[ii, jj], 1.0) + eq(l_x[ii, jj + 1], 1.0) + eq(l_y[ii + 1, jj], 1.0) + if (np.sum(np.array([eq(l_x[ii, jj], 1.0), eq(l_y[ii, jj], 1.0), + eq(l_x[ii, jj + 1], 1.0), eq(l_y[ii + 1, jj], 1.0)])) >= 3): + #if count_cane >=3: + S[ii, jj] = 1.0 + elif (eq(l_x[ii, jj], 0.) and neq(l_y[ii, jj], 0.) + and neq(l_x[ii, jj + 1], 0.) and neq(l_y[ii + 1, jj], 0)): + S[ii, jj] = 0.5 * (l_y[ii, jj] + l_y[ii + 1, jj]) * l_x[ii, jj + 1] + elif (eq(l_x[ii, jj + 1], 0.) and neq(l_y[ii, jj], 0.) + and neq(l_x[ii, jj], 0.) and neq(l_y[ii + 1, jj], 0.)): + S[ii, jj] = 0.5 * (l_y[ii, jj] + l_y[ii + 1, jj]) * l_x[ii, jj] + elif (eq(l_y[ii, jj], 0.) and neq(l_x[ii, jj], 0.) + and neq(l_x[ii, jj + 1], 0.) and neq(l_y[ii + 1, jj], 0.)): + S[ii, jj] = 0.5 * (l_x[ii, jj] + l_x[ii, jj + 1]) * l_y[ii + 1, jj] + elif (eq(l_y[ii + 1, jj], 0.) and neq(l_x[ii, jj], 0.) and neq(l_y[ii, jj], 0.) + and neq(l_x[ii, jj + 1], 0.)): + S[ii, jj] = 0.5 * (l_x[ii, jj] + l_x[ii, jj + 1]) * l_y[ii, jj] + elif (eq(l_x[ii, jj], 0.) and eq(l_y[ii, jj], 0.) + and neq(l_x[ii, jj + 1], 0.) and neq(l_y[ii + 1, jj], 0.)): + S[ii, jj] = 0.5 * l_x[ii, jj + 1] * l_y[ii + 1, jj] + elif (eq(l_x[ii, jj], 0.) and eq(l_y[ii + 1, jj], 0.) + and neq(l_x[ii, jj + 1], 0.) and neq(l_y[ii, jj], 0.)): + S[ii, jj] = 0.5 * l_x[ii, jj + 1] * l_y[ii, jj] + elif (eq(l_x[ii, jj + 1], 0.) and eq(l_y[ii + 1, jj], 0.) + and neq(l_x[ii, jj], 0.) and neq(l_y[ii, jj], 0.)): + S[ii, jj] = 0.5 * l_x[ii, jj] * l_y[ii, jj] + elif (eq(l_x[ii, jj + 1], 0.) and eq(l_y[ii, jj], 0.) + and neq(l_x[ii, jj], 0.) and neq(l_y[ii + 1, jj], 0.)): + S[ii, jj] = 0.5 * l_x[ii, jj] * l_y[ii + 1, jj] + elif (0. < l_x[ii, jj] <= 1. and 0. < l_y[ii, jj] <= 1. + and eq(l_x[ii, jj + 1], 1.) and eq(l_y[ii + 1, jj], 1.)): + S[ii, jj] = 1. - 0.5 * (1. - l_x[ii, jj]) * (1. - l_y[ii, jj]) + elif (0. < l_x[ii, jj] <= 1. and 0. < l_y[ii + 1, jj] <= 1. + and eq(l_x[ii, jj + 1], 1.) and eq(l_y[ii, jj], 1.)): + S[ii, jj] = 1. - 0.5 * (1. - l_x[ii, jj]) * (1. - l_y[ii + 1, jj]) + elif (0. < l_x[ii, jj + 1] <= 1. and 0. < l_y[ii + 1, jj] <= 1. + and eq(l_x[ii, jj], 1.) and eq(l_y[ii, jj], 1.)): + S[ii, jj] = 1. - 0.5 * (1. - l_x[ii, jj + 1]) * (1. - l_y[ii + 1, jj]) + elif (0. < l_x[ii, jj + 1] <= 1. and 0. < l_y[ii, jj] <= 1. + and eq(l_x[ii, jj], 1.) and eq(l_y[ii + 1, jj], 1.)): + S[ii, jj] = 1. - 0.5 * (1. - l_x[ii, jj + 1]) * (1. - l_y[ii, jj]) + + l_x *= dx + l_y *= dy + +""" +Function to mark wich cells are interior (int), require extension (unst), +are on the boundary(bound), are available for intrusion (avail) +""" +@jit('(f8[:,:], f8[:,:], i4, i4, f8, f8, f8[:,:], b1[:,:], f8[:,:], b1[:,:], b1[:,:], b1[:,:])', nopython=True) +def mark_cells(l_x, l_y, nx, ny, dx, dy, S, flag_int_cell, S_stab, + flag_unst_cell, flag_bound_cell, flag_avail_cell): + for ii in range(nx): + for jj in range(ny): + flag_int_cell[ii, jj] = S[ii, jj] > 0 + if flag_int_cell[ii, jj]: + S_stab[ii, jj] = 0.5 * np.max( + np.array([l_x[ii, jj] * dy, l_x[ii, jj + 1] * dy, + l_y[ii, jj] * dx, + l_y[ii + 1, jj] * dx])) + flag_unst_cell[ii, jj] = S[ii, jj] < S_stab[ii, jj] + flag_bound_cell[ii, jj] = ((0 < l_x[ii, jj] < dx) or + (0 < l_y[ii, jj] < dy) or + (0 < l_x[ii, jj + 1] < dx) or + (0 < l_y[ii + 1, jj] < dy)) + flag_avail_cell[ii, jj] = flag_int_cell[ii, jj] and (not flag_unst_cell[ii, jj]) diff --git a/build/lib/wakis/grid3D.py b/build/lib/wakis/grid3D.py new file mode 100644 index 0000000..f08aad4 --- /dev/null +++ b/build/lib/wakis/grid3D.py @@ -0,0 +1,415 @@ +import numpy as np +from grid2D import Grid2D +from grid2D import compute_areas as compute_areas_2D, mark_cells as mark_cells_2D +from numba import jit +from field import Field + +def seg_length(x_1, y_1, z_1, x_2, y_2, z_2): + return np.linalg.norm(np.array([x_1 - x_2, y_1 - y_2, z_1 - z_2])) + + +def eq(a, b, tol=1e-8): + return abs(a - b) < tol + + +def neq(a, b, tol=1e-8): + return not eq(a, b, tol) + + # Undo the normalization + S *= dx * dy + l_x *= dx + l_y *= dy + S_red[:] = S.copy()[:] + +class Grid3D: + """ + Class holding the grid info and the routines for cell extensions. + Constructor arguments: + - xmin, xmax, ymin, ymax, zmin, zmax: extent of the domain. + - nx, ny, nz: number of cells per direction + - conductors: conductor object + - sol_type: type of solver. 'FDTD' for staircased FDTD, 'DM' for Conformal Dey-Mittra FDTD, + 'ECT' for Enlarged Cell Technique conformal FDTD + """ + + def __init__(self, xmin, xmax, ymin, ymax, zmin, zmax, nx, ny, nz, conductors, sol_type): + + self.xmin = xmin + self.xmax = xmax + self.ymin = ymin + self.ymax = ymax + self.zmin = zmin + self.zmax = zmax + self.nx = nx + self.ny = ny + self.nz = nz + self.dx = (xmax - xmin) / nx + self.dy = (ymax - ymin) / ny + self.dz = (ymax - ymin) / nz + self.conductors = conductors + + self.l_x = np.zeros((nx, ny + 1, nz + 1)) + self.l_y = np.zeros((nx + 1, ny, nz + 1)) + self.l_z = np.zeros((nx + 1, ny + 1, nz)) + self.Sxy = np.zeros((nx, ny, nz + 1)) + self.Syz = np.zeros((nx + 1, ny, nz)) + self.Szx = np.zeros((nx, ny + 1, nz)) + self.Sxy_stab = np.zeros_like(self.Sxy) + self.Sxy_enl = np.zeros_like(self.Sxy) + self.Sxy_red = np.zeros_like(self.Sxy) + self.Syz_stab = np.zeros_like(self.Syz) + self.Syz_enl = np.zeros_like(self.Syz) + self.Syz_red = np.zeros_like(self.Syz) + self.Szx_stab = np.zeros_like(self.Szx) + self.Szx_enl = np.zeros_like(self.Szx) + self.Szx_red = np.zeros_like(self.Szx) + self.flag_unst_cell_xy = np.zeros_like(self.Sxy, dtype=bool) + self.flag_intr_cell_xy = np.zeros_like(self.Sxy, dtype=bool) + self.flag_int_cell_xy = np.zeros_like(self.Sxy, dtype=bool) + self.flag_bound_cell_xy = np.zeros_like(self.Sxy, dtype=bool) + self.flag_avail_cell_xy = np.zeros_like(self.Sxy, dtype=bool) + self.flag_ext_cell_xy = np.zeros_like(self.Sxy, dtype=bool) + self.flag_unst_cell_yz = np.zeros_like(self.Syz, dtype=bool) + self.flag_intr_cell_yz = np.zeros_like(self.Syz, dtype=bool) + self.flag_int_cell_yz = np.zeros_like(self.Syz, dtype=bool) + self.flag_bound_cell_yz = np.zeros_like(self.Syz, dtype=bool) + self.flag_avail_cell_yz = np.zeros_like(self.Syz, dtype=bool) + self.flag_ext_cell_yz = np.zeros_like(self.Syz, dtype=bool) + self.flag_unst_cell_zx = np.zeros_like(self.Szx, dtype=bool) + self.flag_intr_cell_zx = np.zeros_like(self.Szx, dtype=bool) + self.flag_int_cell_zx = np.zeros_like(self.Szx, dtype=bool) + self.flag_bound_cell_zx = np.zeros_like(self.Szx, dtype=bool) + self.flag_avail_cell_zx = np.zeros_like(self.Szx, dtype=bool) + self.flag_ext_cell_zx = np.zeros_like(self.Szx, dtype=bool) + self.broken_xy = np.zeros_like(self.Sxy, dtype=bool) + self.broken_yz = np.zeros_like(self.Syz, dtype=bool) + self.broken_zx = np.zeros_like(self.Szx, dtype=bool) + + if (sol_type is not 'FDTD') and (sol_type is not 'DM') and (sol_type is not 'ECT') and (sol_type is not 'FIT'): + raise ValueError("sol_type must be:\n" + + "\t'FDTD' for standard staircased FDTD\n" + + "\t'DM' for Dey-Mittra conformal FDTD\n" + + "\t'ECT' for Enlarged Cell Technique conformal FDTD") + + self.compute_edges() + if sol_type is 'DM' or sol_type is 'FDTD': #or sol_type is 'FIT': + self.compute_areas(self.l_x, self.l_y, self.l_z, self.Sxy, self.Syz, self.Szx, + self.Sxy_red, self.Syz_red, self.Szx_red, + self.nx, self.ny, self.nz, self.dx, self.dy, self.dz) + self.mark_cells(self.l_x, self.l_y, self.l_z, self.nx, self.ny, self.nz, self.dx, self.dy, self.dz, + self.Sxy, self.Syz, self.Szx, self.flag_int_cell_xy, self.flag_int_cell_yz, self.flag_int_cell_zx, + self.Sxy_stab, self.Syz_stab, self.Szx_stab, self.flag_unst_cell_xy, self.flag_unst_cell_yz, + self.flag_unst_cell_zx, self.flag_bound_cell_xy, self.flag_bound_cell_yz, self.flag_bound_cell_zx, + self.flag_avail_cell_xy, self.flag_avail_cell_yz, self.flag_avail_cell_zx) + elif sol_type is 'ECT': + self.compute_areas(self.l_x, self.l_y, self.l_z, self.Sxy, self.Syz, self.Szx, + self.Sxy_red, self.Syz_red, self.Szx_red, + self.nx, self.ny, self.nz, self.dx, self.dy, self.dz) + self.mark_cells(self.l_x, self.l_y, self.l_z, self.nx, self.ny, self.nz, self.dx, self.dy, self.dz, + self.Sxy, self.Syz, self.Szx, self.flag_int_cell_xy, self.flag_int_cell_yz, self.flag_int_cell_zx, + self.Sxy_stab, self.Syz_stab, self.Szx_stab, self.flag_unst_cell_xy, self.flag_unst_cell_yz, + self.flag_unst_cell_zx, self.flag_bound_cell_xy, self.flag_bound_cell_yz, self.flag_bound_cell_zx, + self.flag_avail_cell_xy, self.flag_avail_cell_yz, self.flag_avail_cell_zx) + + + # info about intruded cells (i,j,[(i_borrowing,j_borrowing,area_borrowing, )]) + self.borrowing_xy = np.empty((nx, ny, nz + 1), dtype=object) + self.borrowing_yz = np.empty((nx + 1, ny, nz), dtype=object) + self.borrowing_zx = np.empty((nx, ny + 1, nz), dtype=object) + + for ii in range(nx): + for jj in range(ny): + for kk in range(nz + 1): + self.borrowing_xy[ii, jj, kk] = [] + for ii in range(nx + 1): + for jj in range(ny): + for kk in range(nz): + self.borrowing_yz[ii, jj, kk] = [] + for ii in range(nx): + for jj in range(ny + 1): + for kk in range(nz): + self.borrowing_zx[ii, jj, kk] = [] + + self.flag_ext_cell_xy = self.flag_unst_cell_xy.copy() + self.flag_ext_cell_yz = self.flag_unst_cell_yz.copy() + self.flag_ext_cell_zx = self.flag_unst_cell_zx.copy() + self.mark_cells(self.l_x, self.l_y, self.l_z, self.nx, self.ny, self.nz, self.dx, self.dy, self.dz, + self.Sxy, self.Syz, self.Szx, self.flag_int_cell_xy, self.flag_int_cell_yz, self.flag_int_cell_zx, + self.Sxy_stab, self.Syz_stab, self.Szx_stab, self.flag_unst_cell_xy, self.flag_unst_cell_yz, + self.flag_unst_cell_zx, self.flag_bound_cell_xy, self.flag_bound_cell_yz, self.flag_bound_cell_zx, + self.flag_avail_cell_xy, self.flag_avail_cell_yz, self.flag_avail_cell_zx) + self.compute_extensions() + + + elif sol_type is 'FIT': + + # primal Grid G + self.x = np.linspace(self.xmin, self.xmax, self.nx+1) + self.y = np.linspace(self.ymin, self.ymax, self.ny+1) + self.z = np.linspace(self.zmin, self.zmax, self.nz+1) + + Y, X, Z = np.meshgrid(self.y, self.x, self.z) + + self.L = Field(self.nx, self.ny, self.nz) + self.L.field_x = X[1:, 1:, 1:] - X[:-1, :-1, :-1] + self.L.field_y = Y[1:, 1:, 1:] - Y[:-1, :-1, :-1] + self.L.field_z = Z[1:, 1:, 1:] - Z[:-1, :-1, :-1] + + self.iA = Field(self.nx, self.ny, self.nz) + self.iA.field_x = np.divide(1.0, self.L.field_y * self.L.field_z) + self.iA.field_y = np.divide(1.0, self.L.field_x * self.L.field_z) + self.iA.field_z = np.divide(1.0, self.L.field_x * self.L.field_y) + + # tilde grid ~G + #self.itA = self.iA + #self.tL = self.L + + self.tx = (self.x[1:]+self.x[:-1])/2 + self.ty = (self.y[1:]+self.y[:-1])/2 + self.tz = (self.z[1:]+self.z[:-1])/2 + + self.tx = np.append(self.tx, self.tx[-1]) + self.ty = np.append(self.ty, self.ty[-1]) + self.tz = np.append(self.tz, self.tz[-1]) + + tY, tX, tZ = np.meshgrid(self.ty, self.tx, self.tz) + + self.tL = Field(self.nx, self.ny, self.nz) + self.tL.field_x = tX[1:, 1:, 1:] - tX[:-1, :-1, :-1] + self.tL.field_y = tY[1:, 1:, 1:] - tY[:-1, :-1, :-1] + self.tL.field_z = tZ[1:, 1:, 1:] - tZ[:-1, :-1, :-1] + + self.itA = Field(self.nx, self.ny, self.nz) + aux = self.tL.field_y * self.tL.field_z + self.itA.field_x = np.divide(1.0, aux, out=np.zeros_like(aux), where=aux!=0) + aux = self.tL.field_x * self.tL.field_z + self.itA.field_y = np.divide(1.0, aux, out=np.zeros_like(aux), where=aux!=0) + aux = self.tL.field_x * self.tL.field_y + self.itA.field_z = np.divide(1.0, aux, out=np.zeros_like(aux), where=aux!=0) + del aux + + self.compute_areas(self.l_x, self.l_y, self.l_z, self.Sxy, self.Syz, self.Szx, + self.Sxy_red, self.Syz_red, self.Szx_red, + self.nx, self.ny, self.nz, self.dx, self.dy, self.dz) + self.mark_cells(self.l_x, self.l_y, self.l_z, self.nx, self.ny, self.nz, self.dx, self.dy, self.dz, + self.Sxy, self.Syz, self.Szx, self.flag_int_cell_xy, self.flag_int_cell_yz, self.flag_int_cell_zx, + self.Sxy_stab, self.Syz_stab, self.Szx_stab, self.flag_unst_cell_xy, self.flag_unst_cell_yz, + self.flag_unst_cell_zx, self.flag_bound_cell_xy, self.flag_bound_cell_yz, self.flag_bound_cell_zx, + self.flag_avail_cell_xy, self.flag_avail_cell_yz, self.flag_avail_cell_zx) + + """ + Function to compute the length of the edges of the conformal grid. + Inputs: + - tol: an edge shorter than tol will be considered as of zero length + """ + + def compute_edges(self, tol=1e-8): + # edges xy + for kk in range(self.nz + 1): + for ii in range(self.nx): + for jj in range(self.ny + 1): + x_1 = ii * self.dx + self.xmin + y_1 = jj * self.dy + self.ymin + x_2 = (ii + 1) * self.dx + self.xmin + y_2 = jj * self.dy + self.ymin + z = kk * self.dz + self.zmin + # if point 1 is in conductor + if self.conductors.in_conductor(x_1, y_1, z): + # if point 2 is in conductor, length of the l_x[i, j] is zero + if self.conductors.in_conductor(x_2, y_2, z): + self.l_x[ii, jj, kk] = 0 + # if point 2 is not in conductor, length of l_x[i, j] + # is the fractional length + else: + self.l_x[ii, jj, kk] = seg_length(self.conductors.intersec_x(x_2, + y_2, z), + y_2, z, x_2, y_2, z) + # if point 1 is not in conductor + else: + # if point 2 is in conductor, length of l_x[i, j] is the fractional length + if self.conductors.in_conductor(x_2, y_2, z): + self.l_x[ii, jj, kk] = seg_length(x_1, y_1, z, + self.conductors.intersec_x(x_1, y_1, + z), + y_1, z) + # if point 2 is not in conductor, length of l_x[i, j] is dx + else: + self.l_x[ii, jj, kk] = self.dx + + for kk in range(self.nz + 1): + for ii in range(self.nx + 1): + for jj in range(self.ny): + x_1 = ii * self.dx + self.xmin + y_1 = jj * self.dy + self.ymin + x_3 = ii * self.dx + self.xmin + y_3 = (jj + 1) * self.dy + self.ymin + z = kk * self.dz + self.zmin + # if point 1 is in conductor + if self.conductors.in_conductor(x_1, y_1, z): + # if point 3 to the right is in conductor, length of the l_y[i, j] is zero + if self.conductors.in_conductor(x_3, y_3, z): + self.l_y[ii, jj, kk] = 0 + # if point 3 is not in conductor, length of l_y[i, j] + # is the fractional length + else: + self.l_y[ii, jj, kk] = seg_length(x_3, + self.conductors.intersec_y(x_3, y_3, + z), + z, x_3, y_3, z) + # if point 1 is not in conductor + else: + # if point 3 is in conductor, length of the l_y[i, j] + # is the fractional length + if self.conductors.in_conductor(x_3, y_3, z): + self.l_y[ii, jj, kk] = seg_length(x_1, y_1, z, x_1, + self.conductors.intersec_y(x_1, y_1, + z), + z) + # if point 3 is not in conductor, length of l_y[i, j] is dy + else: + self.l_y[ii, jj, kk] = self.dy + + for jj in range(self.ny + 1): + for ii in range(self.nx + 1): + for kk in range(self.nz): + y = jj*self.dy + self.ymin + x_1 = ii * self.dx + self.xmin + z_1 = kk * self.dz + self.zmin + x_3 = ii * self.dx + self.xmin + z_3 = (kk + 1) * self.dz + self.zmin + # if point 1 is in conductor + if self.conductors.in_conductor(x_1, y, z_1): + # if point 3 to the right is in conductor, length of the l_y[i, j] is zero + if self.conductors.in_conductor(x_3, y, z_3): + self.l_z[ii, jj, kk] = 0 + # if point 3 is not in conductor, length of l_y[i, j] + # is the fractional length + else: + self.l_z[ii, jj, kk] = seg_length(x_3, y, + self.conductors.intersec_z(x_3, + y, z_3), + x_3, y, z_3) + # if point 1 is not in conductor + else: + # if point 3 is in conductor, length of the l_y[i, j] + # is the fractional length + if self.conductors.in_conductor(x_3, y, z_3): + self.l_z[ii, jj, kk] = seg_length(x_1, y, z_1, x_1, y, + self.conductors.intersec_z(x_1, y, + z_1)) + # if point 3 is not in conductor, length of l_y[i, j] is dy + else: + self.l_z[ii, jj, kk] = self.dz + + # set to zero the length of very small cells + if tol > 0.: + self.l_x /= self.dx + self.l_y /= self.dy + self.l_z /= self.dz + + low_values_flags = abs(self.l_x) < tol + high_values_flags = abs(self.l_x - 1.) < tol + self.l_x[low_values_flags] = 0 + self.l_x[high_values_flags] = 1. + + low_values_flags = abs(self.l_y) < tol + high_values_flags = abs(self.l_y - 1.) < tol + self.l_y[low_values_flags] = 0 + self.l_y[high_values_flags] = 1. + + low_values_flags = abs(self.l_z) < tol + high_values_flags = abs(self.l_z - 1.) < tol + self.l_z[low_values_flags] = 0 + self.l_z[high_values_flags] = 1. + + self.l_x *= self.dx + self.l_y *= self.dy + self.l_z *= self.dz + + + """ + Function to compute the area of the cells of the conformal grid. + """ + @staticmethod + @jit('(f8[:,:,:], f8[:,:,:], f8[:,:,:], f8[:,:,:], f8[:,:,:], f8[:,:,:], f8[:,:,:], f8[:,:,:], f8[:,:,:], i4, i4, i4, f8, f8, f8)', nopython=True) + def compute_areas(l_x, l_y, l_z, Sxy, Syz, Szx, Sxy_red, Syz_red, Szx_red, nx, ny, nz, dx, dy, dz): + + for kk in range(nz + 1): + compute_areas_2D(l_x[:, :, kk], l_y[:, :, kk], Sxy[:, :, kk], Sxy_red[:, :, kk], + nx, ny, dx, dy) + + for ii in range(nx + 1): + compute_areas_2D(l_y[ii, :, :], l_z[ii, :, :], Syz[ii, :, :], Syz_red[ii, :, :], + ny, nz, dy, dz) + + for jj in range(ny + 1): + compute_areas_2D(l_x[:, jj, :], l_z[:, jj, :], Szx[:, jj, :], Szx_red[:, jj, :], + nx, nz, dx, dz) + + """ + Function to mark wich cells are interior (int), require extension (unst), + are on the boundary(bound), are available for intrusion (avail) + """ + @staticmethod + @jit('(f8[:,:,:], f8[:,:,:], f8[:,:,:], i4, i4, i4, f8, f8, f8, f8[:,:,:], f8[:,:,:], f8[:,:,:], b1[:,:,:], b1[:,:,:], b1[:,:,:], f8[:,:,:], f8[:,:,:], f8[:,:,:], b1[:,:,:], b1[:,:,:], b1[:,:,:], b1[:,:,:], b1[:,:,:], b1[:,:,:], b1[:,:,:], b1[:,:,:], b1[:,:,:])', nopython=True) + def mark_cells(l_x, l_y, l_z, nx, ny, nz, dx, dy, dz, Sxy, Syz, Szx, flag_int_cell_xy, flag_int_cell_yz, flag_int_cell_zx, + Sxy_stab, Syz_stab, Szx_stab, flag_unst_cell_xy, flag_unst_cell_yz, flag_unst_cell_zx, flag_bound_cell_xy, + flag_bound_cell_yz, flag_bound_cell_zx, flag_avail_cell_xy, flag_avail_cell_yz, flag_avail_cell_zx): + + for kk in range(nz + 1): + mark_cells_2D(l_x[:, :, kk], l_y[:, :, kk], nx, ny, dx, dy, Sxy[:, :, kk], + flag_int_cell_xy[:, :, kk], Sxy_stab[:, :, kk], flag_unst_cell_xy[:, :, kk], + flag_bound_cell_xy[:, :, kk], flag_avail_cell_xy[:, :, kk]) + + for ii in range(nx + 1): + mark_cells_2D(l_y[ii, :, :], l_z[ii, :, :], ny, nz, dy, dz, Syz[ii, :, :], + flag_int_cell_yz[ii, :, :], Syz_stab[ii, :, :], flag_unst_cell_yz[ii, :, :], + flag_bound_cell_yz[ii, :, :], flag_avail_cell_yz[ii, :, :]) + + for jj in range(ny + 1): + mark_cells_2D(l_x[:, jj, :], l_z[:, jj, :], nx, nz, dx, dz, Szx[:, jj, :], + flag_int_cell_zx[:, jj, :], Szx_stab[:, jj, :], flag_unst_cell_zx[:, jj, :], + flag_bound_cell_zx[:, jj, :], flag_avail_cell_zx[:, jj, :]) + + """ + Function to compute the extension of the unstable cells + """ + + def compute_extensions(self): + #breakpoint() + for kk in range(self.nz + 1): + #breakpoint() + Grid2D.compute_extensions(nx=self.nx, ny=self.ny, S=self.Sxy[:, :, kk], + flag_int_cell=self.flag_int_cell_xy[:, :, kk], + S_stab=self.Sxy_stab[:, :, kk], S_enl=self.Sxy_enl[:, :, kk], + S_red=self.Sxy_red[:, :, kk], + flag_unst_cell=self.flag_unst_cell_xy[:, :, kk], + flag_avail_cell=self.flag_avail_cell_xy[:, :, kk], + flag_ext_cell=self.flag_ext_cell_xy[:, :, kk], + flag_intr_cell=self.flag_intr_cell_xy[:,:,kk], + borrowing=self.borrowing_xy[:, :, kk], + kk=kk, l_verbose=False) + + for ii in range(self.nx + 1): + Grid2D.compute_extensions(nx=self.ny, ny=self.nz, S=self.Syz[ii, :, :], + flag_int_cell=self.flag_int_cell_yz[ii, :, :], + S_stab=self.Syz_stab[ii, :, :], S_enl=self.Syz_enl[ii, :, :], + S_red=self.Syz_red[ii, :, :], + flag_unst_cell=self.flag_unst_cell_yz[ii, :, :], + flag_avail_cell=self.flag_avail_cell_yz[ii, :, :], + flag_ext_cell=self.flag_ext_cell_yz[ii, :, :], + borrowing=self.borrowing_yz[ii, :, :], + flag_intr_cell=self.flag_intr_cell_yz[ii,:,:], + kk = ii, l_verbose=False) + + for jj in range(self.ny + 1): + Grid2D.compute_extensions(nx=self.nx, ny=self.nz, S=self.Szx[:, jj, :], + flag_int_cell=self.flag_int_cell_zx[:, jj, :], + S_stab=self.Szx_stab[:, jj, :], S_enl=self.Szx_enl[:, jj, :], + S_red=self.Szx_red[:, jj, :], + flag_unst_cell=self.flag_unst_cell_zx[:, jj, :], + flag_avail_cell=self.flag_avail_cell_zx[:, jj, :], + flag_ext_cell=self.flag_ext_cell_zx[:, jj, :], + flag_intr_cell=self.flag_intr_cell_zx[:,jj,:], + borrowing=self.borrowing_zx[:, jj, :], + kk = jj, l_verbose=False) diff --git a/build/lib/wakis/gridFIT3D.py b/build/lib/wakis/gridFIT3D.py new file mode 100644 index 0000000..d6e447b --- /dev/null +++ b/build/lib/wakis/gridFIT3D.py @@ -0,0 +1,841 @@ +# copyright ################################# # +# This file is part of the wakis Package. # +# Copyright (c) CERN, 2024. # +# ########################################### # + +import numpy as np +import pyvista as pv +from functools import partial +from scipy.optimize import least_squares + +from .field import Field + +try: + from mpi4py import MPI + imported_mpi = True +except ImportError: + imported_mpi = False + +class GridFIT3D: + """ + Class holding the grid information and + stl importing handling using PyVista + + Parameters + ---------- + xmin, xmax, ymin, ymax, zmin, zmax: float + extent of the domain. + Nx, Ny, Nz: int + number of cells per direction + stl_solids: dict, optional + stl files to import in the domain. + {'Solid 1': stl_1, 'Solid 2': stl_2, ...} + If stl files are not in the same folder, + add the path to the file name. + stl_materials: dict, optional + Material properties associated with stl + {'Solid 1': [eps1, mu1], + 'Solid 2': [eps1, mu1], + ...} + stl_rotate: list or dict, optional + Angle of rotation to apply to the stl models: [rot_x, rot_y, rot_z] + - if list, it will be applied to all stls in `stl_solids` + - if dict, it must contain the same keys as `stl_solids`, + indicating the rotation angle per stl + stl_scale: float or dict, optional + Scaling value to apply to the stl model to convert to [m] + - if float, it will be applied to all stl in `stl_solids` + - if dict, it must contain the same keys as `stl_solids` + tol: float, default 1e-3 + Tolerance factor for stl import, used in grid.select_enclosed_points. + Importing tolerance is computed by: tol*min(dx,dy,dz). + verbose: int or bool, default 1 + Enable verbose ouput on the terminal + """ + + def __init__(self, xmin, xmax, ymin, ymax, zmin, zmax, + Nx, Ny, Nz, + x=None, y=None, z=None, + use_mpi=False, + use_mesh_refinement=False, refinement_method='insert', refinement_tol=1e-8, + snap_points=None, snap_tol=1e-5, snap_solids=None, + stl_solids=None, stl_materials=None, + stl_rotate=[0., 0., 0.], stl_translate=[0., 0., 0.], stl_scale=1.0, + stl_colors=None, verbose=1, stl_tol=1e-3): + + if verbose: print('Generating grid...') + self.verbose = verbose + self.use_mpi = use_mpi + self.use_mesh_refinement = use_mesh_refinement + + # domain limits + self.xmin = xmin + self.xmax = xmax + self.ymin = ymin + self.ymax = ymax + self.zmin = zmin + self.zmax = zmax + self.Nx = Nx + self.Ny = Ny + self.Nz = Nz + self.dx = (xmax - xmin) / Nx + self.dy = (ymax - ymin) / Ny + self.dz = (zmax - zmin) / Nz + + # stl info + self.stl_solids = stl_solids + self.stl_materials = stl_materials + self.stl_rotate = stl_rotate + self.stl_translate = stl_translate + self.stl_scale = stl_scale + self.stl_colors = stl_colors + if stl_solids is not None: + self._prepare_stl_dicts() + + # MPI subdivide domain + if self.use_mpi: + self.ZMIN = None + self.ZMAX = None + self.NZ = None + self.Z = None + if imported_mpi: + self.mpi_initialize() + if self.verbose: print(f"MPI initialized for {self.rank} of {self.size}") + else: + raise ImportError("*** mpi4py is required when use_mpi=True but was not found") + + # primal Grid G base axis x, y, z + self.x = x + self.y = y + self.z = z + self.refinement_method = refinement_method + self.snap_points = snap_points + self.snap_tol = snap_tol + self.snap_solids = snap_solids # if None, use all stl_solids + + if self.x is not None and self.y is not None and self.z is not None: + # allow user to set the grid axis manually + self.Nx = len(self.x) - 1 + self.Ny = len(self.y) - 1 + self.Nz = len(self.z) - 1 + self.dx = np.min(np.diff(self.x)) + self.dy = np.min(np.diff(self.y)) + self.dz = np.min(np.diff(self.z)) + + elif self.use_mesh_refinement: + if verbose: print('Applying mesh refinement...') + if self.snap_points is None and stl_solids is not None: + self.compute_snap_points(snap_solids=snap_solids, snap_tol=snap_tol) + self.refine_xyz_axis(method=refinement_method, tol=refinement_tol) # obtain self.x, self.y, self.z + else: + self.x = np.linspace(self.xmin, self.xmax, self.Nx+1) + self.y = np.linspace(self.ymin, self.ymax, self.Ny+1) + self.z = np.linspace(self.zmin, self.zmax, self.Nz+1) + + # grid G and tilde grid ~G, lengths and inverse areas + self.compute_grid() + + # tolerance for stl import tol*min(dx,dy,dz) + if verbose: print('Importing STL solids...') + self.tol = stl_tol + if stl_solids is not None: + self.mark_cells_in_stl() + if stl_colors is None: + self.assign_colors() + + def compute_grid(self): + X, Y, Z = np.meshgrid(self.x, self.y, self.z, indexing='ij') + self.grid = pv.StructuredGrid(X.transpose(), Y.transpose(), Z.transpose()) + + self.L = Field(self.Nx, self.Ny, self.Nz) + self.L.field_x = X[1:, 1:, 1:] - X[:-1, :-1, :-1] + self.L.field_y = Y[1:, 1:, 1:] - Y[:-1, :-1, :-1] + self.L.field_z = Z[1:, 1:, 1:] - Z[:-1, :-1, :-1] + + self.iA = Field(self.Nx, self.Ny, self.Nz) + self.iA.field_x = np.divide(1.0, self.L.field_y * self.L.field_z) + self.iA.field_y = np.divide(1.0, self.L.field_x * self.L.field_z) + self.iA.field_z = np.divide(1.0, self.L.field_x * self.L.field_y) + + # tilde grid ~G + self.tx = (self.x[1:]+self.x[:-1])/2 + self.ty = (self.y[1:]+self.y[:-1])/2 + self.tz = (self.z[1:]+self.z[:-1])/2 + + self.tx = np.append(self.tx, self.tx[-1]) + self.ty = np.append(self.ty, self.ty[-1]) + self.tz = np.append(self.tz, self.tz[-1]) + + tX, tY, tZ = np.meshgrid(self.tx, self.ty, self.tz, indexing='ij') + + self.tL = Field(self.Nx, self.Ny, self.Nz) + self.tL.field_x = tX[1:, 1:, 1:] - tX[:-1, :-1, :-1] + self.tL.field_y = tY[1:, 1:, 1:] - tY[:-1, :-1, :-1] + self.tL.field_z = tZ[1:, 1:, 1:] - tZ[:-1, :-1, :-1] + + self.itA = Field(self.Nx, self.Ny, self.Nz) + aux = self.tL.field_y * self.tL.field_z + self.itA.field_x = np.divide(1.0, aux, out=np.zeros_like(aux), where=aux!=0) + aux = self.tL.field_x * self.tL.field_z + self.itA.field_y = np.divide(1.0, aux, out=np.zeros_like(aux), where=aux!=0) + aux = self.tL.field_x * self.tL.field_y + self.itA.field_z = np.divide(1.0, aux, out=np.zeros_like(aux), where=aux!=0) + del aux + + def mpi_initialize(self): + comm = MPI.COMM_WORLD # Get MPI communicator + self.comm = comm + self.rank = self.comm.Get_rank() + self.size = self.comm.Get_size() + + # Error handling for Nz < size + if self.Nz < self.size: + raise ValueError(f"Nz ({self.Nz}) must be greater than or equal to the number of MPI processes ({self.size}).") + + # global z quantities [ALLCAPS] + self.ZMIN = self.zmin + self.ZMAX = self.zmax + self.NZ = self.Nz - self.Nz%(self.size) # ensure multiple of MPI size + self.Z = np.linspace(self.ZMIN, self.ZMAX, self.NZ+1)[:-1] + self.dz/2 + + if self.verbose and self.rank==0: + print(f"Global grid ZMIN={self.ZMIN}, ZMAX={self.ZMAX}, NZ={self.NZ}") + + # MPI subdomain quantities + self.Nz = self.NZ // (self.size) + self.dz = (self.ZMAX - self.ZMIN) / self.NZ + self.zmin = self.rank * self.Nz * self.dz + self.ZMIN + self.zmax = (self.rank+1) * self.Nz * self.dz + self.ZMIN + + if self.verbose: print(f"MPI rank {self.rank} of {self.size} initialized with \ + zmin={self.zmin}, zmax={self.zmax}, Nz={self.Nz}") + # Add ghost cells + self.n_ghosts = 1 + if self.rank > 0: + self.zmin += - self.n_ghosts * self.dz + self.Nz += self.n_ghosts + if self.rank < (self.size-1): + self.zmax += self.n_ghosts * self.dz + self.Nz += self.n_ghosts + + # Support for single core + if self.rank == 0 and self.size == 1: + self.zmax += self.n_ghosts * self.dz + self.Nz += self.n_ghosts + + def mpi_gather_asGrid(self): + _grid = None + if self.rank == 0: + print(f"Generating global grid from {self.ZMIN} to {self.ZMAX}") + _grid = GridFIT3D(self.xmin, self.xmax, + self.ymin, self.ymax, + self.ZMIN, self.ZMAX, + self.Nx, self.Ny, self.NZ, + use_mpi=False, + stl_solids=self.stl_solids, + stl_materials=self.stl_materials, + stl_scale=self.stl_scale, + stl_rotate=self.stl_rotate, + stl_translate=self.stl_translate, + stl_colors=self.stl_colors, + verbose=self.verbose, + tol=self.tol, + ) + return _grid + + def _prepare_stl_dicts(self): + if type(self.stl_solids) is not dict: + if type(self.stl_solids) is str: + self.stl_solids = {'Solid 1' : self.stl_solids} + else: + raise Exception('Attribute `stl_solids` must contain a string or a dictionary') + + if type(self.stl_rotate) is not dict: + # if not a dict, the same values will be applied to all solids + stl_rotate = {} + for key in self.stl_solids.keys(): + stl_rotate[key] = self.stl_rotate + self.stl_rotate = stl_rotate + + if type(self.stl_scale) is not dict: + # if not a dict, the same values will be applied to all solids + stl_scale = {} + for key in self.stl_solids.keys(): + stl_scale[key] = self.stl_scale + self.stl_scale = stl_scale + + if type(self.stl_translate) is not dict: + # if not a dict, the same values will be applied to all solids + stl_translate = {} + for key in self.stl_solids.keys(): + stl_translate[key] = self.stl_translate + self.stl_translate = stl_translate + + def mark_cells_in_stl(self): + # Obtain masks with grid cells inside each stl solid + tol = np.min([self.dx, self.dy, self.dz])*self.tol + for key in self.stl_solids.keys(): + + surf = self.read_stl(key) + + # mark cells in stl [True == in stl, False == out stl] + try: + select = self.grid.select_enclosed_points(surf, tolerance=tol) + except: + select = self.grid.select_enclosed_points(surf, tolerance=tol, check_surface=False) + self.grid[key] = select.point_data_to_cell_data()['SelectedPoints'] > tol + + def read_stl(self, key): + # import stl + surf = pv.read(self.stl_solids[key]) + + # rotate + surf = surf.rotate_x(self.stl_rotate[key][0]) + surf = surf.rotate_y(self.stl_rotate[key][1]) + surf = surf.rotate_z(self.stl_rotate[key][2]) + + # translate + surf = surf.translate(self.stl_translate[key]) + + # scale + surf = surf.scale(self.stl_scale[key]) + + return surf + + def compute_snap_points(self, snap_solids=None, snap_tol=1e-8): + if self.verbose: print('* Calculating snappy points...') + # Support for user-defined stl_keys as list + if snap_solids is None: + snap_solids = self.stl_solids.keys() + + # Union of all the surfaces + # [TODO]: should use | for union instead or +? + model = None + for key in snap_solids: + solid = self.read_stl(key) + if model is None: + model = solid + else: + model = model + solid + + edges = model.extract_feature_edges(boundary_edges=True, manifold_edges=False) + + # Extract points lying in the X-Z plane (Y ≈ 0) + xz_plane_points = edges.points[np.abs(edges.points[:, 1]) < snap_tol] + # Extract points lying in the Y-Z plane (X ≈ 0) + yz_plane_points = edges.points[np.abs(edges.points[:, 0]) < snap_tol] + # Extract points lying in the X-Y plane (Z ≈ 0) + xy_plane_points = edges.points[np.abs(edges.points[:, 2]) < 1e-5] + + self.snap_points = np.r_[xz_plane_points, yz_plane_points, xy_plane_points] + + # get the unique x, y, z coordinates + x_snaps = np.unique(np.round(self.snap_points[:, 0], 5)) + y_snaps = np.unique(np.round(self.snap_points[:, 1], 5)) + z_snaps = np.unique(np.round(self.snap_points[:, 2], 5)) + + # Include simulation domain bounds + self.x_snaps = np.unique(np.concatenate(([self.xmin], x_snaps, [self.xmax]))) + self.y_snaps = np.unique(np.concatenate(([self.ymin], y_snaps, [self.ymax]))) + self.z_snaps = np.unique(np.concatenate(([self.zmin], z_snaps, [self.zmax]))) + + def plot_snap_points(self, snap_solids=None, snap_tol=1e-8): + # TODO + # Support for user-defined stl_keys as list + if snap_solids is None: + snap_solids = self.stl_solids.keys() + + # Union of all the surfaces + # [TODO]: should use | for union instead or +? + model = None + for key in snap_solids: + solid = self.read_stl(key) + if model is None: + model = solid + else: + model = model + solid + + edges = model.extract_feature_edges(boundary_edges=True, manifold_edges=False) + + # Extract points lying in the X-Z plane (Y ≈ 0) + xz_plane_points = edges.points[np.abs(edges.points[:, 1]) < snap_tol] + # Extract points lying in the Y-Z plane (X ≈ 0) + yz_plane_points = edges.points[np.abs(edges.points[:, 0]) < snap_tol] + # Extract points lying in the X-Y plane (Z ≈ 0) + xy_plane_points = edges.points[np.abs(edges.points[:, 2]) < 1e-5] + + xz_cloud = pv.PolyData(xz_plane_points) + yz_cloud = pv.PolyData(yz_plane_points) + xy_cloud = pv.PolyData(xy_plane_points) + + pl = pv.Plotter() + pl.add_mesh(model, color='white', opacity=0.5, label='base STL') + pl.add_mesh(edges, color='black', line_width=5, opacity=0.8,) + pl.add_mesh(xz_cloud, color='green', point_size=20, render_points_as_spheres=True, label='XZ plane points') + pl.add_mesh(yz_cloud, color='orange', point_size=20, render_points_as_spheres=True, label='YZ plane points') + pl.add_mesh(xy_cloud, color='magenta', point_size=20, render_points_as_spheres=True, label='XY plane points') + pl.add_legend() + pl.show() + + def refine_axis(self, xmin, xmax, Nx, x_snaps, + method='insert', tol=1e-12): + + # Loss function to minimize cell size spread + def loss_function(x, x0, is_snap): + # avoid moving snap points + penalty_snap = np.sum((x[is_snap] - x0[is_snap])**2) * 1000 + # avoid gaps < uniform gap + dx = np.diff(x) + threshold = 1/(len(x)-1) # or a hardcoded `min_spacing` + penalty_small_gaps = np.sum((threshold - dx[dx < threshold])**2) + # avoid large spread in gap length + dx = np.diff(x) + penalty_variance = np.std(dx) * 10 + #return penalty_snap + penalty_small_gaps + penalty_variance + return np.hstack([penalty_snap, penalty_small_gaps, penalty_variance]) + + # Uniformly distributed points as initial guess + x_snaps = (x_snaps-xmin)/(xmax-xmin) # normalize to [0,1] + + if method == 'insert': + x0 = np.unique(np.append(x_snaps, np.linspace(0, 1, Nx - len(x_snaps)))) + + elif method == 'neighbor': + x = np.linspace(0, 1, Nx) + dx = np.diff(x)[0] + mask = np.zeros_like(x, dtype=bool) + i=0 + for s in x_snaps: + m = np.isclose(x, s, rtol=0.0, atol=dx/2) + if np.sum(m)>0: + x[np.argmax(m)] = s + x0 = x.copy() + + elif method == 'subdivision': + # x = snaps + while len(x) < Nx: + #idx of segments sorted min -> max + idx_max_diffs = np.argsort(np.diff(x))[-1] # take bigger + + print(f"Bigger segment starts at {x[idx_max_diffs]}") + # compute new point in the middle of the segment + val = x[idx_max_diffs] + (x[idx_max_diffs + 1] - x[idx_max_diffs]) / 2 + + # insert the new point + x = np.insert(x, idx_max_diffs+1, val) + x = np.unique(x) + print(f"Inserted point {val} at index {idx_max_diffs}") + x0 = x.copy() + else: + raise ValueError(f"Method {method} not supported. Use 'insert', 'neighbor' or 'subdivision'.") + + # minimize segment length spread for the test points + is_snap = np.isin(x0, x_snaps) + result = least_squares(loss_function, + x0=x0.copy(), + bounds=(0,1),#(zmin, zmax), + jac='3-point', + method='dogbox', + loss='arctan', + gtol=tol, + ftol=tol, + xtol=tol, + verbose=1, + args=(x0.copy(), is_snap.copy()), + ) + # transform back to [xmin, xmax] + return result.x*(xmax-xmin)+xmin + + def refine_xyz_axis(self, method='insert', tol=1e-6): + '''Refine the grid in the x, y, z axis + using the snap points extracted from the stl solids. + The snap points are used to refine the grid ''' + + if self.verbose: print(f'* Refining x axis with {len(self.x_snaps)} snaps...') + self.x = self.refine_axis(self.xmin, self.xmax, self.Nx+1, self.x_snaps, + method=method, tol=tol) + + if self.verbose: print(f'* Refining y axis with {len(self.y_snaps)} snaps...') + self.y = self.refine_axis(self.ymin, self.ymax, self.Ny+1, self.y_snaps, + method=method, tol=tol) + + if self.verbose: print(f'* Refining z axis with {len(self.z_snaps)} snaps...') + self.z = self.refine_axis(self.zmin, self.zmax, self.Nz+1, self.z_snaps, + method=method, tol=tol) + + self.Nx = len(self.x) - 1 + self.Ny = len(self.y) - 1 + self.Nz = len(self.z) - 1 + self.dx = np.min(np.diff(self.x)) #TODO: should this be an array? + self.dy = np.min(np.diff(self.y)) + self.dz = np.min(np.diff(self.z)) + + print(f"Refined grid: Nx = {len(self.x)}, Ny ={len(self.y)}, Nz = {len(self.z)}") + + def assign_colors(self): + '''Classify colors assigned to each solid + based on the categories in `material_colors` dict + inside `materials.py` + ''' + self.stl_colors = {} + + for key in self.stl_solids: + mat = self.stl_materials[key] + if type(mat) is str: + self.stl_colors[key] = mat + elif len(mat) == 2: + if mat[0] is np.inf: #eps_r + self.stl_colors[key] = 'pec' + elif mat[0] > 1.0: #eps_r + self.stl_colors[key] = 'dielectric' + else: + self.stl_colors[key] = 'vacuum' + elif len(mat) == 3: + self.stl_colors[key] = 'lossy metal' + else: + self.stl_colors[key] = 'other' + + def plot_solids(self, bounding_box=False, show_grid=False, anti_aliasing=None, + opacity=1.0, specular=0.5, offscreen=False, **kwargs): + """ + Generates a 3D visualization of the imported STL geometries using PyVista. + + Parameters: + ----------- + bounding_box : bool, optional + If True, adds a bounding box around the plotted geometry (default: False). + + show_grid : bool, optional + If True, adds the grid's mesh wireframe to the display (default: False). + + anti_aliasing : str or None, optional + Enables anti-aliasing if provided. Valid values depend on PyVista settings (default: None). + + opacity : float, optional + Controls the transparency of the plotted solids. A value of 1.0 is fully opaque, + while 0.0 is fully transparent (default: 1.0). + + specular : float, optional + Adjusts the specular lighting effect on the surface. Higher values increase shininess (default: 0.5). + + **kwargs : dict + Additional keyword arguments passed to `pyvista.add_mesh()`, allowing customization of the mesh rendering. + + Notes: + ------ + - Colors are determined by the `GridFIT3D.stl_colors` attribute dictionary if not None + - Solids labeled as 'vacuum' are rendered with a default opacity of 0.3 for visibility.e. + - The camera is positioned at an angle to provide better depth perception. + - If `bounding_box=True`, a bounding box is drawn around the model. + - If `anti_aliasing` is specified, it is enabled to improve rendering quality. + + """ + + from .materials import material_colors + + pl = pv.Plotter() + pl.add_mesh(self.grid, opacity=0., name='grid', show_scalar_bar=False) + for key in self.stl_solids: + try: + color = material_colors[self.stl_colors[key]] # match library e.g. 'vacuum' + except: + color = self.stl_colors[key] # specifies color e.g. 'tab:red' + + if self.stl_colors[key] == 'vacuum' or self.stl_materials[key] == 'vacuum': + _opacity = 0.3 + else: + _opacity = opacity + pl.add_mesh(self.read_stl(key), color=color, + opacity=_opacity, specular=specular, smooth_shading=True, + **kwargs) + + pl.set_background('mistyrose', top='white') + try: pl.add_logo_widget('../docs/img/wakis-logo-pink.png') + except: pass + pl.camera_position = 'zx' + pl.camera.azimuth += 30 + pl.camera.elevation += 30 + pl.add_axes() + + if anti_aliasing is not None: + pl.enable_anti_aliasing(anti_aliasing) + + if bounding_box: + pl.add_bounding_box() + + if show_grid: + pl.add_mesh(self.grid, style='wireframe', color='grey', opacity=0.3, name='grid') + + if offscreen: + pl.export_html('grid_plot_solids.html') + else: + pl.show() + + def plot_stl_mask(self, stl_solid, cmap='viridis', bounding_box=True, show_grid=True, + add_stl='all', stl_opacity=0., stl_colors=None, + xmax=None, ymax=None, zmax=None, + anti_aliasing='ssaa', offscreen=False): + + """ + Interactive 3D visualization of the structured grid mask and imported STL geometries. + + This routine uses PyVista to display the grid scalar field corresponding to a + chosen STL mask. It provides interactive slider widgets to clip the domain + along the X, Y, and Z directions. At each slider position, the clipped + scalar field is shown with a colormap while the grid structure is shown + as a 2D slice in wireframe. Optionally, one or more STL geometries can + be added to the scene, along with a bounding box of the simulation domain. + + Parameters + ---------- + stl_solid : str + Key name of the `stl_solids` dictionary to retrieve the mask for + visualization (used as the scalar field). + cmap : str, default 'viridis' + Colormap used to visualize the clipped scalar values. + bounding_box : bool, default True + If True, add a static wireframe bounding box of the simulation domain. + show_grid : bool, default True + If True, adds the computational grid overlay on the clipped slice + add_stl : {'all', str, list[str]}, default 'all' + STL geometries to add: + * 'all' → add all STL solids found in `self.stl_solids` + * str → add a single STL solid by key + * list → add a list of STL solids by key + If None, no STL surfaces are added. + stl_opacity : float, default 0.0 + Opacity of the STL surfaces (0 = fully transparent, 1 = fully opaque). + stl_colors : str, list[str], dict, or None, default None + Color(s) of the STL surfaces: + * str → single color for all STL surfaces + * list → per-solid colors, in order + * dict → mapping from STL key to color (using `material_colors`) + * None → use default colors from `self.stl_colors` + xmax, ymax, zmax : float, optional + Initial clipping positions along each axis. If None, use the + maximum domain extent. + anti_aliasing : {'ssaa', 'fxaa', None}, default 'ssaa' + Anti-aliasing mode passed to `pl.enable_anti_aliasing`. + offscreen : bool, default False + If True, render offscreen and export the scene to + ``grid_stl_mask_.html``. If False, open an interactive window. + + Notes + ----- + * Three sliders (X, Y, Z) control clipping of the scalar field by a box + along the respective axis. The clipped scalar field is shown with the + given colormap. A simultaneous 2D slice of the grid is displayed in + wireframe at the clip location. + * STL solids can be visualized in transparent mode to show the relation + between the structured grid and the geometry. + * A static domain bounding box can be added for reference. + * Camera, background, and lighting are pre-configured for clarity + """ + from .materials import material_colors + if stl_colors is None: + stl_colors = self.stl_colors + + if xmax is None: xmax = self.xmax + if ymax is None: ymax = self.ymax + if zmax is None: zmax = self.zmax + + pv.global_theme.allow_empty_mesh = True + pl = pv.Plotter() + vals = {'x':xmax, 'y':ymax, 'z':zmax} + + # --- Update function --- + def update_clip(val, axis="x"): + vals[axis] = val + # define bounds dynamically + if axis == "x": + slice_obj = self.grid.slice(normal="x", origin=(val, 0, 0)) + elif axis == "y": + slice_obj = self.grid.slice(normal="y", origin=(0, val, 0)) + else: # z + slice_obj = self.grid.slice(normal="z", origin=(0, 0, val)) + + # add clipped volume (scalars) + pl.add_mesh( + self.grid.clip_box(bounds=(self.xmin, vals['x'], + self.ymin, vals['y'], + self.zmin, vals['z']), invert=False), + scalars=stl_solid, + cmap=cmap, + name="clip", + ) + + # add slice wireframe (grid structure) + if show_grid: + pl.add_mesh(slice_obj, style="wireframe", color="grey", name="slice") + + # Plot stl surface(s) + if add_stl is not None: + if type(add_stl) is str: #add all stl solids + if add_stl.lower() == 'all': + for i, key in enumerate(self.stl_solids): + surf = self.read_stl(key) + if type(stl_colors) is dict: + if type(stl_colors[key]) is str: + pl.add_mesh(surf, color=material_colors[stl_colors[key]], opacity=stl_opacity, silhouette=dict(color=material_colors[stl_colors[key]]), name=key) + else: + pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=dict(color=stl_colors[key]), name=key) + elif type(stl_colors) is list: + pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, silhouette=dict(color=stl_colors[i]), name=key) + else: + pl.add_mesh(surf, color='white', opacity=stl_opacity, silhouette=True, name=key) + else: #add 1 selected stl solid + key = add_stl + surf = self.read_stl(key) + if type(stl_colors[key]) is str: + pl.add_mesh(surf, color=material_colors[stl_colors[key]], opacity=stl_opacity, silhouette=dict(color=material_colors[stl_colors[key]]), name=key) + else: + pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=dict(color=stl_colors[key]), name=key) + + elif type(add_stl) is list: #add selected list of stl solids + for i, key in enumerate(add_stl): + surf = self.read_stl(key) + if type(stl_colors[key]) is dict: + if type(stl_colors) is str: + pl.add_mesh(surf, color=material_colors[stl_colors[key]], opacity=stl_opacity, silhouette=dict(color=material_colors[stl_colors[key]]), name=key) + else: + pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=dict(color=stl_colors[key]), name=key) + elif type(stl_colors) is list: + pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, silhouette=dict(color=stl_colors[i]), name=key) + else: + pl.add_mesh(surf, color='white', opacity=stl_opacity, silhouette=True, name=key) + + # --- Sliders (placed side-by-side vertically) --- + pl.add_slider_widget( + lambda val: update_clip(val, "x"), + [self.xmin, self.xmax], + value=xmax, title="X Clip", + pointa=(0.8, 0.8), pointb=(0.95, 0.8), # top-right + style='modern', + ) + + pl.add_slider_widget( + lambda val: update_clip(val, "y"), + [self.ymin, self.ymax], + value=ymax, title="Y Clip", + pointa=(0.8, 0.6), pointb=(0.95, 0.6), # middle-right + style='modern', + ) + + pl.add_slider_widget( + lambda val: update_clip(val, "z"), + [self.zmin, self.zmax], + value=zmax, title="Z Clip", + pointa=(0.8, 0.4), pointb=(0.95, 0.4), # lower-right + style='modern', + ) + + # Camera orientation + pl.camera_position = 'zx' + pl.camera.azimuth += 30 + pl.camera.elevation += 30 + pl.set_background('mistyrose', top='white') + try: pl.add_logo_widget('../docs/img/wakis-logo-pink.png') + except: pass + pl.add_axes() + pl.enable_3_lights() + pl.enable_anti_aliasing(anti_aliasing) + + if bounding_box: + pl.add_mesh(pv.Box(bounds=(self.xmin, self.xmax, self.ymin, self.ymax, self.zmin, self.zmax)), + style="wireframe", color="black", line_width=2, name="domain_box") + + if offscreen: + pl.export_html(f'grid_stl_mask_{stl_solid}.html') + else: + pl.show() + + def inspect(self, add_stl=None, stl_opacity=0.5, stl_colors=None, + anti_aliasing='ssaa', offscreen=False): + + '''3D plot using pyvista to visualize + the structured grid and + the imported stl geometries + + Parameters + --- + add_stl: str or list, optional + List or str of stl solids to add to the plot by `pv.add_mesh` + stl_opacity: float, default 0.1 + Opacity of the stl surfaces (0 - Transparent, 1 - Opaque) + stl_colors: str or list of str, default 'white' + Color of the stl surfaces + ''' + from .materials import material_colors + if stl_colors is None: + stl_colors = self.stl_colors + + pv.global_theme.allow_empty_mesh = True + pl = pv.Plotter() + pl.add_mesh(self.grid, show_edges=True, cmap=['white', 'white'], name='grid') + def clip(widget): + # Plot structured grid + b = widget.bounds + x = self.x[np.logical_and(self.x>=b[0], self.x<=b[1])] + y = self.y[np.logical_and(self.y>=b[2], self.y<=b[3])] + z = self.z[np.logical_and(self.z>=b[4], self.z<=b[5])] + X, Y, Z = np.meshgrid(x, y, z, indexing='ij') + grid = pv.StructuredGrid(X.transpose(), Y.transpose(), Z.transpose()) + + pl.add_mesh(grid, show_edges=True, cmap=['white', 'white'], name='grid') + # Plot stl surface(s) + if add_stl is not None: #add 1 selected stl solid + if type(add_stl) is str: + key = add_stl + surf = self.read_stl(key) + surf = surf.clip_box(widget.bounds, invert=False) + if type(stl_colors[key]) is str: + pl.add_mesh(surf, color=material_colors[self.stl_colors[key]], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) + else: + pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) + + elif type(add_stl) is list: #add selected list of stl solids + for i, key in enumerate(add_stl): + surf = self.read_stl(key) + surf = surf.clip_box(widget.bounds, invert=False) + if type(stl_colors) is dict: + if type(stl_colors[key]) is str: + pl.add_mesh(surf, color=material_colors[stl_colors[key]], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) + else: + pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) + elif type(stl_colors) is list: + pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) + else: + pl.add_mesh(surf, color='white', opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) + + else: #add all stl solids + for i, key in enumerate(self.stl_solids): + surf = self.read_stl(key) + surf = surf.clip_box(widget.bounds, invert=False) + if type(stl_colors) is dict: + if type(stl_colors[key]) is str: + pl.add_mesh(surf, color=material_colors[stl_colors[key]], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) + else: + pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) + elif type(stl_colors) is list: + pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) + else: + pl.add_mesh(surf, color='white', opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) + + _ = pl.add_box_widget(callback=clip, rotation_enabled=False) + + # Camera orientation + pl.camera_position = 'zx' + pl.camera.azimuth += 30 + pl.camera.elevation += 30 + pl.set_background('mistyrose', top='white') + try: pl.add_logo_widget('../docs/img/wakis-logo-pink.png') + except: pass + #pl.camera.zoom(zoom) + pl.add_axes() + pl.enable_3_lights() + pl.enable_anti_aliasing(anti_aliasing) + + if offscreen: + pl.export_html('grid_inspect.html') + else: + pl.show() \ No newline at end of file diff --git a/build/lib/wakis/materials.py b/build/lib/wakis/materials.py new file mode 100644 index 0000000..155d8e2 --- /dev/null +++ b/build/lib/wakis/materials.py @@ -0,0 +1,50 @@ +# copyright ################################# # +# This file is part of the wakis Package. # +# Copyright (c) CERN, 2024. # +# ########################################### # + +''' +Material library dictionary + +Format (non-conductive): +{ + 'material key' : [eps_r, mu_r], +} + +Format (conductive): +{ + 'material key' : [eps_r, mu_r, sigma[S/m]], +} + +! Note: +* 'material key' in lower case only +* eps = eps_r*eps_0 and mu = mu_r*mu_0 +''' + +import numpy as np + +material_lib = { + 'pec' : [np.inf, 1.], + + 'vacuum' : [1.0, 1.0], + + 'dielectric' : [10., 1.0], + + 'lossy metal' : [10, 1.0, 10], + + 'copper' : [5.8e+07, 1.0, 5.8e+07], + + 'berillium' : [2.5e+07, 1.0, 2.5e+07], +} + +material_colors = { + 'pec' : 'silver', + + 'vacuum' : 'tab:blue', + + 'dielectric' : 'tab:green', + + 'lossy metal' : 'tab:orange', + + 'other' : 'white', +} \ No newline at end of file diff --git a/build/lib/wakis/plotting.py b/build/lib/wakis/plotting.py new file mode 100644 index 0000000..cdf2ef6 --- /dev/null +++ b/build/lib/wakis/plotting.py @@ -0,0 +1,820 @@ +# copyright ################################# # +# This file is part of the wakis Package. # +# Copyright (c) CERN, 2024. # +# ########################################### # + +import numpy as np +import matplotlib.pyplot as plt + +class PlotMixin: + + def plot3D(self, field='E', component='z', clim=None, hide_solids=None, + show_solids=None, add_stl=None, stl_opacity=0.1, stl_colors='white', + title=None, cmap='jet', + clip_interactive=False, clip_normal='-y', + clip_box=False, clip_bounds=None, + off_screen=False, zoom=0.5, camera_position=None, + nan_opacity=1.0, n=None): + ''' + Built-in 3D plotting using PyVista + + Parameters: + ----------- + field: str, default 'E' + 3D field magnitude ('E', 'H', or 'J') to plot + To plot a component 'Ex', 'Hy' is also accepted + component: str, default 'z' + 3D field compoonent ('x', 'y', 'z', 'Abs') to plot. It will be overriden + if a component is defined in field + clim: list, optional + Colorbar limits for the field plot [min, max] + hide_solids: bool, optional + Mask the values inside solid to np.nan. NaNs will be shown in gray, + since there is a bug with the nan_opacity parameter + show_solids: bool, optional + Mask the values outside solid to np.nan. + add_stl: str or list, optional + List or str of stl solids to add to the plot by `pv.add_mesh` + stl_opacity: float, default 0.1 + Opacity of the stl surfaces (0 - Transparent, 1 - Opaque) + stl_colors: str or list of str, default 'white' + Color of the stl surfaces + title: str, optional + Title used to save the screenshot of the 3D plot (Path+Name) if off_screen=True + cmap: str, default 'jet' + Colormap name to use in the field display + clip_interactive: bool, default False + Enable an interactive widget to clip out part of the domain, plane normal is defined by + `clip_normal` parameter + clip_normal: str, default '-y' + Normal direction of the clip_volume interactive plane + clip_box: bool, default False + Enable a static box clipping of the domain. The box bounds are defined by `clip_bounds` parameter + clip_bounds: Default None + List of bounds [xmin, xmax, ymin, ymax, zmin, zmax] of the box to clip if clip_box is active. + field_on_stl : bool, default False + Samples the field on the stl file specified in `add_stl`. + field_opacity : optional, default 1.0 + Sets de opacity of the `field_on_stl` plot + off_screen: bool, default False + Enable plot rendering off screen, for gif frames generation. + Plot will not be rendered if set to True. + n: int, optional + Timestep number to be added to the plot title and figsave title. + ''' + if self.use_mpi: + print('*** plot3D is not supported when `use_mpi=True`') + return + + import pyvista as pv + + if len(field) == 2: #support for e.g. field='Ex' + component = field[1] + field = field[0] + + if title is None: + title = field + component +'3d' + + if self.plotter_active and not off_screen: + self.plotter_active = False + + if not self.plotter_active: + + pl = pv.Plotter(off_screen=off_screen) + + # Plot stl surface(s) + if add_stl is not None: + if type(add_stl) is str: + key = add_stl + surf = self.grid.read_stl(key) + pl.add_mesh(surf, color=stl_colors, opacity=stl_opacity, smooth_shading=True) + + elif type(add_stl) is list: + for i, key in enumerate(add_stl): + surf = self.grid.read_stl(key) + if type(stl_colors) is list: + pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, smooth_shading=True) + else: + pl.add_mesh(surf, color=stl_colors, opacity=stl_opacity, smooth_shading=True) + else: + key = self.grid.stl_solids.keys()[0] + surf = self.grid.read_stl(key) + pl.add_mesh(surf, color=stl_colors, opacity=stl_opacity, smooth_shading=True) + + if camera_position is None: + pl.camera_position = 'zx' + pl.camera.azimuth += 30 + pl.camera.elevation += 30 + else: + pl.camera_position = camera_position + + pl.set_background('mistyrose', top='white') + try: pl.add_logo_widget('docs/img/wakis-logo-pink.png') + except: pass + pl.camera.zoom(zoom) + pl.add_axes() + pl.enable_3_lights() + + if off_screen: + self.plotter_active = True + else: + pl = self.pl + + + # Plot field + if field == 'E': + self.grid.grid.cell_data[field+component] = np.reshape(self.E[:, :, :, component], self.N) + elif field == 'H': + self.grid.grid.cell_data[field+component] = np.reshape(self.H[:, :, :, component], self.N) + elif field == 'J': + self.grid.grid.cell_data[field+component] = np.reshape(self.J[:, :, :, component], self.N) + else: + print("`field` value not valid") + + points = self.grid.grid.cell_data_to_point_data() #interpolate + + # Mask the values inside solid to np.nan + if hide_solids is not None: + tol = np.min([self.dx, self.dy, self.dz])*1e-3 + if type(hide_solids) is str: + surf = self.grid.read_stl(hide_solids) + select = self.grid.grid.select_enclosed_points(surf, tolerance=tol) + mask = select['SelectedPoints'] > 0 + + elif type(hide_solids) is list: + for i, solid in enumerate(hide_solids): + surf = self.grid.read_stl(solid) + select = self.grid.grid.select_enclosed_points(surf, tolerance=tol) + if i == 0: + mask = select['SelectedPoints'] > 0 + else: + mask += select['SelectedPoints'] > 0 + + points[field+component][mask] = np.nan + + # Mask the values outside solid to np.nan + if show_solids is not None: + tol = np.min([self.dx, self.dy, self.dz])*1e-3 + if type(show_solids) is str: + surf = self.grid.read_stl(show_solids) + select = self.grid.grid.select_enclosed_points(surf, tolerance=tol) + mask = select['SelectedPoints'] > 0 + + elif type(show_solids) is list: + for solid in show_solids: + surf = self.grid.read_stl(solid) + select = self.grid.grid.select_enclosed_points(surf, tolerance=tol) + if i == 0: + mask = select['SelectedPoints'] > 0 + else: + mask += select['SelectedPoints'] > 0 + + points[field+component][np.logical_not(mask)] = np.nan + + # Clip a rectangle of the domain + if clip_box: + if clip_bounds is None: + Lx, Ly = (self.grid.xmax-self.grid.xmin), (self.grid.ymax-self.grid.ymin) + clip_bounds = [self.grid.xmax-Lx/2, self.grid.xmax, + self.grid.ymax-Ly/2, self.grid.ymax, + self.grid.zmin, self.grid.zmax] + + ac1 = pl.add_mesh(points.clip_box(bounds=clip_bounds), opacity=nan_opacity, + scalars=field+component, cmap=cmap, clim=clim) + + # Enable an interactive widget to clip out part of the domain with a plane, with clip_normal + elif clip_interactive: + ac1 = pl.add_mesh_clip_plane(points, normal=clip_normal, opacity=1.0, + scalars=field+component, cmap=cmap, clim=clim, + normal_rotation=False, nan_opacity=nan_opacity) + else: + print('Plotting option inconsistent') + + # Save + if n is not None: + pl.add_title(field+component+f' field, timestep={n}', font='times', font_size=12) + title += '_'+str(n).zfill(6) + if off_screen: + pl.screenshot(title+'.png') + pl.remove_actor(ac1) + self.pl = pl + else: + pl.show(full_screen=False) + + def plot3DonSTL(self, field='E', component='z', clim=None, cmap='jet', log_scale=False, + stl_with_field=None, field_opacity=1.0, tolerance=None, + stl_transparent=None, stl_opacity=0.1, stl_colors='white', + clip_plane = False, clip_interactive=False, + clip_normal='-x', clip_origin=[0,0,0], + clip_box=False, clip_bounds=None, + title=None, off_screen=False, n=None, + zoom=0.5, camera_position=None, **kwargs): + ''' + Built-in 3D plotting using PyVista + + Parameters: + ----------- + field: str, default 'E' + 3D field magnitude ('E', 'H', or 'J') to plot + To plot a component 'Ex', 'Hy' is also accepted + component: str, default 'z' + 3D field compoonent ('x', 'y', 'z', 'Abs') to plot. It will be overriden + if a component is defined in field + clim: list, optional + Colorbar limits for the field plot [min, max] + cmap: str or cmap obj, default 'jet' + Colormap to use for the field plot + log_scale: bool, default False + Turns on logarithmic scale colorbar + stl_with_field : list or str + STL str name or list of names to samples the selected field on + field_opacity : optional, default 1.0 + Sets de opacity of the `field_on_stl` plot + tolerance : float, default None + Tolerance to apply to PyVista's sampling algorithm + stl_transparent: list or str, default None + STL name or list of names to add to the scene with the selected transparency and color + stl_opacity: float, default 0.1 + Opacity of the STL solids without field + stl_colors: list or str, default 'white' + str or list of colors to use for each STL solid + clip_interactive: bool, default False + Enable an interactive widget to clip out part of the domain, plane normal is defined by + `clip_normal` parameter + clip_plane: bool, default False + Clip stl_with_field surface with a plane and show field on such plane + clip_normal: str, default '-y' + Normal direction of the clip_volume interactive plane and the clip_plane + clip_origin: list, default [0,0,0] + Origin of the clipping plane for the clip_plane option + clip_box: bool, default False + Enable a static box clipping of the domain. The box bounds are defined by `clip_bounds` parameter + clip_bounds: Default None + List of bounds [xmin, xmax, ymin, ymax, zmin, zmax] of the box to clip if clip_box is active. + off_screen: bool, default False + Enable plot rendering off screen, for gif frames generation. + Plot will not be rendered if set to True. + title: str + Name to use in the .png savefile with off_screen is True + n: int, optional + Timestep number to be added to the plot title and figsave title. + **kwargs: optional + PyVista's add_mesh optional arguments: + https://docs.pyvista.org/api/plotting/_autosummary/pyvista.plotter.add_mesh + ''' + if self.use_mpi: + print('*** plot3D is not supported when `use_mpi=True`') + return + + import pyvista as pv + + if len(field) == 2: #support for e.g. field='Ex' + component = field[1] + field = field[0] + + if title is None: + title = field + component +'3d' + + if self.plotter_active and not off_screen: + self.plotter_active = False + + if not self.plotter_active: + pl = pv.Plotter(off_screen=off_screen, lighting='none') + light = pv.Light(light_type='headlight') + pl.add_light(light) + + # Plot stl surface(s) + if stl_transparent is not None: + if type(stl_transparent) is str: + key = stl_transparent + surf = self.grid.read_stl(key) + pl.add_mesh(surf, color=stl_colors, opacity=stl_opacity, smooth_shading=True) + + elif type(stl_transparent) is list: + for i, key in enumerate(stl_transparent): + surf = self.grid.read_stl(key) + if type(stl_colors) is list: + pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, smooth_shading=True) + else: + pl.add_mesh(surf, color=stl_colors, opacity=stl_opacity, smooth_shading=True) + + if off_screen: + self.plotter_active = True + else: + pl = self.pl + + # Plot field + if field == 'E': + self.grid.grid.cell_data[field+component] = np.reshape(self.E[:, :, :, component], self.N) + + elif field == 'H': + self.grid.grid.cell_data[field+component] = np.reshape(self.H[:, :, :, component], self.N) + + elif field == 'J': + self.grid.grid.cell_data[field+component] = np.reshape(self.J[:, :, :, component], self.N) + else: + print("`field` value not valid") + + points = self.grid.grid.cell_data_to_point_data() #interpolate + + # Interpolate fields on stl + if stl_with_field is not None: + if type(stl_with_field) is str: + key = stl_with_field + surf = self.grid.read_stl(key) + if clip_plane: + try: surf = surf.clip_closed_surface(normal=clip_normal, origin=clip_origin).subdivide_adaptive(max_edge_len=3*self.dz) + except: print("Surface non-manifold, clip with plane skipped") + + fieldonsurf = surf.sample(points, tolerance) + + if clip_interactive: # interactive plotting with a plane + ac1 = pl.add_mesh_clip_plane(fieldonsurf, normal=clip_normal, normal_rotation=False, + scalars=field+component, opacity=field_opacity, + cmap=cmap, clim=clim, log_scale=log_scale, + **kwargs) + + elif clip_box: # Clip a rectangle of the domain + if clip_bounds is None: + Lx, Ly = (self.grid.xmax-self.grid.xmin), (self.grid.ymax-self.grid.ymin) + clip_bounds = [self.grid.xmax-Lx/2, self.grid.xmax, + self.grid.ymax-Ly/2, self.grid.ymax, + self.grid.zmin, self.grid.zmax] + + ac1 = pl.add_mesh(fieldonsurf.clip_box(bounds=clip_bounds), cmap=cmap, clim=clim, + scalars=field+component, opacity=field_opacity, + log_scale=log_scale, + **kwargs) + + else: + ac1 = pl.add_mesh(fieldonsurf, cmap=cmap, clim=clim, + scalars=field+component, opacity=field_opacity, + log_scale=log_scale, + **kwargs) + + elif type(stl_with_field) is list: + for i, key in enumerate(stl_with_field): + surf = self.grid.read_stl(key) + if clip_plane: + try: surf = surf.clip_closed_surface(normal=clip_normal, origin=clip_origin) + except: print("Surface non-manifold, clip with plane skipped") + + fieldonsurf = surf.sample(points) + + if clip_interactive: # interactive plotting with a plane + ac1 = pl.add_mesh_clip_plane(fieldonsurf, normal=clip_normal, normal_rotation=False, + scalars=field+component, opacity=field_opacity, + cmap=cmap, clim=clim, log_scale=log_scale, + **kwargs) + elif clip_box: # Clip a rectangle of the domain + if clip_bounds is None: + Lx, Ly = (self.grid.xmax-self.grid.xmin), (self.grid.ymax-self.grid.ymin) + clip_bounds = [self.grid.xmax-Lx/2, self.grid.xmax, + self.grid.ymax-Ly/2, self.grid.ymax, + self.grid.zmin, self.grid.zmax] + ac1 = pl.add_mesh(fieldonsurf.clip_box(bounds=clip_bounds), cmap=cmap, clim=clim, + scalars=field+component, opacity=field_opacity, + log_scale=log_scale, **kwargs) + else: + ac1 = pl.add_mesh(fieldonsurf, cmap=cmap, clim=clim, + scalars=field+component, opacity=field_opacity, + log_scale=log_scale, + **kwargs) + if camera_position is None: + pl.camera_position = 'zx' + pl.camera.azimuth += 30 + pl.camera.elevation += 30 + else: + pl.camera_position = camera_position + + pl.set_background('mistyrose', top='white') + try: pl.add_logo_widget('docs/img/wakis-logo-pink.png') + except: pass + pl.camera.zoom(zoom) + pl.add_axes() + pl.enable_anti_aliasing() + #pl.enable_3_lights() + + if n is not None: + pl.add_title(field+component+f' field, timestep={n}', font='times', font_size=12) + title += '_'+str(n).zfill(6) + + # Save + if off_screen: + pl.screenshot(title+'.png') + pl.remove_actor(ac1) + self.pl = pl + else: + pl.show(full_screen=False) + + def plot2D(self, field='E', component='z', plane='ZY', pos=0.5, norm=None, + vmin=None, vmax=None, figsize=[8,4], cmap='jet', patch_alpha=0.1, + patch_reverse=False, add_patch=False, title=None, off_screen=False, + n=None, interpolation='antialiased', dpi=100, return_handles=False): + ''' + Built-in 2D plotting of a field slice using matplotlib + + Parameters: + ---------- + field: str, default 'E' + Field magnitude ('E', 'H', or 'J') to plot + To plot a component 'Ex', 'Hy' is also accepted + component: str, default 'z' + Field compoonent ('x', 'y', 'z', 'Abs') to plot. It will be overriden + if a component is defined in field + plane: arr or str, default 'XZ' + Plane where to plot the 2d field cut: array of 2 slices() and 1 int [x,y,z] + or a str 'XY', 'ZY' or 'ZX' + pos: float, default 0.5 + Position of the cutting plane, as a franction of the plane's normal dimension + e.g. plane 'XZ' wil be sitting at y=pos*(ymax-ymin) + norm: str, default None + Plotting scale to pass to matplotlib imshow: 'linear', 'log', 'symlog' + ** Only for matplotlib version >= 3.8 + vmin: list, optional + Colorbar min limit for the field plot + vmax: list, optional + Colorbar max limit for the field plot + figsize: list, default [8,4] + Figure size to pass to the plot initialization + add_patch: str or list, optional + List or str of stl solids to add to the plot by `pv.add_mesh` + patch_alpha: float, default 0.1 + Value for the transparency of the patch if `add_patch = True` + title: str, optional + Title used to save the screenshot of the 3D plot (Path+Name) if off_screen=True. + If n is provided, 'str(n).zfill(6)' will be added to the title. + cmap: str, default 'jet' + Colormap name to use in the field display + off_screen: bool, default False + Enable plot rendering off screen, for gif frames generation. + Plot will not be rendered if set to True. + n: int, optional + Timestep number to be added to the plot title and figsave title. + interpolation: str, default 'antialiased' + Interpolation method to pass to matplotlib imshow e.g., 'none', + 'antialiased', 'nearest', 'bilinear', 'bicubic', 'spline16', 'spline36', + ''' + from mpl_toolkits.axes_grid1 import make_axes_locatable + + Nx, Ny, Nz = self.Nx, self.Ny, self.Nz + xmin, xmax = self.grid.xmin, self.grid.xmax + ymin, ymax = self.grid.ymin, self.grid.ymax + zmin, zmax = self.grid.zmin, self.grid.zmax + _z = self.z + + if self.use_mpi: + zmin, zmax = self.grid.ZMIN, self.grid.ZMAX + Nz = self.grid.NZ + _z = self.grid.Z + + if type(field) is str: + if len(field) == 2: #support for e.g. field='Ex' + component = field[1] + field = field[0] + elif len(field) == 4: #support for e.g. field='EAbs' + component = field[1:] + field = field[0] + + if title is None: + title = field + component +'2d' + + if type(plane) is not str and len(plane) == 3: + x, y, z = plane[0], plane[1], plane[2] + + if type(plane[2]) is int: + cut = f'(x,y,a) a={round(self.z[z],3)}' + xax, yax = 'y', 'x' + extent = [self.y[y].min(), self.y[y].max(), + self.x[x].min(), self.x[x].max()] + + if type(plane[0]) is int: + cut = f'(a,y,z) a={round(self.x[x],3)}' + xax, yax = 'z', 'y' + extent = [_z[z].min(), _z[z].max(), + self.y[y].min(), self.y[y].max()] + + if type(plane[1]) is int: + cut = f'(x,a,z) a={round(self.y[y],3)}' + xax, yax = 'z', 'x' + extent = [_z[z].min(), _z[z].max(), + self.x[x].min(), self.x[x].max()] + + elif plane == 'XY': + x, y, z = slice(0,Nx), slice(0,Ny), int(Nz*pos) #plane XY + cut = f'(x,y,a) a={round(pos*(zmax-zmin)+zmin,3)}' + xax, yax = 'y', 'x' + extent = [ymin, ymax, xmin, xmax] + + elif plane == 'ZY' or plane == 'YZ': + x, y, z = int(Nx*pos), slice(0,Ny), slice(0,Nz) #plane ZY + cut = f'(a,y,z) a={round(pos*(xmax-xmin)+xmin,3)}' + xax, yax = 'z', 'y' + extent = [zmin, zmax, ymin, ymax] + + elif plane == 'ZX' or plane == 'XZ': + x, y, z = slice(0,Nx), int(Ny*pos), slice(0,Nz) #plane XZ + cut = f'(x,a,z) a={round(pos*(ymax-ymin)+ymin,3)}' + xax, yax = 'z', 'x' + extent = [zmin, zmax, xmin, xmax] + + else: + print("Plane needs to be an array of slices [x,y,z] or a str 'XY', 'ZY', 'ZX'") + + if self.use_mpi: # only in rank=0 + _field = self.mpi_gather(field, x=x, y=y, z=z, component=component) + + if self.rank == 0: + fig, ax = plt.subplots(1,1, figsize=figsize, dpi=dpi) + im = ax.imshow(_field, cmap=cmap, norm=norm, + extent=extent, origin='lower', vmin=vmin, vmax=vmax, + interpolation=interpolation) + + fig.colorbar(im, cax=make_axes_locatable(ax).append_axes('right', size='5%', pad=0.05)) + ax.set_title(f'Wakis {field}{component}{cut}') + ax.set_xlabel(xax) + ax.set_ylabel(yax) + + if n is not None: + fig.suptitle('$'+str(field)+'_{'+str(component)+'}$ field, timestep='+str(n)) + title += '_'+str(n).zfill(6) + + fig.tight_layout() + + if off_screen: + fig.savefig(title+'.png') + plt.clf() + plt.close(fig) + elif return_handles: + return fig, ax + else: + plt.show(block=False) + else: + fig, ax = plt.subplots(1,1, figsize=figsize, dpi=dpi) + if field == 'E': + _field = self.E[x, y, z, component] + if field == 'H': + _field = self.H[x, y, z, component] + if field == 'J': + _field = self.J[x, y, z, component] + + im = ax.imshow(_field, cmap=cmap, norm=norm, + extent=extent, origin='lower', vmin=vmin, vmax=vmax, + interpolation=interpolation) + + fig.colorbar(im, cax=make_axes_locatable(ax).append_axes('right', size='5%', pad=0.05)) + ax.set_title(f'Wakis {field}{component}{cut}') + ax.set_xlabel(xax) + ax.set_ylabel(yax) + + # Patch stl - not supported when running MPI + if add_patch is not None: + if type(add_patch) is str: + mask = np.reshape(self.grid.grid[add_patch], (Nx, Ny, Nz)) + patch = np.ones((Nx, Ny, Nz)) + if patch_reverse: + patch[mask] = np.nan + else: + patch[np.logical_not(mask)] = np.nan + ax.imshow(patch[x,y,z], cmap='Greys', extent=extent, origin='lower', alpha=patch_alpha) + + elif type(add_patch) is list: + for solid in add_patch: + mask = np.reshape(self.grid.grid[solid], (Nx, Ny, Nz)) + patch = np.ones((Nx, Ny, Nz)) + if patch_reverse: + patch[mask] = np.nan + else: + patch[np.logical_not(mask)] = np.nan + ax.imshow(patch[x,y,z], cmap='Greys', extent=extent, origin='lower', alpha=patch_alpha) + + if n is not None: + fig.suptitle('$'+str(field)+'_{'+str(component)+'}$ field, timestep='+str(n)) + title += '_'+str(n).zfill(6) + + fig.tight_layout() + + if off_screen: + fig.savefig(title+'.png') + plt.clf() + plt.close(fig) + elif return_handles: + return fig, ax + else: + plt.show(block=False) + + def plot1D(self, field='E', component='z', line='z', pos=[0.5], + xscale='linear', yscale='linear', xlim=None, ylim=None, + figsize=[8,4], title=None, off_screen=False, n=None, + colors=None, **kwargs): + ''' + Built-in 1D plotting of a field line using matplotlib + + Parameters: + ---------- + field: str, default 'E' + Field magnitude ('E', 'H', or 'J') to plot + To plot a component 'Ex', 'Hy' is also accepted + component: str, default 'z' + Field compoonent ('x', 'y', 'z', 'Abs') to plot. It will be overriden + if a component is defined in field + line: str or list, default 'z' + line of indexes to plot. E.g. line=[0, slice(10,Ny-10), 0] + pos: float or list, default 0.5 + Float or list of floats betwwen 0-1 indicating the cut position. + Only used if line is str. + xlim, ylim: tupple + limits for x and y axis (see matplotlib.ax.set_xlim for more) + xscale, yscale: str + scale to use in x and y axes (see matplotlib.ax.set_xscale for more) + figsize: list, default [8,4] + Figure size to pass to the plot initialization + title: str, optional + Title used to save the screenshot of the 3D plot (Path+Name) if off_screen=True. + If n is provided, 'str(n).zfill(6)' will be added to the title. + cmap: str, default 'jet' + Colormap name to use in the field display + off_screen: bool, default False + Enable plot rendering off screen, for gif frames generation. + Plot will not be rendered if set to True. + n: int, optional + Timestep number to be added to the plot title and figsave title. + colors: list, optional + List of matplotlib-compatible colors. len(colors) >= len(pos) + **kwargs: + Keyword arguments to be passed to the `matplotlib.plot` function. + Default kwargs used: + kwargs = {'color':'g', 'lw':1.2, 'ls':'-'} + ''' + Nx, Ny, Nz = self.Nx, self.Ny, self.Nz + xmin, xmax = self.grid.xmin, self.grid.xmax + ymin, ymax = self.grid.ymin, self.grid.ymax + zmin, zmax = self.grid.zmin, self.grid.zmax + _z = self.z + + if self.use_mpi: + zmin, zmax = self.grid.ZMIN, self.grid.ZMAX + Nz = self.grid.NZ + _z = self.grid.Z + if self.rank == 0: + fig, ax = plt.subplots(1,1, figsize=figsize) + else: + fig, ax = plt.subplots(1,1, figsize=figsize) + + + plotkw = {'lw':1.2, 'ls':'-'} + + if colors is None: + colors = ['k', 'tab:red', 'tab:blue', 'tab:green', + 'tab:orange', 'tab:purple', 'tab:pink'] + plotkw.update(kwargs) + + if len(field) == 2: #support for e.g. field='Ex' + component = field[1] + field = field[0] + + if title is None: + title = field + component +'1d' + + if type(pos) is not list: #support for a list of cut positions + pos_arr = [pos] + else: + pos_arr = pos + + for i, pos in enumerate(pos_arr): + + if type(line) is not str and len(line) == 3: + x, y, z = line[0], line[1], line[2] + + #z-axis + if type(line[2]) is slice: + cut = f'(a,b,z) a={round(self.x[x],3)}, b={round(self.y[y],3)}' + xax = 'z' + xx = _z[z] + xlims = (_z[z].min(), _z[z].max()) + + #x-axis + elif type(line[0]) is slice: + cut = f'(x,a,b) a={round(self.y[y],3)}, b={round(_z[z],3)}' + xax = 'x' + xx = self.x[x] + xlims = (self.x[x].min(), self.x[x].max()) + + #y-axis + elif type(line[1]) is slice: + cut = f'(a,y,b) a={round(self.x[x],3)}, b={round(_z[z],3)}' + xax = 'y' + xx = self.y[y] + xlims = (self.y[y].min(), self.y[y].max()) + + elif line.lower() == 'x': + x, y, z = slice(0,Nx), int(Ny*pos), int(Nz*pos) #x-axis + cut = f'(x,a,b) a={round(self.y[y],3)}, b={round(_z[z],3)}' + xax = 'x' + xx = self.x[x] + xlims = (xmin, xmax) + + elif line.lower() == 'y': + x, y, z = int(Nx*pos), slice(0,Ny), int(Nz*pos) #y-axis + cut = f'(a,y,b) a={round(self.x[x],3)}, b={round(_z[z],3)}' + xax = 'y' + xx = self.y[y] + xlims = (ymin, ymax) + + elif line.lower() == 'z': + x, y, z = int(Nx*pos), int(Ny*pos), slice(0,Nz) #z-axis + cut = f'(a,b,z) a={round(self.x[x],3)}, b={round(self.y[y],3)}' + xax = 'z' + xx = _z[z] + xlims = (zmin, zmax) + + else: + print("line needs to be an array of slices [x,y,z] or a str 'x', 'y', 'z'") + + if i == 0: # first one on top + zorder = 10 + else: zorder = i + + if self.use_mpi: # only in rank=0 + _field = self.mpi_gather(field, x=x, y=y, z=z, component=component) + + if self.rank == 0: + ax.plot(xx, _field, color=colors[i], zorder=zorder, + label=f'{field}{component}{cut}', **plotkw) + else: + if field == 'E': + _field = self.E[x, y, z, component] + if field == 'H': + _field = self.H[x, y, z, component] + if field == 'J': + _field = self.J[x, y, z, component] + + ax.plot(xx, _field, color=colors[i], zorder=zorder, + label=f'{field}{component}{cut}', **plotkw) + + if not self.use_mpi: + yax = f'{field}{component} amplitude' + + ax.set_title(f'Wakis {field}{component}'+(len(pos_arr)==1)*f'{cut}') + ax.set_xlabel(xax) + ax.set_ylabel(yax, color=colors[0]) + ax.set_xlim(xlims) + + ax.set_xscale(xscale) + ax.set_yscale(yscale) + + if len(pos_arr) > 1: + ax.legend(loc=1) + + if xlim is not None: + ax.set_xlim(xlim) + if ylim is not None: + ax.set_ylim(ylim) + + if n is not None: + fig.suptitle('$'+field+'_{'+component+'}$ field, timestep='+str(n)) + title += '_'+str(n).zfill(6) + + fig.tight_layout() + + if off_screen: + fig.savefig(title+'.png') + plt.clf() + plt.close(fig) + + else: + plt.show() + + elif self.use_mpi and self.rank == 0: + yax = f'{field}{component} amplitude' + + ax.set_title(f'Wakis {field}{component}'+(len(pos_arr)==1)*f'{cut}') + ax.set_xlabel(xax) + ax.set_ylabel(yax, color=colors[0]) + ax.set_xlim(xlims) + + ax.set_xscale(xscale) + ax.set_yscale(yscale) + + if len(pos_arr) > 1: + ax.legend(loc=1) + + if xlim is not None: + ax.set_xlim(xlim) + if ylim is not None: + ax.set_ylim(ylim) + + if n is not None: + fig.suptitle('$'+field+'_{'+component+'}$ field, timestep='+str(n)) + title += '_'+str(n).zfill(6) + + fig.tight_layout() + + if off_screen: + fig.savefig(title+'.png') + plt.clf() + plt.close(fig) + + else: + plt.show() + diff --git a/build/lib/wakis/pmlBlock2D.py b/build/lib/wakis/pmlBlock2D.py new file mode 100644 index 0000000..073a82d --- /dev/null +++ b/build/lib/wakis/pmlBlock2D.py @@ -0,0 +1,135 @@ +import numpy as np +from scipy.constants import c as c_light, epsilon_0 as eps_0, mu_0 as mu_0 + + +class PmlBlock2D: + def __init__(self, Nx, Ny, dt, dx, dy, lx_block=None, ly_block=None, rx_block=None, ry_block=None): + self.Nx = Nx + self.Ny = Ny + self.dt = dt + self.dx = dx + self.dy = dy + + self.Ex = np.zeros((Nx, Ny + 1)) + self.Ey = np.zeros((Nx + 1, Ny)) + self.Exy = np.zeros((Nx, Ny + 1)) + self.Exz = np.zeros((Nx, Ny + 1)) + self.Eyx = np.zeros((Nx + 1, Ny)) + self.Eyz = np.zeros((Nx + 1, Ny)) + self.Hz = np.zeros((Nx, Ny)) + self.Hzx = np.zeros((Nx, Ny)) + self.Hzy = np.zeros((Nx, Ny)) + self.Jx = np.zeros((self.Nx, self.Ny + 1)) + self.Jy = np.zeros((self.Nx + 1, self.Ny)) + + # the sigmas must be assembled by the solver + self.sigma_x = np.zeros_like(self.Ex) + self.sigma_y = np.zeros_like(self.Ey) + self.sigma_z = np.zeros_like(self.Ex) + self.sigma_star_x = np.zeros_like(self.Hz) + self.sigma_star_y = np.zeros_like(self.Hz) + self.sigma_star_z = np.zeros_like(self.Hz) + + # we assemble these after the sigmas + self.Ax = np.zeros_like(self.sigma_x) + self.Ay = np.zeros_like(self.sigma_y) + self.Az = np.zeros_like(self.sigma_z) + self.Bx = np.zeros_like(self.sigma_x) + self.By = np.zeros_like(self.sigma_y) + self.Bz = np.zeros_like(self.sigma_z) + self.Cx = np.zeros_like(self.sigma_star_x) + self.Cy = np.zeros_like(self.sigma_star_y) + self.Cz = np.zeros_like(self.sigma_star_z) + self.Dx = np.zeros_like(self.sigma_star_x) + self.Dy = np.zeros_like(self.sigma_star_y) + self.Dz = np.zeros_like(self.sigma_star_z) + + self.lx_block = lx_block + self.ly_block = ly_block + self.rx_block = rx_block + self.ry_block = ry_block + + self.C1 = self.dt / (self.dx * mu_0) + self.C2 = self.dt / (self.dy * mu_0) + self.C4 = self.dt / (self.dy * eps_0) + self.C5 = self.dt / (self.dx * eps_0) + self.C3 = self.dt / eps_0 + self.C6 = self.dt / eps_0 + + def assemble_coeffs(self): + self.Ax = (2 * eps_0 - self.dt * self.sigma_x) / (2 * eps_0 + self.dt * self.sigma_x) + self.Ay = (2 * eps_0 - self.dt * self.sigma_y) / (2 * eps_0 + self.dt * self.sigma_y) + self.Az = (2 * eps_0 - self.dt * self.sigma_z) / (2 * eps_0 + self.dt * self.sigma_z) + self.Bx = 2 * self.dt / (2 * eps_0 + self.dt * self.sigma_x) + self.By = 2 * self.dt / (2 * eps_0 + self.dt * self.sigma_y) + self.Bz = 2 * self.dt / (2 * eps_0 + self.dt * self.sigma_z) + self.Cx = (2 * mu_0 - self.dt * self.sigma_star_x) / (2 * mu_0 + self.dt * self.sigma_star_x) + self.Cy = (2 * mu_0 - self.dt * self.sigma_star_y) / (2 * mu_0 + self.dt * self.sigma_star_y) + self.Cz = (2 * mu_0 - self.dt * self.sigma_star_z) / (2 * mu_0 + self.dt * self.sigma_star_z) + self.Dx = 2 * self.dt / (2 * mu_0 + self.dt * self.sigma_star_x) + self.Dy = 2 * self.dt / (2 * mu_0 + self.dt * self.sigma_star_y) + self.Dz = 2 * self.dt / (2 * mu_0 + self.dt * self.sigma_star_z) + + def advance_h_fdtd(self): + Hz = self.Hz + Ex = self.Ex + Ey = self.Ey + + for ii in range(self.Nx): + for jj in range(self.Ny): + self.Hzx[ii, jj] = self.Cx[ii, jj] * self.Hzx[ii, jj] - self.Dx[ii, jj] / self.dx * (self.Ey[ii + 1, jj] + - self.Ey[ii, jj]) + self.Hzy[ii, jj] = self.Cy[ii, jj] * self.Hzy[ii, jj] + self.Dy[ii, jj] / self.dy * (self.Ex[ii, jj + 1] + - self.Ex[ii, jj]) + + self.Hz = self.Hzx + self.Hzy + + def advance_e_fdtd(self): + Hz = self.Hz + Ex = self.Ex + Ey = self.Ey + + for ii in range(self.Nx): + for jj in range(1, self.Ny): + self.Exy[ii, jj] = self.Ay[ii, jj] * self.Exy[ii, jj] + self.By[ii, jj] / self.dy * ( + self.Hz[ii, jj] - self.Hz[ii, jj - 1]) + self.Exz[ii, jj] = self.Az[ii, jj] * self.Exz[ii, jj] + + for ii in range(1, self.Nx): + for jj in range(self.Ny): + self.Eyz[ii, jj] = self.Az[ii, jj] * self.Eyz[ii, jj] + self.Eyx[ii, jj] = self.Ax[ii, jj] * self.Eyx[ii, jj] - self.Bx[ii, jj] / self.dx * ( + self.Hz[ii, jj] - self.Hz[ii - 1, jj]) + + self.Ex = self.Exy + self.Exz + self.Ey = self.Eyx + self.Eyz + + def update_e_boundary(self): + + if self.lx_block is not None: + for jj in range(self.Ny): + self.Eyz[0, jj] = self.Az[0, jj] * self.Eyz[0, jj] + self.Eyx[0, jj] = self.Ax[0, jj] * self.Eyx[0, jj] - self.Bx[0, jj] / self.dx * ( + self.Hz[0, jj] - self.lx_block.Hz[-1, jj]) + + if self.rx_block is not None: + for jj in range(self.Ny): + self.Eyz[-1, jj] = self.Az[-1, jj] * self.Eyz[-1, jj] + self.Eyx[-1, jj] = self.Ax[-1, jj] * self.Eyx[-1, jj] - self.Bx[-1, jj] / self.dx * ( + self.rx_block.Hz[0, jj] - self.Hz[-1, jj]) + + self.Ey = self.Eyx + self.Eyz + + if self.ly_block is not None: + for ii in range(self.Nx): + self.Exy[ii, 0] = self.Ay[ii, 0] * self.Exy[ii, 0] + self.By[ii, 0] / self.dy * ( + self.Hz[ii, 0] - self.ly_block.Hz[ii, - 1]) + self.Exz[ii, 0] = self.Az[ii, 0] * self.Exz[ii, 0] + + if self.ry_block is not None: + for ii in range(self.Nx): + self.Exy[ii, -1] = self.Ay[ii, -1] * self.Exy[ii, -1] + self.By[ii, -1] / self.dy * ( + self.ry_block.Hz[ii, 0] - self.Hz[ii, - 1]) + self.Exz[ii, -1] = self.Az[ii, -1] * self.Exz[ii, -1] + + self.Ex = self.Exy + self.Exz diff --git a/build/lib/wakis/pmlBlock3D.py b/build/lib/wakis/pmlBlock3D.py new file mode 100644 index 0000000..ae98d0b --- /dev/null +++ b/build/lib/wakis/pmlBlock3D.py @@ -0,0 +1,265 @@ +import numpy as np +from scipy.constants import c as c_light, epsilon_0 as eps_0, mu_0 as mu_0 + + +class PmlBlock3D: + def __init__(self, Nx, Ny, Nz, dt, dx, dy, dz, i_block, j_block, k_block, lx_block=None, ly_block=None, lz_block=None, rx_block=None, + ry_block=None, rz_block=None): + self.Nx = Nx + self.Ny = Ny + self.Nz = Nz + self.dt = dt + self.dx = dx + self.dy = dy + self.dz = dz + + self.i_block = i_block + self.j_block = j_block + self.k_block = k_block + + # Solver3D constructor takes care of populating this matrix (should we ensure this somehow?) + self.blocks_mat = np.full((3, 3, 3), None) + + self.Ex = np.zeros((Nx, Ny + 1, Nz + 1)) + self.Ey = np.zeros((Nx + 1, Ny, Nz + 1)) + self.Ez = np.zeros((Nx + 1, Ny + 1, Nz)) + self.Exy = np.zeros((Nx, Ny + 1, Nz + 1)) + self.Exz = np.zeros((Nx, Ny + 1, Nz + 1)) + self.Eyx = np.zeros((Nx + 1, Ny, Nz + 1)) + self.Eyz = np.zeros((Nx + 1, Ny, Nz + 1)) + self.Ezx = np.zeros((Nx + 1, Ny + 1, Nz)) + self.Ezy = np.zeros((Nx + 1, Ny + 1, Nz)) + self.Hx = np.zeros((Nx + 1, Ny, Nz)) + self.Hy = np.zeros((Nx, Ny + 1, Nz)) + self.Hz = np.zeros((Nx, Ny, Nz + 1)) + self.Hzx = np.zeros((Nx, Ny, Nz + 1)) + self.Hzy = np.zeros((Nx, Ny, Nz + 1)) + self.Hxy = np.zeros((Nx + 1, Ny, Nz)) + self.Hxz = np.zeros((Nx + 1, Ny, Nz)) + self.Hyx = np.zeros((Nx, Ny + 1, Nz)) + self.Hyz = np.zeros((Nx, Ny + 1, Nz)) + self.Jx = np.zeros((self.Nx, self.Ny + 1, self.Nz + 1)) + self.Jy = np.zeros((self.Nx + 1, self.Ny, self.Nz + 1)) + self.Jz = np.zeros((self.Nx + 1, self.Ny + 1, self.Nz)) + + # the sigmas must be assembled by the solver + self.sigma_x = np.zeros((Nx + 1, Ny + 1, Nz + 1)) + self.sigma_y = np.zeros((Nx + 1, Ny + 1, Nz + 1)) + self.sigma_z = np.zeros((Nx + 1, Ny + 1, Nz + 1)) + self.sigma_star_x = np.zeros((Nx, Ny + 1, Nz + 1)) + self.sigma_star_y = np.zeros((Nx + 1, Ny, Nz + 1)) + self.sigma_star_z = np.zeros((Nx + 1, Ny + 1, Nz)) + + # we assemble these after the sigmas + self.Ax = np.zeros_like(self.sigma_x) + self.Ay = np.zeros_like(self.sigma_y) + self.Az = np.zeros_like(self.sigma_z) + self.Bx = np.zeros_like(self.sigma_x) + self.By = np.zeros_like(self.sigma_y) + self.Bz = np.zeros_like(self.sigma_z) + self.Cx = np.zeros_like(self.sigma_star_x) + self.Cy = np.zeros_like(self.sigma_star_y) + self.Cz = np.zeros_like(self.sigma_star_z) + self.Dx = np.zeros_like(self.sigma_star_x) + self.Dy = np.zeros_like(self.sigma_star_y) + self.Dz = np.zeros_like(self.sigma_star_z) + + self.lx_block = lx_block + self.ly_block = ly_block + self.lz_block = lz_block + self.rx_block = rx_block + self.ry_block = ry_block + self.rz_block = rz_block + + self.C1 = self.dt / (self.dx * mu_0) + self.C2 = self.dt / (self.dy * mu_0) + self.C4 = self.dt / (self.dy * eps_0) + self.C5 = self.dt / (self.dx * eps_0) + self.C3 = self.dt / eps_0 + self.C6 = self.dt / eps_0 + + def assemble_coeffs(self): + self.Ax = (2 * eps_0 - self.dt * self.sigma_x) / (2 * eps_0 + self.dt * self.sigma_x) + self.Ay = (2 * eps_0 - self.dt * self.sigma_y) / (2 * eps_0 + self.dt * self.sigma_y) + self.Az = (2 * eps_0 - self.dt * self.sigma_z) / (2 * eps_0 + self.dt * self.sigma_z) + self.Bx = 2 * self.dt / (2 * eps_0 + self.dt * self.sigma_x) + self.By = 2 * self.dt / (2 * eps_0 + self.dt * self.sigma_y) + self.Bz = 2 * self.dt / (2 * eps_0 + self.dt * self.sigma_z) + self.Cx = (2 * mu_0 - self.dt * self.sigma_star_x) / (2 * mu_0 + self.dt * self.sigma_star_x) + self.Cy = (2 * mu_0 - self.dt * self.sigma_star_y) / (2 * mu_0 + self.dt * self.sigma_star_y) + self.Cz = (2 * mu_0 - self.dt * self.sigma_star_z) / (2 * mu_0 + self.dt * self.sigma_star_z) + self.Dx = 2 * self.dt / (2 * mu_0 + self.dt * self.sigma_star_x) + self.Dy = 2 * self.dt / (2 * mu_0 + self.dt * self.sigma_star_y) + self.Dz = 2 * self.dt / (2 * mu_0 + self.dt * self.sigma_star_z) + + def advance_h_fdtd(self): + Hz = self.Hz + Ex = self.Ex + Ey = self.Ey + + for ii in range(self.Nx): + for jj in range(self.Ny): + for kk in range(self.Nz + 1): + self.Hzx[ii, jj, kk] = self.Cx[ii, jj, kk] * self.Hzx[ii, jj, kk] - self.Dx[ii, jj, kk] / self.dx * (self.Ey[ii + 1, jj, kk] + - self.Ey[ii, jj, kk]) + self.Hzy[ii, jj, kk] = self.Cy[ii, jj, kk] * self.Hzy[ii, jj, kk] + self.Dy[ii, jj, kk] / self.dy * (self.Ex[ii, jj + 1, kk] + - self.Ex[ii, jj, kk]) + + for ii in range(self.Nx + 1): + for jj in range(self.Ny): + for kk in range(self.Nz): + self.Hxy[ii, jj, kk] = self.Cy[ii, jj, kk] * self.Hxy[ii, jj, kk] - self.Dy[ii, jj, kk] / self.dy * (self.Ez[ii, jj + 1, kk] + - self.Ez[ii, jj, kk]) + self.Hxz[ii, jj, kk] = self.Cz[ii, jj, kk] * self.Hxz[ii, jj, kk] + self.Dz[ii, jj, kk] / self.dz * (self.Ey[ii, jj, kk + 1] + - self.Ey[ii, jj, kk]) + + for ii in range(self.Nx): + for jj in range(self.Ny + 1): + for kk in range(self.Nz): + self.Hyx[ii, jj, kk] = self.Cx[ii, jj, kk] * self.Hyx[ii, jj, kk] + self.Dx[ii, jj, kk] / self.dx * (self.Ez[ii + 1, jj, kk] + - self.Ez[ii, jj, kk]) + self.Hyz[ii, jj, kk] = self.Cz[ii, jj, kk] * self.Hyz[ii, jj, kk] - self.Dz[ii, jj, kk] / self.dz * (self.Ex[ii, jj, kk + 1] + - self.Ex[ii, jj, kk]) + + def advance_e_fdtd(self): + Hz = self.Hz + Ex = self.Ex + Ey = self.Ey + + for ii in range(self.Nx): + for jj in range(1, self.Ny): + for kk in range(self.Nz + 1): + self.Exy[ii, jj, kk] = self.Ay[ii, jj, kk] * self.Exy[ii, jj, kk] + self.By[ii, jj, kk] / self.dy * ( + self.Hz[ii, jj, kk] - self.Hz[ii, jj - 1, kk]) + + for ii in range(self.Nx): + for jj in range(self.Ny + 1): + for kk in range(1, self.Nz): + self.Exz[ii, jj, kk] = self.Az[ii, jj, kk] * self.Exz[ii, jj, kk] - self.Bz[ii, jj, kk] / self.dz * ( + self.Hy[ii, jj, kk] - self.Hy[ii, jj, kk - 1]) + + for ii in range(1, self.Nx): + for jj in range(self.Ny): + for kk in range(self.Nz + 1): + self.Eyx[ii, jj, kk] = self.Ax[ii, jj, kk] * self.Eyx[ii, jj, kk] - self.Bx[ii, jj, kk] / self.dx * ( + self.Hz[ii, jj, kk] - self.Hz[ii - 1, jj, kk]) + + for ii in range(self.Nx + 1): + for jj in range(self.Ny): + for kk in range(1, self.Nz): + self.Eyz[ii, jj, kk] = self.Az[ii, jj, kk] * self.Eyz[ii, jj, kk] + self.Bz[ii, jj, kk] / self.dz * ( + self.Hx[ii, jj, kk] - self.Hx[ii, jj, kk - 1]) + + for ii in range(1, self.Nx): + for jj in range(self.Ny + 1): + for kk in range(self.Nz): + self.Ezx[ii, jj, kk] = self.Ax[ii, jj, kk]*self.Ezx[ii, jj, kk] + self.Bx[ii, jj, kk] / self.dx * ( + self.Hy[ii, jj, kk] - self.Hy[ii - 1, jj, kk]) + + for ii in range(self.Nx + 1): + for jj in range(1, self.Ny): + for kk in range(self.Nz): + self.Ezy[ii, jj, kk] = self.Ay[ii, jj, kk]*self.Ezy[ii, jj, kk] - self.By[ii, jj, kk] / self.dy * ( + self.Hx[ii, jj, kk] - self.Hx[ii, jj - 1, kk]) + + def sum_e_fields(self): + self.Ex = self.Exy + self.Exz + self.Ey = self.Eyx + self.Eyz + self.Ez = self.Ezx + self.Ezy + + def sum_h_fields(self): + self.Hx = self.Hxy + self.Hxz + self.Hy = self.Hyx + self.Hyz + self.Hz = self.Hzx + self.Hzy + + def update_e_boundary(self): + + i_block = self.i_block + j_block = self.j_block + k_block = self.k_block + blocks_mat = self.blocks_mat + Nx = self.Nx + Ny = self.Ny + Nz = self.Nz + Ax = self.Ax + Ay = self.Ay + Az = self.Az + Bx = self.Bx + By = self.By + Bz = self.Bz + Exy = self.Exy + Exz = self.Exz + Eyx = self.Eyx + Eyz = self.Eyz + Ezx = self.Ezx + Ezy = self.Ezy + Hx = self.Hx + Hy = self.Hy + Hz = self.Hz + dx = self.dx + dy = self.dy + dz = self.dz + + # Separate update on edges doesn't seem to be needed as derivatives are taken in one direction per time + # Update E on "lower" faces + if i_block > 0 and blocks_mat[i_block - 1, j_block, k_block] is not None: + for jj in range(Ny): + for kk in range(Nz + 1): + Eyx[0, jj, kk] = Ax[0, jj, kk] * Eyx[0, jj, kk] - Bx[0, jj, kk] / dx * (Hz[0, jj, kk] + - blocks_mat[i_block - 1, j_block, k_block].Hz[-1, jj, kk]) + for jj in range(Ny + 1): + for kk in range(Nz): + Ezx[0, jj, kk] = Ax[0, jj, kk] * Ezx[0, jj, kk] + Bx[0, jj, kk] / dx * (Hy[0, jj, kk] + - blocks_mat[i_block - 1, j_block, k_block].Hy[-1, jj, kk]) + + if j_block > 0 and blocks_mat[i_block, j_block - 1, k_block] is not None: + for ii in range(Nx): + for kk in range(Nz + 1): + Exy[ii, 0, kk] = Ay[ii, 0, kk] * Exy[ii, 0, kk] + By[ii, 0, kk] / dy * (Hz[ii, 0, kk] + - blocks_mat[i_block, j_block - 1, k_block].Hz[ii, -1, kk]) + for ii in range(Nx + 1): + for kk in range(Nz): + Ezy[ii, 0, kk] = Ay[ii, 0, kk] * Ezy[ii, 0, kk] - By[ii, 0, kk] / dy * (Hx[ii, 0, kk] + - blocks_mat[i_block, j_block - 1, k_block].Hx[ii, -1, kk]) + + if k_block > 0 and blocks_mat[i_block, j_block, k_block - 1] is not None: + for ii in range(Nx + 1): + for jj in range(Ny): + Eyz[ii, jj, 0] = Az[ii, jj, 0] * Eyz[ii, jj, 0] + Bz[ii, jj, 0] / dz * (Hx[ii, jj, 0] + - blocks_mat[i_block, j_block, k_block - 1].Hx[ii, jj, - 1]) + for ii in range(Nx): + for jj in range(Ny + 1): + Exz[ii, jj, 0] = Az[ii, jj, 0] * Exz[ii, jj, 0] - Bz[ii, jj, 0] / dz * (Hy[ii, jj, 0] + - blocks_mat[i_block, j_block, k_block - 1].Hy[ii, jj, - 1]) + + # Update E on "upper" faces + if i_block < 2 and blocks_mat[i_block + 1, j_block, k_block] is not None: + for jj in range(Ny): + for kk in range(Nz + 1): + Eyx[Nx, jj, kk] = Ax[Nx, jj, kk] * Eyx[Nx, jj, kk] - Bx[Nx, jj, kk] / dx * ( + blocks_mat[i_block + 1, j_block, k_block].Hz[0, jj, kk] - Hz[-1, jj, kk]) + for jj in range(Ny + 1): + for kk in range(Nz): + Ezx[Nx, jj, kk] = Ax[Nx, jj, kk] * Ezx[Nx, jj, kk] + Bx[Nx, jj, kk] / dx * ( + blocks_mat[i_block + 1, j_block, k_block].Hy[0, jj, kk] - Hy[-1, jj, kk]) + + if j_block < 2 and blocks_mat[i_block, j_block + 1, k_block] is not None: + for ii in range(Nx): + for kk in range(Nz + 1): + Exy[ii, Ny, kk] = Ay[ii, Ny, kk] * Exy[ii, Ny, kk] + By[ii, Ny, kk] / dy * ( + blocks_mat[i_block, j_block + 1, k_block].Hz[ii, 0, kk] - Hz[ii, -1, kk]) + for ii in range(Nx + 1): + for kk in range(Nz): + Ezy[ii, Ny, kk] = Ay[ii, Ny, kk] * Ezy[ii, Ny, kk] - By[ii, Ny, kk] / dy * ( + blocks_mat[i_block, j_block + 1, k_block].Hx[ii, 0, kk] - Hx[ii, -1, kk]) + + if k_block < 2 and blocks_mat[i_block, j_block, k_block + 1] is not None: + for ii in range(Nx + 1): + for jj in range(Ny): + Eyz[ii, jj, Nz] = Az[ii, jj, Nz] * Eyz[ii, jj, Nz] + Bz[ii, jj, Nz] / dz * ( + blocks_mat[i_block, j_block, k_block + 1].Hx[ii, jj, 0] - Hx[ii, jj, -1]) + for ii in range(Nx): + for jj in range(Ny + 1): + Exz[ii, jj, Nz] = Az[ii, jj, Nz] * Exz[ii, jj, Nz] - Bz[ii, jj, Nz] / dz * ( + blocks_mat[i_block, j_block, k_block + 1].Hy[ii, jj, 0] - Hy[ii, jj, -1]) + diff --git a/build/lib/wakis/routines.py b/build/lib/wakis/routines.py new file mode 100644 index 0000000..46da44f --- /dev/null +++ b/build/lib/wakis/routines.py @@ -0,0 +1,336 @@ +# copyright ################################# # +# This file is part of the wakis Package. # +# Copyright (c) CERN, 2024. # +# ########################################### # + +import numpy as np +import h5py +from tqdm import tqdm +from scipy.constants import c as c_light + +from wakis.sources import Beam + +class RoutinesMixin(): + + def emsolve(self, Nt, source=None, callback=None, + save=False, fields=['E'], components=['Abs'], save_every=1, subdomain=None, + plot=False, plot_every=1, use_etd=False, + plot3d=False, **kwargs): + ''' + Run the simulation and save the selected field components in HDF5 files + for every timestep. Each field will be saved in a separate HDF5 file 'Xy.h5' + where X is the field and y the component. + + Parameters: + ---------- + Nt: int + Number of timesteps to run + source: source object + source object from `sources.py` defining the time-dependednt source. + It should have an update function `source.update(solver, t)` + save: bool + Flag to enable saving the field in HDF5 format + fields: list, default ['E'] + 3D field magnitude ('E', 'H', or 'J') to save + 'Ex', 'Hy', etc., is also accepted and will override + the `components` parameter. + components: list, default ['z'] + Field compoonent ('x', 'y', 'z', 'Abs') to save. It will be overriden + if a component is specified in the`field` parameter + save_every: int, default 1 + Number of timesteps between saves + subdomain: list, default None + Slice [x,y,z] of the domain to be saved + plot: bool, default False + Flag to enable 2D plotting + plot3d: bool, default False + Flag to enable 3D plotting + plot_every: int + Number of timesteps between consecutive plots + **kwargs: + Keyword arguments to be passed to the Plot2D function. + * Default kwargs used for 2D plotting: + {'field':'E', 'component':'z', + 'plane':'ZY', 'pos':0.5, 'title':'Ez', + 'cmap':'rainbow', 'patch_reverse':True, + 'off_screen': True, 'interpolation':'spline36'} + * Default kwargs used for 3D plotting: + {'field':'E', 'component':'z', + 'add_stl':None, 'stl_opacity':0.0, 'stl_colors':'white', + 'title':'Ez', 'cmap':'jet', 'clip_volume':False, 'clip_normal':'-y', + 'field_on_stl':True, 'field_opacity':1.0, + 'off_screen':True, 'zoom':1.0, 'nan_opacity':1.0} + + Raises: + ------- + ImportError: + If the hdf5 dependency cannot be imported + + Dependencies: + ------------- + h5py + ''' + self.Nt = Nt + if source is not None: self.source = source + + if save: + + hfs = {} + for field in fields: + if len(field) == 1: + for component in components: + hfs[field+component] = h5py.File(field+component+'.h5', 'w') + else: + hfs[field] = h5py.File(field+'.h5', 'w') + + for hf in hfs: + hf['x'], hf['y'], hf['z'] = self.x, self.y, self.z + hf['t'] = np.arange(0, Nt*self.dt, save_every*self.dt) + + if subdomain is not None: + xx, yy, zz = subdomain + else: + xx, yy, zz = slice(0,self.Nx), slice(0,self.Ny), slice(0,self.Nz) + + if plot: + plotkw = {'field':'E', 'component':'z', + 'plane':'ZY', 'pos':0.5, 'cmap':'rainbow', + 'patch_reverse':True, 'title':'Ez', + 'off_screen': True, 'interpolation':'spline36'} + plotkw.update(kwargs) + + if plot3d: + plotkw = {'field':'E', 'component':'z', + 'add_stl':None, 'stl_opacity':0.0, 'stl_colors':'white', + 'title':'Ez', 'cmap':'jet', 'clip_volume':False, 'clip_normal':'-y', + 'field_on_stl':True, 'field_opacity':1.0, + 'off_screen':True, 'zoom':1.0, 'nan_opacity':1.0} + + plotkw.update(kwargs) + + # get ABC values + if self.activate_abc: + E_abc_2, H_abc_2 = self.get_abc() + E_abc_1, H_abc_1 = self.get_abc() + + # Time loop + for n in tqdm(range(Nt)): + + if source is not None: + source.update(self, n*self.dt) + + if save and n%save_every == 0: + for field in hfs.keys(): + try: + d = getattr(self, field[0])[xx,yy,zz,field[1:]] + except: + raise(f'Component {field} not valid. Input must have a \ + field ["E", "H", "J"] and a component ["x", "y", "z", "Abs"]') + + # Save timestep in HDF5 + hfs[field]['#'+str(n).zfill(5)] = d + + # Advance + self.one_step() + + # Plot + if plot and n%plot_every == 0: + self.plot2D(n=n, **plotkw) + + if plot3d and n%plot_every == 0: + self.plot3D(n=n, **plotkw) + + # ABC BCs + if self.activate_abc: + self.update_abc(E_abc_2, H_abc_2) # n-2 + E_abc_2, H_abc_2 = E_abc_1, H_abc_1 # n-1 + E_abc_1, H_abc_1 = self.get_abc() # n + + # Callback func(solver, t) + if callback is not None: + callback(self, n*self.dt) + + # End + if save: + for hf in hfs: + hf.close() + + def wakesolve(self, wakelength, + wake=None, + callback=None, + compute_plane='both', + plot=False, plot_from=None, plot_every=1, plot_until=None, + save_J=False, + add_space=None, #for legacy + use_edt=None, #deprecated + **kwargs): + ''' + Run the EM simulation and compute the longitudinal (z) and transverse (x,y) + wake potential WP(s) and impedance Z(s). + + The `Ez` field is saved every timestep in a subdomain (xtest, ytest, z) around + the beam trajectory in HDF5 format file `Ez.h5`. + + The computed results are available as Solver class attributes: + - wake potential: WP (longitudinal), WPx, WPy (transverse) [V/pC] + - impedance: Z (longitudinal), Zx, Zy (transverse) [Ohm] + - beam charge distribution: lambdas (distance) [C/m] lambdaf (spectrum) [C] + + Parameters: + ----------- + wakelength: float + Desired length of the wake in [m] to be computed + + Maximum simulation time in [s] can be computed from the wakelength parameter as: + .. math:: t_{max} = t_{inj} + (wakelength + (z_{max}-z_{min}))/c + wake: Wake obj, default None + `Wake()` object containing the information needed to run + the wake solver calculation. See Wake() docstring for more information. + Can be passed at `Solver()` instantiation as parameter too. + save_J: bool, default False + Flag to enable saving the current J in a diferent HDF5 file 'Jz.h5' + plot: bool, default False + Flag to enable 2D plotting + plot_every: int + Number of timesteps between consecutive plots + **kwargs: + Keyword arguments to be passed to the Plot2D function. + Default kwargs used: + {'plane':'ZY', 'pos':0.5, 'title':'Ez', + 'cmap':'rainbow', 'patch_reverse':True, + 'off_screen': True, 'interpolation':'spline36'} + + Raises: + ------- + AttributeError: + If the Wake object is not provided + ImportError: + If the hdf5 dependency cannot be imported + + Dependencies: + ------------- + h5py + ''' + + if wake is not None: self.wake = wake + if self.wake is None: + raise('Wake solver information not passed to the solver instantiation') + + if add_space is not None: #legacy support + self.wake.skip_cells = add_space + + # plot params defaults + if plot: + plotkw = {'plane':'ZY', 'pos':0.5, 'title':'Ez', + 'cmap':'rainbow', 'patch_reverse':True, + 'off_screen': True, 'interpolation':'spline36'} + plotkw.update(kwargs) + + # integration path (test position) + self.xtest, self.ytest = self.wake.xtest, self.wake.ytest + self.ixt, self.iyt = np.abs(self.x-self.xtest).argmin(), np.abs(self.y-self.ytest).argmin() + if compute_plane.lower() == 'longitudinal': + xx, yy = self.ixt, self.iyt + else: + xx, yy = slice(self.ixt-1, self.ixt+2), slice(self.iyt-1, self.iyt+2) + + # Compute simulation time + self.wake.wakelength = wakelength + self.ti = self.wake.ti + self.v = self.wake.v + if self.use_mpi: #E- should it be zmin, zmax instead? + z = self.Z # use global coords + zz = slice(0, self.NZ) + else: + z = self.z + zz = slice(0, self.Nz) + + tmax = (wakelength + self.ti*self.v + (z.max()-z.min()))/self.v #[s] + Nt = int(tmax/self.dt) + self.tmax, self.Nt = tmax, Nt + + # Add beam source + beam = Beam(q=self.wake.q, + sigmaz=self.wake.sigmaz, + beta=self.wake.beta, + xsource=self.wake.xsource, ysource=self.wake.ysource, + ) + + # hdf5 + self.Ez_file = self.wake.Ez_file + hf = None # needed for MPI + if self.use_mpi: + if self.rank==0: + hf = h5py.File(self.Ez_file, 'w') + hf['x'], hf['y'], hf['z'] = self.x[xx], self.y[yy], z[zz] + hf['t'] = np.arange(0, Nt*self.dt, self.dt) + + if save_J: + hfJ = h5py.File('Jz.h5', 'w') + hfJ['x'], hfJ['y'], hfJ['z'] = self.x[xx], self.y[yy], z[zz] + hfJ['t'] = np.arange(0, Nt*self.dt, self.dt) + else: + hf = h5py.File(self.Ez_file, 'w') + hf['x'], hf['y'], hf['z'] = self.x[xx], self.y[yy], z[zz] + hf['t'] = np.arange(0, Nt*self.dt, self.dt) + + if save_J: + hfJ = h5py.File('Jz.h5', 'w') + hfJ['x'], hfJ['y'], hfJ['z'] = self.x[xx], self.y[yy], z[zz] + hfJ['t'] = np.arange(0, Nt*self.dt, self.dt) + + def save_to_h5(self, hf, field, x, y, z, comp, n): + if self.use_mpi: + _field = self.mpi_gather(field, x, y, z, comp) + if self.rank == 0: + hf['#'+str(n).zfill(5)] = _field + else: + hf['#'+str(n).zfill(5)] = getattr(self, field)[x, y, z, comp] + + if plot_until is None: plot_until = Nt + if plot_from is None: plot_from = int(self.ti/self.dt) + + print('Running electromagnetic time-domain simulation...') + for n in tqdm(range(Nt)): + + # Initial condition + beam.update(self, n*self.dt) + + # Save + save_to_h5(self, hf, 'E', xx, yy, zz, 'z', n) + if save_J: + save_to_h5(self, hfJ, 'J', xx, yy, zz, 'z', n) + + # Advance + self.one_step() + + # Plot + if plot: + if n%plot_every == 0 and nplot_from: + self.plot2D(field='E', component='z', n=n, **plotkw) + else: + pass + + # Callback func(solver, t) + if callback is not None: + callback(self, n*self.dt) + + # End of time loop + if self.use_mpi: + if self.rank==0: + hf.close() + if save_J: + hfJ.close() + + # Compute wakefield magnitudes is done inside WakeSolver + self.wake.solve(compute_plane=compute_plane) + else: + hf.close() + if save_J: + hfJ.close() + + # Compute wakefield magnitudes is done inside WakeSolver + self.wake.solve(compute_plane=compute_plane) + + + \ No newline at end of file diff --git a/build/lib/wakis/solver2D.py b/build/lib/wakis/solver2D.py new file mode 100644 index 0000000..2b89458 --- /dev/null +++ b/build/lib/wakis/solver2D.py @@ -0,0 +1,484 @@ +import numpy as np +from scipy.constants import c as c_light, epsilon_0 as eps_0, mu_0 as mu_0 +from pmlBlock2D import PmlBlock2D + + +def eq(a, b, tol=1e-8): + return abs(a - b) < tol + + +def neq(a, b, tol=1e-8): + return not eq(a, b, tol) + + +class EMSolver2D: + def __init__(self, grid, sol_type, cfln, i_s, j_s, bc_low, bc_high, + N_pml_low=None, N_pml_high=None): + self.grid = grid + self.type = type + self.cfln = cfln + + self.dt = cfln / (c_light * np.sqrt(1 / self.grid.dx ** 2 + 1 / self.grid.dy ** 2)) + self.dx = self.grid.dx + self.dy = self.grid.dy + self.Nx = self.grid.nx + self.Ny = self.grid.ny + self.sol_type = sol_type + + self.N_pml_low = np.zeros(2, dtype=int) + self.N_pml_high = np.zeros(2, dtype=int) + self.bc_low = bc_low + self.bc_high = bc_high + + if bc_low[0] == 'pml': + self.N_pml_low[0] = 10 if N_pml_low is None else N_pml_low[0] + if bc_low[1] == 'pml': + self.N_pml_low[1] = 10 if N_pml_low is None else N_pml_low[1] + if bc_high[0] == 'pml': + self.N_pml_high[0] = 10 if N_pml_high is None else N_pml_high[0] + if bc_high[1] == 'pml': + self.N_pml_high[1] = 10 if N_pml_high is None else N_pml_high[1] + + self.blocks = [] + self.pml_ly = None + self.pml_lx = None + self.pml_rx = None + self.pml_ry = None + self.pml_lxly = None + self.pml_rxly = None + self.pml_lxry = None + self.pml_rxry = None + + if bc_low[0] is 'pml': + self.pml_lx = PmlBlock2D(self.N_pml_low[0], self.Ny, self.dt, self.dx, self.dy) + self.blocks.append(self.pml_lx) + if bc_low[1] is 'pml': + self.pml_lxly = PmlBlock2D(self.N_pml_low[0], self.N_pml_low[1], self.dt, self.dx, self.dy) + self.blocks.append(self.pml_lxly) + if bc_high[1] is 'pml': + self.pml_lxry = PmlBlock2D(self.N_pml_low[0], self.N_pml_high[1], self.dt, self.dx, self.dy) + self.blocks.append(self.pml_lxry) + + if bc_high[0] is 'pml': + self.pml_rx = PmlBlock2D(self.N_pml_high[0], self.Ny, self.dt, self.dx, self.dy) + self.blocks.append(self.pml_rx) + if bc_low[1] is 'pml': + self.pml_rxry = PmlBlock2D(self.N_pml_high[0], self.N_pml_high[1], self.dt, self.dx, self.dy) + self.blocks.append(self.pml_rxry) + if bc_high[1] is 'pml': + self.pml_rxly = PmlBlock2D(self.N_pml_high[0], self.N_pml_low[1], self.dt, self.dx, self.dy) + self.blocks.append(self.pml_rxly) + + if bc_low[1] is 'pml': + self.pml_ly = PmlBlock2D(self.Nx, self.N_pml_low[1], self.dt, self.dx, self.dy) + self.blocks.append(self.pml_ly) + + if bc_high[1] is 'pml': + self.pml_ry = PmlBlock2D(self.Nx, self.N_pml_high[1], self.dt, self.dx, self.dy) + self.blocks.append(self.pml_ry) + + self.organize_pmls() + self.alpha_pml = 3 + self.R0_pml = 0.001 + + self.assemble_conductivities_pmls() + + self.assemble_coeffs_pmls() + + self.N_tot_x = self.Nx + self.N_pml_low[0] + self.N_pml_high[0] + self.N_tot_y = self.Ny + self.N_pml_low[0] + self.N_pml_high[0] + self.sol_type = sol_type + + self.Ex = np.zeros((self.Nx, self.Ny + 1)) + self.Ey = np.zeros((self.Nx + 1, self.Ny)) + self.Hz = np.zeros((self.Nx, self.Ny)) + self.V_new = np.zeros((self.Nx, self.Ny)) + self.Jx = np.zeros((self.Nx, self.Ny + 1)) + self.Jy = np.zeros((self.Nx + 1, self.Ny)) + + self.rho = np.zeros((self.Nx, self.Ny)) + + if (sol_type is not 'FDTD') and (sol_type is not 'DM') and (sol_type is not 'ECT'): + raise ValueError("sol_type must be:\n" + + "\t'FDTD' for standard staircased FDTD\n" + + "\t'DM' for Dey-Mittra conformal FDTD\n" + + "\t'ECT' for Enlarged Cell Technique conformal FDTD") + + if sol_type is 'DM' or sol_type is 'ECT': + self.Vxy = np.zeros((self.Nx, self.Ny)) + if sol_type is 'ECT': + self.V_enl = np.zeros((self.Nx, self.Ny)) + + if sol_type is 'ECT' or sol_type is 'DM': + self.C1 = self.dt / mu_0 + self.C4 = self.dt / (eps_0 * self.dy) + self.C5 = self.dt / (eps_0 * self.dx) + self.C3 = self.dt / eps_0 + self.C6 = self.dt / eps_0 + + if sol_type is 'FDTD': + Z_0 = np.sqrt(mu_0 / eps_0) + + self.C1 = self.dt / (self.dx * mu_0) + self.C2 = self.dt / (self.dy * mu_0) + self.C4 = self.dt / (self.dy * eps_0) + self.C5 = self.dt / (self.dx * eps_0) + self.C3 = self.dt / eps_0 + self.C6 = self.dt / eps_0 + + # indices for the source + self.i_s = i_s + self.j_s = j_s + + self.time = 0 + + def organize_pmls(self): + if self.bc_low[0] == 'pml': + self.pml_lx.rx_block = self + if self.bc_low[1] is 'pml': + self.pml_lx.ly_block = self.pml_lxly + self.pml_lxly.ry_block = self.pml_lx + self.pml_lxly.rx_block = self.pml_ly + if self.bc_high[1] is 'pml': + self.pml_lx.ry_block = self.pml_lxry + self.pml_lxry.ly_block = self.pml_lx + self.pml_lxry.rx_block = self.pml_ry + + if self.bc_high[0] is 'pml': + self.pml_rx.lx_block = self + if self.bc_high[1] is 'pml': + self.pml_rx.ry_block = self.pml_rxry + self.pml_rxly.lx_block = self.pml_ly + self.pml_rxly.ry_block = self.pml_rx + if self.bc_low[1] is 'pml': + self.pml_rx.ly_block = self.pml_rxly + self.pml_rxry.lx_block = self.pml_ry + self.pml_rxry.ly_block = self.pml_rx + + + if self.bc_low[1] is 'pml': + self.pml_ly.ry_block = self + if self.bc_low[0] is 'pml': + self.pml_ly.lx_block = self.pml_lxly + if self.bc_low[0] is 'pml': + self.pml_ly.rx_block = self.pml_rxly + + if self.bc_high[1] is 'pml': + self.pml_ry.ly_block = self + if self.bc_low[0] is 'pml': + self.pml_ry.lx_block = self.pml_lxry + if self.bc_high[0] is 'pml': + self.pml_ry.rx_block = self.pml_rxry + + def assemble_conductivities_pmls(self): + sigma_m_low_x = 0 + sigma_m_high_x = 0 + sigma_m_low_y = 0 + sigma_m_high_y = 0 + if self.bc_low[0] is 'pml': + sigma_m_low_x = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_low[0]-1)*self.dx) * np.log(self.R0_pml) + if self.bc_low[1] is 'pml': + sigma_m_low_y = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_low[1]-1)*self.dy) * np.log(self.R0_pml) + if self.bc_high[0] is 'pml': + sigma_m_high_x = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_high[0]-1)*self.dy) * np.log(self.R0_pml) + if self.bc_high[1] is 'pml': + sigma_m_high_y = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_high[1]-1)*self.dy) * np.log(self.R0_pml) + + if self.bc_low[0] is 'pml': + for n in range(self.N_pml_low[0]): + self.pml_lx.sigma_x[-(n+1), :] = sigma_m_low_x * (n / (self.N_pml_low[0])) ** self.alpha_pml + if self.bc_low[1] is 'pml': + for n in range((self.N_pml_low[1])): + self.pml_lxly.sigma_y[:, -(n+1)] = sigma_m_low_y * (n / (self.N_pml_low[1])) ** self.alpha_pml + for n in range((self.N_pml_low[0])): + self.pml_lxly.sigma_x[-(n+1), :] = sigma_m_low_x * (n / (self.N_pml_low[0])) ** self.alpha_pml + if self.bc_high[1] is 'pml': + for n in range(self.N_pml_high[1]): + self.pml_lxry.sigma_y[:, n] = sigma_m_high_y * (n / (self.N_pml_high[1])) ** self.alpha_pml + for n in range(self.N_pml_low[0]): + self.pml_lxry.sigma_x[-(n+1), :] = sigma_m_low_x * (n / (self.N_pml_low[0])) ** self.alpha_pml + + if self.bc_high[0] is 'pml': + for n in range(self.N_pml_high[0]): + self.pml_rx.sigma_x[n, :] = sigma_m_high_x * (n / (self.N_pml_high[0])) ** self.alpha_pml + if self.bc_high[1] is 'pml': + for n in range(self.N_pml_high[0]): + self.pml_rxry.sigma_x[n, :] = sigma_m_high_x * (n / (self.N_pml_high[0])) ** self.alpha_pml + for n in range(self.N_pml_high[1]): + self.pml_rxry.sigma_y[:, n] = sigma_m_high_y * (n / (self.N_pml_high[1])) ** self.alpha_pml + if self.bc_low[1] == 'pml': + for n in range(self.N_pml_low[1]): + self.pml_rxly.sigma_y[:, -(n+1)] = sigma_m_low_y * (n / (self.N_pml_low[1])) ** self.alpha_pml + for n in range(self.N_pml_high[0]): + self.pml_rxly.sigma_x[n, :] = sigma_m_high_x * (n / (self.N_pml_high[0])) ** self.alpha_pml + + if self.bc_low[1] is 'pml': + for n in range(self.N_pml_low[1]): + self.pml_ly.sigma_y[:, -(n+1)] = sigma_m_low_y * (n / (self.N_pml_low[1])) ** self.alpha_pml + + if self.bc_high[1] is 'pml': + for n in range(self.N_pml_high[1]): + self.pml_ry.sigma_y[:, n] = sigma_m_high_y * (n / (self.N_pml_high[1])) ** self.alpha_pml + + if self.pml_lx is not None: + self.pml_lx.sigma_star_x = self.pml_lx.sigma_x * mu_0 / eps_0 + self.pml_lx.sigma_star_y = self.pml_lx.sigma_y * mu_0 / eps_0 + if self.pml_ly is not None: + self.pml_ly.sigma_star_x = self.pml_ly.sigma_x * mu_0 / eps_0 + self.pml_ly.sigma_star_y = self.pml_ly.sigma_y * mu_0 / eps_0 + if self.pml_rx is not None: + self.pml_rx.sigma_star_x = self.pml_rx.sigma_x * mu_0 / eps_0 + self.pml_rx.sigma_star_y = self.pml_rx.sigma_y * mu_0 / eps_0 + if self.pml_ry is not None: + self.pml_ry.sigma_star_x = self.pml_ry.sigma_x * mu_0 / eps_0 + self.pml_ry.sigma_star_y = self.pml_ry.sigma_y * mu_0 / eps_0 + if self.pml_lxly is not None: + self.pml_lxly.sigma_star_x = self.pml_lxly.sigma_x * mu_0 / eps_0 + self.pml_lxly.sigma_star_y = self.pml_lxly.sigma_y * mu_0 / eps_0 + if self.pml_lxry is not None: + self.pml_lxry.sigma_star_x = self.pml_lxry.sigma_x * mu_0 / eps_0 + self.pml_lxry.sigma_star_y = self.pml_lxry.sigma_y * mu_0 / eps_0 + if self.pml_rxry is not None: + self.pml_rxry.sigma_star_x = self.pml_rxry.sigma_x * mu_0 / eps_0 + self.pml_rxry.sigma_star_y = self.pml_rxry.sigma_y * mu_0 / eps_0 + if self.pml_rxly is not None: + self.pml_rxly.sigma_star_x = self.pml_rxly.sigma_x * mu_0 / eps_0 + self.pml_rxly.sigma_star_y = self.pml_rxly.sigma_y * mu_0 / eps_0 + + def assemble_coeffs_pmls(self): + if self.bc_low[0] is 'pml': + self.pml_lx.assemble_coeffs() + if self.bc_low[1] is 'pml': + self.pml_lxly.assemble_coeffs() + if self.bc_high[1] is 'pml': + self.pml_lxry.assemble_coeffs() + + if self.bc_high[0] is 'pml': + self.pml_rx.assemble_coeffs() + if self.bc_low[0] is 'pml': + self.pml_rxry.assemble_coeffs() + if self.bc_high[1] is 'pml': + self.pml_rxly.assemble_coeffs() + + if self.bc_low[1] is 'pml': + self.pml_ly.assemble_coeffs() + + if self.bc_high[1] is 'pml': + self.pml_ry.assemble_coeffs() + + def update_e_boundary(self): + Ex = self.Ex + Ey = self.Ey + Hz = self.Hz + if self.pml_lx is not None: + for jj in range(self.Ny): + Ey[0, jj] = Ey[0, jj] - self.C3 * self.Jy[0, jj] - self.C5 * (Hz[0, jj] - self.pml_lx.Hz[-1, jj]) + + if self.pml_rx is not None: + for jj in range(self.Ny): + Ey[-1, jj] = Ey[-1, jj] - self.C3 * self.Jy[-1, jj] - self.C5 * (self.pml_rx.Hz[0, jj] - Hz[-1, jj]) + + if self.pml_ly is not None: + for ii in range(self.Nx): + Ex[ii, 0] = Ex[ii, 0] - self.C3 * self.Jx[ii, 0] + self.C4 * (Hz[ii, 0] - self.pml_ly.Hz[ii, -1]) + + if self.pml_ry is not None: + for ii in range(self.Nx): + Ex[ii, -1] = Ex[ii, -1] - self.C3 * self.Jx[ii, -1] + self.C4 * (self.pml_ry.Hz[ii, 0] - Hz[ii, -1]) + + for block in self.blocks: + block.update_e_boundary() + + def gauss(self, t): + tau = 10 * self.dt + if t < 6 * tau: + return 100 * np.exp(-(t - 3 * tau) ** 2 / tau ** 2) + else: + return 0. + + def one_step(self): + if self.sol_type == 'ECT': + self.compute_v_and_rho() + self.one_step_ect(Nx=self.Nx, Ny=self.Ny, V_enl=self.V_enl, + rho=self.rho, Hz=self.Hz, C1=self.C1, + flag_int_cell=self.grid.flag_int_cell, + flag_unst_cell=self.grid.flag_unst_cell, + flag_intr_cell=self.grid.flag_intr_cell, + S=self.grid.S, + borrowing=self.grid.borrowing, S_enl=self.grid.S_enl, + S_red=self.grid.S_red, V_new = self.V_new, dt=self.dt) + for block in self.blocks: + block.advance_h_fdtd() + + self.advance_e_dm() + + for block in self.blocks: + block.advance_e_fdtd() + + self.update_e_boundary() + + self.time += self.dt + + if self.sol_type == 'FDTD': + self.one_step_fdtd() + if self.sol_type == 'DM': + self.one_step_dm() + + def one_step_fdtd(self): + Z_0 = np.sqrt(mu_0 / eps_0) + Ex = self.Ex + Ey = self.Ey + Hz = self.Hz + + self.advance_h_fdtd() + for block in self.blocks: + block.advance_h_fdtd() + self.advance_e_fdtd() + for block in self.blocks: + block.advance_e_fdtd() + self.update_e_boundary() + + self.time += self.dt + + def one_step_dm(self): + self.compute_v_and_rho() + + for i in range(self.Nx): + for j in range(self.Ny): + if self.grid.flag_int_cell[i, j]: + self.Hz[i, j] = self.Hz[i, j] - self.dt / (mu_0 * self.grid.S[i, j]) * self.Vxy[i, j] + + for block in self.blocks: + block.advance_h_fdtd() + + self.advance_e_dm() + + for block in self.blocks: + block.advance_e_fdtd() + + self.update_e_boundary_dm() + + self.time += self.dt + + def advance_h_fdtd(self): + Ex = self.Ex + Ey = self.Ey + Hz = self.Hz + for ii in range(self.Nx): + for jj in range(self.Ny): + if self.grid.flag_int_cell[ii, jj]: + Hz[ii, jj] = (Hz[ii, jj] - self.C1 * (Ey[ii + 1, jj] - Ey[ii, jj]) + + self.C2 * (Ex[ii, jj + 1]- Ex[ii, jj])) + + def advance_e_fdtd(self): + Z_0 = np.sqrt(mu_0 / eps_0) + Ex = self.Ex + Ey = self.Ey + Hz = self.Hz + for ii in range(self.Nx): + for jj in range(1, self.Ny): + if self.grid.flag_int_cell[ii, jj]: + if self.grid.l_x[ii, jj] > 0: + Ex[ii, jj] = Ex[ii, jj] - self.C3 * self.Jx[ii, jj] + self.C4 * ( + Hz[ii, jj] - Hz[ii, jj - 1]) + + for ii in range(1, self.Nx): + for jj in range(self.Ny): + if self.grid.flag_int_cell[ii, jj]: + if self.grid.l_y[ii, jj] > 0: + Ey[ii, jj] = Ey[ii, jj] - self.C3 * self.Jy[ii, jj] - self.C5 * ( + Hz[ii, jj] - Hz[ii - 1, jj]) + + @staticmethod + def one_step_ect(Nx=None, Ny=None, V_enl=None, rho=None, Hz=None, C1=None, flag_int_cell=None, + flag_unst_cell=None, flag_intr_cell=None, S=None, borrowing=None, S_enl=None, + S_red=None, V_new=None, dt = None, comp=None, kk=None): + + #if dt==None: dt = self.dt + + V_enl = np.zeros((Nx, Ny)) + + # take care of unstable cells + for ii in range(Nx): + for jj in range(Ny): + if flag_int_cell[ii, jj] and flag_unst_cell[ii, jj]: + + V_enl[ii, jj] = rho[ii, jj] * S[ii, jj] + + if len(borrowing[ii, jj]) == 0: + print('error in one_step_ect') + for (ip, jp, patch, _) in borrowing[ii, jj]: + + V_enl[ii, jj] += rho[ip, jp] * patch + + rho_enl = V_enl[ii, jj] / S_enl[ii, jj] + + # communicate to the intruded cell the intruding rho + for (ip, jp, patch, _) in borrowing[ii, jj]: + V_enl[ip, jp] += rho_enl * patch + + Hz[ii, jj] = Hz[ii, jj] - dt/mu_0 * rho_enl + + + # take care of stable cells + for ii in range(Nx): + for jj in range(Ny): + if flag_int_cell[ii, jj] and not flag_unst_cell[ii, jj]: + # stable cell which hasn't been intruded + if not flag_intr_cell[ii, jj]: + Hz[ii, jj] = Hz[ii, jj] - dt/mu_0 * rho[ii, jj] + # stable cell which has been intruded + else: + V_enl[ii, jj] += rho[ii, jj] * S_red[ii, jj] + Hz[ii, jj] = Hz[ii, jj] - dt/mu_0 * V_enl[ii, jj] / S[ii, jj] + + + def compute_v_and_rho(self): + l_y = self.grid.l_y + l_x = self.grid.l_x + for ii in range(self.Nx): + for jj in range(self.Ny): + if self.grid.flag_int_cell[ii, jj]: + self.Vxy[ii, jj] = ( + self.Ey[ii + 1, jj] * l_y[ii + 1, jj] - self.Ey[ii, jj] * l_y[ + ii, jj] + - self.Ex[ii, jj + 1] * l_x[ii, jj + 1] + self.Ex[ii, jj] * l_x[ + ii, jj]) + if self.sol_type != 'DM': + self.rho[ii, jj] = self.Vxy[ii, jj] / self.grid.S[ii, jj] + + def advance_e_dm(self): + for ii in range(self.Nx): + for jj in range(1, self.Ny): + if self.grid.l_x[ii, jj] > 0: + self.Ex[ii, jj] = self.Ex[ii, jj] + self.dt / (eps_0 * self.dy) * ( + self.Hz[ii, jj] - self.Hz[ii, jj - 1]) - self.C3 * self.Jx[ii, jj] + + for ii in range(1, self.Nx): + for jj in range(self.Ny): + if self.grid.l_y[ii, jj] > 0: + self.Ey[ii, jj] = self.Ey[ii, jj] - self.dt / (eps_0 * self.dx) * ( + self.Hz[ii, jj] - self.Hz[ii - 1, jj]) - self.C3 * self.Jy[ii, jj] + + def update_e_boundary_dm(self): + Ex = self.Ex + Ey = self.Ey + Hz = self.Hz + if self.pml_lx is not None: + for jj in range(self.Ny): + Ey[0, jj] = Ey[0, jj] - self.C3 * self.Jy[0, jj] - self.dt / (eps_0 * self.dy) * (Hz[0, jj] - self.pml_lx.Hz[-1, jj]) + + if self.pml_rx is not None: + for jj in range(self.Ny): + Ey[-1, jj] = Ey[-1, jj] - self.C3 * self.Jy[-1, jj] - self.dt / (eps_0 * self.dy) * (self.pml_rx.Hz[0, jj] - Hz[-1, jj]) + + if self.pml_ly is not None: + for ii in range(self.Nx): + Ex[ii, 0] = Ex[ii, 0] - self.C3 * self.Jx[ii, 0] + self.dt / (eps_0 * self.dy) * (Hz[ii, 0] - self.pml_ly.Hz[ii, -1]) + + if self.pml_ry is not None: + for ii in range(self.Nx): + Ex[ii, -1] = Ex[ii, -1] - self.C3 * self.Jx[ii, -1] + self.dt / (eps_0 * self.dy) * (self.pml_ry.Hz[ii, 0] - Hz[ii, -1]) + + for block in self.blocks: + block.update_e_boundary() diff --git a/build/lib/wakis/solver3D.py b/build/lib/wakis/solver3D.py new file mode 100644 index 0000000..1d9e42b --- /dev/null +++ b/build/lib/wakis/solver3D.py @@ -0,0 +1,1126 @@ +import numpy as np +from scipy.constants import c as c_light, epsilon_0 as eps_0, mu_0 as mu_0 +from solver2D import EMSolver2D +from pmlBlock3D import PmlBlock3D + +from numba import jit + +def eq(a, b, tol=1e-8): + return abs(a - b) < tol + + +def neq(a, b, tol=1e-8): + return not eq(a, b, tol) + + +class EMSolver3D: + def __init__(self, grid, sol_type, cfln=0.5, + bc_low=['Dirichlet', 'Dirichlet', 'Dirichlet'], + bc_high=['Dirichlet', 'Dirichlet', 'Dirichlet'], + i_s=0, j_s=0, k_s=0, N_pml_low=None, N_pml_high=None): + + self.grid = grid + self.type = type + self.cfln = cfln + + self.dt = cfln / (c_light * np.sqrt(1 / self.grid.dx ** 2 + 1 / self.grid.dy ** 2 + + 1 / self.grid.dz ** 2)) + self.dx = self.grid.dx + self.dy = self.grid.dy + self.dz = self.grid.dz + self.Nx = self.grid.nx + self.Ny = self.grid.ny + self.Nz = self.grid.nz + self.sol_type = sol_type + + self.N_pml_low = np.zeros(3, dtype=int) + self.N_pml_high = np.zeros(3, dtype=int) + self.bc_low = bc_low + self.bc_high = bc_high + + Nx = self.Nx + Ny = self.Ny + Nz = self.Nz + + self.sigma_x = np.zeros((Nx + 1, Ny + 1, Nz + 1)) + self.sigma_y = np.zeros((Nx + 1, Ny + 1, Nz + 1)) + self.sigma_z = np.zeros((Nx + 1, Ny + 1, Nz + 1)) + self.sigma_star_x = np.zeros((Nx, Ny + 1, Nz + 1)) + self.sigma_star_y = np.zeros((Nx + 1, Ny, Nz + 1)) + self.sigma_star_z = np.zeros((Nx + 1, Ny + 1, Nz)) + + if bc_low[0] == 'pml': + self.N_pml_low[0] = 10 if N_pml_low is None else N_pml_low[0] + if bc_low[1] == 'pml': + self.N_pml_low[1] = 10 if N_pml_low is None else N_pml_low[1] + if bc_low[2] == 'pml': + self.N_pml_low[2] = 10 if N_pml_low is None else N_pml_low[2] + if bc_high[0] == 'pml': + self.N_pml_high[0] = 10 if N_pml_high is None else N_pml_high[0] + if bc_high[1] == 'pml': + self.N_pml_high[1] = 10 if N_pml_high is None else N_pml_high[1] + if bc_high[2] == 'pml': + self.N_pml_high[2] = 10 if N_pml_high is None else N_pml_high[2] + + self.blocks = [] + + self.blocks_mat = np.full((3, 3, 3), None) + self.blocks_mat[1, 1, 1] = self + + self.connect_pmls() + + self.alpha_pml = 3 + self.R0_pml = 0.001 + + self.assemble_conductivities_pmls() + self.assemble_coeffs_pmls() + + self.Ex = np.zeros((self.Nx, self.Ny + 1, self.Nz + 1)) + self.Ey = np.zeros((self.Nx + 1, self.Ny, self.Nz + 1)) + self.Ez = np.zeros((self.Nx + 1, self.Ny + 1, self.Nz)) + self.Hx = np.zeros((self.Nx + 1, self.Ny, self.Nz)) + self.Hy = np.zeros((self.Nx, self.Ny + 1, self.Nz)) + self.Hz = np.zeros((self.Nx, self.Ny, self.Nz + 1)) + self.Jx = np.zeros((self.Nx, self.Ny + 1, self.Nz + 1)) + self.Jy = np.zeros((self.Nx + 1, self.Ny, self.Nz + 1)) + self.Jz = np.zeros((self.Nx + 1, self.Ny + 1, self.Nz)) + self.rho_xy = np.zeros((self.Nx, self.Ny, self.Nz + 1)) + self.rho_yz = np.zeros((self.Nx + 1, self.Ny, self.Nz)) + self.rho_zx = np.zeros((self.Nx, self.Ny + 1, self.Nz)) + + if (sol_type is not 'FDTD') and (sol_type is not 'DM') and (sol_type is not 'ECT'): + raise ValueError("sol_type must be:\n" + + "\t'FDTD' for standard staircased FDTD\n" + + "\t'DM' for Dey-Mittra conformal FDTD\n" + + "\t'ECT' for Enlarged Cell Technique conformal FDTD") + + if sol_type is 'DM' or sol_type is 'ECT': + self.Vxy = np.zeros((self.Nx, self.Ny, self.Nz + 1)) + self.Vyz = np.zeros((self.Nx + 1, self.Ny, self.Nz)) + self.Vzx = np.zeros((self.Nx, self.Ny + 1, self.Nz)) + if sol_type is 'ECT': + self.Vxy_enl = np.zeros((self.Nx, self.Ny, self.Nz + 1)) + self.Vyz_enl = np.zeros((self.Nx + 1, self.Ny, self.Nz)) + self.Vzx_enl = np.zeros((self.Nx, self.Ny + 1, self.Nz)) + + self.C1 = self.dt / (self.dx * mu_0) + self.C2 = self.dt / (self.dy * mu_0) + self.C7 = self.dt / (self.dz * mu_0) + self.C4 = self.dt / (self.dy * eps_0) + self.C5 = self.dt / (self.dx * eps_0) + self.C8 = self.dt / (self.dz * eps_0) + self.C3 = self.dt / eps_0 + self.C6 = self.dt / eps_0 + + self.CN = self.dt / mu_0 + + # indices for the source + self.i_s = i_s + self.j_s = j_s + self.k_s = k_s + + self.time = 0 + + def connect_pmls(self): + bc_low = self.bc_low + bc_high = self.bc_high + if bc_low[0] is 'pml': + i_block = 0 + j_block = 1 + k_block = 1 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.Ny, self.Nz, self.dt, + self.dx, + self.dy, self.dz, i_block, j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_low[1] is 'pml': + i_block = 0 + j_block = 0 + k_block = 1 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.N_pml_low[1], self.Nz, + self.dt, self.dx, self.dy, self.dz, i_block, + j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_low[2] is 'pml': + i_block = 0 + j_block = 0 + k_block = 0 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.N_pml_low[1], + self.N_pml_low[2], self.dt, self.dx, + self.dy, + self.dz, i_block, j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_high[2] is 'pml': + i_block = 0 + j_block = 0 + k_block = 2 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.N_pml_low[1], + self.N_pml_high[2], self.dt, self.dx, + self.dy, + self.dz, i_block, j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_high[1] is 'pml': + i_block = 0 + j_block = 2 + k_block = 1 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.N_pml_high[1], self.Nz, + self.dt, self.dx, self.dy, self.dz, i_block, + j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_low[2] is 'pml': + i_block = 0 + j_block = 2 + k_block = 0 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.N_pml_high[1], + self.N_pml_low[2], self.dt, self.dx, + self.dy, + self.dz, i_block, j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_high[2] is 'pml': + i_block = 0 + j_block = 2 + k_block = 2 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.N_pml_high[1], + self.N_pml_high[2], self.dt, self.dx, + self.dy, + self.dz, i_block, j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_low[2] is 'pml': + i_block = 0 + j_block = 1 + k_block = 0 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.Ny, self.N_pml_low[2], + self.dt, self.dx, self.dy, self.dz, i_block, + j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_high[2] is 'pml': + i_block = 0 + j_block = 1 + k_block = 2 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.Ny, self.N_pml_high[2], + self.dt, self.dx, self.dy, self.dz, i_block, + j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + + if bc_high[0] is 'pml': + i_block = 2 + j_block = 1 + k_block = 1 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.Ny, self.Nz, self.dt, + self.dx, self.dy, self.dz, i_block, j_block, + k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_low[1] is 'pml': + i_block = 2 + j_block = 0 + k_block = 1 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.N_pml_low[1], self.Nz, + self.dt, self.dx, self.dy, self.dz, i_block, + j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_low[2] is 'pml': + i_block = 2 + j_block = 0 + k_block = 0 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.N_pml_low[1], + self.N_pml_low[2], self.dt, self.dx, + self.dy, + self.dz, i_block, j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_high[2] is 'pml': + i_block = 2 + j_block = 0 + k_block = 2 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.N_pml_low[1], + self.N_pml_high[2], self.dt, self.dx, + self.dy, + self.dz, i_block, j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_high[1] is 'pml': + i_block = 2 + j_block = 2 + k_block = 1 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.N_pml_high[1], self.Nz, + self.dt, self.dx, self.dy, self.dz, i_block, + j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_low[2] is 'pml': + i_block = 2 + j_block = 2 + k_block = 0 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.N_pml_high[1], + self.N_pml_low[2], self.dt, self.dx, + self.dy, + self.dz, i_block, j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_high[2] is 'pml': + i_block = 2 + j_block = 2 + k_block = 2 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.N_pml_high[1], + self.N_pml_high[2], self.dt, self.dx, + self.dy, + self.dz, i_block, j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_low[2] is 'pml': + i_block = 2 + j_block = 1 + k_block = 0 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.Ny, self.N_pml_low[2], + self.dt, self.dx, self.dy, self.dz, i_block, + j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_high[2] is 'pml': + i_block = 2 + j_block = 1 + k_block = 2 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.Ny, self.N_pml_high[2], + self.dt, self.dx, self.dy, self.dz, i_block, + j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + + if bc_low[1] is 'pml': + i_block = 1 + j_block = 0 + k_block = 1 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.Nx, self.N_pml_low[1], self.Nz, self.dt, + self.dx, + self.dy, self.dz, i_block, j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_low[2] is 'pml': + i_block = 1 + j_block = 0 + k_block = 0 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.Nx, self.N_pml_low[1], self.N_pml_low[2], + self.dt, self.dx, self.dy, self.dz, i_block, + j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_high[2] is 'pml': + i_block = 1 + j_block = 0 + k_block = 2 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.Nx, self.N_pml_low[1], self.N_pml_high[2], + self.dt, self.dx, self.dy, self.dz, i_block, + j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + + if bc_high[1] is 'pml': + i_block = 1 + j_block = 2 + k_block = 1 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.Nx, self.N_pml_high[1], self.Nz, self.dt, + self.dx, self.dy, self.dz, i_block, j_block, + k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_low[2] is 'pml': + i_block = 1 + j_block = 2 + k_block = 0 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.Nx, self.N_pml_low[1], self.N_pml_low[2], + self.dt, self.dx, self.dy, self.dz, i_block, + j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + if bc_high[2] is 'pml': + i_block = 1 + j_block = 2 + k_block = 2 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.Nx, self.N_pml_low[1], self.N_pml_high[2], + self.dt, self.dx, self.dy, self.dz, i_block, + j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + + if bc_low[2] is 'pml': + i_block = 1 + j_block = 1 + k_block = 0 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.Nx, self.Ny, self.N_pml_low[2], self.dt, + self.dx, + self.dy, self.dz, i_block, j_block, k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + + if bc_high[2] is 'pml': + i_block = 1 + j_block = 1 + k_block = 2 + self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.Nx, self.Ny, self.N_pml_high[2], self.dt, + self.dx, self.dy, self.dz, i_block, j_block, + k_block) + self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) + + for block in self.blocks: + block.blocks_mat = self.blocks_mat + + def assemble_conductivities_pmls(self): + sigma_m_low_x = 0 + sigma_m_high_x = 0 + sigma_m_low_y = 0 + sigma_m_high_y = 0 + sigma_m_low_z = 0 + sigma_m_high_z = 0 + Z_0_2 = mu_0 / eps_0 + + if self.bc_low[0] is 'pml': + sigma_m_low_x = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_low[0] - 1) * self.dx) * np.log(self.R0_pml) + if self.bc_low[1] is 'pml': + sigma_m_low_y = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_low[1] - 1) * self.dy) * np.log(self.R0_pml) + if self.bc_low[2] is 'pml': + sigma_m_low_z = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_low[2] - 1) * self.dy) * np.log(self.R0_pml) + if self.bc_high[0] is 'pml': + sigma_m_high_x = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_high[0] - 1) * self.dy) * np.log(self.R0_pml) + if self.bc_high[1] is 'pml': + sigma_m_high_y = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_high[1] - 1) * self.dy) * np.log(self.R0_pml) + if self.bc_high[2] is 'pml': + sigma_m_high_z = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_high[2] - 1) * self.dy) * np.log(self.R0_pml) + + + if self.bc_low[0] is 'pml': + (i_block, j_block, k_block) = (0, 1, 1) + for n in range(self.N_pml_low[0]): + self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 + if self.bc_low[1] is 'pml': + (i_block, j_block, k_block) = (0, 0, 1) + for n in range((self.N_pml_low[0])): + self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_low[1])): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 + if self.bc_low[2] is 'pml': + (i_block, j_block, k_block) = (0, 0, 0) + for n in range((self.N_pml_low[0])): + self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_low[1])): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_low[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 + if self.bc_low[2] is 'pml': + (i_block, j_block, k_block) = (0, 0, 2) + for n in range((self.N_pml_low[0])): + self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_low[1])): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_high[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 + if self.bc_high[1] is 'pml': + (i_block, j_block, k_block) = (0, 2, 1) + for n in range(self.N_pml_low[0]): + self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 + for n in range(self.N_pml_high[1]): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 + + if self.bc_low[2] is 'pml': + (i_block, j_block, k_block) = (0, 2, 0) + for n in range((self.N_pml_low[0])): + self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_high[1])): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_low[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 + if self.bc_low[2] is 'pml': + (i_block, j_block, k_block) = (0, 2, 2) + for n in range((self.N_pml_low[0])): + self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_high[1])): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_high[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 + if self.bc_low[2] is 'pml': + (i_block, j_block, k_block) = (0, 1, 0) + for n in range((self.N_pml_low[0])): + self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_low[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 + if self.bc_high[2] is 'pml': + (i_block, j_block, k_block) = (0, 1, 2) + for n in range((self.N_pml_low[0])): + self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( + n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_high[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 + + if self.bc_high[0] is 'pml': + (i_block, j_block, k_block) = (2, 1, 1) + for n in range(self.N_pml_high[0]): + self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 + if self.bc_low[1] is 'pml': + (i_block, j_block, k_block) = (2, 0, 1) + for n in range(self.N_pml_high[0]): + self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_low[1])): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 + if self.bc_low[2] is 'pml': + (i_block, j_block, k_block) = (2, 0, 0) + for n in range(self.N_pml_high[0]): + self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_low[1])): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_low[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 + if self.bc_low[2] is 'pml': + (i_block, j_block, k_block) = (2, 0, 2) + for n in range(self.N_pml_high[0]): + self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_low[1])): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_high[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 + if self.bc_high[1] is 'pml': + (i_block, j_block, k_block) = (2, 2, 1) + for n in range(self.N_pml_high[0]): + self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 + for n in range(self.N_pml_high[1]): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 + if self.bc_low[2] is 'pml': + (i_block, j_block, k_block) = (2, 2, 0) + for n in range(self.N_pml_high[0]): + self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_high[1])): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_low[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 + if self.bc_low[2] is 'pml': + (i_block, j_block, k_block) = (2, 2, 2) + for n in range(self.N_pml_high[0]): + self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_high[1])): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_high[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 + if self.bc_low[2] is 'pml': + (i_block, j_block, k_block) = (2, 1, 0) + for n in range(self.N_pml_high[0]): + self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_low[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 + if self.bc_high[2] is 'pml': + (i_block, j_block, k_block) = (2, 1, 2) + for n in range(self.N_pml_high[0]): + self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( + n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_high[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 + + if self.bc_low[1] is 'pml': + (i_block, j_block, k_block) = (1, 0, 1) + for n in range(self.N_pml_low[1]): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 + if self.bc_low[2] is 'pml': + (i_block, j_block, k_block) = (1, 0, 0) + for n in range(self.N_pml_low[1]): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_low[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 + if self.bc_high[2] is 'pml': + (i_block, j_block, k_block) = (1, 0, 2) + for n in range(self.N_pml_low[1]): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( + n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_high[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 + + if self.bc_high[1] is 'pml': + (i_block, j_block, k_block) = (1, 2, 1) + for n in range(self.N_pml_high[1]): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 + if self.bc_low[2] is 'pml': + (i_block, j_block, k_block) = (1, 2, 0) + for n in range(self.N_pml_high[1]): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_low[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 + if self.bc_high[2] is 'pml': + (i_block, j_block, k_block) = (1, 2, 2) + for n in range(self.N_pml_high[1]): + self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( + n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 + for n in range((self.N_pml_high[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 + + if self.bc_low[2] is 'pml': + (i_block, j_block, k_block) = (1, 1, 0) + for n in range((self.N_pml_low[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( + n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 + + if self.bc_high[2] is 'pml': + (i_block, j_block, k_block) = (1, 1, 2) + for n in range((self.N_pml_high[2])): + self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml + self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( + n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 + + #for i_block in range(3): + # for j_block in range(3): + # for k_block in range(3): + # block = self.blocks_mat[i_block, j_block, k_block] + # if block is not None: + # if not (i_block == 1 and j_block == 1 and k_block == 1): + # block.sigma_star_x = block.sigma_x * mu_0 / eps_0 + # block.sigma_star_y = block.sigma_y * mu_0 / eps_0 + # block.sigma_star_z = block.sigma_z * mu_0 / eps_0 + + def assemble_coeffs_pmls(self): + for i_block in range(3): + for j_block in range(3): + for k_block in range(3): + if self.blocks_mat[i_block, j_block, k_block] is not None: + if not (i_block == 1 and j_block == 1 and k_block == 1): + self.blocks_mat[i_block, j_block, k_block].assemble_coeffs() + + def update_e_boundary(self): + Ex = self.Ex + Ey = self.Ey + Ez = self.Ez + Hx = self.Hx + Hy = self.Hy + Hz = self.Hz + + Nx = self.Nx + Ny = self.Ny + Nz = self.Nz + + # Update E on "lower" faces + if self.blocks_mat[0, 1, 1] is not None: + for jj in range(self.Ny): + for kk in range(1, self.Nz): + Ey[0, jj, kk] = (Ey[0, jj, kk] - self.C3 * self.Jy[0, jj, kk] + + self.C8 * (Hx[0, jj, kk] - Hx[0, jj, kk - 1]) - + self.C5 * (Hz[0, jj, kk] - self.blocks_mat[0, 1, 1].Hz[-1, jj, kk])) + + for jj in range(1, self.Ny): + for kk in range(self.Nz): + Ez[0, jj, kk] = (Ez[0, jj, kk] - self.C3 * self.Jz[0, jj, kk] + + self.C5 * (Hy[0, jj, kk] - self.blocks_mat[0, 1, 1].Hy[-1, jj, kk]) - + self.C4 * (Hx[0, jj, kk] - Hx[0, jj - 1, kk])) + + if self.blocks_mat[1, 0, 1] is not None: + for ii in range(self.Nx): + for kk in range(1, self.Nz): + Ex[ii, 0, kk] = (Ex[ii, 0, kk] - self.C3 * self.Jx[ii, 0, kk] + + self.C4 * (Hz[ii, 0, kk] - self.blocks_mat[1, 0, 1].Hz[ii, -1, kk]) - + self.C8 * (Hy[ii, 0, kk] - Hy[ii, 0, kk - 1])) + for ii in range(1, self.Nx): + for kk in range(self.Nz): + Ez[ii, 0, kk] = (Ez[ii, 0, kk] - self.C3 * self.Jz[ii, 0, kk] + + self.C5 * (Hy[ii, 0, kk] - Hy[ii - 1, 0, kk]) - + self.C4 * (Hx[ii, 0, kk] - self.blocks_mat[1, 0, 1].Hx[ii, -1, kk])) + + if self.blocks_mat[1, 1, 0] is not None: + for ii in range(self.Nx): + for jj in range(1, self.Ny): + Ex[ii, jj, 0] = (Ex[ii, jj, 0] - self.C3 * self.Jx[ii, jj, 0] + + self.C4 * (Hz[ii, jj, 0] - Hz[ii, jj - 1, 0]) - + self.C8 * (Hy[ii, jj, 0] - self.blocks_mat[1, 1, 0].Hy[ii, jj, -1])) + for ii in range(1, self.Nx): + for jj in range(self.Ny): + Ey[ii, jj, 0] = (Ey[ii, jj, 0] - self.C3 * self.Jy[ii, jj, 0] + + self.C8 * (Hx[ii, jj, 0] - self.blocks_mat[1, 1, 0].Hx[ii, jj, -1]) - + self.C5 * (Hz[ii, jj, 0] - Hz[ii - 1, jj, 0])) + + # Update E on "upper" faces + if self.blocks_mat[2, 1, 1] is not None: + for jj in range(self.Ny): + for kk in range(1, self.Nz): + Ey[Nx, jj, kk] = (Ey[Nx, jj, kk] - self.C3 * self.Jy[Nx, jj, kk] + + self.C8 * (Hx[Nx, jj, kk] - Hx[Nx, jj, kk - 1]) - + self.C5 * (self.blocks_mat[2, 1, 1].Hz[0, jj, kk] - Hz[Nx - 1, jj, kk])) + for jj in range(1, self.Ny): + for kk in range(self.Nz): + Ez[Nx, jj, kk] = (Ez[Nx, jj, kk] - self.C3 * self.Jz[Nx, jj, kk] + + self.C5 * (self.blocks_mat[2, 1, 1].Hy[0, jj, kk] - Hy[Nx - 1, jj, kk]) - + self.C4 * (Hx[Nx, jj, kk] - Hx[Nx, jj - 1, kk])) + + if self.blocks_mat[1, 2, 1] is not None: + for ii in range(self.Nx): + for kk in range(1, self.Nz): + Ex[ii, Ny, kk] = (Ex[ii, Ny, kk] - self.C3 * self.Jx[ii, Ny, kk] + + self.C4 * (self.blocks_mat[1, 2, 1].Hz[ii, 0, kk] - Hz[ii, Ny - 1, kk]) - + self.C8 * (Hy[ii, Ny, kk] - Hy[ii, Ny, kk - 1])) + for ii in range(1, self.Nx): + for kk in range(self.Nz): + Ez[ii, Ny, kk] = (Ez[ii, Ny, kk] - self.C3 * self.Jz[ii, Ny, kk] + + self.C5 * (Hy[ii, Ny, kk] - Hy[ii - 1, Ny, kk]) - + self.C4 * (self.blocks_mat[1, 2, 1].Hx[ii, 0, kk] - Hx[ii, Ny - 1, kk])) + + if self.blocks_mat[1, 1, 2] is not None: + for ii in range(self.Nx): + for jj in range(1, self.Ny): + Ex[ii, jj, Nz] = (Ex[ii, jj, Nz] - self.C3 * self.Jx[ii, jj, Nz] + + self.C4 * (Hz[ii, jj, Nz] - Hz[ii, jj - 1, Nz]) - + self.C8 * (self.blocks_mat[1, 1, 2].Hy[ii, jj, 0] - Hy[ii, jj, Nz - 1])) + for ii in range(1, self.Nx): + for jj in range(self.Ny): + Ey[ii, jj, Nz] = (Ey[ii, jj, Nz] - self.C3 * self.Jy[ii, jj, Nz] + + self.C8 * (self.blocks_mat[1, 1, 2].Hx[ii, jj, 0] - Hx[ii, jj, Nz - 1]) - + self.C5 * (Hz[ii, jj, Nz] - Hz[ii - 1, jj, Nz])) + # Update Ez on edges (xy) + if self.blocks_mat[0, 1, 1] is not None and self.blocks_mat[1, 0, 1] is not None: + for kk in range(Nz): + Ez[0, 0, kk] = (Ez[0, 0, kk] - self.C3 * self.Jz[0, 0, kk] + + self.C5 * (Hy[0, 0, kk] - self.blocks_mat[0, 1, 1].Hy[-1, 0, kk]) - + self.C4 * (Hx[0, 0, kk] - self.blocks_mat[1, 0, 1].Hx[0, -1, kk])) + if self.blocks_mat[0, 1, 1] is not None and self.blocks_mat[1, 2, 1] is not None: + for kk in range(Nz): + Ez[0, Ny, kk] = (Ez[0, Ny, kk] - self.C3 * self.Jz[0, Ny, kk] + + self.C5 * (Hy[0, Ny, kk] - self.blocks_mat[0, 1, 1].Hy[-1, Ny, kk]) - + self.C4 * (self.blocks_mat[1, 2, 1].Hx[0, 0, kk] - Hx[0, Ny - 1, kk])) + if self.blocks_mat[2, 1, 1] is not None and self.blocks_mat[1, 0, 1] is not None: + for kk in range(Nz): + Ez[Nx, 0, kk] = (Ez[Nx, 0, kk] - self.C3 * self.Jz[Nx, 0, kk] + + self.C5 * (self.blocks_mat[2, 1, 1].Hy[0, 0, kk] - Hy[Nx - 1, 0, kk]) - + self.C4 * (Hx[Nx, 0, kk] - self.blocks_mat[1, 0, 1].Hx[Nx, -1, kk])) + if self.blocks_mat[2, 1, 1] is not None and self.blocks_mat[1, 2, 1] is not None: + for kk in range(Nz): + Ez[Nx, Ny, kk] = (Ez[Nx, Ny, kk] - self.C3 * self.Jz[Nx, Ny, kk] + + self.C5 * (self.blocks_mat[2, 1, 1].Hy[0, Ny, kk] - Hy[Nx - 1, Ny, kk]) - + self.C4 * (self.blocks_mat[1, 2, 1].Hx[Nx, 0, kk] - Hx[Nx, Ny - 1, kk])) + # Update Ex on edges (yz) + if self.blocks_mat[1, 0, 1] is not None and self.blocks_mat[1, 1, 0] is not None: + for ii in range(Nx): + Ex[ii, 0, 0] = (Ex[ii, 0, 0] - self.C3 * self.Jx[ii, 0, 0] + + self.C4 * (Hz[ii, 0, 0] - self.blocks_mat[1, 0, 1].Hz[ii, -1, 0]) - + self.C8 * (Hy[ii, 0, 0] - self.blocks_mat[1, 1, 0].Hy[ii, 0, -1])) + if self.blocks_mat[1, 0, 1] is not None and self.blocks_mat[1, 1, 2] is not None: + for ii in range(Nx): + Ex[ii, 0, Nz] = (Ex[ii, 0, Nz] - self.C3 * self.Jx[ii, 0, Nz] + + self.C4 * (Hz[ii, 0, Nz] - self.blocks_mat[1, 0, 1].Hz[ii, -1, Nz]) - + self.C8 * (self.blocks_mat[1, 1, 2].Hy[ii, 0, 0] - Hy[ii, 0, Nz - 1])) + if self.blocks_mat[1, 2, 1] is not None and self.blocks_mat[1, 1, 0] is not None: + for ii in range(Nx): + Ex[ii, Ny, 0] = (Ex[ii, Ny, 0] - self.C3 * self.Jx[ii, Ny, 0] + + self.C4 * (self.blocks_mat[1, 2, 1].Hz[ii, 0, 0] - Hz[ii, Ny - 1, 0]) - + self.C8 * (Hy[ii, Ny, 0] - self.blocks_mat[1, 1, 0].Hy[ii, Ny, -1])) + if self.blocks_mat[1, 2, 1] is not None and self.blocks_mat[1, 1, 2] is not None: + for ii in range(Nx): + Ex[ii, Ny, Nz] = (Ex[ii, Ny, Nz] - self.C3 * self.Jx[ii, Ny, Nz] + + self.C4 * (self.blocks_mat[1, 2, 1].Hz[ii, 0, Nz] - Hz[ii, Ny - 1, Nz]) - + self.C8 * (self.blocks_mat[1, 1, 2].Hy[ii, Ny, 0] - Hy[ii, Ny, Nz - 1])) + + # Update Ey on edges (xz) + if self.blocks_mat[0, 1, 1] is not None and self.blocks_mat[1, 1, 0] is not None: + for jj in range(Ny): + Ey[0, jj, 0] = (Ey[0, jj, 0] - self.C3 * self.Jy[0, jj, 0] + + self.C8 * (Hx[0, jj, 0] - self.blocks_mat[1, 1, 0].Hx[0, jj, -1]) - + self.C5 * (Hz[0, jj, 0] - self.blocks_mat[0, 1, 1].Hz[-1, jj, 0])) + if self.blocks_mat[0, 1, 1] is not None and self.blocks_mat[1, 1, 2] is not None: + for jj in range(Ny): + Ey[0, jj, Nz] = (Ey[0, jj, Nz] - self.C3 * self.Jy[0, jj, Nz] + + self.C8 * (self.blocks_mat[1, 1, 2].Hx[0, jj, 0] - Hx[0, jj, Nz - 1]) - + self.C5 * (Hz[0, jj, Nz] - self.blocks_mat[0, 1, 1].Hz[-1, jj, Nz])) + if self.blocks_mat[2, 1, 1] is not None and self.blocks_mat[1, 1, 0] is not None: + for jj in range(Ny): + Ey[Nx, jj, 0] = (Ey[Nx, jj, 0] - self.C3 * self.Jy[Nx, jj, 0] + + self.C8 * (Hx[Nx, jj, 0] - self.blocks_mat[1, 1, 0].Hx[Nx, jj, -1]) - + self.C5 * (self.blocks_mat[2, 1, 1].Hz[0, jj, 0] - Hz[Nx - 1, jj, 0])) + if self.blocks_mat[2, 1, 1] is not None and self.blocks_mat[1, 1, 2] is not None: + for jj in range(Ny): + Ey[Nx, jj, Nz] = (Ey[Nx, jj, Nz] - self.C3 * self.Jy[Nx, jj, Nz] + + self.C8 * (self.blocks_mat[1, 1, 2].Hx[Nx, jj, 0] - Hx[Nx, jj, Nz - 1]) - + self.C5 * (self.blocks_mat[2, 1, 1].Hz[0, jj, Nz] - Hz[Nx - 1, jj, Nz])) + + def gauss(self, t): + tau = 10 * self.dt + if t < 6 * tau: + return 100 * np.exp(-(t - 3 * tau) ** 2 / tau ** 2) + else: + return 0. + + def one_step(self): + if self.sol_type == 'ECT': + self.one_step_ect() + if self.sol_type == 'FDTD': + self.one_step_fdtd() + if self.sol_type == 'DM': + self.one_step_dm() + + self.time += self.dt + + def one_step_ect(self): + self.compute_v_and_rho() + self.advance_h_ect() + for block in self.blocks: + block.advance_h_fdtd() + block.sum_h_fields() + self.advance_e_dm() + self.update_e_boundary() + for block in self.blocks: + block.advance_e_fdtd() + block.update_e_boundary() + block.sum_e_fields() + + def one_step_fdtd(self): + self.advance_h_fdtd(self.grid.Sxy, self.grid.Syz, self.grid.Szx, self.Ex, self.Ey, + self.Ez, self.Hx, self.Hy, self.Hz, self.Nx, self.Ny, self.Nz, + self.C1, self.C2, self.C7) + + for block in self.blocks: + block.advance_h_fdtd() + + block.sum_h_fields() + self.advance_e_fdtd(self.grid.l_x, self.grid.l_y, self.grid.l_z, self.Ex, self.Ey, self.Ez, + self.Hx, self.Hy, self.Hz, self.Jx, self.Jy, self.Jz, self.Nx, self.Ny, + self.Nz, self.C3, self.C4, self.C5, self.C8) + self.update_e_boundary() + for block in self.blocks: + block.advance_e_fdtd() + block.update_e_boundary() + block.sum_e_fields() + + @staticmethod + @jit(nopython=True) + def advance_h_fdtd(Sxy, Syz, Szx, Ex, Ey, Ez, Hx, Hy, Hz, Nx, Ny, Nz, C1, C2, C7): + + # Compute cell voltages + for ii in range(Nx + 1): + for jj in range(Ny): + for kk in range(Nz): + if Syz[ii, jj, kk] > 0: + Hx[ii, jj, kk] = (Hx[ii, jj, kk] - + C2 * (Ez[ii, jj + 1, kk] - Ez[ii, jj, kk]) + + C7 * (Ey[ii, jj, kk + 1] - Ey[ii, jj, kk])) + + for ii in range(Nx): + for jj in range(Ny + 1): + for kk in range(Nz): + if Szx[ii, jj, kk] > 0: + Hy[ii, jj, kk] = (Hy[ii, jj, kk] - + C7 * (Ex[ii, jj, kk + 1] - Ex[ii, jj, kk]) + + C1 * (Ez[ii + 1, jj, kk] - Ez[ii, jj, kk])) + + for ii in range(Nx): + for jj in range(Ny): + for kk in range(Nz + 1): + if Sxy[ii, jj, kk] > 0: + Hz[ii, jj, kk] = (Hz[ii, jj, kk] - + C1 * (Ey[ii + 1, jj, kk] - Ey[ii, jj, kk]) + + C2 * (Ex[ii, jj + 1, kk] - Ex[ii, jj, kk])) + + @staticmethod + @jit(nopython=True) + def advance_e_fdtd(l_x, l_y, l_z, Ex, Ey, Ez, Hx, Hy, Hz, Jx, Jy, Jz, Nx, Ny, Nz, C3, C4, C5, C8): + + for ii in range(Nx): + for jj in range(1, Ny): + for kk in range(1, Nz): + if l_x[ii, jj, kk] > 0: + Ex[ii, jj, kk] = (Ex[ii, jj, kk] - C3 * Jx[ii, jj, kk] + + C4 * (Hz[ii, jj, kk] - Hz[ii, jj - 1, kk]) - + C8 * (Hy[ii, jj, kk] - Hy[ii, jj, kk - 1])) + + for ii in range(1, Nx): + for jj in range(Ny): + for kk in range(1, Nz): + if l_y[ii, jj, kk] > 0: + Ey[ii, jj, kk] = (Ey[ii, jj, kk] - C3 * Jy[ii, jj, kk] + + C8 * (Hx[ii, jj, kk] - Hx[ii, jj, kk - 1]) - + C5 * (Hz[ii, jj, kk] - Hz[ii - 1, jj, kk])) + + for ii in range(1, Nx): + for jj in range(1, Ny): + for kk in range(Nz): + if l_z[ii, jj, kk] > 0: + Ez[ii, jj, kk] = (Ez[ii, jj, kk] - C3 * Jz[ii, jj, kk] + + C5 * (Hy[ii, jj, kk] - Hy[ii - 1, jj, kk]) - + C4 * (Hx[ii, jj, kk] - Hx[ii, jj - 1, kk])) + + def one_step_dm(self): + self.compute_v_and_rho() + + for i in range(self.Nx): + for j in range(self.Ny): + for k in range(self.Nz + 1): + if self.grid.flag_int_cell_xy[i, j, k]: + self.Hz[i, j, k] = (self.Hz[i, j, k] - + self.dt / (mu_0 * self.grid.Sxy[i, j, k]) * + self.Vxy[i, j, k]) + + for i in range(self.Nx + 1): + for j in range(self.Ny): + for k in range(self.Nz): + if self.grid.flag_int_cell_yz[i, j, k]: + self.Hx[i, j, k] = (self.Hx[i, j, k] - + self.dt / (mu_0 * self.grid.Syz[i, j, k]) * + self.Vyz[i, j, k]) + + for i in range(self.Nx): + for j in range(self.Ny + 1): + for k in range(self.Nz): + if self.grid.flag_int_cell_zx[i, j, k]: + self.Hy[i, j, k] = (self.Hy[i, j, k] - + self.dt / (mu_0 * self.grid.Szx[i, j, k]) * + self.Vzx[i, j, k]) + + self.advance_e_dm() + + def advance_h_ect(self, dt=None): + + for ii in range(self.Nx + 1): + EMSolver2D.advance_h_ect(Nx=self.Ny, Ny=self.Nz, V_enl=self.Vyz_enl[ii, :, :], + rho=self.rho_yz[ii, :, :], Hz=self.Hx[ii, :, :], C1=self.CN, + flag_int_cell=self.grid.flag_int_cell_yz[ii, :, :], + flag_unst_cell=self.grid.flag_unst_cell_yz[ii, :, :], + flag_intr_cell=self.grid.flag_intr_cell_yz[ii,:,:], + S=self.grid.Syz[ii, :, :], + borrowing=self.grid.borrowing_yz[ii, :, :], + S_enl=self.grid.Syz_enl[ii, :, :], + S_red=self.grid.Syz_red[ii, :, :], dt=dt, comp='x', kk=ii) + + for jj in range(self.Ny + 1): + EMSolver2D.advance_h_ect(Nx=self.Nx, Ny=self.Nz, V_enl=self.Vzx_enl[:, jj, :], + rho=self.rho_zx[:, jj, :], Hz=self.Hy[:, jj, :], C1=self.CN, + flag_int_cell=self.grid.flag_int_cell_zx[:, jj, :], + flag_unst_cell=self.grid.flag_unst_cell_zx[:, jj, :], + flag_intr_cell=self.grid.flag_intr_cell_zx[:,jj,:], + S=self.grid.Szx[:, jj, :], + borrowing=self.grid.borrowing_zx[:, jj, :], + S_enl=self.grid.Szx_enl[:, jj, :], + S_red=self.grid.Szx_red[:, jj, :], dt=dt, comp='y', kk=jj) + + for kk in range(self.Nz + 1): + EMSolver2D.advance_h_ect(Nx=self.Nx, Ny=self.Ny, V_enl=self.Vxy_enl[:, :, kk], + rho=self.rho_xy[:, :, kk], Hz=self.Hz[:, :, kk], C1=self.CN, + flag_int_cell=self.grid.flag_int_cell_xy[:, :, kk], + flag_unst_cell=self.grid.flag_unst_cell_xy[:, :, kk], + flag_intr_cell=self.grid.flag_intr_cell_xy[:,:,kk], + S=self.grid.Sxy[:, :, kk], + borrowing=self.grid.borrowing_xy[:, :, kk], + S_enl=self.grid.Sxy_enl[:, :, kk], + S_red=self.grid.Sxy_red[:, :, kk], dt=dt, comp='z', kk=kk) + + def compute_v_and_rho(self): + l_x = self.grid.l_x + l_y = self.grid.l_y + l_z = self.grid.l_z + + for ii in range(self.Nx): + for jj in range(self.Ny): + for kk in range(self.Nz + 1): + if self.grid.flag_int_cell_xy[ii, jj, kk]: + self.Vxy[ii, jj, kk] = (self.Ex[ii, jj, kk] * l_x[ii, jj, kk] - + self.Ex[ii, jj + 1, kk] * l_x[ii, jj + 1, kk] + + self.Ey[ii + 1, jj, kk] * l_y[ii + 1, jj, kk] - + self.Ey[ii, jj, kk] * l_y[ii, jj, kk]) + + if self.sol_type != 'DM': + self.rho_xy[ii, jj, kk] = (self.Vxy[ii, jj, kk] / + self.grid.Sxy[ii, jj, kk]) + + + for ii in range(self.Nx + 1): + for jj in range(self.Ny): + for kk in range(self.Nz): + if self.grid.flag_int_cell_yz[ii, jj, kk]: + self.Vyz[ii, jj, kk] = (self.Ey[ii, jj, kk] * l_y[ii, jj, kk] - + self.Ey[ii, jj, kk + 1] * l_y[ii, jj, kk + 1] + + self.Ez[ii, jj + 1, kk] * l_z[ii, jj + 1, kk] - + self.Ez[ii, jj, kk] * l_z[ii, jj, kk]) + + if self.sol_type != 'DM': + self.rho_yz[ii, jj, kk] = (self.Vyz[ii, jj, kk] / + self.grid.Syz[ii, jj, kk]) + + + for ii in range(self.Nx): + for jj in range(self.Ny + 1): + for kk in range(self.Nz): + if self.grid.flag_int_cell_zx[ii, jj, kk]: + self.Vzx[ii, jj, kk] = (self.Ez[ii, jj, kk] * l_z[ii, jj, kk] - + self.Ez[ii + 1, jj, kk] * l_z[ii + 1, jj, kk] + + self.Ex[ii, jj, kk + 1] * l_x[ii, jj, kk + 1] - + self.Ex[ii, jj, kk] * l_x[ii, jj, kk]) + + if self.sol_type != 'DM': + self.rho_zx[ii, jj, kk] = (self.Vzx[ii, jj, kk] / + self.grid.Szx[ii, jj, kk]) + + def advance_e_dm(self, dt=None): + + if dt==None: dt = self.dt + + Ex = self.Ex + Ey = self.Ey + Ez = self.Ez + Hx = self.Hx + Hy = self.Hy + Hz = self.Hz + + C4 = dt / (self.dy * eps_0) + C5 = dt / (self.dx * eps_0) + C8 = dt / (self.dz * eps_0) + C3 = dt / eps_0 + + for ii in range(self.Nx): + for jj in range(1, self.Ny): + for kk in range(1, self.Nz): + if self.grid.l_x[ii, jj, kk] > 0: + Ex[ii, jj, kk] = (Ex[ii, jj, kk] - self.C3 * self.Jx[ii, jj, kk] + + self.C4 * (Hz[ii, jj, kk] - Hz[ii, jj - 1, kk]) - + self.C8 * (Hy[ii, jj, kk] - Hy[ii, jj, kk - 1])) + + for ii in range(1, self.Nx): + for jj in range(self.Ny): + for kk in range(1, self.Nz): + if self.grid.l_y[ii, jj, kk] > 0: + Ey[ii, jj, kk] = (Ey[ii, jj, kk] - self.C3 * self.Jy[ii, jj, kk] + + self.C8 * (Hx[ii, jj, kk] - Hx[ii, jj, kk - 1]) - + self.C5 * (Hz[ii, jj, kk] - Hz[ii - 1, jj, kk])) + + for ii in range(1, self.Nx): + for jj in range(1, self.Ny): + for kk in range(self.Nz): + if self.grid.l_z[ii, jj, kk] > 0: + Ez[ii, jj, kk] = (Ez[ii, jj, kk] - self.C3 * self.Jz[ii, jj, kk] + + self.C5 * (Hy[ii, jj, kk] - Hy[ii - 1, jj, kk]) - + self.C4 * (Hx[ii, jj, kk] - Hx[ii, jj - 1, kk])) diff --git a/build/lib/wakis/solverFIT3D.py b/build/lib/wakis/solverFIT3D.py new file mode 100644 index 0000000..c142264 --- /dev/null +++ b/build/lib/wakis/solverFIT3D.py @@ -0,0 +1,1110 @@ +# copyright ################################# # +# This file is part of the wakis Package. # +# Copyright (c) CERN, 2024. # +# ########################################### # + +from tqdm import tqdm + +import numpy as np +import time +import h5py + +from scipy.constants import c as c_light, epsilon_0 as eps_0, mu_0 as mu_0 +from scipy.sparse import csc_matrix as sparse_mat +from scipy.sparse import diags, hstack, vstack + +from .field import Field +from .materials import material_lib +from .plotting import PlotMixin +from .routines import RoutinesMixin + +try: + from cupyx.scipy.sparse import csc_matrix as gpu_sparse_mat + imported_cupyx = True +except ImportError: + imported_cupyx = False + +try: + from sparse_dot_mkl import csr_matrix as mkl_sparse_mat, dot_product_mkl + imported_mkl = True +except ImportError: + imported_mkl = False + + +class SolverFIT3D(PlotMixin, RoutinesMixin): + + def __init__(self, grid, wake=None, cfln=0.5, dt=None, + bc_low=['Periodic', 'Periodic', 'Periodic'], + bc_high=['Periodic', 'Periodic', 'Periodic'], + use_stl=False, use_conductors=False, + use_gpu=False, use_mpi=False, dtype=np.float64, + n_pml=10, bg=[1.0, 1.0], verbose=1): + ''' + Class holding the 3D time-domain electromagnetic solver + algorithm based on the Finite Integration Technique (FIT) + + Parameters: + ----------- + grid: GridFIT3D object + Instance of GridFIT3D class containing the simulation mesh and the + imported geometry + wake: WakeSolver object, optional + Instance of WakeSolver class containing the beam parameters. Needed to + run a wakefield simulation to compute wake potential and impedance + cfln: float, default 0.5 + Convergence condition by Courant–Friedrichs–Lewy, used to compute the + simulation timestep + dt: float, optional + Simulation timestep. If not None, it overrides the cfln-based timestep + bc_low: list, default ['Periodic', 'Periodic', 'Periodic'] + Domain box boundary conditions for X-, Y-, Z- + bc_high: list, default ['Periodic', 'Periodic', 'Periodic'] + Domain box boundary conditions for X+, Y+, Z+ + use_conductors: bool, default False + If true, enables geometry import based on elements from `conductors.py` + use_stl: bool, default False + If true, activates all the solids and materials passed to the `grid` object + use_gpu: bool, default False, + Using cupyx, enables GPU accelerated computation of every timestep + bg: list, default [1.0, 1.0] + Background material for the simulation box [eps_r, mu_r, sigma]. Default is vacuum. + It supports any material from the material library in `materials.py`, of a + custom list of floats. If conductivity (sigma) is passed, + it enables flag: use_conductivity + verbose: int or bool, default True + Enable verbose ouput on the terminal if 1 or True + + Attributes + ---------- + E: Field object + Object to access the Electric field data in [V/m]. + E.g.: solver.E[:,:,n,'z'] gives a 2D numpy.ndarray fieldmap of Ez component, located at the n-th cell in z + H: Field object + Object to access the Magnetic field data in [A/m]. + E.g.: solver.H[i,j,k,'x'] gives a point value of Hx component, located at the i,j,k cell + J: Field object + Object to access the Current density field data in [A/m^2]. + ieps: Field object + Object to access the ε^-1 tensor containing 1/permittivity values in the 3 dimensions. + imu: Field object + Object to access the μ^-1 tensor containing 1/permeability values in the 3 dimensions. + sigma: Field object + Object to access the condutcity σ tensor in the 3 dimensions. + ''' + + self.verbose = verbose + if verbose: t0 = time.time() + + # Flags + self.step_0 = True + self.nstep = int(0) + self.plotter_active = False + self.use_conductors = use_conductors + self.use_stl = use_stl + self.use_gpu = use_gpu + self.use_mpi = use_mpi + self.activate_abc = False # Will turn true if abc BCs are chosen + self.activate_pml = False # Will turn true if pml BCs are chosen + self.use_conductivity = False # Will turn true if conductive material or pml is added + self.imported_mkl = imported_mkl # Use MKL backend when available + self.one_step = self._one_step + if use_stl: + self.use_conductors = False + + # Grid + self.grid = grid + + self.Nx = self.grid.Nx + self.Ny = self.grid.Ny + self.Nz = self.grid.Nz + self.N = self.Nx*self.Ny*self.Nz + + self.dx = self.grid.dx + self.dy = self.grid.dy + self.dz = self.grid.dz + + self.x = self.grid.x[:-1]+self.dx/2 + self.y = self.grid.y[:-1]+self.dy/2 + self.z = self.grid.z[:-1]+self.dz/2 + + self.L = self.grid.L + self.iA = self.grid.iA + self.tL = self.grid.tL + self.itA = self.grid.itA + + # Wake computation + self.wake = wake + + # Fields + self.dtype = dtype + self.E = Field(self.Nx, self.Ny, self.Nz, use_gpu=self.use_gpu, dtype=self.dtype) + self.H = Field(self.Nx, self.Ny, self.Nz, use_gpu=self.use_gpu, dtype=self.dtype) + self.J = Field(self.Nx, self.Ny, self.Nz, use_gpu=self.use_gpu, dtype=self.dtype) + + # MPI init + if self.use_mpi: + if self.grid.use_mpi: + self.mpi_initialize() + self.one_step = self.mpi_one_step + else: + print('*** Grid not subdivided for MPI, set `use_mpi`=True also in `GridFIT3D` to enable MPI') + + # Matrices + if verbose: print('Assembling operator matrices...') + N = self.N + self.Px = diags([-1, 1], [0, 1], shape=(N, N), dtype=np.int8) + self.Py = diags([-1, 1], [0, self.Nx], shape=(N, N), dtype=np.int8) + self.Pz = diags([-1, 1], [0, self.Nx*self.Ny], shape=(N, N), dtype=np.int8) + + # original grid + self.Ds = diags(self.L.toarray(), shape=(3*N, 3*N), dtype=self.dtype) + self.iDa = diags(self.iA.toarray(), shape=(3*N, 3*N), dtype=self.dtype) + + # tilde grid + self.tDs = diags(self.tL.toarray(), shape=(3*N, 3*N), dtype=self.dtype) + self.itDa = diags(self.itA.toarray(), shape=(3*N, 3*N), dtype=self.dtype) + + # Curl matrix + self.C = vstack([ + hstack([sparse_mat((N,N)), -self.Pz, self.Py]), + hstack([self.Pz, sparse_mat((N,N)), -self.Px]), + hstack([-self.Py, self.Px, sparse_mat((N,N))]) + ], dtype=np.int8) + + # Boundaries + if verbose: print('Applying boundary conditions...') + self.bc_low = bc_low + self.bc_high = bc_high + self.apply_bc_to_C() + + # Materials + if verbose: print('Adding material tensors...') + if type(bg) is str: + bg = material_lib[bg.lower()] + + if len(bg) == 3: + self.eps_bg, self.mu_bg, self.sigma_bg = bg[0]*eps_0, bg[1]*mu_0, bg[2] + self.use_conductivity = True + else: + self.eps_bg, self.mu_bg, self.sigma_bg = bg[0]*eps_0, bg[1]*mu_0, 0.0 + + self.ieps = Field(self.Nx, self.Ny, self.Nz, use_ones=True, dtype=self.dtype)*(1./self.eps_bg) + self.imu = Field(self.Nx, self.Ny, self.Nz, use_ones=True, dtype=self.dtype)*(1./self.mu_bg) + self.sigma = Field(self.Nx, self.Ny, self.Nz, use_ones=True, dtype=self.dtype)*self.sigma_bg + + if self.use_stl: + self.apply_stl() + + # Fill PML BCs + if self.activate_pml: + if verbose: print('Filling PML sigmas...') + self.n_pml = n_pml + self.pml_lo = 5e-3 + self.pml_hi = 1.e-1 + self.pml_func = np.geomspace + self.fill_pml_sigmas() + + # Timestep calculation + if verbose: print('Calculating maximal stable timestep...') + self.cfln = cfln + if dt is None: + self.dt = cfln / (c_light * np.sqrt(1 / self.grid.dx ** 2 + 1 / self.grid.dy ** 2 + 1 / self.grid.dz ** 2)) + else: + self.dt = dt + self.dt = dtype(self.dt) + + if self.use_conductivity: # relaxation time criterion tau + + mask = np.logical_and(self.sigma.toarray()!=0, #for non-conductive + self.ieps.toarray()!=0) #for PEC eps=inf + + self.tau = (1/self.ieps.toarray()[mask]) / \ + self.sigma.toarray()[mask] + + if self.dt > self.tau.min(): + self.dt = self.tau.min() + + # Pre-computing + if verbose: print('Pre-computing...') + self.iDeps = diags(self.ieps.toarray(), shape=(3*N, 3*N), dtype=self.dtype) + self.iDmu = diags(self.imu.toarray(), shape=(3*N, 3*N), dtype=self.dtype) + self.Dsigma = diags(self.sigma.toarray(), shape=(3*N, 3*N), dtype=self.dtype) + + self.tDsiDmuiDaC = self.tDs * self.iDmu * self.iDa * self.C + self.itDaiDepsDstC = self.itDa * self.iDeps * self.Ds * self.C.transpose() + + if imported_mkl and not self.use_gpu: # MKL backend for CPU + print('Using MKL backend for time-stepping...') + self.tDsiDmuiDaC = mkl_sparse_mat(self.tDsiDmuiDaC) + self.itDaiDepsDstC = mkl_sparse_mat(self.itDaiDepsDstC) + self.one_step = self.one_step_mkl + + # Move to GPU + if use_gpu: + if verbose: print('Moving to GPU...') + if imported_cupyx: + self.tDsiDmuiDaC = gpu_sparse_mat(self.tDsiDmuiDaC) + self.itDaiDepsDstC = gpu_sparse_mat(self.itDaiDepsDstC) + self.ieps.to_gpu() + self.sigma.to_gpu() + else: + raise ImportError('*** cupyx could not be imported, please check CUDA installation') + + if verbose: print(f'Total initialization time: {time.time() - t0} s') + + def update_tensors(self, tensor='all'): + '''Update tensor matrices after + Field ieps, imu or sigma have been modified + and pre-compute the time-stepping matrices + + Parameters: + ----------- + tensor : str, default 'all' + Name of the tensor to update: 'ieps', 'imu', 'sigma' + for permitivity, permeability and conductivity, respectively. + If left to default 'all', all thre tensors will be recomputed. + ''' + if self.verbose: print(f'Re-computing tensor "{tensor}"...') + + if tensor == 'ieps': + self.iDeps = diags(self.ieps.toarray(), shape=(3*self.N, 3*self.N), dtype=self.dtype) + elif tensor =='imu': + self.iDmu = diags(self.imu.toarray(), shape=(3*self.N, 3*self.N), dtype=self.dtype) + elif tensor == 'sigma': + self.Dsigma = diags(self.sigma.toarray(), shape=(3*self.N, 3*self.N), dtype=self.dtype) + elif tensor == 'all': + self.iDeps = diags(self.ieps.toarray(), shape=(3*self.N, 3*self.N), dtype=self.dtype) + self.iDmu = diags(self.imu.toarray(), shape=(3*self.N, 3*self.N), dtype=self.dtype) + self.Dsigma = diags(self.sigma.toarray(), shape=(3*self.N, 3*self.N), dtype=self.dtype) + + if self.verbose: print('Re-Pre-computing ...') + self.tDsiDmuiDaC = self.tDs * self.iDmu * self.iDa * self.C + self.itDaiDepsDstC = self.itDa * self.iDeps * self.Ds * self.C.transpose() + self.step_0 = False + + def _one_step(self): + if self.step_0: + self.set_ghosts_to_0() + self.step_0 = False + self.attrcleanup() + + self.H.fromarray(self.H.toarray() - + self.dt*self.tDsiDmuiDaC*self.E.toarray() + ) + + self.E.fromarray(self.E.toarray() + + self.dt*(self.itDaiDepsDstC*self.H.toarray() + - self.ieps.toarray()*self.J.toarray() + ) + ) + + #include current computation + if self.use_conductivity: + self.J.fromarray(self.sigma.toarray()*self.E.toarray()) + + def one_step_mkl(self): + if self.step_0: + self.set_ghosts_to_0() + self.step_0 = False + self.attrcleanup() + + self.H.fromarray(self.H.toarray() - + self.dt*dot_product_mkl(self.tDsiDmuiDaC,self.E.toarray()) + ) + + self.E.fromarray(self.E.toarray() + + self.dt*(dot_product_mkl(self.itDaiDepsDstC,self.H.toarray()) + - self.ieps.toarray()*self.J.toarray() + ) + ) + + #include current computation + if self.use_conductivity: + self.J.fromarray(self.sigma.toarray()*self.E.toarray()) + + def mpi_initialize(self): + self.comm = self.grid.comm + self.rank = self.grid.rank + self.size = self.grid.size + + self.NZ = self.grid.NZ + self.ZMIN = self.grid.ZMIN + self.ZMAX = self.grid.ZMAX + self.Z = self.grid.Z + + def mpi_one_step(self): + if self.step_0: + self.set_ghosts_to_0() + self.step_0 = False + self.attrcleanup() + + self.H.fromarray(self.H.toarray() - + self.dt*self.tDsiDmuiDaC*self.E.toarray() + ) + + self.mpi_communicate(self.H) + self.mpi_communicate(self.J) + self.E.fromarray(self.E.toarray() + + self.dt*(self.itDaiDepsDstC * self.H.toarray() + - self.ieps.toarray()*self.J.toarray() + ) + ) + + self.mpi_communicate(self.E) + # include current computation + if self.use_conductivity: + self.J.fromarray(self.sigma.toarray()*self.E.toarray()) + + def mpi_communicate(self, field): + if self.use_gpu: + field.from_gpu() + + # ghosts lo + if self.rank > 0: + for d in ['x','y','z']: + self.comm.Sendrecv(field[:, :, 1, d], + recvbuf=field[:, :, 0, d], + dest=self.rank-1, sendtag=0, + source=self.rank-1, recvtag=1) + # ghosts hi + if self.rank < self.size - 1: + for d in ['x','y','z']: + self.comm.Sendrecv(field[:, :, -2, d], + recvbuf=field[:, :, -1, d], + dest=self.rank+1, sendtag=1, + source=self.rank+1, recvtag=0) + + if self.use_gpu: + field.to_gpu() + + def mpi_gather(self, field, x=None, y=None, z=None, component=None): + ''' + Gather a specific component or slice of a distributed field from all MPI ranks. + + This function collects a selected component of a field (E, H, J, or custom) + from all MPI processes along the z-axis and reconstructs the global field data + on the root rank (rank 0). The user can specify slices or single indices + along x, y, and z to control the subset of data gathered. + + Parameters + ---------- + field : str or Field obj + The field to gather. If a string, it must begin with one of: + - `'E'`, `'H'`, or `'J'` followed optionally by a component label + (e.g., `'Ex'`, `'Hz'`, `'JAbs'`). + If no component is specified, defaults to `'z'`. + + x : int or slice, optional + Range of x-indices to gather. If None, defaults to the full x-range. + + y : int or slice, optional + Range of y-indices to gather. If None, defaults to the full y-range. + + z : int or slice, optional + Range of z-indices to gather. If None, defaults to the full z-range. + + component : str or slice, optional + Component of the field to gather ('x', 'y', 'z', or a slice). + If None and not inferred from `field`, defaults to `'z'`. + + Returns + ------- + numpy.ndarray or None + The gathered field values assembled on rank 0 with shape depending on + the selected slices along (x, y, z). Returns `None` on non-root ranks. + + Notes + ----- + - Assumes field data is distributed along the z-dimension. + - Automatically handles removal of ghost cells in reconstruction. + - Field components are inferred from the input `field` string or the + `component` argument. + - This method supports full 3D subvolume extraction or 1D/2D slices + for performance diagnostics and visualization. + + Examples + -------- + >>> # Gather Ex component on full domain on rank 0 + >>> global_Ex = solver.mpi_gather('Ex') + + >>> # Gather a 2D yz-slice at x=10 of the J field + >>> yz_J = solver.mpi_gather('J', x=10) + ''' + + if x is None: + x = slice(0, self.Nx) + if y is None: + y = slice(0, self.Ny) + if z is None: + z = slice(0, self.NZ) + + if type(field) is str: + if len(field) == 2: #support for e.g. field='Ex' + component = field[1] + field = field[0] + elif len(field) == 4: #support for Abs + component = field[1:] + field = field[0] + elif component is None: + component = 'z' + print("[!] `component` not specified, using default component='z'") + + if field == 'E': + local = self.E[x, y, :, component].ravel() + elif field == 'H': + local = self.H[x, y, :, component].ravel() + elif field == 'J': + local = self.J[x, y, :, component].ravel() + else: + if component is None: + component = 'z' + print("[!] `component` not specified, using default component='z'") + local = field[x, y, :, component].ravel() + + buffer = self.comm.gather(local, root=0) + _field = None + + if self.rank == 0: + if type(x) is int and type(y) is int: # 1d array at x=a, y=b + nz = self.NZ//self.size + _field = np.zeros((self.NZ)) + for r in range(self.size): + zz = np.s_[r*nz:(r+1)*nz] + if r == 0: + _field[zz] = np.reshape(buffer[r], (nz+self.grid.n_ghosts))[:-1] + elif r == (self.size-1): + _field[zz] = np.reshape(buffer[r], (nz+self.grid.n_ghosts))[1:] + else: + _field[zz] = np.reshape(buffer[r], (nz+2*self.grid.n_ghosts))[1:-1] + _field = _field[z] + + elif type(x) is int: # 2d slice at x=a + ny = y.stop-y.start + nz = self.NZ//self.size + _field = np.zeros((ny, self.NZ)) + for r in range(self.size): + zz = np.s_[r*nz:(r+1)*nz] + if r == 0: + _field[:, zz] = np.reshape(buffer[r], (ny, nz+self.grid.n_ghosts))[:, :-1] + elif r == (self.size-1): + _field[:, zz] = np.reshape(buffer[r], (ny, nz+self.grid.n_ghosts))[:, 1:] + else: + _field[:, zz] = np.reshape(buffer[r], (ny, nz+2*self.grid.n_ghosts))[:, 1:-1] + _field = _field[:, z] + + elif type(y) is int: # 2d slice at y=a + nx = x.stop-x.start + nz = self.NZ//self.size + _field = np.zeros((nx, self.NZ)) + for r in range(self.size): + zz = np.s_[r*nz:(r+1)*nz] + if r == 0: + _field[:, zz] = np.reshape(buffer[r], (nx, nz+self.grid.n_ghosts))[:, :-1] + elif r == (self.size-1): + _field[:, zz] = np.reshape(buffer[r], (nx, nz+self.grid.n_ghosts))[:, 1:] + else: + _field[:, zz] = np.reshape(buffer[r], (nx, nz+2*self.grid.n_ghosts))[:, 1:-1] + _field = _field[:, z] + + else: # both type slice -> 3d array + nx = x.stop-x.start + ny = y.stop-y.start + nz = self.NZ//self.size + _field = np.zeros((nx, ny, self.NZ)) + for r in range(self.size): + zz = np.s_[r*nz:(r+1)*nz] + if r == 0: + _field[:, :, zz] = np.reshape(buffer[r], (nx, ny, nz+self.grid.n_ghosts))[:, :, :-1] + elif r == (self.size-1): + _field[:, :, zz] = np.reshape(buffer[r], (nx, ny, nz+self.grid.n_ghosts))[:, :, 1:] + else: + _field[:, :, zz] = np.reshape(buffer[r], (nx, ny, nz+2*self.grid.n_ghosts))[:, :, 1:-1] + _field = _field[:, :, z] + + return _field + + def mpi_gather_asField(self, field): + ''' + Gather distributed field data from all MPI ranks and return a global Field object. + + This method collects the specified electromagnetic field data (E, H, or J) + from all MPI processes and assembles it into a single global `Field` object + on the root rank (rank 0). The field data can be specified as a string + ('E', 'H', or 'J') or as a `wakis.Field` object. + + Parameters + ---------- + field : str or Field obj + The field to gather. If a string, it must be one of: + - `'E'` for the electric field + - `'H'` for the magnetic field + - `'J'` for the current density + + Passing a `wakis.Field` is also supported (e.g. ieps, imu, sigma) + + Returns + ------- + Field + A `wakis.Field` object containing the gathered global field data + with shape (Nx, Ny, NZ, 3). Only returned on rank 0. On other + ranks, the returned value is undefined and should not be used. + + Notes + ----- + - This method assumes the field is distributed along the `z`-axis. + - Ghost cells are removed appropriately when reassembling the global field. + ''' + + _field = Field(self.Nx, self.Ny, self.NZ) + + for d in ['x','y','z']: + if type(field) is str: + if field == 'E': + local = self.E[:, :, :,d].ravel() + elif field == 'H': + local = self.H[:, :, :,d].ravel() + elif field == 'J': + local = self.J[:, :, :,d].ravel() + else: + local = field[:, :, :, d].ravel() + + buffer = self.comm.gather(local, root=0) + if self.rank == 0: + nz = self.NZ//self.size + for r in range(self.size): + zz = np.s_[r*nz:(r+1)*nz] + if r == 0: + _field[:, :, zz, d] = np.reshape(buffer[r], (self.Nx, self.Ny, nz+self.grid.n_ghosts))[:, :, :-1] + elif r == (self.size-1): + _field[:, :, zz, d] = np.reshape(buffer[r], (self.Nx, self.Ny, nz+self.grid.n_ghosts))[:, :, 1:] + else: + _field[:, :, zz, d] = np.reshape(buffer[r], (self.Nx, self.Ny, nz+2*self.grid.n_ghosts))[:, :, 1:-1] + + return _field + + def apply_bc_to_C(self): + ''' + Modifies rows or columns of C and tDs and itDa matrices + according to bc_low and bc_high + ''' + xlo, ylo, zlo = 1., 1., 1. + xhi, yhi, zhi = 1., 1., 1. + + # Check BCs for internal MPI subdomains + if self.use_mpi and self.grid.use_mpi: + if self.rank > 0: + self.bc_low=['pec', 'pec', 'mpi'] + + if self.rank < self.size - 1: + self.bc_high=['pec', 'pec', 'mpi'] + + # Perodic: out == in + if any(True for x in self.bc_low if x.lower() == 'periodic'): + if self.bc_low[0].lower() == 'periodic' and self.bc_high[0].lower() == 'periodic': + self.tL[-1, :, :, 'x'] = self.L[0, :, :, 'x'] + self.itA[-1, :, :, 'y'] = self.iA[0, :, :, 'y'] + self.itA[-1, :, :, 'z'] = self.iA[0, :, :, 'z'] + + if self.bc_low[1].lower() == 'periodic' and self.bc_high[1].lower() == 'periodic': + self.tL[:, -1, :, 'y'] = self.L[:, 0, :, 'y'] + self.itA[:, -1, :, 'x'] = self.iA[:, 0, :, 'x'] + self.itA[:, -1, :, 'z'] = self.iA[:, 0, :, 'z'] + + if self.bc_low[2].lower() == 'periodic' and self.bc_high[2].lower() == 'periodic': + self.tL[:, :, -1, 'z'] = self.L[:, :, 0, 'z'] + self.itA[:, :, -1, 'x'] = self.iA[:, :, 0, 'x'] + self.itA[:, :, -1, 'y'] = self.iA[:, :, 0, 'y'] + + self.tDs = diags(self.tL.toarray(), shape=(3*self.N, 3*self.N), dtype=float) + self.itDa = diags(self.itA.toarray(), shape=(3*self.N, 3*self.N), dtype=float) + + # Dirichlet PEC: tangential E field = 0 at boundary + if any(True for x in self.bc_low if x.lower() in ('electric','pec','pml')) \ + or any(True for x in self.bc_high if x.lower() in ('electric','pec','pml')): + if self.bc_low[0].lower() in ('electric','pec', 'pml'): + xlo = 0 + if self.bc_low[1].lower() in ('electric','pec', 'pml'): + ylo = 0 + if self.bc_low[2].lower() in ('electric','pec', 'pml'): + zlo = 0 + if self.bc_high[0].lower() in ('electric','pec', 'pml'): + xhi = 0 + if self.bc_high[1].lower() in ('electric','pec', 'pml'): + yhi = 0 + if self.bc_high[2].lower() in ('electric','pec', 'pml'): + zhi = 0 + + # Assemble matrix + self.BC = Field(self.Nx, self.Ny, self.Nz, dtype=np.int8, use_ones=True) + + for d in ['x', 'y', 'z']: #tangential to zero + if d != 'x': + self.BC[0, :, :, d] = xlo + self.BC[-1, :, :, d] = xhi + if d != 'y': + self.BC[:, 0, :, d] = ylo + self.BC[:, -1, :, d] = yhi + if d != 'z': + self.BC[:, :, 0, d] = zlo + self.BC[:, :, -1, d] = zhi + + self.Dbc = diags(self.BC.toarray(), + shape=(3*self.N, 3*self.N), + dtype=np.int8 + ) + + # Update C (columns) + self.C = self.C*self.Dbc + + + # Dirichlet PMC: tangential H field = 0 at boundary + if any(True for x in self.bc_low if x.lower() in ('magnetic','pmc')) \ + or any(True for x in self.bc_high if x.lower() in ('magnetic','pmc')): + if self.bc_low[0].lower() == 'magnetic' or self.bc_low[0] == 'pmc': + xlo = 0 + if self.bc_low[1].lower() == 'magnetic' or self.bc_low[1] == 'pmc': + ylo = 0 + if self.bc_low[2].lower() == 'magnetic' or self.bc_low[2] == 'pmc': + zlo = 0 + if self.bc_high[0].lower() == 'magnetic' or self.bc_high[0] == 'pmc': + xhi = 0 + if self.bc_high[1].lower() == 'magnetic' or self.bc_high[1] == 'pmc': + yhi = 0 + if self.bc_high[2].lower() == 'magnetic' or self.bc_high[2] == 'pmc': + zhi = 0 + + # Assemble matrix + self.BC = Field(self.Nx, self.Ny, self.Nz, dtype=np.int8, use_ones=True) + + for d in ['x', 'y', 'z']: #tangential to zero + if d != 'x': + self.BC[0, :, :, d] = xlo + self.BC[-1, :, :, d] = xhi + if d != 'y': + self.BC[:, 0, :, d] = ylo + self.BC[:, -1, :, d] = yhi + if d != 'z': + self.BC[:, :, 0, d] = zlo + self.BC[:, :, -1, d] = zhi + + self.Dbc = diags(self.BC.toarray(), + shape=(3*self.N, 3*self.N), + dtype=np.int8 + ) + + # Update C (rows) + self.C = self.Dbc*self.C + + # Absorbing boundary conditions ABC + if any(True for x in self.bc_low if x.lower() == 'abc') \ + or any(True for x in self.bc_high if x.lower() == 'abc'): + if self.bc_high[0].lower() == 'abc': + self.tL[-1, :, :, 'x'] = self.L[0, :, :, 'x'] + self.itA[-1, :, :, 'y'] = self.iA[0, :, :, 'y'] + self.itA[-1, :, :, 'z'] = self.iA[0, :, :, 'z'] + + if self.bc_high[1].lower() == 'abc': + self.tL[:, -1, :, 'y'] = self.L[:, 0, :, 'y'] + self.itA[:, -1, :, 'x'] = self.iA[:, 0, :, 'x'] + self.itA[:, -1, :, 'z'] = self.iA[:, 0, :, 'z'] + + if self.bc_high[2].lower() == 'abc': + self.tL[:, :, -1, 'z'] = self.L[:, :, 0, 'z'] + self.itA[:, :, -1, 'x'] = self.iA[:, :, 0, 'x'] + self.itA[:, :, -1, 'y'] = self.iA[:, :, 0, 'y'] + + self.tDs = diags(self.tL.toarray(), shape=(3*self.N, 3*self.N), dtype=float) + self.itDa = diags(self.itA.toarray(), shape=(3*self.N, 3*self.N), dtype=float) + self.activate_abc = True + + # Perfect Matching Layers (PML) + if any(True for x in self.bc_low if x.lower() == 'pml') \ + or any(True for x in self.bc_high if x.lower() == 'pml'): + self.activate_pml = True + self.use_conductivity = True + + def fill_pml_sigmas(self): + ''' + Routine to calculate pml sigmas and apply them + to the conductivity tensor sigma + + [IN-PROGRESS] + ''' + + # Initialize + sx, sy, sz = np.zeros(self.Nx), np.zeros(self.Ny), np.zeros(self.Nz) + pml_exp = 2 + + # Fill + if self.bc_low[0].lower() == 'pml': + #sx[0:self.n_pml] = eps_0/(2*self.dt)*((self.x[self.n_pml] - self.x[:self.n_pml])/(self.n_pml*self.dx))**pml_exp + sx[0:self.n_pml] = np.linspace( self.pml_hi, self.pml_lo, self.n_pml) + for d in ['x', 'y', 'z']: + for i in range(self.n_pml): + self.ieps[i, :, :, d] = 1./eps_0 + self.sigma[i, :, :, d] = sx[i] + #if sx[i] > 0 : self.ieps[i, :, :, d] = 1/(eps_0+sx[i]*(2*self.dt)) + + if self.bc_low[1].lower() == 'pml': + #sy[0:self.n_pml] = 1/(2*self.dt)*((self.y[self.n_pml] - self.y[:self.n_pml])/(self.n_pml*self.dy))**pml_exp + sy[0:self.n_pml] = self.pml_func( self.pml_hi, self.pml_lo, self.n_pml) + for d in ['x', 'y', 'z']: + for j in range(self.n_pml): + self.ieps[:, j, :, d] = 1./eps_0 + self.sigma[:, j, :, d] = sy[j] + #if sy[j] > 0 : self.ieps[:, j, :, d] = 1/(eps_0+sy[j]*(2*self.dt)) + + if self.bc_low[2].lower() == 'pml': + #sz[0:self.n_pml] = eps_0/(2*self.dt)*((self.z[self.n_pml] - self.z[:self.n_pml])/(self.n_pml*self.dz))**pml_exp + sz[0:self.n_pml] = self.pml_func( self.pml_hi, self.pml_lo, self.n_pml) + for d in ['x', 'y', 'z']: + for k in range(self.n_pml): + self.ieps[:, :, k, d] = 1./eps_0 + self.sigma[:, :, k, d] = sz[k] + #if sz[k] > 0. : self.ieps[:, :, k, d] = 1/(np.mean(sz[:self.n_pml])*eps_0) + + if self.bc_high[0].lower() == 'pml': + #sx[-self.n_pml:] = 1/(2*self.dt)*((self.x[-self.n_pml:] - self.x[-self.n_pml])/(self.n_pml*self.dx))**pml_exp + sx[-self.n_pml:] = self.pml_func( self.pml_lo, self.pml_hi, self.n_pml) + for d in ['x', 'y', 'z']: + for i in range(self.n_pml): + i +=1 + self.ieps[-i, :, :, d] = 1./eps_0 + self.sigma[-i, :, :, d] = sx[-i] + #if sx[-i] > 0 : self.ieps[-i, :, :, d] = 1/(eps_0+sx[-i]*(2*self.dt)) + + if self.bc_high[1].lower() == 'pml': + #sy[-self.n_pml:] = 1/(2*self.dt)*((self.y[-self.n_pml:] - self.y[-self.n_pml])/(self.n_pml*self.dy))**pml_exp + sy[-self.n_pml:] = self.pml_func( self.pml_lo, self.pml_hi, self.n_pml) + for d in ['x', 'y', 'z']: + for j in range(self.n_pml): + j +=1 + self.ieps[:, -j, :, d] = 1./eps_0 + self.sigma[:, -j, :, d] = sy[-j] + #if sy[-j] > 0 : self.ieps[:, -j, :, d] = 1/(eps_0+sy[-j]*(2*self.dt)) + + if self.bc_high[2].lower() == 'pml': + #sz[-self.n_pml:] = eps_0/(2*self.dt)*((self.z[-self.n_pml:] - self.z[-self.n_pml])/(self.n_pml*self.dz))**pml_exp + sz[-self.n_pml:] = self.pml_func( self.pml_lo, self.pml_hi, self.n_pml) + for d in ['x', 'y', 'z']: + for k in range(self.n_pml): + k +=1 + self.ieps[:, :, -k, d] = 1./eps_0 + self.sigma[:, :, -k, d] = sz[-k] + #self.ieps[:, :, -k, d] = 1/(np.mean(sz[-self.n_pml:])*eps_0) + + def get_abc(self): + ''' + Save the n-2 timestep to apply ABC + ''' + E_abc, H_abc = {}, {} + + if self.bc_low[0].lower() == 'abc': + E_abc[0] = {} + H_abc[0] = {} + for d in ['x', 'y', 'z']: + E_abc[0][d+'lo'] = self.E[1, :, :, d] + H_abc[0][d+'lo'] = self.H[1, :, :, d] + + if self.bc_low[1].lower() == 'abc': + E_abc[1] = {} + H_abc[1] = {} + for d in ['x', 'y', 'z']: + E_abc[1][d+'lo'] = self.E[:, 1, :, d] + H_abc[1][d+'lo'] = self.H[:, 1, :, d] + + if self.bc_low[2].lower() == 'abc': + E_abc[2] = {} + H_abc[2] = {} + for d in ['x', 'y', 'z']: + E_abc[2][d+'lo'] = self.E[:, :, 1, d] + H_abc[2][d+'lo'] = self.H[:, :, 1, d] + + if self.bc_high[0].lower() == 'abc': + E_abc[0] = {} + H_abc[0] = {} + for d in ['x', 'y', 'z']: + E_abc[0][d+'hi'] = self.E[-1, :, :, d] + H_abc[0][d+'hi'] = self.H[-1, :, :, d] + + if self.bc_high[1].lower() == 'abc': + E_abc[1] = {} + H_abc[1] = {} + for d in ['x', 'y', 'z']: + E_abc[1][d+'hi'] = self.E[:, -1, :, d] + H_abc[1][d+'hi'] = self.H[:, -1, :, d] + + if self.bc_high[2].lower() == 'abc': + E_abc[2] = {} + H_abc[2] = {} + for d in ['x', 'y', 'z']: + E_abc[2][d+'hi'] = self.E[:, :, -1, d] + H_abc[2][d+'hi'] = self.H[:, :, -1, d] + + return E_abc, H_abc + + def update_abc(self, E_abc, H_abc): + ''' + Apply ABC algo to the selected BC, + to be applied after each timestep + ''' + + if self.bc_low[0].lower() == 'abc': + for d in ['x', 'y', 'z']: + self.E[0, :, :, d] = E_abc[0][d+'lo'] + self.H[0, :, :, d] = H_abc[0][d+'lo'] + + if self.bc_low[1].lower() == 'abc': + for d in ['x', 'y', 'z']: + self.E[:, 0, :, d] = E_abc[1][d+'lo'] + self.H[:, 0, :, d] = H_abc[1][d+'lo'] + + if self.bc_low[2].lower() == 'abc': + for d in ['x', 'y', 'z']: + self.E[:, :, 0, d] = E_abc[2][d+'lo'] + self.H[:, :, 0, d] = H_abc[2][d+'lo'] + + if self.bc_high[0].lower() == 'abc': + for d in ['x', 'y', 'z']: + self.E[-1, :, :, d] = E_abc[0][d+'hi'] + self.H[-1, :, :, d] = H_abc[0][d+'hi'] + + if self.bc_high[1].lower() == 'abc': + for d in ['x', 'y', 'z']: + self.E[:, -1, :, d] = E_abc[1][d+'hi'] + self.H[:, -1, :, d] = H_abc[1][d+'hi'] + + if self.bc_high[2].lower() == 'abc': + for d in ['x', 'y', 'z']: + self.E[:, :, -1, d] = E_abc[2][d+'hi'] + self.H[:, :, -1, d] = H_abc[2][d+'hi'] + + def set_ghosts_to_0(self): + ''' + Cleanup for initial conditions if they are + accidentally applied to the ghost cells + ''' + # Set H ghost quantities to 0 + for d in ['x', 'y', 'z']: #tangential to zero + if d != 'x': + self.H[-1, :, :, d] = 0. + if d != 'y': + self.H[:, -1, :, d] = 0. + if d != 'z': + self.H[:, :, -1, d] = 0. + + # Set E ghost quantities to 0 + self.E[-1, :, :, 'x'] = 0. + self.E[:, -1, :, 'y'] = 0. + self.E[:, :, -1, 'z'] = 0. + + def apply_conductors(self): + ''' + Set the 1/epsilon values inside the PEC conductors to zero + ''' + self.flag_in_conductors = self.grid.flag_int_cell_yz[:-1,:,:] \ + + self.grid.flag_int_cell_zx[:,:-1,:] \ + + self.grid.flag_int_cell_xy[:,:,:-1] + + self.ieps *= self.flag_in_conductors + + def set_field_in_conductors_to_0(self): + ''' + Cleanup for initial conditions if they are + accidentally applied to the conductors + ''' + self.flag_cleanup = self.grid.flag_int_cell_yz[:-1,:,:] \ + + self.grid.flag_int_cell_zx[:,:-1,:] \ + + self.grid.flag_int_cell_xy[:,:,:-1] + + self.H *= self.flag_cleanup + self.E *= self.flag_cleanup + + def apply_stl(self): + ''' + Mask the cells inside the stl and assing the material + defined by the user + + * Note: stl material should contain **relative** epsilon and mu + ** Note 2: when assigning the stl material, the default values + 1./eps_0 and 1./mu_0 are substracted + ''' + grid = self.grid.grid + self.stl_solids = self.grid.stl_solids + self.stl_materials = self.grid.stl_materials + + for key in self.stl_solids.keys(): + + mask = np.reshape(grid[key], (self.Nx, self.Ny, self.Nz)).astype(int) + + if type(self.stl_materials[key]) is str: + # Retrieve from material library + mat_key = self.stl_materials[key].lower() + + eps = material_lib[mat_key][0]*eps_0 + mu = material_lib[mat_key][1]*mu_0 + + # Setting to zero + self.ieps += self.ieps * (-1.0*mask) + self.imu += self.imu * (-1.0*mask) + + # Adding new values + self.ieps += mask * 1./eps + self.imu += mask * 1./mu + + # Conductivity + if len(material_lib[mat_key]) == 3: + sigma = material_lib[mat_key][2] + self.sigma += self.sigma * (-1.0*mask) + self.sigma += mask * sigma + self.use_conductivity = True + + elif self.sigma_bg > 0.0: # assumed sigma = 0 + self.sigma += self.sigma * (-1.0*mask) + + else: + # From input + eps = self.stl_materials[key][0]*eps_0 + mu = self.stl_materials[key][1]*mu_0 + + # Setting to zero + self.ieps += self.ieps * (-1.0*mask) + self.imu += self.imu * (-1.0*mask) + + # Adding new values + self.ieps += mask * 1./eps + self.imu += mask * 1./mu + + # Conductivity + if len(self.stl_materials[key]) == 3: + sigma = self.stl_materials[key][2] + self.sigma += self.sigma * (-1.0*mask) + self.sigma += mask * sigma + self.use_conductivity = True + + elif self.sigma_bg > 0.0: # assumed sigma = 0 + self.sigma += self.sigma * (-1.0*mask) + + def attrcleanup(self): + # Fields + del self.L, self.tL, self.iA, self.itA + if hasattr(self, 'BC'): + del self.BC + del self.Dbc + + # Matrices + del self.Px, self.Py, self.Pz + del self.Ds, self.iDa, self.tDs, self.itDa + del self.C + + def save_state(self, filename="solver_state.h5", close=True): + """Save the solver state to an HDF5 file. + + This function saves only the key state variables (`H`, `E`, `J`) that are updated + in `one_step()`, storing them as datasets in an HDF5 file. + + Parameters: + ----------- + filename : str, optional + The name of the HDF5 file where the state will be stored. Default is "solver_state.h5". + close : bool, optional (default=True) + - If True, the HDF5 file is closed after saving, and the function returns nothing. + - If False, the function returns an open HDF5 file object, which must be closed manually. + + Returns: + -------- + h5py.File or None + - If `close=True`, nothing is returned. + - If `close=False`, returns an open `h5py.File` object for further manipulation. + """ + + if self.use_mpi: # MPI savestate + E = self.mpi_gather_asField('E') + H = self.mpi_gather_asField('H') + J = self.mpi_gather_asField('J') + + if self.rank == 0: + state = h5py.File(filename, "w") + state.create_dataset("H", data=self.H.toarray()) + state.create_dataset("E", data=self.E.toarray()) + state.create_dataset("J", data=self.J.toarray()) + # TODO: check for MPI-GPU + + elif self.use_gpu: # GPU savestate + state = h5py.File(filename, "w") + state.create_dataset("H", data=self.H.toarray().get()) + state.create_dataset("E", data=self.E.toarray().get()) + state.create_dataset("J", data=self.J.toarray().get()) + + else: # CPU savestate + state = h5py.File(filename, "w") + state.create_dataset("H", data=self.H.toarray()) + state.create_dataset("E", data=self.E.toarray()) + state.create_dataset("J", data=self.J.toarray()) + + if close: + state.close() + else: + return state # Caller must close this manually + + def load_state(self, filename="solver_state.h5"): + """Load the solver state from an HDF5 file. + + Reads the saved state variables (`H`, `E`, `J`) from the specified HDF5 file + and restores them to the solver. + + Parameters: + ----------- + filename : str, optional + The name of the HDF5 file to load the solver state from. Default is "solver_state.h5". + + Returns: + -------- + None + """ + state = h5py.File(filename, "r") + + self.E.fromarray(state["E"][:]) + self.H.fromarray(state["H"][:]) + self.J.fromarray(state["J"][:]) + + # TODO: support MPI loadstate + + state.close() + + def read_state(self, filename="solver_state.h5"): + """Open an HDF5 file for reading without loading its contents. + + This function returns an open `h5py.File` object, allowing the caller + to manually inspect or extract data as needed. The file must be closed + by the caller after use. + + Parameters: + ----------- + filename : str, optional + The name of the HDF5 file to open. Default is "solver_state.h5". + + Returns: + -------- + h5py.File + An open HDF5 file object in read mode. + """ + return h5py.File(filename, "r") + + def reset_fields(self): + """ + Resets the electromagnetic field components to zero. + + This function clears the electric field (E), magnetic field (H), and + current density (J) by setting all their components to zero in the + simulation domain. It ensures a clean restart for a new simulation. + + Notes + ----- + - This method is useful when reusing an existing simulation object + without reinitializing all attributes. + """ + for d in ['x', 'y', 'z']: + self.E[:, :, :, d] = 0.0 + self.H[:, :, :, d] = 0.0 + self.J[:, :, :, d] = 0.0 \ No newline at end of file diff --git a/build/lib/wakis/sources.py b/build/lib/wakis/sources.py new file mode 100644 index 0000000..ad2f205 --- /dev/null +++ b/build/lib/wakis/sources.py @@ -0,0 +1,418 @@ +# copyright ################################# # +# This file is part of the wakis Package. # +# Copyright (c) CERN, 2024. # +# ########################################### # + +''' +The `sources.py` script containts different classes +defining a time-dependent sources to be installed +in the electromagnetic simulation. + +All sources need an update function that will be called +every simulation timestep, e.g.: + def update(self, t, *args, **kwargs)` +''' + +import numpy as np +import matplotlib.pyplot as plt +from scipy.constants import mu_0, c as c_light + +class Beam: + def __init__(self, xsource=0., ysource=0., beta=1.0, + q=1e-9, sigmaz=None, ti=None): + ''' + Updates the current J every timestep + to introduce a gaussian beam + moving in +z direction + + Parameters + --- + xsource, ysource: float, default 0. + Transverse position of the source [m] + beta: float, default 1.0 + Relativistic beta of the beam [0-1.0] + q: float, default 1e-9 + Beam charge [C] + sigmaz: float, default None + Beam longitudinal sigma [m] + ti: float, default 9.55*sigmaz + Injection time [s] + ''' + + self.xsource, self.ysource = xsource, ysource + self.sigmaz = sigmaz + self.q = q + self.beta = beta + self.v = c_light*beta + if ti is not None: + self.ti = ti + else: self.ti = 8.548921333333334*self.sigmaz/self.v + self.is_first_update = True + + def update(self, solver, t): + if self.is_first_update: + self.ixs, self.iys = np.abs(solver.x-self.xsource).argmin(), np.abs(solver.y-self.ysource).argmin() + self.is_first_update = False + if hasattr(solver, 'ZMIN'): # support for MPI + self.zmin = solver.ZMIN + solver.dz/2 + else: + self.zmin = solver.z.min() + # reference shift + s0 = self.zmin - self.v*self.ti + s = solver.z - self.v*t + # gaussian + profile = 1/np.sqrt(2*np.pi*self.sigmaz**2)*np.exp(-(s-s0)**2/(2*self.sigmaz**2)) + # update + solver.J[self.ixs,self.iys,:,'z'] = self.q*self.v*profile/solver.dx/solver.dy + + def plot(self, t): + # reference shift + s0 = - self.v*self.ti + s = - self.v*t + # gaussian + profile = 1/np.sqrt(2*np.pi*self.sigmaz**2)*np.exp(-(s-s0)**2/(2*self.sigmaz**2)) + source = self.q*self.v*profile + + fig, ax = plt.subplots() + ax.plot(t, source, 'darkorange') + ax.set_xlabel('Time [s]') + ax.set_ylabel('Current Jz [Cm/s]', color='darkorange') + ax.set_ylim(0., +np.abs(source).max()*1.3) + + fig.tight_layout() + plt.show() + +class PlaneWave: + def __init__(self, xs=None, ys=None, zs=0, nodes=None, f=None, + amplitude=1.0, beta=1.0, phase=0): + ''' + Updates the fields E and H every timestep + to introduce a planewave excitation at the given + xs, ys slice, moving in z+ direction + + Parameters + --- + xs, ys: slice, default 0:N + Transverse positions of the source (indexes) + zs: int or slice, default 0 + Injection position in z + nodes: float, default 15 + Number of nodes between z.min and z.max + f: float, default nodes/T + Frequency of the plane wave [Hz]. It overrides nodes param. + beta: float, default 1. + Relativistic beta + + # TODO: support different directions + ''' + # Check inputs and update self + self.nodes = nodes + self.beta = beta + self.xs, self.ys = xs, ys + self.zs = zs + self.f = f + self.amplitude = amplitude + self.is_first_update = True + self.phase = phase + + self.vp = self.beta*c_light # wavefront velocity beta*c + self.w = 2*np.pi*self.f # ang. frequency + self.kz = self.w/c_light # wave number + self.tmax = np.inf + + if self.nodes is not None: + self.tmax = self.nodes/self.f + + def update(self, solver, t): + if self.is_first_update: + if self.xs is None: + self.xs = slice(0, solver.Nx) + if self.ys is None: + self.ys = slice(0, solver.Ny) + + self.is_first_update = False + + if t <= self.tmax: + solver.H[self.xs,self.ys,self.zs,'y'] = -self.amplitude * np.cos(self.w*t+self.phase) + solver.E[self.xs,self.ys,self.zs,'x'] = self.amplitude * np.cos(self.w*t+self.phase) /(self.kz/(mu_0*self.vp)) + else: + solver.H[self.xs,self.ys,self.zs,'y'] = 0. + solver.E[self.xs,self.ys,self.zs,'x'] = 0. + + def plot(self, t): + fig, ax = plt.subplots() + + sourceH = -self.amplitude * np.cos(self.w*t+self.phase) + sourceE = self.amplitude * np.cos(self.w*t+self.phase) /(self.kz/(mu_0*self.vp)) + + sourceH[t>self.tmax] = 0. + sourceE[t>self.tmax] = 0. + + ax.plot(t, sourceH, 'b') + ax.set_xlabel('Time [s]') + ax.set_ylabel('Magnetic field Hy [A/m]', color='b') + ax.set_ylim(-np.abs(sourceH).max(), +np.abs(sourceH).max()) + + axx = ax.twinx() + axx.plot(t, sourceE, 'r') + axx.set_ylabel('Electric field Ex [V/m]', color='r') + axx.set_ylim(-np.abs(sourceE).max(), +np.abs(sourceE).max()) + + fig.tight_layout() + plt.show() + +class WavePacket: + def __init__(self, xs=None, ys=None, zs=0, + sigmaz=None, sigmaxy=None, tinj=None, + wavelength=None, f=None, + amplitude=1.0, beta=1.0, phase=0): + ''' + Updates E and H fields every timestep to + introduce a 2d gaussian wave packetat the + given xs, ys slice travelling in z+ direction + + Parameters + ---- + xs, ys: slice, default 0:N + Transverse positions of the source [index] + zs: int, default 0 + Longitudinal position of the source [index] + sigmaz: float, default 10*dz + Longitudinal gaussian sigma [m] + sigmaxy: float, default 5*dx + Longitudinal gaussian sigma [m] + tinj: + Injection time delay [m], default 6*sigmaz + wavelength: float, default 10*dz + Wave packet wavelength [m] f=c/wavelength + f: float, default None + Wave packet frequency [Hz], overrides wavelength + beta: float, default 1. + Relativistic beta + ''' + + self.beta = beta + self.xs, self.ys = xs, ys + self.zs = zs + self.f = f + self.wavelength = wavelength + self.sigmaxy = sigmaxy + self.sigmaz = sigmaz + self.tinj = tinj + self.amplitude = amplitude + self.phase = phase + + if self.f is None: + self.f = c_light/self.wavelength + self.w = 2*np.pi*self.f + + if self.tinj is None: + self.tinj = 6*self.sigmaz + + self.is_first_update = True + + def update(self, solver, t): + if self.is_first_update: + if self.xs is None: + self.xs = slice(0, solver.Nx) + if self.ys is None: + self.ys = slice(0, solver.Ny) + if self.sigmaz is None: + self.sigmaz = 10*solver.dz + if self.sigmaxy is None: + self.sigmaxy = 5*solver.dx + if self.tinj is None: + self.tinj = 6*self.sigmaz + + self.is_first_update = False + + # reference shift + s0 = solver.z.min()-self.tinj + s = solver.z.min()-self.beta*c_light*t + + # 2d gaussian + X, Y = np.meshgrid(solver.x[self.xs], solver.y[self.ys]) + gaussxy = np.exp(-(X**2+Y**2)/(2*self.sigmaxy**2)) + gausst = np.exp(-(s-s0)**2/(2*self.sigmaz**2)) + + # Update + solver.H[self.xs,self.ys,self.zs,'y'] = -self.amplitude*np.cos(self.w*t+self.phase)*gaussxy*gausst + solver.E[self.xs,self.ys,self.zs,'x'] = self.amplitude*mu_0*c_light*np.cos(self.w*t+self.phase)*gaussxy*gausst + + def plot(self, t, zmin=0): + fig, ax = plt.subplots() + + # compute source evolution + s0 = zmin-self.tinj + s = zmin-self.beta*c_light*t + gausst = np.exp(-(s-s0)**2/(2*self.sigmaz**2)) + + sourceH = -self.amplitude*np.cos(self.w*t+self.phase)*gausst + ax.plot(t, sourceH, 'b') + ax.set_xlabel('Time [s]') + ax.set_ylabel('Magnetic field Hy [A/m]', color='b') + ax.set_ylim(-np.abs(sourceH).max(), +np.abs(sourceH).max()) + + sourceE = self.amplitude*mu_0*c_light*np.cos(self.w*t+self.phase)*gausst + axx = ax.twinx() + axx.plot(t, sourceE, 'r') + axx.set_ylabel('Electric field Ex [V/m]', color='r') + axx.set_ylim(-np.abs(sourceE).max(), +np.abs(sourceE).max()) + + fig.tight_layout() + plt.show() + +class Dipole: + def __init__(self, field='E', component='z', + xs=None, ys=None, zs=None, + nodes=10, f=None, amplitude=1.0, + phase=0): + ''' + Updates the given field and component every timestep + to introduce a dipole-like sinusoidal excitation + + Parameters + --- + field: str, default 'E' + Field to add to source to. Supports component e.g. 'Ex' + component: str, default 'z' + If not specified in field, component of the field to add the source to + xs, ys, zs: int or slice, default N/2 + Positions of the source (indexes) + nodes: float, default 15 + Number of nodes between z.min and z.max + f: float, default nodes/T + Frequency of the plane wave [Hz]. It overrides nodes param. + ''' + # Check inputs and update self + self.nodes = nodes + self.xs, self.ys, self.zs = xs, ys, zs + self.f = f + self.field = field + self.component = component + self.amplitude = amplitude + self.phase = phase + + if len(field) == 2: #support for e.g. field='Ex' + self.component = field[1] + self.field = field[0] + + self.is_first_update = True + + + def update(self, solver, t): + if self.is_first_update: + if self.xs is None: + self.xs = int(solver.Nx/2) + if self.ys is None: + self.ys = int(solver.Ny/2) + if self.zs is None: + self.zs = int(solver.Nz/2) + if self.f is None: + T = (solver.z.max()-solver.z.min())/c_light + self.f = self.nodes/T + + self.w = 2*np.pi*self.f + self.is_first_update = False + + if self.field == 'E': + solver.E[self.xs,self.ys,self.zs,self.component] = self.amplitude*np.sin(self.w*t+self.phase) + elif self.field == 'H': + solver.H[self.xs,self.ys,self.zs,self.component] = self.amplitude*np.sin(self.w*t+self.phase) + elif self.field == 'J': + solver.J[self.xs,self.ys,self.zs,self.component] = self.amplitude*np.sin(self.w*t+self.phase) + else: + print(f'Field "{self.field}" not valid, should be "E", "H" or "J"]') + +class Pulse: + def __init__(self, field='E', component='z', + xs=None, ys=None, zs=None, + shape='Harris', L=None, amplitude=1.0, + delay=0.): + ''' + Injects an electromagnetic pulse at the given + source point (xs, ys, zs), with the selected + shape, length and amplitude. + + Parameters + ---- + field: str, default 'E' + Field to add to source to. Supports component e.g. 'Ex' + component: str, default 'z' + If not specified in field, component of the field to add the source to + xs, ys, zs: int or slice, default N/2 + Positions of the source (indexes) + shape: str, default 'Harris' + Profile of the pulse in time: ['Harris', 'Gaussian'] + L: float, default 50*dt + width of the pulse (~10*sigma) + + Note: injection time for the gaussian pulse t0=5*L to ensure smooth derivative. + ''' + # Check inputs and update self + + self.xs, self.ys, self.zs = xs, ys, zs + self.field = field + self.component = component + self.amplitude = amplitude + self.shape = shape + self.L = L + self.delay = delay + + if len(field) == 2: #support for e.g. field='Ex' + self.component = field[1] + self.field = field[0] + + if shape.lower() == 'harris': + self.tprofile = self.harris_pulse + elif shape.lower() == 'gaussian': + self.tprofile = self.gaussian_pulse + elif shape.lower() == 'rectangular': + self.tprofile = self.rectangular_pulse + else: + print('** shape does not, match available types: "Harris", "Gaussian", "Rectangular"') + + self.is_first_update = True + + def harris_pulse(self, t): + t = t*c_light - self.delay + try: + if t0.: + return 1.0 + else: + return 0.0 + + def update(self, solver, t): + if self.is_first_update: + if self.xs is None: + self.xs = int(solver.Nx/2) + if self.ys is None: + self.ys = int(solver.Ny/2) + if self.zs is None: + self.zs = int(solver.Nz/2) + if self.L is None: + self.L = 50*solver.dt + + self.is_first_update = False + + if self.field == 'E': + solver.E[self.xs,self.ys,self.zs,self.component] = self.amplitude*self.tprofile(t) + elif self.field == 'H': + solver.H[self.xs,self.ys,self.zs,self.component] = self.amplitude*self.tprofile(t) + elif self.field == 'J': + solver.J[self.xs,self.ys,self.zs,self.component] = self.amplitude*self.tprofile(t) + else: + print(f'Field "{self.field}" not valid, should be "E", "H" or "J"]') \ No newline at end of file diff --git a/build/lib/wakis/wakeSolver.py b/build/lib/wakis/wakeSolver.py new file mode 100644 index 0000000..328dedb --- /dev/null +++ b/build/lib/wakis/wakeSolver.py @@ -0,0 +1,1288 @@ +# copyright ################################# # +# This file is part of the wakis Package. # +# Copyright (c) CERN, 2024. # +# ########################################### # + +import time +import os +import glob +import shutil +import h5py +import numpy as np +from tqdm import tqdm +from scipy.constants import c as c_light + +class WakeSolver(): + ''' Class for wake potential and impedance + calculation from 3D time domain E fields + ''' + + def __init__(self, wakelength=None, q=1e-9, sigmaz=1e-3, beta=1.0, + xsource=0., ysource=0., xtest=0., ytest=0., + chargedist=None, ti=None, + compute_plane='both', skip_cells=0, add_space=None, + Ez_file='Ez.h5', save=True, results_folder='results/', + verbose=0, logfile=False): + ''' + Parameters + ---------- + wakelength : float, optional + Wakelength to be simulated. If not provided, it will be calculated from the Ez field data. + q : float + Beam total charge in [C] + sigmaz : float + Beam sigma in the longitudinal direction [m] + beta : float, deafult 1.0 + Ratio of beam's velocity to the speed of light c [a.u.] + xsource : float, default 0. + Beam center in the transverse plane, x-dir [m] + ysource : float, default 0. + Beam center in the transverse plane, y-dir [m] + xtest : float, default 0. + Integration path center in the transverse plane, x-dir [m] + ytest : float, default 0. + Integration path center in the transverse plane, y-dir [m] + ti : float, optional + Injection time, when beam enters domain [s]. If not provided, + the default value ti=8.53*sigmaz will be used + chargedist : dict or str, default None + If not provided, an analytic gaussian with sigmaz and q will be used. + When str, specifies the filename containing the charge distribution data + When dict, should contain the charge distribution data in keys (e.g.): {'X','Y'} + 'X' : longitudinal coordinate [m] + 'Y' : charge distribution in [C/m] + Ez_file : str, default 'Ez.h5' + hdf5 file containing Ez(x,y,z) data for every timestep + save: bool, default True + Flag to enable saving the wake potential, impedance and charge distribution + results in `.txt` files. + - Longitudinal: WP.txt, Z.txt. + - Transverse: WPx.txt, WPy.txt, Zx.txt, Zy.txt + - Charge distribution: lambda.txt, spectrum.txt + verbose: bool, default 0 + Controls the level of verbose in the terminal output + logfile: bool, default False + Creates a `wake.log` file with the summary of the input parameters + and calculations performed + + Attributes + ---------- + Ezt : ndarray + Matrix (nz x nt) containing Ez(x_test, y_test, z, t) + where nz = len(z), nt = len(t) + s : ndarray + Wakelegth vector s=c_light*t-z [m] representing the distance between + the source and the integration point. Goes from -ti*c_light + to the simulated wakelength where ti is the beam injection time. + WP : ndarray + Longitudinal wake potential WP(s) [V/pC] + WP_3d : ndarray + Longitudinal wake potential in 3d WP(x,y,s). Shape = (2*n+1, 2*n+1, len(s)) + where n = n_transverse_cells and s the wakelength array [V/pC] + Z : ndarray + Longitudinal impedance [Ohm] computed by the fourier-transformation of the + longitudinal component of the wake potential, which is divided by the + fourier-transformed charge distribution line function lambda(s) using a + single-sided DFT with 1000 samples. + WPx : ndarray + Trasnverse wake potential in x direction WPx(s) [V/pC] + WPy : ndarray + Transverse wake potential in y direction WP(s) [V/pC] + Zx : ndarray + Trasnverse impedance in x-dir Zx(f) [Ohm] + Zy : ndarray + Transverse impedance in y-dir Zy(f) [Ohm] + lambdas : ndarray + Linear charge distribution of the passing beam λ(s) [C/m] + lambdaf : ndarray + Charge distribution spectrum λ(f) [C] + dx : float + Ez field mesh step in transverse plane, x-dir [m] + dy : float + Ez field mesh step in transverse plane, y-dir [m] + x : ndarray + vector containing x-coordinates for field monitor [m] + y : ndarray + vector containing y-coordinates for field monitor [m] + n_transverse_cells : int, default 1 + Number of transverse cells used for the 3d calculation: 2*n+1 + This determines de size of the 3d wake potential + + ''' + + #beam + self.q = q + self.sigmaz = sigmaz + self.beta = beta + self.v = self.beta*c_light + self.xsource, self.ysource = xsource, ysource + self.xtest, self.ytest = xtest, ytest + self.chargedist = chargedist + self.ti = ti + self.skip_cells = skip_cells + self.compute_plane = compute_plane + self.DE_model = None + + if add_space is not None: #legacy support for add_space + self.skip_cells = add_space + + # Injection time + if ti is not None: + self.ti = ti + else: + #ti = 8.548921333333334*self.sigmaz/self.v #injection time as in CST for beta = 1 + ti = 8.548921333333334*self.sigmaz/(np.sqrt(self.beta)*self.v) #injection time as in CST for beta <=1 + self.ti = ti + + #field + self.Ez_file = Ez_file + self.Ez_hf = None + self.Ezt = None #Ez(x_t, y_t, z, t) + self.t = None + self.xf, self.yf, self.zf = None, None, None #field subdomain + self.x, self.y, self.z = None, None, None #full simulation domain + + #solver init + self.wakelength = wakelength + self.s = None + self.lambdas = None + self.WP = None + self.WP_3d = None + self.n_transverse_cells = 1 + self.WPx, self.WPy = None, None + self.f = None + self.Z = None + self.Zx, self.Zy = None, None + self.lambdaf = None + + #user + self.verbose = verbose + self.save = save + self.logfile = logfile + self.folder = results_folder + + if self.save: + if not os.path.exists(self.folder): + os.mkdir(self.folder) + + # create log + if self.log: + self.params_to_log() + + def solve(self, compute_plane=None, **kwargs): + ''' + Perform the wake potential and impedance for + longitudinal and transverse plane + ''' + if compute_plane is None: + compute_plane = self.compute_plane + + for key, val in kwargs.items(): + setattr(self, key, val) + + t0 = time.time() + + if compute_plane.lower() == 'both' or 'transverse': + # Obtain longitudinal Wake potential + self.calc_long_WP_3d() + + #Obtain transverse Wake potential + self.calc_trans_WP() + + #Obtain the longitudinal impedance + self.calc_long_Z() + + #Obtain transverse impedance + self.calc_trans_Z() + + elif compute_plane == 'longitudinal': + # Obtain longitudinal Wake potential + self.calc_long_WP() + + #Obtain the longitudinal impedance + self.calc_long_Z() + + #Elapsed time + t1 = time.time() + totalt = t1-t0 + self.log('Calculation terminated in %ds' %totalt) + + def update_long_WP(self, t): + '''WIP + calculation of wake potential on the fly + TODO: simplify logic, add transverse WP + ''' + + it = int(t/self.dt) + if it == 0: + # --- setup once before time loop --- + # self.s already computed as in your calc_long_WP + self.s = np.arange(-self.ti*self.v, self.wake.wakelength, self.dt*self.v) # 1D array of s-values (m) + self.WP = np.zeros_like(self.s, dtype=np.float64) # accumulator + + # --- inside your time loop, after Ez is updated for current timestep `it` --- + # get Ez at the probe (shape nz,) + # if you store Ez as Ezt[k, it] style, do: + Ez_curr = self.E[self.Nx//2, self.Ny//2, :, 'z'] # or extract from self.E at (xmid, ymid, :) if not using Ezt + + # compute s values for each z where current Ez contributes + s_vals = self.v * t - self.v * self.ti + np.min(self.z) - self.z + # convert to fractional index in s-array + idxf = (s_vals - self.s[0]) / (self.v*self.dt) # float indices + + # mask in-range contributions + mask = (idxf >= 0.0) & (idxf <= (len(self.s) - 1)) + if not np.any(mask): + # no contributions fall into s range this timestep + pass + else: + idxf_m = idxf[mask] + Ez_m = Ez_curr[mask] + + # integer floor indices + i0 = np.floor(idxf_m).astype(int) + frac = idxf_m - i0 + + # handle points that land exactly on last bin: we add all to last bin + last_bin_mask = (i0 >= len(self.s) - 1) + if np.any(last_bin_mask): + # assign entirely to last bin (no i0+1 available) + self.WP[-1] += np.sum(Ez_m[last_bin_mask]) * self.dz + # drop those from other accumulation + keep = ~last_bin_mask + i0 = i0[keep] + frac = frac[keep] + Ez_m = Ez_m[keep] + + # accumulate with linear interpolation weights + if i0.size: + # faster: use bincount + left = np.bincount(i0, weights=Ez_m * (1.0 - frac) * self.dz, minlength=len(self.s)) + right = np.bincount(i0 + 1, weights=Ez_m * frac * self.dz, minlength=len(self.s)) + self.WP += left + right + + if it == self.Nt-1: + WP = WP/(self.q*1e12) + + def calc_long_WP(self, Ezt=None,**kwargs): + ''' + Obtains the wake potential from the pre-computed longitudinal + Ez(z,t) field from the specified solver. + Parameters can be passed as **kwargs. + + Parameters + ---------- + t : ndarray + vector containing time values [s] + z : ndarray + vector containint z-coordinates [m] + sigmaz : float + Beam longitudinal sigma, to calculate injection time [m] + q : float + Beam charge, to normalize wake potential + ti : float, default 8.53*sigmaz/c + Injection time needed to set the negative part of s vector + and wakelength + Ezt : ndarray, default None + Matrix (nz x nt) containing Ez(x_test, y_test, z, t) + where nz = len(z), nt = len(t) + Ez_file : str, default 'Ez.h5' + HDF5 file containing the Ez(x, y, z) field data + for every timestep. Needed only if Ezt is not provided. + ''' + for key, val in kwargs.items(): + setattr(self, key, val) + + # Read h5 + if Ezt is not None: + self.Ezt = Ezt + + elif self.Ez_hf is None: + self.read_Ez() + + # time variables + nt = len(self.t) + dt = self.t[2]-self.t[1] + ti = self.ti + + # longitudinal variables + if self.zf is None: self.zf = self.z + dz = self.zf[2]-self.zf[1] + zmax = np.max(self.zf) #should it be domain's edge instead? + zmin = np.min(self.zf) + + if self.skip_cells !=0: + zz = slice(self.skip_cells, -self.skip_cells) + else: + zz = np.s_[:] + z = self.zf[zz] + nz = len(z) + + # Set Wake length and s + if self.wakelength is not None: + wakelength = self.wakelength + else: + wakelength = nt*dt*self.v - (zmax-zmin) - ti*self.v + self.wakelength = wakelength + + s = np.arange(-self.ti*self.v, wakelength, dt*self.v) + + self.log('Max simulated time = '+str(round(self.t[-1]*1.0e9,4))+' ns') + self.log('Wakelength = '+str(round(wakelength,3))+'m') + + # Initialize + WP = np.zeros_like(s) + keys = list(self.Ez_hf.keys()) + + # check for rounding errors + if nt > len(keys)-4: + nt = len(keys)-4 + self.log('*** rounding error in number of timesteps') + + # Assembly Ez field + self.log('Assembling Ez field...') + Ezt = np.zeros((nz,nt)) #Assembly Ez field + Ez = self.Ez_hf[keys[0]] + + if len(Ez.shape) == 3: + for n in range(nt): + Ez = self.Ez_hf[keys[n]] + Ezt[:, n] = Ez[Ez.shape[0]//2+1,Ez.shape[1]//2+1,zz] + + elif len(Ez.shape) == 1: + for n in range(nt): + Ezt[:, n] = self.Ez_hf[keys[n]] + self.Ezt = Ezt + + # integral of (Ez(xtest, ytest, z, t=(s+z)/c))dz + print('Calculating longitudinal wake potential WP(s)...') + with tqdm(total=len(s)*len(z)) as pbar: + for n in range(len(s)): + for k in range(nz): + ts = (z[k]+s[n])/self.v-zmin/self.v-self.t[0]+ti + it = int(ts/dt) #find index for t + if it < nt: + WP[n] = WP[n]+(Ezt[k, it])*dz #compute integral + pbar.update(1) + + WP = WP/(self.q*1e12) # [V/pC] + + self.s = s + self.WP = WP + + if self.save: + np.savetxt(self.folder+'WP.txt', np.c_[self.s,self.WP], header=' s [m]'+' '*20+'WP [V/pC]'+'\n'+'-'*48) + + def calc_long_WP_3d(self, **kwargs): + ''' + Obtains the 3d wake potential from the pre-computed Ez(x,y,z) + field from the specified solver. The calculation + Parameters can be passed as **kwargs. + + Parameters + ---------- + Ez_file : str, default 'Ez.h5' + HDF5 file containing the Ez(x,y,z) field data for every timestep + t : ndarray + vector containing time values [s] + z : ndarray + vector containing z-coordinates [m] + q : float + Total beam charge in [C]. Default is 1e9 C + n_transverse_cells : int, default 1 + Number of transverse cells used for the 3d calculation: 2*n+1 + This determines de size of the 3d wake potential + ''' + self.log('\n') + self.log('Longitudinal wake potential') + self.log('-'*24) + + for key, val in kwargs.items(): + setattr(self, key, val) + + # Read h5 + if self.Ez_hf is None: + self.read_Ez() + + # time variables + nt = len(self.t) + dt = self.t[2]-self.t[1] + ti = self.ti + + # longitudinal varianles + if self.zf is None: self.zf = self.z + dz = self.zf[2]-self.zf[1] + zmax = np.max(self.zf) + zmin = np.min(self.zf) + + if self.skip_cells !=0: + zz = slice(self.skip_cells, -self.skip_cells) + else: + zz = np.s_[:] + z = self.zf[zz] + nz = len(z) + + # Set Wake length and s + if self.wakelength is not None: + wakelength = self.wakelength + else: + wakelength = nt*dt*self.v - (zmax-zmin) - ti*self.v + self.wakelength = wakelength + + s = np.arange(-self.ti*self.v, wakelength, dt*self.v) + + self.log(f'* Max simulated time = {np.max(self.t)} s') + self.log(f'* Wakelength = {wakelength} m') + + #field subvolume in No.cells for x, y + i0, j0 = self.n_transverse_cells, self.n_transverse_cells + WP = np.zeros_like(s) + WP_3d = np.zeros((i0*2+1,j0*2+1,len(s))) + Ezt = np.zeros((nz,nt)) + keys = list(self.Ez_hf.keys()) + + # check for rounding errors + if nt > len(keys)-4: + nt = len(keys)-4 + self.log('*** rounding error in number of timesteps') + + print('Calculating longitudinal wake potential WP(s)') + with tqdm(total=len(s)*(i0*2+1)*(j0*2+1)) as pbar: + for i in range(-i0,i0+1,1): + for j in range(-j0,j0+1,1): + + # Assembly Ez field + for n in range(nt): + Ez = self.Ez_hf[keys[n]] + Ezt[:, n] = Ez[Ez.shape[0]//2+i,Ez.shape[1]//2+j, zz] + + # integral of (Ez(xtest, ytest, z, t=(s+z)/c))dz + for n in range(len(s)): + for k in range(nz): + ts = (z[k]+s[n])/self.v-zmin/self.v-self.t[0]+ti + it = int(ts/dt) #find index for t + if it < nt: + WP[n] = WP[n]+(Ezt[k, it])*dz #compute integral + + pbar.update(1) + + WP = WP/(self.q*1e12) # [V/pC] + WP_3d[i0+i,j0+j,:] = WP + + self.s = s + self.WP = WP_3d[i0,j0,:] + self.WP_3d = WP_3d + + self.log(f'Elapsed time {pbar.format_dict["elapsed"]} s') + + if self.save: + np.savetxt(self.folder+'WP.txt', np.c_[self.s, self.WP], header=' s [m]'+' '*20+'WP [V/pC]'+'\n'+'-'*48) + + def calc_trans_WP(self, **kwargs): + ''' + Obtains the transverse wake potential from the longitudinal + wake potential in 3d using the Panofsky-Wenzel theorem using a + second-order scheme for the gradient calculation + + Parameters + ---------- + WP_3d : ndarray + Longitudinal wake potential in 3d WP(x,y,s). Shape = (2*n+1, 2*n+1, len(s)) + where n = n_transverse_cells and s the wakelength array + s : ndarray + Wakelegth vector s=c*t-z representing the distance between + the source and the integration point. Goes from -8.53*sigmat to WL + where sigmat = sigmaz/c and WL is the Wakelength + dx : float + Ez field mesh step in transverse plane, x-dir [m] + dy : float + Ez field mesh step in transverse plane, y-dir [m] + x : ndarray, optional + vector containing x-coordinates [m] + y : ndarray, optional + vector containing y-coordinates [m] + n_transverse_cells : int, default 1 + Number of transverse cells used for the 3d calculation: 2*n+1 + This determines de size of the 3d wake potential + ''' + + for key, val in kwargs.items(): + setattr(self, key, val) + + self.log('\n') + self.log('Transverse wake potential') + self.log('-'*24) + self.log(f'* No. transverse cells = {self.n_transverse_cells}') + + # Obtain dx, dy, ds + if 'dx' in kwargs.keys() and 'dy' in kwargs.keys(): + dx = kwargs['dx'] + dy = kwargs['dy'] + else: + dx=self.xf[2]-self.xf[1] + dy=self.yf[2]-self.yf[1] + + ds = self.s[2]-self.s[1] + i0, j0 = self.n_transverse_cells, self.n_transverse_cells + + # Initialize variables + WPx = np.zeros_like(self.s) + WPy = np.zeros_like(self.s) + int_WP = np.zeros_like(self.WP_3d) + + print('Calculating transverse wake potential WPx, WPy...') + # Obtain the transverse wake potential + with tqdm(total=len(self.s)*(i0*2+1)*(j0*2+1)) as pbar: + for n in range(len(self.s)): + for i in range(-i0,i0+1,1): + for j in range(-j0,j0+1,1): + # Perform the integral + int_WP[i0+i,j0+j,n]=np.sum(self.WP_3d[i0+i,j0+j,0:n])*ds + pbar.update(1) + + # Perform the gradient (second order scheme) + WPx[n] = - (int_WP[i0+1,j0,n]-int_WP[i0-1,j0,n])/(2*dx) + WPy[n] = - (int_WP[i0,j0+1,n]-int_WP[i0,j0-1,n])/(2*dy) + + self.WPx = WPx + self.WPy = WPy + + self.log(f'Elapsed time {pbar.format_dict["elapsed"]} s') + + if self.save: + np.savetxt(self.folder+'WPx.txt', np.c_[self.s,self.WPx], header=' s [m]'+' '*20+'WP [V/pC]'+'\n'+'-'*48) + np.savetxt(self.folder+'WPy.txt', np.c_[self.s,self.WPx], header=' s [m]'+' '*20+'WP [V/pC]'+'\n'+'-'*48) + + def calc_long_Z(self, samples=1001, fmax=None, **kwargs): + ''' + Obtains the longitudinal impedance from the longitudinal + wake potential and the beam charge distribution using a + single-sided DFT with 1000 samples. + Parameters can be passed as **kwargs + + Parameters + ---------- + WP : ndarray + Longitudinal wake potential WP(s) + s : ndarray + Wakelegth vector s=c*t-z representing the distance between + the source and the integration point. Goes from -8.53*sigmat to WL + where sigmat = sigmaz/c and WL is the Wakelength + lambdas : ndarray + Charge distribution λ(s) interpolated to s axis, normalized by the beam charge + chargedist : ndarray, optional + Charge distribution λ(z). Not needed if lambdas is specified + q : float, optional + Total beam charge in [C]. Not needed if lambdas is specified + z : ndarray + vector containing z-coordinates [m]. Not needed if lambdas is specified + sigmaz : float + Beam sigma in the longitudinal direction [m]. + Used to calculate maximum frequency of interest fmax=c/(3*sigmaz) + ''' + self.log('\n') + self.log('Longitudinal impedance') + self.log('-'*24) + + for key, val in kwargs.items(): + setattr(self, key, val) + + print('Calculating longitudinal impedance Z...') + self.log(f'Single sided DFT with number of samples = {samples}') + + # setup charge distribution in s + if self.lambdas is None and self.chargedist is not None: + self.calc_lambdas() + elif self.lambdas is None and self.chargedist is None: + self.calc_lambdas_analytic() + try: + self.log('! Using analytic charge distribution λ(s) since no data was provided') + except: #ascii encoder error handling + self.log('! Using analytic charge distribution since no data was provided') + + # Set up the DFT computation + ds = np.mean(self.s[1:]-self.s[:-1]) + if fmax is None: + fmax = self.v/self.sigmaz/3 #max frequency of interest + N = int((self.v/ds)//fmax*samples) #to obtain a 1000 sample single-sided DFT + + # Obtain DFTs - is it v or c? + lambdafft = np.fft.fft(self.lambdas*self.v, n=N) + WPfft = np.fft.fft(self.WP*1e12, n=N) + ffft = np.fft.fftfreq(len(WPfft), ds/self.v) + + # Mask invalid frequencies + mask = np.logical_and(ffft >= 0 , ffft < fmax) + WPf = WPfft[mask]*ds + lambdaf = lambdafft[mask]*ds + self.f = ffft[mask] # Positive frequencies + + # Compute the impedance + self.Z = - WPf / lambdaf + self.lambdaf = lambdaf + + if self.save: + np.savetxt(self.folder+'Z.txt', np.c_[self.f, self.Z], header=' f [Hz]'+' '*20+'Z [Ohm]'+'\n'+'-'*48) + np.savetxt(self.folder+'spectrum.txt', np.c_[self.f, self.lambdaf], header=' f [Hz]'+' '*20+'Charge distribution spectrum [C/s]'+'\n'+'-'*48) + + def calc_trans_Z(self, samples=1001): + ''' + Obtains the transverse impedance from the transverse + wake potential and the beam charge distribution using a + single-sided DFT with 1000 samples + Parameters can be passed as **kwargs + ''' + self.log('\n') + self.log('Transverse impedance') + self.log('-'*24) + + print('Calculating transverse impedance Zx, Zy...') + self.log(f'Single sided DFT with number of samples = {samples}') + + # Set up the DFT computation + ds = self.s[2]-self.s[1] + fmax=1*self.v/self.sigmaz/3 + N=int((self.v/ds)//fmax*samples) #to obtain a 1000 sample single-sided DFT + + # Obtain DFTs + + # Normalized charge distribution λ(w) + lambdafft = np.fft.fft(self.lambdas*self.v, n=N) + ffft=np.fft.fftfreq(len(lambdafft), ds/self.v) + mask = np.logical_and(ffft >= 0 , ffft < fmax) + lambdaf = lambdafft[mask]*ds + + # Horizontal impedance Zx⊥(w) + WPxfft = np.fft.fft(self.WPx*1e12, n=N) + WPxf = WPxfft[mask]*ds + + self.Zx = 1j * WPxf / lambdaf + + # Vertical impedance Zy⊥(w) + WPyfft = np.fft.fft(self.WPy*1e12, n=N) + WPyf = WPyfft[mask]*ds + + self.Zy = 1j * WPyf / lambdaf + + self.fx = ffft[mask] + self.fy = ffft[mask] + + if self.save: + np.savetxt(self.folder+'Zx.txt', np.c_[self.fx, self.Zx], header=' f [Hz]'+' '*20+'Zx [Ohm]'+'\n'+'-'*48) + np.savetxt(self.folder+'Zy.txt', np.c_[self.fy, self.Zy], header=' f [Hz]'+' '*20+'Zy [Ohm]'+'\n'+'-'*48) + + def calc_lambdas(self, **kwargs): + '''Obtains normalized charge distribution in terms of s + λ(s) to use in the Impedance calculation + + Parameters + ---------- + s : ndarray + Wakelegth vector s=c*t-z representing the distance between + the source and the integration point. Goes from -8.53*sigmat to WL + where sigmat = sigmaz/c and WL is the Wakelength + chargedist : ndarray, optional + Charge distribution λ(z) + q : float, optional + Total beam charge in [C] + z : ndarray, optional + vector containing z-coordinates of the domain [m] + zf : ndarray, optional + vector containing z-coordinates of the field monitor [m]. N + ''' + for key, val in kwargs.items(): + setattr(self, key, val) + + if type(self.chargedist) is str: + d = self.read_txt(self.chargedist) + keys = list(d.keys()) + z = d[keys[0]] + chargedist = d[keys[1]] + + elif (self.chargedist) is dict: + keys = list(self.chargedist.keys()) + z = self.chargedist[keys[0]] + chargedist = self.chargedist[keys[1]] + + else: + chargedist = self.chargedist + if len(self.z) == len(self.chargedist): + z = self.z + elif len(self.zf) == len(self.chargedist): + z = self.zf + else: + self.log('Dimension error: check input dimensions') + + self.lambdas = np.interp(self.s, z, chargedist/self.q) + + if self.save: + np.savetxt(self.folder+'lambda.txt', np.c_[self.s, self.lambdas], header=' s [Hz]'+' '*20+'Charge distribution [C/m]'+'\n'+'-'*48) + + def calc_lambdas_analytic(self, **kwargs): + '''Obtains normalized charge distribution in s λ(z) + as an analytical gaussian centered in s=0 and std + equal sigmaz + + Parameters + ---------- + s : ndarray + Wakelegth vector s=c*t-z representing the distance between + the source and the integration point. Goes from -8.53*sigmat to WL + where sigmat = sigmaz/c and WL is the Wakelength + sigmaz : float + Beam sigma in the longitudinal direction [m] + ''' + + for key, val in kwargs.items(): + setattr(self, key, val) + + self.lambdas = 1/(self.sigmaz*np.sqrt(2*np.pi))*np.exp(-(self.s**2)/(2*self.sigmaz**2)) + + if self.save: + np.savetxt(self.folder+'lambda.txt', np.c_[self.s, self.lambdas], header=' s [Hz]'+' '*20+'Charge distribution [C/m]'+'\n'+'-'*48) + + def get_SmartBounds(self, freq_data=None, impedance_data=None, + minimum_peak_height=1.0, distance=3, inspect_bounds=True, + Rs_bounds=[0.8, 10], Q_bounds=[0.5, 5], fres_bounds=[-0.01e9, +0.01e9] + ): + import iddefix + + self.log('\nCalculating bounds using the Smart Bound Determination...') + # Smart bounds + # Find the main resonators and estimate the bounds -courtesy of Malthe Raschke! + bounds = iddefix.SmartBoundDetermination(freq_data, np.real(impedance_data), minimum_peak_height=minimum_peak_height, + Rs_bounds=Rs_bounds, Q_bounds=Q_bounds, fres_bounds=fres_bounds) + + bounds.find(minimum_peak_height=minimum_peak_height, distance=distance) + + if inspect_bounds: + bounds.inspect() + bounds.to_table() + return bounds + + bounds.to_table() + return bounds + + + def get_DEmodel_fitting(self, freq_data=None, impedance_data=None, + plane='longitudinal', dim='z', + parameterBounds=None, N_resonators=None, DE_kernel='DE', + maxiter=1e5, cmaes_sigma=0.01, popsize=150, tol=1e-3, + use_minimization=True, minimization_margin=[0.3, 0.2, 0.01], + minimum_peak_height=1.0, distance=3, inspect_bounds=False, + Rs_bounds=[0.8, 10], Q_bounds=[0.5, 5], fres_bounds=[-0.01e9, +0.01e9], + ): + import iddefix + + if freq_data is None or impedance_data is None: + if plane == 'longitudinal' and dim == 'z': + freq_data = self.f + impedance_data = self.Z + elif plane == 'transverse': + if dim == 'x': + freq_data = self.fx + impedance_data = self.Zx + elif dim == 'y': + freq_data = self.fy + impedance_data = self.Zy + else: + raise ValueError('Invalid dimension. Use dim = "x" or "y".') + else: + raise ValueError('Invalid plane or dimension. Use plane = "longitudinal" or "transverse" and choose the dimension dim = "z", "x" or "y".') + + if parameterBounds is None or N_resonators is None: + bounds = self.get_SmartBounds(parameterBounds=parameterBounds, + N_resonators=N_resonators, + minimum_peak_height=minimum_peak_height, + distance=distance, + inspect_bounds=inspect_bounds, + Rs_bounds=Rs_bounds, Q_bounds=Q_bounds, fres_bounds=fres_bounds) + N_resonators = bounds.N_resonators + parameterBounds = bounds.parameterBounds + + # Build the differential evolution model + print('Fitting the impedance using Differential Evolution...') + self.log('\nExtrapolating wake potential using Differential Evolution...') + + objectiveFunction=iddefix.ObjectiveFunctions.sumOfSquaredErrorReal + DE_model = iddefix.EvolutionaryAlgorithm(freq_data, + np.real(impedance_data), + N_resonators=N_resonators, + parameterBounds=parameterBounds, + plane=plane, + fitFunction='impedance', + wake_length=self.wakelength, # in [m] + objectiveFunction=objectiveFunction, + ) + + if DE_kernel == 'DE': + DE_model.run_differential_evolution(maxiter=int(maxiter), + popsize=popsize, + tol=tol, + mutation=(0.3, 0.8), + crossover_rate=0.5) + + elif DE_kernel == 'CMAES': #TODO: fix UnboundLocalError + DE_model.run_cmaes(maxiter=int(maxiter), + popsize=popsize, + sigma=cmaes_sigma,) + + if use_minimization: + self.log('Running minimization algorithm...') + DE_model.run_minimization_algorithm(minimization_margin) + + self.DE_model = DE_model + self.log(DE_model.warning) + if self.verbose: + print(DE_model.warning) + + return DE_model + + def get_extrapolated_wake(self, wakelength, sigma=None, use_minimization=True): + ''' + Get the extrapolated wake potential [V/pC] from the DE model + ''' + if self.DE_model is None: + raise AttributeError('Run get_DEmodel() first to obtain the DE model') + + if sigma is None: + sigma = self.sigmaz/c_light + + # Get the extrapolated wake potential + # TODO: add beta + t = np.arange(self.s[0]/c_light, wakelength/c_light, (self.s[2]-self.s[1])/c_light) + wake_potential = self.DE_model.get_wake_potential(t, sigma=sigma, + use_minimization=use_minimization) + + s = t * c_light # Convert time to distance [m] + return s, -wake_potential*1e-12 # in [V/pC] + CST convention + + def get_extrapolated_wake_function(self, wakelength, use_minimization=True): + ''' + Get the extrapolated wake function (a.k.a. Green function) from the DE model + ''' + if self.DE_model is None: + raise AttributeError('Run get_DEmodel() first to obtain the DE model') + + t = np.arange(self.s[0]/c_light, wakelength/c_light, (self.s[2]-self.s[1])/c_light) + wake_function = self.DE_model.get_wake(t, use_minimization=use_minimization) + return t, wake_function + + def get_extrapolated_impedance(self, f=None, use_minimization=True, + wakelength=None): + ''' + Get the extrapolated impedance [Ohm] from the DE model + ''' + if self.DE_model is None: + raise AttributeError('Run get_DEmodel() first to obtain the DE model') + + if f is None: + f = self.DE_model.frequency_data + + impedance = self.DE_model.get_impedance(frequency_data=f, + use_minimization=use_minimization, + wakelength=wakelength) + return f, impedance + + @staticmethod + def calc_impedance_from_wake(wake, s=None, t=None, fmax=None, + samples=None, verbose=True): + if type(wake) == list: + t = wake[0] + wake = wake[1] + if s is not None: + t = s/c_light + elif s is None and t is None: + raise AttributeError('Provide time data through parameter "t" [s] or "s" [m]') + dt = np.mean(t[1:]-t[:-1]) + + # Maximum frequency: fmax = 1/dt + if fmax is not None: + aux = np.arange(t.min(), t.max(), 1/fmax/2) + wake = np.interp(aux, t, wake) + dt = np.mean(aux[1:]-aux[:-1]); del aux + else: fmax = 1/dt + # Time resolution: fres=(1/len(wake)/dt/2) + + # Obtain DFTs + if samples is not None: + Wfft = np.fft.fft(wake, n=2*samples) + else: + Wfft = np.fft.fft(wake) + + ffft = np.fft.fftfreq(len(Wfft), dt) + + # Mask invalid frequencies + mask = np.logical_and(ffft >= 0 , ffft < fmax) + Z = Wfft[mask]/len(wake)*2 + f = ffft[mask] # Positive frequencies + + if verbose: + print(f'* Number of samples = {len(f)}') + print(f'* Maximum frequency = {f.max()} Hz') + print(f'* Maximum resolution = {np.mean(f[1:]-f[:-1])} Hz') + + return [f, Z] + + @staticmethod + def calc_wake_from_impedance(impedance, f=None, tmax=None, + samples=None, pad=0, verbose=True): + + if len(impedance) == 2: + f = impedance[0] + Z = impedance[1] + elif f is None: + raise AttributeError('Provide frequency data through parameter "f"') + else: + Z = impedance + df = np.mean(f[1:]-f[:-1]) + + # Maximum time: tmax = 1/(f[2]-f[1]) + if tmax is not None: + aux = np.arange(f.min(), f.max(), 1/tmax) + Z = np.interp(aux, f, Z) + df = np.mean(aux[1:]-aux[:-1]); del aux + else: tmax = 1/df + + # Time resolution: tres=(1/len(Z)/(f[2]-f[1])) + # pad = int(1/df/tres - len(Z)) + # wake = np.real(np.fft.ifft(np.pad(Z, pad))) + wake = np.real(-1*np.fft.fft(Z, n=samples)) + wake = np.roll(wake,-1) + # Inverse fourier transform of impedance + t = np.linspace(0, tmax, len(wake)) + + if verbose: + print(f'* Number of samples = {len(t)}') + print(f'* Maximum time = {t.max()} s') + print(f'* Maximum resolution = {np.mean(t[1:]-t[:-1])} s') + + return [t, wake] + + def read_Ez(self, filename=None, return_value=False): + ''' + Read the Ez.h5 file containing the Ez field information + ''' + + if filename is None: + filename = self.Ez_file + + hf = h5py.File(filename, 'r') + print(f'Reading h5 file {filename}' ) + self.log('Size of the h5 file: ' + str(round((os.path.getsize(filename)/10**9),2))+' Gb') + + #Set attributes + self.Ez_hf = hf + self.Ez_file = filename + if 'x' in hf.keys(): + self.xf = np.array(hf['x']) + if 'y' in hf.keys(): + self.yf = np.array(hf['y']) + if 'z' in hf.keys(): + self.zf = np.array(hf['z']) + if 't' in hf.keys(): + self.t = np.array(hf['t']) + + if return_value: + return hf + + def read_txt(self, txt, skiprows=2, delimiter=None, usecols=None): + ''' + Reads txt variables from ascii files and + returns data in a dictionary. Header should + be the first line. + ''' + + try: + load = np.loadtxt(txt, skiprows=skiprows, delimiter=delimiter, usecols=usecols) + except: + load = np.loadtxt(txt, skiprows=skiprows, delimiter=delimiter, + usecols=usecols, dtype=np.complex128) + + try: # keys == header names + with open(txt) as f: + header = f.readline() + + header = header.replace(' ', '') + header = header.replace('#', '') + header = header.replace('\n', '') + header = header.split(']') + + d = {} + for i in range(len(load[0,:])): + d[header[i]+']'] = load[:, i] + + except: #keys == int 0, 1, ... + d = {} + for i in range(len(load[0,:])): + d[i] = load[:, i] + + return d + + def save_txt(self, f_name, x_data=None, y_data=None, x_name='X [-]', y_name='Y [-]'): + """ + Saves x and y data to a text file in a two-column format. + + This function exports the provided `x_data` and `y_data` to a `.txt` file, + formatting the output with a header that includes custom column names. + + Parameters + ---------- + f_name : str + Name of the output file (without the `.txt` extension). + x_data : numpy.ndarray, optional + Array containing x-axis data. If None, the file is not saved. + y_data : numpy.ndarray, optional + Array containing y-axis data. If None, the file is not saved. + x_name : str, optional + Label for the x-axis column in the output file. Default is `"X [-]"`. + y_name : str, optional + Label for the y-axis column in the output file. Default is `"Y [-]"`. + + Notes + ----- + - The data is saved in a two-column format where `x_data` and `y_data` + are combined column-wise. + - If `x_data` or `y_data` is missing, the function prints a warning and does not save a file. + + Examples + -------- + Save two NumPy arrays to `data.txt`: + + >>> x = np.linspace(0, 10, 5) + >>> y = np.sin(x) + >>> save_txt("data", x, y, x_name="Time [s]", y_name="Amplitude") + + The saved file will look like: + + Time [s] Amplitude + -------------------------------- + 0.00 0.00 + 2.50 0.59 + 5.00 -0.99 + 7.50 0.94 + 10.00 -0.54 + """ + if x_data is not None and y_data is not None: + np.savetxt(f_name+'.txt', np.c_[x_data, y_data], header=' '+x_name+' '*20+y_name+'\n'+'-'*48) + else: + print('txt not saved, please provide x_data and y_data') + + def load_results(self, folder): + '''Load all txt from a given folder + + The txt files are generated when + the attribute`save = True` is used + ''' + if not folder.endswith('/'): + folder = folder + '/' + + _, self.lambdas = self.read_txt(folder+'lambda.txt').values() + _, self.WPx = self.read_txt(folder+'WPx.txt').values() + _, self.WPy = self.read_txt(folder+'WPy.txt').values() + self.s, self.WP = self.read_txt(folder+'WP.txt').values() + + _, self.lambdaf = self.read_txt(folder+'spectrum.txt').values() + _, self.Zx = self.read_txt(folder+'Zx.txt').values() + _, self.Zy = self.read_txt(folder+'Zy.txt').values() + self.f, self.Z = self.read_txt(folder+'Z.txt').values() + + self.f = np.abs(self.f) + self.wakelength = self.s[-1] + + def copy(self): + obj = type(self).__new__(self.__class__) + obj.__dict__.update(self.__dict__) + return obj + + def log(self, txt): + + if self.verbose: + print('\x1b[2;37m'+txt+'\x1b[0m') + + if not self.logfile: + return + + title = 'wake' + f = open(title + '.log', "a") + f.write(txt + '\r\n') + f.close() + + def params_to_log(self): + self.log(time.asctime()) + self.log('Wake computation') + self.log('='*24) + self.log(f'* Charge q = {self.q} [C]') + self.log(f'* Beam sigmaz = {self.sigmaz} [m]') + self.log(f'* xsource, ysource = {self.xsource}, {self.ysource} [m]') + self.log(f'* xtest, ytest = {self.xtest}, {self.ytest} [m]') + self.log(f'* Beam injection time ti= {self.ti} [s]') + + if self.chargedist is not None: + if type(self.chargedist) is str: + self.log(f'* Charge distribution file: {self.chargedist}') + else: + self.log(f'* Charge distribution data is provided') + else: + self.log(f'* Charge distribution analytic') + + self.log('\n') + + def read_cst_3d(self, path=None, folder='3d', filename='Ez.h5', units=1e-3): + ''' + Read CST 3d exports folder and store the + Ez field information into a matrix Ez(x,y,z) + for every timestep into a single `.h5` file + compatible with wakis. + + Parameters + ---------- + path: str, default None + Path to the field data + folder: str, default '3d' + Folder containing the CST field data .txt files + filename: str, default 'Ez.h5' + Name of the h5 file that will be generated + ''' + + self.log('Reading 3d CST field exports') + self.log('-'*24) + + if path is None: + path = folder + '/' + + # Rename files with E-02, E-03 + for file in glob.glob(path +'*E-02.txt'): + file=file.split(path) + title=file[1].split('_') + num=title[1].split('E') + num[0]=float(num[0])/100 + + ntitle=title[0]+'_'+str(num[0])+'.txt' + shutil.copy(path+file[1], path+file[1]+'.old') + os.rename(path+file[1], path+ntitle) + + for file in glob.glob(path +'*E-03.txt'): + file=file.split(path) + title=file[1].split('_') + num=title[1].split('E') + num[0]=float(num[0])/1000 + + ntitle=title[0]+'_'+str(num[0])+'.txt' + shutil.copy(path+file[1], path+file[1]+'.old') + os.rename(path+file[1], path+ntitle) + + for file in glob.glob(path +'*_0.txt'): + file=file.split(path) + title=file[1].split('_') + num=title[1].split('.') + num[0]=float(num[0]) + + ntitle=title[0]+'_'+str(num[0])+'.txt' + shutil.copy(path+file[1], path+file[1]+'.old') + os.rename(path+file[1], path+ntitle) + + #sort + try: + def sorter(item): + num=item.split(path)[1].split('_')[1].split('.txt')[0] + return float(num) + fnames = sorted(glob.glob(path+'*.txt'), key=sorter) + except: + fnames = sorted(glob.glob(path+'*.txt')) + + #Get the number of longitudinal and transverse cells used for Ez + i=0 + with open(fnames[0]) as f: + lines=f.readlines() + n_rows = len(lines)-3 #n of rows minus the header + x1=lines[3].split()[0] + + while True: + i+=1 + x2=lines[i+3].split()[0] + if x1==x2: + break + + n_transverse_cells=i + n_longitudinal_cells=int(n_rows/(n_transverse_cells**2)) + + # Create h5 file + if os.path.exists(path+filename): + os.remove(path+filename) + + hf = h5py.File(path+filename, 'w') + + # Initialize variables + Ez=np.zeros((n_transverse_cells, n_transverse_cells, n_longitudinal_cells)) + x=np.zeros((n_transverse_cells)) + y=np.zeros((n_transverse_cells)) + z=np.zeros((n_longitudinal_cells)) + t=[] + + nsteps, i, j, k = 0, 0, 0, 0 + skip=-4 #number of rows to skip + rows=skip + + # Start scan + self.log(f'Scanning files in {path}:') + for file in tqdm(fnames): + #self.log.debug('Scanning file '+ file + '...') + title=file.split(path) + title2=title[1].split('_') + + try: + num=title2[1].split('.txt') + t.append(float(num[0])*1e-9) + except: + t.append(nsteps) + + with open(file) as f: + for line in f: + rows+=1 + columns = line.split() + + if rows>=0 and len(columns)>1: + k=int(rows/n_transverse_cells**2) + j=int(rows/n_transverse_cells-n_transverse_cells*k) + i=int(rows-j*n_transverse_cells-k*n_transverse_cells**2) + + if k>= n_longitudinal_cells: + k = int(n_longitudinal_cells-1) + + Ez[i,j,k]=float(columns[5]) + + x[i]=float(columns[0])*units + y[j]=float(columns[1])*units + z[k]=float(columns[2])*units + + if nsteps == 0: + prefix='0'*5 + hf.create_dataset('Ez_'+prefix+str(nsteps), data=Ez) + else: + prefix='0'*(5-int(np.log10(nsteps))) + hf.create_dataset('Ez_'+prefix+str(nsteps), data=Ez) + + i, j, k = 0, 0, 0 + rows=skip + nsteps+=1 + + #close file + f.close() + + hf['x'] = x + hf['y'] = y + hf['z'] = z + hf['t'] = t + + hf.close() + + #set field info + self.log('Ez field is stored in a matrix with shape '+str(Ez.shape)+' in '+str(int(nsteps))+' datasets') + self.log(f'Finished scanning files - hdf5 file {filename} succesfully generated') + + #Update self + self.xf = x + self.yf = y + self.zf = z + self.t = np.array(t) diff --git a/wakis/__init__.py b/wakis/__init__.py index 16ebb49..7bd1b23 100644 --- a/wakis/__init__.py +++ b/wakis/__init__.py @@ -10,10 +10,12 @@ from . import materials from . import wakeSolver from . import geometry +from . import logger from .field import Field from .gridFIT3D import GridFIT3D from .solverFIT3D import SolverFIT3D from .wakeSolver import WakeSolver +from .logger import Logger from ._version import __version__ \ No newline at end of file diff --git a/wakis/logger.py b/wakis/logger.py new file mode 100644 index 0000000..c6730ec --- /dev/null +++ b/wakis/logger.py @@ -0,0 +1,35 @@ +# copyright ################################# # +# This file is part of the wakis Package. # +# Copyright (c) CERN, 2024. # +# ########################################### # + +from tqdm import tqdm + +import numpy as np +import time +import h5py + +from scipy.constants import c as c_light, epsilon_0 as eps_0, mu_0 as mu_0 +from scipy.sparse import csc_matrix as sparse_mat +from scipy.sparse import diags, hstack, vstack + + + +class Logger(): + + def __init__(self): + self.grid_logs = None + self.solver_logs = None + self.wakeSolver_logs = None + + def assign_grid_logs(self, grid_logs): + self.grid_logs = grid_logs + + def assign_solver_logs(self, solver_logs): + self.solver_logs = solver_logs + + def assign_wakeSolver_logs(self, wakeSolver_logs): + self.wakeSolver_logs = wakeSolver_logs + + def save_logs(self, filename): + print('Here I would save the logs to file:' + self.grid_logs + self.solver_logs + self.wakeSolver_logs) \ No newline at end of file diff --git a/wakis/solverFIT3D.py b/wakis/solverFIT3D.py index c142264..0887491 100644 --- a/wakis/solverFIT3D.py +++ b/wakis/solverFIT3D.py @@ -17,6 +17,7 @@ from .materials import material_lib from .plotting import PlotMixin from .routines import RoutinesMixin +from .logger import Logger try: from cupyx.scipy.sparse import csc_matrix as gpu_sparse_mat @@ -33,7 +34,7 @@ class SolverFIT3D(PlotMixin, RoutinesMixin): - def __init__(self, grid, wake=None, cfln=0.5, dt=None, + def __init__(self, grid, wake=None, logger=None, cfln=0.5, dt=None, bc_low=['Periodic', 'Periodic', 'Periodic'], bc_high=['Periodic', 'Periodic', 'Periodic'], use_stl=False, use_conductors=False, @@ -93,7 +94,15 @@ def __init__(self, grid, wake=None, cfln=0.5, dt=None, ''' self.verbose = verbose - if verbose: t0 = time.time() + + t0 = time.time() + solver_logs = {} + solver_logs["use_gpu"] = use_gpu + solver_logs["use_mpi"] = use_mpi + solver_logs["bc_low"] = bc_low + solver_logs["bc_high"] = bc_high + solver_logs["n_pml"] = n_pml + solver_logs["bg"] = bg # Flags self.step_0 = True @@ -213,6 +222,8 @@ def __init__(self, grid, wake=None, cfln=0.5, dt=None, self.dt = dt self.dt = dtype(self.dt) + solver_logs["dt"] = self.dt + if self.use_conductivity: # relaxation time criterion tau mask = np.logical_and(self.sigma.toarray()!=0, #for non-conductive @@ -252,6 +263,11 @@ def __init__(self, grid, wake=None, cfln=0.5, dt=None, if verbose: print(f'Total initialization time: {time.time() - t0} s') + solver_logs["solverInitializationTime"] = time.time() - t0 + + if logger is not None: + logger.assign_solver_logs(solver_logs) + def update_tensors(self, tensor='all'): '''Update tensor matrices after Field ieps, imu or sigma have been modified From 28d4c7443ec2c448ef76f0de80450c8297ebd22b Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Mon, 10 Nov 2025 18:05:47 +0100 Subject: [PATCH 02/39] Created Logfile and amended the test --- tests/test_001_pec_cubic_cavity.py | 2 +- tests/test_007_mpi_lossy_cavity.py | 24 ++++++++++++- wakis/gridFIT3D.py | 21 +++++++++++ wakis/logger.py | 56 +++++++++++++++++++++++------- wakis/routines.py | 9 ++++- wakis/solverFIT3D.py | 29 +++++++--------- wakis/wakeSolver.py | 30 +++++++++------- 7 files changed, 127 insertions(+), 44 deletions(-) diff --git a/tests/test_001_pec_cubic_cavity.py b/tests/test_001_pec_cubic_cavity.py index e015de7..aacf49e 100644 --- a/tests/test_001_pec_cubic_cavity.py +++ b/tests/test_001_pec_cubic_cavity.py @@ -99,7 +99,7 @@ def test_simulation(self): skip_cells = 12 # no. cells to skip in WP integration wake = WakeSolver(q=q, sigmaz=sigmaz, beta=beta, xsource=xs, ysource=ys, xtest=xt, ytest=yt, - save=False, logfile=False, Ez_file='tests/001_Ez.h5', + save=False, Ez_file='tests/001_Ez.h5', skip_cells=skip_cells, ) diff --git a/tests/test_007_mpi_lossy_cavity.py b/tests/test_007_mpi_lossy_cavity.py index 0650d68..9769884 100644 --- a/tests/test_007_mpi_lossy_cavity.py +++ b/tests/test_007_mpi_lossy_cavity.py @@ -82,6 +82,18 @@ class TestMPILossyCavity: -6.04105997e+01 ,-3.06532160e+01 ,-1.17749936e+01 ,-3.12574866e+00, -7.35339521e-01 ,-1.13085658e-01 , 7.18247535e-01 , 8.73829036e-02]) + grid_logs = {'Nx': 60, 'Ny': 60, 'Nz': 140, 'dx': 0.008666666348775227, 'dy': 0.008666666348775227, + 'dz': 0.005714285799435207, 'stl_solids': ['tests/stl/007_vacuum_cavity.stl', 'tests/stl/007_lossymetal_shell.stl'], + 'stl_materials': ['vacuum', [30, 1.0, 30]], 'gridInitializationTime': 0} + + solver_logs = {'use_gpu': False, 'use_mpi': False, 'bc_low': ['pec', 'pec', 'pec'], + 'bc_high': ['pec', 'pec', 'pec'], 'n_pml': 10, 'bg': 'pec', + 'dt': 6.970326728398968e-12, 'solverInitializationTime': 0} + + wakeSolver_logs = {'ti': 2.8516132094735135e-09, 'q': 1e-09, 'sigmaz': 0.1, 'beta': 1.0, + 'xsource': 0.0, 'ysource': 0.0, 'xtest': 0.0, 'ytest': 0.0, 'chargedist': None, + 'skip_cells': 10, 'results_folder': 'tests/007_results/', 'wakelength': 10.0, 'simulationTime': 0} + img_folder = 'tests/007_img/' def test_mpi_import(self): @@ -314,4 +326,14 @@ def test_long_impedance(self): assert np.allclose(np.real(wake.Z)[::20], np.real(self.Z), rtol=0.1), "Real Impedance samples failed" assert np.allclose(np.imag(wake.Z)[::20], np.imag(self.Z), rtol=0.1), "Imag Impedance samples failed" assert np.cumsum(np.abs(wake.Z))[-1] == pytest.approx(250910.51090497518, 0.1), "Abs Impedance cumsum failed" - \ No newline at end of file + + def test_log_file(self): + global solver + solver.logger.grid_logs["gridInitializationTime"] = 0 #times can vary + solver.logger.solver_logs["solverInitializationTime"] = 0 + solver.logger.wakeSolver_logs["simulationTime"] = 0 + logfile = os.path.join(solver.logger.wakeSolver_logs["results_folder"], "Simulation_Parameters.log") + assert os.path.exists(logfile), "Log file not created" + assert solver.logger.grid_logs == self.grid_logs, "Grid logs do not match expected values" + assert solver.logger.solver_logs == self.solver_logs, "Solver logs do not match expected values" + assert solver.logger.wakeSolver_logs == self.wakeSolver_logs, "WakeSolver logs do not match expected values" \ No newline at end of file diff --git a/wakis/gridFIT3D.py b/wakis/gridFIT3D.py index d6e447b..0213823 100644 --- a/wakis/gridFIT3D.py +++ b/wakis/gridFIT3D.py @@ -7,8 +7,10 @@ import pyvista as pv from functools import partial from scipy.optimize import least_squares +import time from .field import Field +from .logger import Logger try: from mpi4py import MPI @@ -63,6 +65,7 @@ def __init__(self, xmin, xmax, ymin, ymax, zmin, zmax, stl_rotate=[0., 0., 0.], stl_translate=[0., 0., 0.], stl_scale=1.0, stl_colors=None, verbose=1, stl_tol=1e-3): + t0 = time.time() if verbose: print('Generating grid...') self.verbose = verbose self.use_mpi = use_mpi @@ -143,6 +146,24 @@ def __init__(self, xmin, xmax, ymin, ymax, zmin, zmax, if stl_colors is None: self.assign_colors() + # Forward Parameters to logger + self.logger = Logger() + self.logger.grid_logs["Nx"] = self.Nx + self.logger.grid_logs["Ny"] = self.Ny + self.logger.grid_logs["Nz"] = self.Nz + self.logger.grid_logs["dx"] = self.dx + self.logger.grid_logs["dy"] = self.dy + self.logger.grid_logs["dz"] = self.dz + self.logger.grid_logs["stl_solids"] = list(self.stl_solids.values()) if self.stl_solids is not None else [] + self.logger.grid_logs["stl_materials"] = list(self.stl_materials.values()) if self.stl_materials is not None else [] + if stl_rotate != [0., 0., 0.]: + self.logger.grid_logs["stl_rotate"] = self.stl_rotate + if stl_translate != [0., 0., 0.]: + self.logger.grid_logs["stl_translate"] = self.stl_translate + if stl_scale != 1.0: + self.logger.grid_logs["stl_scale"] = self.stl_scale + self.logger.grid_logs["gridInitializationTime"] = time.time()-t0 + def compute_grid(self): X, Y, Z = np.meshgrid(self.x, self.y, self.z, indexing='ij') self.grid = pv.StructuredGrid(X.transpose(), Y.transpose(), Z.transpose()) diff --git a/wakis/logger.py b/wakis/logger.py index c6730ec..d7fad6f 100644 --- a/wakis/logger.py +++ b/wakis/logger.py @@ -1,6 +1,6 @@ # copyright ################################# # # This file is part of the wakis Package. # -# Copyright (c) CERN, 2024. # +# Copyright (c) CERN, 2025. # # ########################################### # from tqdm import tqdm @@ -8,6 +8,8 @@ import numpy as np import time import h5py +import os +import json from scipy.constants import c as c_light, epsilon_0 as eps_0, mu_0 as mu_0 from scipy.sparse import csc_matrix as sparse_mat @@ -18,18 +20,48 @@ class Logger(): def __init__(self): - self.grid_logs = None - self.solver_logs = None - self.wakeSolver_logs = None + self.grid_logs = {} + self.solver_logs = {} + self.wakeSolver_logs = {} - def assign_grid_logs(self, grid_logs): - self.grid_logs = grid_logs + def save_logs(self): + """ + Save all logs (grid, solver, wakeSolver) into log-file inside the results folder. + """ + logfile = os.path.join(self.wakeSolver_logs["results_folder"], "Simulation_Parameters.log") - def assign_solver_logs(self, solver_logs): - self.solver_logs = solver_logs + # Write sections + if not os.path.exists(self.wakeSolver_logs["results_folder"]): + os.mkdir(self.wakeSolver_logs["results_folder"]) + + with open(logfile, "w", encoding="utf-8") as fh: + fh.write("Simulation Parameters\n") + fh.write("""=====================\n\n""") - def assign_wakeSolver_logs(self, wakeSolver_logs): - self.wakeSolver_logs = wakeSolver_logs + sections = [ + ("WakeSolver Logs", self.wakeSolver_logs), + ("Solver Logs", self.solver_logs), + ("Grid Logs", self.grid_logs), + ] - def save_logs(self, filename): - print('Here I would save the logs to file:' + self.grid_logs + self.solver_logs + self.wakeSolver_logs) \ No newline at end of file + for title, data in sections: + fh.write(f"\n## {title} ##\n") + if not data: + fh.write("(empty)\n") + continue + + # convert non-serializable values to strings recursively + def _convert(obj): + if isinstance(obj, dict): + return {k: _convert(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_convert(v) for v in obj] + try: + json.dumps(obj) + return obj + except Exception: + return str(obj) + + clean = _convert(data) + fh.write(json.dumps(clean, indent=2, ensure_ascii=False)) + fh.write("\n") \ No newline at end of file diff --git a/wakis/routines.py b/wakis/routines.py index b525a41..5936e66 100644 --- a/wakis/routines.py +++ b/wakis/routines.py @@ -5,6 +5,7 @@ import numpy as np import h5py +import time from tqdm import tqdm from scipy.constants import c as c_light from wakis.sources import Beam @@ -297,6 +298,7 @@ def save_to_h5(self, hf, field, x, y, z, comp, n): if plot_from is None: plot_from = int(self.ti/self.dt) print('Running electromagnetic time-domain simulation...') + t0 = time.time() for n in tqdm(range(Nt)): # Initial condition @@ -340,4 +342,9 @@ def save_to_h5(self, hf, field, x, y, z, comp, n): # Compute wakefield magnitudes is done inside WakeSolver self.wake.solve(compute_plane=compute_plane) - + + # Forward parameters to logger + self.logger.wakeSolver_logs=self.wake.logger.wakeSolver_logs + self.logger.wakeSolver_logs["wakelength"]=wakelength + self.logger.wakeSolver_logs["simulationTime"]=time.time()-t0 + self.logger.save_logs() \ No newline at end of file diff --git a/wakis/solverFIT3D.py b/wakis/solverFIT3D.py index 0887491..cec8db8 100644 --- a/wakis/solverFIT3D.py +++ b/wakis/solverFIT3D.py @@ -34,7 +34,7 @@ class SolverFIT3D(PlotMixin, RoutinesMixin): - def __init__(self, grid, wake=None, logger=None, cfln=0.5, dt=None, + def __init__(self, grid, wake=None, cfln=0.5, dt=None, bc_low=['Periodic', 'Periodic', 'Periodic'], bc_high=['Periodic', 'Periodic', 'Periodic'], use_stl=False, use_conductors=False, @@ -94,15 +94,7 @@ def __init__(self, grid, wake=None, logger=None, cfln=0.5, dt=None, ''' self.verbose = verbose - t0 = time.time() - solver_logs = {} - solver_logs["use_gpu"] = use_gpu - solver_logs["use_mpi"] = use_mpi - solver_logs["bc_low"] = bc_low - solver_logs["bc_high"] = bc_high - solver_logs["n_pml"] = n_pml - solver_logs["bg"] = bg # Flags self.step_0 = True @@ -122,7 +114,7 @@ def __init__(self, grid, wake=None, logger=None, cfln=0.5, dt=None, # Grid self.grid = grid - + bg_log = bg self.Nx = self.grid.Nx self.Ny = self.grid.Ny self.Nz = self.grid.Nz @@ -222,8 +214,6 @@ def __init__(self, grid, wake=None, logger=None, cfln=0.5, dt=None, self.dt = dt self.dt = dtype(self.dt) - solver_logs["dt"] = self.dt - if self.use_conductivity: # relaxation time criterion tau mask = np.logical_and(self.sigma.toarray()!=0, #for non-conductive @@ -263,10 +253,17 @@ def __init__(self, grid, wake=None, logger=None, cfln=0.5, dt=None, if verbose: print(f'Total initialization time: {time.time() - t0} s') - solver_logs["solverInitializationTime"] = time.time() - t0 - - if logger is not None: - logger.assign_solver_logs(solver_logs) + # assign logs + self.logger = Logger() + self.logger.grid_logs = self.grid.logger.grid_logs + self.logger.solver_logs["use_gpu"] = use_gpu + self.logger.solver_logs["use_mpi"] = use_mpi + self.logger.solver_logs["bc_low"] = bc_low + self.logger.solver_logs["bc_high"] = bc_high + self.logger.solver_logs["n_pml"] = n_pml + self.logger.solver_logs["bg"] = bg_log + self.logger.solver_logs["dt"] = self.dt + self.logger.solver_logs["solverInitializationTime"] = time.time() - t0 def update_tensors(self, tensor='all'): '''Update tensor matrices after diff --git a/wakis/wakeSolver.py b/wakis/wakeSolver.py index 0078ac6..0e20525 100644 --- a/wakis/wakeSolver.py +++ b/wakis/wakeSolver.py @@ -12,6 +12,8 @@ from tqdm import tqdm from scipy.constants import c as c_light +from .logger import Logger + class WakeSolver(): ''' Class for wake potential and impedance calculation from 3D time domain E fields @@ -22,7 +24,7 @@ def __init__(self, wakelength=None, q=1e-9, sigmaz=1e-3, beta=1.0, chargedist=None, ti=None, compute_plane='both', skip_cells=0, add_space=None, Ez_file='Ez.h5', save=True, results_folder='results/', - verbose=0, logfile=False, counter_moving=False): + verbose=0, counter_moving=False): ''' Parameters ---------- @@ -61,9 +63,6 @@ def __init__(self, wakelength=None, q=1e-9, sigmaz=1e-3, beta=1.0, - Charge distribution: lambda.txt, spectrum.txt verbose: bool, default 0 Controls the level of verbose in the terminal output - logfile: bool, default False - Creates a `wake.log` file with the summary of the input parameters - and calculations performed counter_moving: bool, default False If the test charge is moving in the same or opposite direction to the source @@ -161,8 +160,8 @@ def __init__(self, wakelength=None, q=1e-9, sigmaz=1e-3, beta=1.0, #user self.verbose = verbose + self.logger = Logger() self.save = save - self.logfile = logfile self.folder = results_folder if self.save: @@ -173,6 +172,19 @@ def __init__(self, wakelength=None, q=1e-9, sigmaz=1e-3, beta=1.0, if self.log: self.params_to_log() + # Forward parameters to logger + self.logger.wakeSolver_logs["ti"]=self.ti + self.logger.wakeSolver_logs["q"]=self.q + self.logger.wakeSolver_logs["sigmaz"]=self.sigmaz + self.logger.wakeSolver_logs["beta"]=self.beta + self.logger.wakeSolver_logs["xsource"]=self.xsource + self.logger.wakeSolver_logs["ysource"]=self.ysource + self.logger.wakeSolver_logs["xtest"]=self.xtest + self.logger.wakeSolver_logs["ytest"]=self.ytest + self.logger.wakeSolver_logs["chargedist"]=self.chargedist + self.logger.wakeSolver_logs["skip_cells"]=self.skip_cells + self.logger.wakeSolver_logs["results_folder"]=self.folder + def solve(self, compute_plane=None, **kwargs): ''' Perform the wake potential and impedance for @@ -1114,14 +1126,6 @@ def log(self, txt): if self.verbose: print('\x1b[2;37m'+txt+'\x1b[0m') - - if not self.logfile: - return - - title = 'wake' - f = open(title + '.log', "a") - f.write(txt + '\r\n') - f.close() def params_to_log(self): self.log(time.asctime()) From d04601b3a50ef164e5a78240e05e602bd39f8032 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:52:37 +0100 Subject: [PATCH 03/39] Delete build/lib/wakis/__init__.py --- build/lib/wakis/__init__.py | 19 ------------------- 1 file changed, 19 deletions(-) delete mode 100644 build/lib/wakis/__init__.py diff --git a/build/lib/wakis/__init__.py b/build/lib/wakis/__init__.py deleted file mode 100644 index 16ebb49..0000000 --- a/build/lib/wakis/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -# copyright ################################# # -# This file is part of the wakis Package. # -# Copyright (c) CERN, 2024. # -# ########################################### # - -from . import field -from . import gridFIT3D -from . import solverFIT3D -from . import sources -from . import materials -from . import wakeSolver -from . import geometry - -from .field import Field -from .gridFIT3D import GridFIT3D -from .solverFIT3D import SolverFIT3D -from .wakeSolver import WakeSolver - -from ._version import __version__ \ No newline at end of file From 35f509049f35f1dc3bd3eeca356321ab7977d37e Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:53:45 +0100 Subject: [PATCH 04/39] Delete build/lib/wakis/_version.py --- build/lib/wakis/_version.py | 1 - 1 file changed, 1 deletion(-) delete mode 100644 build/lib/wakis/_version.py diff --git a/build/lib/wakis/_version.py b/build/lib/wakis/_version.py deleted file mode 100644 index 3966a5f..0000000 --- a/build/lib/wakis/_version.py +++ /dev/null @@ -1 +0,0 @@ -__version__ = '0.6.1' \ No newline at end of file From 70e5f61536ade351105cd86272171df3feed72ad Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:53:54 +0100 Subject: [PATCH 05/39] Delete build/lib/wakis/conductors.py --- build/lib/wakis/conductors.py | 204 ---------------------------------- 1 file changed, 204 deletions(-) delete mode 100644 build/lib/wakis/conductors.py diff --git a/build/lib/wakis/conductors.py b/build/lib/wakis/conductors.py deleted file mode 100644 index 1f02a2d..0000000 --- a/build/lib/wakis/conductors.py +++ /dev/null @@ -1,204 +0,0 @@ -import numpy as np -import scipy.optimize - -class OutRect: - def __init__(self, Lx, Ly, x_cent, y_cent): - self.Lx = Lx - self.Ly = Ly - self.x_cent = x_cent - self.y_cent = y_cent - # self.theta = theta - # self.R = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]) - # self.mR = np.array([[np.cos(-theta), -np.sin(-theta)], [np.sin(-theta), np.cos(-theta)]]) - - def out_conductor(self, x, y): - # [xx, yy] = np.dot(self.mR, np.array([x, y])) - return (-0.5 * self.Lx + self.x_cent < x < 0.5 * self.Lx + self.x_cent) and ( - -0.5 * self.Ly + self.y_cent < y < 0.5 * self.Ly + self.y_cent) - - def in_conductor(self, x, y): - return not self.out_conductor(x, y) - - def intersec_x(self, x, y): - # [xx, yy] = np.dot(self.mR, np.array([x, y])) - inters_1 = -0.5 * self.Lx + self.x_cent - inters_2 = inters_1 + self.Lx - if abs(x - inters_1) < abs(x - inters_2): - return inters_1 - else: - return inters_2 - - # [xxx, _] = np.dot(self.R, np.array([inters, yy])) - # return xxx - - def intersec_y(self, x, y): - # [xx, yy] = np.dot(self.mR, np.array([x, y])) - inters_1 = -0.5 * self.Ly + self.y_cent - inters_2 = inters_1 + self.Ly - if abs(y - inters_1) < abs(y - inters_2): - return inters_1 - else: - return inters_2 - - # [_, yyy] = np.dot(self.R, np.array([xx, inters])) - # return yyy - -class ImpFunc: - def __init__(self, func): - self.func = func - - def out_conductor(self, x, y): - return self.func(x, y) < 0 - - def in_conductor(self, x, y): - return self.func(x, y) > 0 - - def intersec_x(self, x, y): - func_x = lambda t : self.func(t, y) - - return scipy.optimize.newton_krylov(func_x, x) - - def intersec_y(self, x, y): - func_y = lambda t : self.func(x, t) - - return scipy.optimize.newton_krylov(func_y, y) - -class Plane: - def __init__(self, m_plane, q_plane, tol=0, sign=1): - self.tol = tol # 1e-16 - self.m_plane = m_plane - self.q_plane = q_plane - self.sign = sign - - def in_conductor(self, x, y): - if self.sign == 1: - return y - self.m_plane * x - self.q_plane >= self.tol - elif self.sign == -1: - return y - self.m_plane * x - self.q_plane <= self.tol - else: - print('sign must be + or - 1') - - def out_conductor(self, x, y): - return not self.in_conductor(x, y) - - def intersec_x(self, x, y): - return y / self.m_plane - self.q_plane / self.m_plane - - def intersec_y(self, x, y): - return self.m_plane * x + self.q_plane - - -class InCircle: - def __init__(self, radius, x_cent, y_cent): - self.radius = radius - self.x_cent = x_cent - self.y_cent = y_cent - - def in_conductor(self, x, y): - return np.square(x - self.x_cent) + np.square(y - self.y_cent) <= np.square(self.radius) - - def out_conductor(self, x, y): - return not self.in_conductor(x, y) - - def intersec_x(self, x, y): - if abs(y - self.y_cent) <= self.radius: - inters_1 = np.sqrt(np.square(self.radius) - np.square(y - self.y_cent)) + self.y_cent - inters_2 = -np.sqrt(np.square(self.radius) - np.square(y - self.y_cent)) + self.y_cent - if abs(x - inters_1) < abs(x - inters_2): - return inters_1 - else: - return inters_2 - else: - return np.inf - - def intersec_y(self, x, y): - if abs(x - self.x_cent) <= self.radius: - inters_1 = np.sqrt(np.square(self.radius) - np.square(x - self.x_cent)) + self.x_cent - inters_2 = -np.sqrt(np.square(self.radius) - np.square(x - self.x_cent)) + self.x_cent - if abs(y - inters_1) < abs(y - inters_2): - return inters_1 - else: - return inters_2 - else: - return np.inf - - -class OutCircle: - def __init__(self, radius, x_cent, y_cent): - self.radius = radius - self.x_cent = x_cent - self.y_cent = y_cent - - def in_conductor(self, x, y): - return np.square(x - self.x_cent) + np.square(y - self.y_cent) >= np.square(self.radius) - - def out_conductor(self, x, y): - return not self.in_conductor(x, y) - - def intersec_x(self, x, y): - # if abs(y - self.y_cent) > self.radius: - inters_1 = np.sqrt(np.square(self.radius) - np.square(y - self.y_cent)) + self.x_cent - inters_2 = -np.sqrt(np.square(self.radius) - np.square(y - self.y_cent)) + self.x_cent - if abs(x - inters_1) < abs(x - inters_2): - return inters_1 - else: - return inters_2 - # else: - # return np.inf - - def intersec_y(self, x, y): - # if abs(x - self.x_cent) > self.radius: - inters_1 = np.sqrt(np.square(self.radius) - np.square(x - self.x_cent)) + self.y_cent - inters_2 = -np.sqrt(np.square(self.radius) - np.square(x - self.x_cent)) + self.y_cent - if abs(y - inters_1) < abs(y - inters_2): - return inters_1 - else: - return inters_2 - # else: - # return np.inf - - -class ConductorsAssembly: - def __init__(self, conductors): - self.conductors = conductors - - def in_conductor(self, x, y): - for conductor in self.conductors: - if conductor.in_conductor(x, y): - return True - - return False - - def out_conductor(self, x, y): - return not self.in_conductor(x, y) - - def intersec_x(self, x, y): - list_inters = np.zeros_like(self.conductors) - for ii, conductor in enumerate(self.conductors): - list_inters[ii] = conductor.intersec_x(x, y) - - dist_inters = abs(list_inters - x) - return list_inters[np.argmin(dist_inters)] - - def intersec_y(self, x, y): - list_inters = np.zeros_like(self.conductors) - for ii, conductor in enumerate(self.conductors): - list_inters[ii] = conductor.intersec_y(x, y) - - dist_inters = abs(list_inters - y) - return list_inters[np.argmin(dist_inters)] - - -class noConductor: - def out_conductor(self, x, y): - return True - - def in_conductor(self, x, y): - return False - - def intersec_x(self, x, y): - return 1000 - - def intersec_y(self, x, y): - return 1000 - From 8ba80b501e6e9cddf2585876377edd49110a421d Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:54:02 +0100 Subject: [PATCH 06/39] Delete build/lib/wakis/conductors3d.py --- build/lib/wakis/conductors3d.py | 223 -------------------------------- 1 file changed, 223 deletions(-) delete mode 100644 build/lib/wakis/conductors3d.py diff --git a/build/lib/wakis/conductors3d.py b/build/lib/wakis/conductors3d.py deleted file mode 100644 index cbdaf24..0000000 --- a/build/lib/wakis/conductors3d.py +++ /dev/null @@ -1,223 +0,0 @@ -import numpy as np - - -class ConductorsAssembly: - def __init__(self, conductors): - self.conductors = conductors - - def in_conductor(self, x, y, z): - for conductor in self.conductors: - if conductor.in_conductor(x, y, z): - return True - - return False - - def out_conductor(self, x, y, z): - return not self.in_conductor(x, y, z) - - def intersec_x(self, x, y, z): - list_inters = np.zeros_like(self.conductors) - for ii, conductor in enumerate(self.conductors): - list_inters[ii] = conductor.intersec_x(x, y, z) - - dist_inters = abs(list_inters - x) - return list_inters[np.argmin(dist_inters)] - - def intersec_y(self, x, y, z): - list_inters = np.zeros_like(self.conductors) - for ii, conductor in enumerate(self.conductors): - list_inters[ii] = conductor.intersec_y(x, y, z) - - dist_inters = abs(list_inters - y) - return list_inters[np.argmin(dist_inters)] - - def intersec_z(self, x, y, z): - list_inters = np.zeros_like(self.conductors) - for ii, conductor in enumerate(self.conductors): - list_inters[ii] = conductor.intersec_z(x, y, z) - - dist_inters = abs(list_inters - z) - return list_inters[np.argmin(dist_inters)] - - -class InCube: - - def __init__(self, lx, ly, lz, x_cent, y_cent, z_cent): - self.lx = lx - self.ly = ly - self.lz = lz - self.x_cent = x_cent - self.y_cent = y_cent - self.z_cent = z_cent - - def out_conductor(self, x, y, z): - return (-0.5 * self.lx + self.x_cent < x < 0.5 * self.lx + self.x_cent and - -0.5 * self.ly + self.y_cent < y < 0.5 * self.ly + self.y_cent and - -0.5 * self.lz + self.z_cent < z < 0.5 * self.lz + self.z_cent) - - def in_conductor(self, x, y, z): - return not self.out_conductor(x, y, z) - - def intersec_x(self, x, y, z): - inters_1 = -0.5 * self.lx + self.x_cent - inters_2 = inters_1 + self.lx - if abs(x - inters_1) < abs(x - inters_2): - return inters_1 - else: - return inters_2 - - def intersec_y(self, x, y, z): - inters_1 = -0.5 * self.ly + self.y_cent - inters_2 = inters_1 + self.ly - if abs(y - inters_1) < abs(y - inters_2): - return inters_1 - else: - return inters_2 - - def intersec_z(self, x, y, z): - inters_1 = -0.5 * self.lz + self.z_cent - inters_2 = inters_1 + self.lz - if abs(z - inters_1) < abs(z - inters_2): - return inters_1 - else: - return inters_2 - - -class InSphere: - - def __init__(self, radius, x_cent, y_cent, z_cent): - self.radius = radius - self.x_cent = x_cent - self.y_cent = y_cent - self.z_cent = z_cent - - def in_conductor(self, x, y, z): - return (np.square(x - self.x_cent) + np.square(y - self.y_cent) + np.square(z - self.z_cent) - <= np.square(self.radius)) - - def out_conductor(self, x, y, z): - return not self.out_conductor(x, y, z) - - def intersec_x(self, x, y, z): - inters_1 = (np.sqrt(np.square(self.radius) - np.square(y - self.y_cent) - - np.square(z - self.z_cent)) + self.x_cent) - inters_2 = -(np.sqrt(np.square(self.radius) - np.square(y - self.y_cent) - - np.square(z - self.z_cent)) + self.x_cent) - - if abs(x - inters_1) < abs(x - inters_2): - return inters_1 - else: - return inters_2 - - def intersec_y(self, x, y, z): - inters_1 = (np.sqrt(np.square(self.radius) - np.square(x - self.x_cent) - - np.square(z - self.z_cent)) + self.y_cent) - inters_2 = -(np.sqrt(np.square(self.radius) - np.square(x - self.x_cent) - - np.square(z - self.z_cent)) + self.y_cent) - if abs(y - inters_1) < abs(y - inters_2): - return inters_1 - else: - return inters_2 - - def intersec_z(self, x, y, z): - inters_1 = (np.sqrt(np.square(self.radius) - np.square(x - self.x_cent) - - np.square(y - self.y_cent)) + self.z_cent) - inters_2 = -(np.sqrt(np.square(self.radius) - np.square(x - self.x_cent) - - np.square(y - self.y_cent)) + self.z_cent) - if abs(z - inters_1) < abs(z - inters_2): - return inters_1 - else: - return inters_2 - - -class noConductor: - def out_conductor(self, x, y, z): - return True - - def in_conductor(self, x, y, z): - return False - - def intersec_x(self, x, y, z): - return 1000 - - def intersec_y(self, x, y, z): - return 1000 - - def intersec_z(self, x, y, z): - return 1000 - - -class Plane: - def __init__(self, p, n): - self.p = p - self.n = n - - def in_conductor(self, x, y, z): - return self.n[0] * (x - self.p[0]) + self.n[1] * (y - self.p[1]) + self.n[2] * (z - self.p[2]) <= 0 - - def out_conductor(self, x, y, z): - return not self.in_conductor(x, y, z) - - def intersec_x(self, x, y, z): - if self.n[0] == 0: - return 1000 - else: - return -(self.n[1] * (y - self.p[1]) + self.n[2] * (z - self.p[2])) / self.n[0] + self.p[0] - - def intersec_y(self, x, y, z): - if self.n[1] == 0: - return 1000 - else: - return -(self.n[0] * (x - self.p[0]) + self.n[2] * (z - self.p[2])) / self.n[1] + self.p[1] - - def intersec_z(self, x, y, z): - if self.n[2] == 0: - return 1000 - else: - return -(self.n[0] * (x - self.p[0]) + self.n[1] * (y - self.p[1])) / self.n[2] + self.p[2] - - -class OutSphere: - def __init__(self, radius, x_cent=0, y_cent=0, z_cent=0): - self.radius = radius - self.x_cent = x_cent - self.y_cent = y_cent - self.z_cent = z_cent - - def in_conductor(self, x, y, z): - return (np.square(x - self.x_cent) + np.square(y - self.y_cent) + np.square(z - self.z_cent) - >= np.square(self.radius)) - - def out_conductor(self, x, y, z): - return not self.in_conductor(x, y, z) - - def intersec_x(self, x, y, z): - inters_1 = (np.sqrt(np.square(self.radius) - np.square(y - self.y_cent) - - np.square(z - self.z_cent)) + self.x_cent) - inters_2 = -(np.sqrt(np.square(self.radius) - np.square(y - self.y_cent) - - np.square(z - self.z_cent)) + self.x_cent) - - if abs(x - inters_1) < abs(x - inters_2): - return inters_1 - else: - return inters_2 - - def intersec_y(self, x, y, z): - inters_1 = (np.sqrt(np.square(self.radius) - np.square(x - self.x_cent) - - np.square(z - self.z_cent)) + self.y_cent) - inters_2 = -(np.sqrt(np.square(self.radius) - np.square(x - self.x_cent) - - np.square(z - self.z_cent)) + self.y_cent) - if abs(y - inters_1) < abs(y - inters_2): - return inters_1 - else: - return inters_2 - - def intersec_z(self, x, y, z): - inters_1 = (np.sqrt(np.square(self.radius) - np.square(x - self.x_cent) - - np.square(y - self.y_cent)) + self.z_cent) - inters_2 = -(np.sqrt(np.square(self.radius) - np.square(x - self.x_cent) - - np.square(y - self.y_cent)) + self.z_cent) - if abs(z - inters_1) < abs(z - inters_2): - return inters_1 - else: - return inters_2 From 57d969583b76d867c27af0254eb5bfc7abfd536b Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:54:12 +0100 Subject: [PATCH 07/39] Delete build/lib/wakis/field.py --- build/lib/wakis/field.py | 667 --------------------------------------- 1 file changed, 667 deletions(-) delete mode 100644 build/lib/wakis/field.py diff --git a/build/lib/wakis/field.py b/build/lib/wakis/field.py deleted file mode 100644 index 3dce54b..0000000 --- a/build/lib/wakis/field.py +++ /dev/null @@ -1,667 +0,0 @@ -# copyright ################################# # -# This file is part of the wakis Package. # -# Copyright (c) CERN, 2024. # -# ########################################### # - - -import numpy as xp -import copy - -try: - import cupy as xp_gpu - imported_cupy = True -except ImportError: - imported_cupy = False - -class Field: - ''' - Class to switch from 3D to collapsed notation by - defining the __getitem__ magic method - - linear numbering: - n = 1 + (i-1) + (j-1)*Nx + (k-1)*Nx*Ny - len(n) = Nx*Ny*Nz - ''' - def __init__(self, Nx, Ny, Nz, dtype=float, - use_ones=False, use_gpu=False): - - self.Nx = Nx - self.Ny = Ny - self.Nz = Nz - self.N = Nx*Ny*Nz - self.dtype = dtype - self.on_gpu = use_gpu - - if self.on_gpu: - if imported_cupy: - self.xp = xp_gpu - else: - print('*** cupy could not be imported, please CUDA check installation') - else: - self.xp = xp - - if use_ones: - self.array = self.xp.ones(self.N*3, dtype=self.dtype, order='F') - else: - self.array = self.xp.zeros(self.N*3, dtype=self.dtype, order='F') - - @property - def field_x(self): - return self.array[0:self.N] - - @property - def field_y(self): - return self.array[self.N: 2*self.N] - - @property - def field_z(self): - return self.array[2*self.N:3*self.N] - - @field_x.setter - def field_x(self, value): - if len(value.shape) > 1: - self.from_matrix(value, 'x') - else: - self.array[0:self.N] = value - - @field_y.setter - def field_y(self, value): - if len(value.shape) > 1: - self.from_matrix(value, 'y') - else: - self.array[self.N: 2*self.N] = value - - @field_z.setter - def field_z(self, value): - if len(value.shape) > 1: - self.from_matrix(value, 'z') - else: - self.array[2*self.N:3*self.N] = value - - def toarray(self): - return self.array - - def fromarray(self, array): - self.array[:] = array - - def to_matrix(self, key): - if key == 0 or key == 'x': - return self.xp.reshape(self.array[0:self.N], (self.Nx, self.Ny, self.Nz), order='F') - if key == 1 or key == 'y': - return self.xp.reshape(self.array[self.N: 2*self.N], (self.Nx, self.Ny, self.Nz), order='F') - if key == 2 or key == 'z': - return self.xp.reshape(self.array[2*self.N:3*self.N], (self.Nx, self.Ny, self.Nz), order='F') - - def from_matrix(self, mat, key): - if key == 0 or key == 'x': - self.array[0:self.N] = self.xp.reshape(mat, self.N, order='F') - elif key == 1 or key == 'y': - self.array[self.N: 2*self.N] = self.xp.reshape(mat, self.N, order='F') - elif key == 2 or key == 'z': - self.array[2*self.N:3*self.N] = self.xp.reshape(mat, self.N, order='F') - else: - raise IndexError('Component id not valid') - - def to_gpu(self): - if imported_cupy: - self.xp = xp_gpu - self.array = self.xp.asarray(self.array) # to cupy arr - self.on_gpu = True - else: - print('*** CuPy is not imported') - pass - - def from_gpu(self): - if self.on_gpu: - self.array = self.array.get() # to numpy arr - self.on_gpu = False - else: - print('*** GPU is not enabled') - pass - - def __getitem__(self, key): - - if type(key) is tuple: - if len(key) != 4: - raise IndexError('Need 3 indexes and component to access the field') - if key[3] == 0 or key[3] == 'x': - if self.on_gpu: - field = self.xp.reshape(self.array[0:self.N], (self.Nx, self.Ny, self.Nz), order='F') - return field[key[0], key[1], key[2]].get() - else: - field = self.xp.reshape(self.array[0:self.N], (self.Nx, self.Ny, self.Nz), order='F') - return field[key[0], key[1], key[2]] - elif key[3] == 1 or key[3] == 'y': - if self.on_gpu: - field = self.xp.reshape(self.array[self.N: 2*self.N], (self.Nx, self.Ny, self.Nz), order='F') - return field[key[0], key[1], key[2]].get() - else: - field = self.xp.reshape(self.array[self.N: 2*self.N], (self.Nx, self.Ny, self.Nz), order='F') - return field[key[0], key[1], key[2]] - elif key[3] == 2 or key[3] == 'z': - if self.on_gpu: - field = self.xp.reshape(self.array[2*self.N:3*self.N], (self.Nx, self.Ny, self.Nz), order='F') - return field[key[0], key[1], key[2]].get() - else: - field = self.xp.reshape(self.array[2*self.N:3*self.N], (self.Nx, self.Ny, self.Nz), order='F') - return field[key[0], key[1], key[2]] - elif type(key[3]) is str and key[3].lower() == 'abs': - field = self.get_abs() - return field[key[0], key[1], key[2]] - else: - raise IndexError('Component id not valid') - - elif type(key) is int: - if key <= self.N: - if self.on_gpu: - return self.array[key].get() - else: - return self.array[key] - else: - raise IndexError('Lexico-graphic index cannot be higher than product of dimensions') - - elif type(key) is slice: - if self.on_gpu: - return self.array[key].get() - else: - return self.array[key] - - else: - raise ValueError('key must be a 3-tuple or an integer') - - def __setitem__(self, key, value): - - if self.on_gpu: - value = self.xp.asarray(value) - - if type(key) is tuple: - if len(key) != 4: - raise IndexError('Need 3 indexes and component to access the field') - else: - field = self.to_matrix(key[3]) - field[key[0], key[1], key[2]] = value - self.from_matrix(field, key[3]) - - elif type(key) is int: - if key <= self.N: - self.array[key] = value - else: - raise IndexError('Lexico-graphic index cannot be higher than product of dimensions') - - elif type(key) is slice: - self.array[key] = value - else: - raise IndexError('key must be a 3-tuple or an integer') - - def __mul__(self, other, dtype=None): - - if dtype is None: - dtype = self.dtype - - # other is number - if type(other) is float or type(other) is int: - mulField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) - mulField.array = self.array * other - - # other is matrix - elif len(other.shape) > 1: - mulField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) - for d in ['x', 'y', 'z']: - mulField.from_matrix(self.to_matrix(d) * other, d) - - # other is 1d array - else: - mulField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) - mulField.array = self.array * other - - return mulField - - def __div__(self, other, dtype=None): - - if dtype is None: - dtype = self.dtype - - # other is number - if type(other) is float or type(other) is int: - divField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) - divField.array = self.array / other - - # other is matrix - if len(other.shape) > 1: - divField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) - for d in ['x', 'y', 'z']: - divField.from_matrix(self.to_matrix(d) / other, d) - - # other is constant or 1d array - else: - divField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) - divField.array = self.array / other - - return divField - - def __add__(self, other, dtype=None): - - if dtype is None: - dtype = self.dtype - - if type(other) is Field: - - addField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) - addField.field_x = self.field_x + other.field_x - addField.field_y = self.field_y + other.field_y - addField.field_z = self.field_z + other.field_z - - # other is matrix - elif len(other.shape) > 1: - addField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) - for d in ['x', 'y', 'z']: - addField.from_matrix(self.to_matrix(d) + other, d) - - # other is constant or 1d array - else: - addField = Field(self.Nx, self.Ny, self.Nz, dtype=dtype) - addField.array = self.array + other - - return addField - - def __repr__(self): - return 'x:\n' + self.field_x.__repr__() + '\n'+ \ - 'y:\n' + self.field_y.__repr__() + '\n'+ \ - 'z:\n' + self.field_z.__repr__() - - def __str__(self): - return 'x:\n' + self.field_x.__str__() + '\n'+ \ - 'y:\n' + self.field_y.__str__() + '\n'+ \ - 'z:\n' + self.field_z.__str__() - - def copy(self): - import copy - obj = type(self).__new__(self.__class__) # Create empty instance - - for key, value in self.__dict__.items(): - if key == "xp": - obj.xp = self.xp # Just copy reference, no need for deepcopy - elif key == "array" and self.on_gpu: - obj.array = self.xp.array(self.array) # Ensure CuPy array is copied properly - else: - obj.__dict__[key] = copy.deepcopy(value) - - return obj - - def compute_ijk(self, n): - if n > (self.N): - raise IndexError('Lexico-graphic index cannot be higher than product of dimensions') - - k = n//(self.Nx*self.Ny) - i = (n-k*self.Nx*self.Ny)%self.Nx - j = (n-k*self.Nx*self.Ny)//self.Nx - - return i, j, k - - def get_abs(self, as_matrix=True): - '''Computes the absolute or magnitude - out of the field components - ''' - if as_matrix: - if self.on_gpu: - return xp.sqrt(self.to_matrix('x')**2 + - self.to_matrix('y')**2 + - self.to_matrix('z')**2 ).get() - else: - return xp.sqrt(self.to_matrix('x')**2 + - self.to_matrix('y')**2 + - self.to_matrix('z')**2 ) - - else: # 1d array - if self.on_gpu: - return xp.sqrt(self.field_x**2 + self.field_y**2, self.field_z**2).get() - else: - return xp.sqrt(self.field_x**2 + self.field_y**2, self.field_z**2) - - def inspect(self, plane='YZ', cmap='bwr', dpi=100, figsize=[8,6], x=None, y=None, z=None, show=True, handles=False, **kwargs): - import matplotlib.pyplot as plt - from mpl_toolkits.axes_grid1 import make_axes_locatable - - if None not in (x,y,z): - pass - elif plane == 'XY': - key=[slice(0,self.Nx), slice(0,self.Ny), int(self.Nz//2)] - x, y, z = key[0], key[1], key[2] - extent = (0, self.Nx, 0, self.Ny) - xax, yax = 'nx', 'ny' - transpose = True - - elif plane == 'XZ': - key=[slice(0,self.Nx), int(self.Ny//2), slice(0,self.Nz)] - x, y, z = key[0], key[1], key[2] - extent = (0, self.Nz, 0, self.Nx) - xax, yax = 'nz', 'nx' - transpose = False - - elif plane == 'YZ': - key=[int(self.Nx//2), slice(0,self.Ny), slice(0,self.Nz)] - x, y, z = key[0], key[1], key[2] - extent = (0, self.Nz, 0, self.Ny) - xax, yax = 'nz', 'ny' - transpose = False - - fig, axs = plt.subplots(1, 3, tight_layout=True, figsize=figsize, dpi=dpi) - dims = {0:'x', 1:'y', 2:'z'} - - im = {} - - for d in [0,1,2]: - field = self.to_matrix(d) - - if self.on_gpu and hasattr(field, 'get'): - field = field.get() - - if transpose: - im[d] = axs[d].imshow(field[x,y,z].T, cmap=cmap, vmin=-field.max(), vmax=field.max(), extent=extent, origin='lower', **kwargs) - - else: - im[d] = axs[d].imshow(field[x,y,z], cmap=cmap, vmin=-field.max(), vmax=field.max(), extent=extent, origin='lower', **kwargs) - - for i, ax in enumerate(axs): - ax.set_title(f'Field {dims[i]}, plane {plane}') - fig.colorbar(im[i], cax=make_axes_locatable(ax).append_axes('right', size='5%', pad=0.1)) - ax.set_xlabel(xax) - ax.set_ylabel(yax) - - if handles: - return fig, axs - - if show: - plt.show() - - def inspect3D(self, field='all', backend='pyista', grid=None, - xmax=None, ymax=None, zmax=None, - bounding_box=True, show_grid=True, - cmap='viridis', dpi=100, show=True, handles=False): - """ - Visualize 3D field data on the structured grid using either Matplotlib - (voxel rendering) or PyVista (interactive clipping and slicing). - - This method provides two complementary visualization backends: - - **Matplotlib**: static voxel plots of the field components (x, y, z) - or all combined, useful for quick inspection, but memory intensive. - - **PyVista**: interactive 3D visualization with sliders to dynamically - clip the volume along X, Y, and Z, and optional wireframe slices. - - Parameters - ---------- - field : {'x', 'y', 'z', 'all'}, default 'all' - Which field component(s) to visualize. - - 'x', 'y', 'z': single component - - 'all': shows all three components - backend : {'matplotlib', 'pyvista'}, default 'pyvista' - Visualization backend to use: - - 'matplotlib': static voxel rendering - - 'pyvista': interactive 3D rendering with clipping sliders - grid : GridFIT3D or pyvista.StructuredGrid, optional - Structured grid object to use for visualization. If None, - a grid is constructed from the solver's internal dimensions. - xmax, ymax, zmax : int or float, optional - Maximum extents in each direction for visualization. Defaults - to the full grid dimensions if not specified. - bounding_box : bool, default True - If True, draw a wireframe bounding box of the simulation domain - (only used in PyVista backend). - show_grid : bool, default True - If True, show wireframe slice planes of the grid during - interactive visualization (PyVista backend). - cmap : str, default 'viridis' - Colormap to apply to the scalar field. - dpi : int, default 100 - Resolution of Matplotlib figures (only for Matplotlib backend). - show : bool, default True - Whether to display the figure/plot immediately. - - If False in PyVista, exports to `field.html` instead. - handles : bool, default False - If True, return figure/axes (Matplotlib) or the Plotter object - (PyVista) for further customization instead of showing directly. - - Returns - ------- - fig, axs : tuple, optional - Returned when `backend='matplotlib'` and `handles=True`. - pl : pyvista.Plotter, optional - Returned when `backend='pyvista'` and `handles=True`. - - Notes - ----- - - The PyVista backend provides interactive sliders to clip the - volume along each axis independently and inspect internal - structures of the 3D field. - - The Matplotlib backend provides a quick static voxel rendering - but is limited in interactivity and scalability. - - """ - - field = field.lower() - - # ---------- matplotlib backend --------------- - if backend.lower() == 'matplotlib': - import matplotlib.pyplot as plt - import matplotlib as mpl - from mpl_toolkits.axes_grid1 import make_axes_locatable - - fig = plt.figure(tight_layout=True, dpi=dpi, figsize=[12,6]) - - plot_x, plot_y, plot_z = False, False, False - - if field == 'all': - plot_x = True - plot_y = True - plot_z = True - - elif field.lower() == 'x': plot_x = True - elif field.lower() == 'y': plot_y = True - elif field.lower() == 'z': plot_z = True - - if xmax is None: xmax = self.Nx - if ymax is None: ymax = self.Ny - if zmax is None: zmax = self.Nz - - x,y,z = self.xp.mgrid[0:xmax+1,0:ymax+1,0:zmax+1] - axs = [] - - # field x - if plot_x: - arr = self.to_matrix('x')[0:int(xmax),0:int(ymax),0:int(zmax)] - if field == 'all': - ax = fig.add_subplot(1, 3, 1, projection='3d') - else: - ax = fig.add_subplot(1, 1, 1, projection='3d') - - vmin, vmax = -self.xp.max(self.xp.abs(arr)), +self.xp.max(self.xp.abs(arr)) - norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax) - colors = mpl.colormaps[cmap](norm(arr)) - vox = ax.voxels(x, y, z, filled=self.xp.ones_like(arr), facecolors=colors) - - m = mpl.cm.ScalarMappable(cmap=cmap, norm=norm) - m.set_array([]) - fig.colorbar(m, ax=ax, shrink=0.5, aspect=10) - ax.set_title(f'Field x') - axs.append(ax) - - # field y - if plot_y: - arr = self.to_matrix('y')[0:int(xmax),0:int(ymax),0:int(zmax)] - if field == 'all': - ax = fig.add_subplot(1, 3, 2, projection='3d') - else: - ax = fig.add_subplot(1, 1, 1, projection='3d') - - vmin, vmax = -self.xp.max(self.xp.abs(arr)), +self.xp.max(self.xp.abs(arr)) - norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax) - colors = mpl.colormaps[cmap](norm(arr)) - vox = ax.voxels(x, y, z, filled=self.xp.ones_like(arr), facecolors=colors) - - m = mpl.cm.ScalarMappable(cmap=cmap, norm=norm) - m.set_array([]) - fig.colorbar(m, ax=ax, shrink=0.5, aspect=10) - ax.set_title(f'Field y') - axs.append(ax) - - # field z - if plot_z: - arr = self.to_matrix('z')[0:int(xmax),0:int(ymax),0:int(zmax)] - if field == 'all': - ax = fig.add_subplot(1, 3, 3, projection='3d') - else: - ax = fig.add_subplot(1, 1, 1, projection='3d') - - vmin, vmax = -self.xp.max(self.xp.abs(arr)), +self.xp.max(self.xp.abs(arr)) - norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax) - colors = mpl.colormaps[cmap](norm(arr)) - vox = ax.voxels(x, y, z, filled=self.xp.ones_like(arr), facecolors=colors) - - m = mpl.cm.ScalarMappable(cmap=cmap, norm=norm) - m.set_array([]) - fig.colorbar(m, ax=ax, shrink=0.5, aspect=10) - ax.set_title(f'Field z') - axs.append(ax) - - dims = {0:'x', 1:'y', 2:'z'} - for i, ax in enumerate(axs): - ax.set_xlabel('Nx') - ax.set_ylabel('Ny') - ax.set_zlabel('Nz') - ax.xaxis.set_pane_color((1.0, 1.0, 1.0, 0.0)) - ax.yaxis.set_pane_color((1.0, 1.0, 1.0, 0.0)) - ax.zaxis.set_pane_color((1.0, 1.0, 1.0, 0.0)) - ax.set_xlim(self.Nx, 0) - ax.set_ylim(self.Ny, 0) - ax.set_zlim(self.Nz, 0) - - if handles: - return fig, axs - - if show: - plt.show() - - # ----------- pyvista backend --------------- - else: - import pyvista as pv - - if grid is not None and hasattr(grid, 'grid'): - xlo, xhi, ylo, yhi, zlo, zhi = grid.xmin, grid.xmax, grid.ymin, grid.ymax, grid.zmin, grid.zmax - grid = grid.grid - if field == 'x': - scalars = 'Field '+field - grid[scalars] = xp.reshape(self.to_matrix(field), self.N) - elif field == 'y': - scalars = 'Field '+field - grid[scalars] = xp.reshape(self.to_matrix(field), self.N) - elif field == 'z': - scalars = 'Field '+field - grid[scalars] = xp.reshape(self.to_matrix(field), self.N) - else: # for all or abs - scalars = 'Field '+'Abs' - grid[scalars] = xp.reshape(self.get_abs(), self.N) - - if xmax is None: xmax = xhi - if ymax is None: ymax = yhi - if zmax is None: zmax = zhi - - else: - print('[!] `grid` is not passed or is not a GridFIT3D object -> Using #N cells instead ') - x = xp.linspace(0, self.Nx, self.Nx+1) - y = xp.linspace(0, self.Ny, self.Ny+1) - z = xp.linspace(0, self.Nz, self.Nz+1) - xlo, xhi, ylo, yhi, zlo, zhi = x.min(), x.max(), y.min(), y.max(), z.min(), z.max() - if xmax is None: xmax = self.Nx - if ymax is None: ymax = self.Ny - if zmax is None: zmax = self.Nz - X, Y, Z = xp.meshgrid(x, y, z, indexing='ij') - grid = pv.StructuredGrid(X.transpose(), Y.transpose(), Z.transpose()) - - if field == 'x': - scalars = 'Field '+field - grid[scalars] = xp.reshape(self.to_matrix(field), self.N) - elif field == 'y': - scalars = 'Field '+field - grid[scalars] = xp.reshape(self.to_matrix(field), self.N) - elif field == 'z': - scalars = 'Field '+field - grid[scalars] = xp.reshape(self.to_matrix(field), self.N) - else: # for all or abs - scalars = 'Field '+'Abs' - grid[scalars] = xp.reshape(self.get_abs(), self.N) - - - pv.global_theme.allow_empty_mesh = True - pl = pv.Plotter() - vals = {'x':xmax, 'y':ymax, 'z':zmax} - - # --- Update function --- - def update_clip(val, axis="x"): - vals[axis] = val - # define bounds dynamically - if axis == "x": - slice_obj = grid.slice(normal="x", origin=(val, 0, 0)) - elif axis == "y": - slice_obj = grid.slice(normal="y", origin=(0, val, 0)) - else: # z - slice_obj = grid.slice(normal="z", origin=(0, 0, val)) - - # add clipped volume (scalars) - pl.add_mesh( - grid.clip_box(bounds=(xlo, vals['x'], ylo, vals['y'], zlo, vals['z']), invert=False), - scalars=scalars, - cmap=cmap, - name="clip", - ) - - # add slice wireframe (grid structure) - if show_grid: - pl.add_mesh(slice_obj, style="wireframe", color="grey", name="slice") - - # --- Sliders (placed side-by-side vertically) --- - pl.add_slider_widget( - lambda value: update_clip(value, "x"), - [xlo, xhi], - value=xmax, title="X Clip", - pointa=(0.8, 0.8), pointb=(0.95, 0.8), # top-right - style='modern', - ) - - pl.add_slider_widget( - lambda value: update_clip(value, "y"), - [ylo, yhi], - value=ymax, title="Y Clip", - pointa=(0.8, 0.6), pointb=(0.95, 0.6), # middle-right - style='modern', - ) - - pl.add_slider_widget( - lambda value: update_clip(value, "z"), - [zlo, zhi], - value=zmax, title="Z Clip", - pointa=(0.8, 0.4), pointb=(0.95, 0.4), # lower-right - style='modern', - ) - - # Camera orientation - pl.camera_position = 'zx' - pl.camera.azimuth += 30 - pl.camera.elevation += 30 - pl.set_background('mistyrose', top='white') - try: pl.add_logo_widget('../docs/img/wakis-logo-pink.png') - except: pass - pl.add_axes() - pl.enable_3_lights() - pl.enable_anti_aliasing() - - if bounding_box: - pl.add_mesh(pv.Box(bounds=(xlo, xhi, ylo, yhi, zlo, zhi)), - style="wireframe", color="black", line_width=2, name="domain_box") - - if handles: - return pl - - if not show: - pl.export_html(f'field.html') - else: - pl.show() From 869fdb9b60aa2f91863464034b784a46b344b664 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:54:21 +0100 Subject: [PATCH 08/39] Delete build/lib/wakis/pmlBlock2D.py --- build/lib/wakis/pmlBlock2D.py | 135 ---------------------------------- 1 file changed, 135 deletions(-) delete mode 100644 build/lib/wakis/pmlBlock2D.py diff --git a/build/lib/wakis/pmlBlock2D.py b/build/lib/wakis/pmlBlock2D.py deleted file mode 100644 index 073a82d..0000000 --- a/build/lib/wakis/pmlBlock2D.py +++ /dev/null @@ -1,135 +0,0 @@ -import numpy as np -from scipy.constants import c as c_light, epsilon_0 as eps_0, mu_0 as mu_0 - - -class PmlBlock2D: - def __init__(self, Nx, Ny, dt, dx, dy, lx_block=None, ly_block=None, rx_block=None, ry_block=None): - self.Nx = Nx - self.Ny = Ny - self.dt = dt - self.dx = dx - self.dy = dy - - self.Ex = np.zeros((Nx, Ny + 1)) - self.Ey = np.zeros((Nx + 1, Ny)) - self.Exy = np.zeros((Nx, Ny + 1)) - self.Exz = np.zeros((Nx, Ny + 1)) - self.Eyx = np.zeros((Nx + 1, Ny)) - self.Eyz = np.zeros((Nx + 1, Ny)) - self.Hz = np.zeros((Nx, Ny)) - self.Hzx = np.zeros((Nx, Ny)) - self.Hzy = np.zeros((Nx, Ny)) - self.Jx = np.zeros((self.Nx, self.Ny + 1)) - self.Jy = np.zeros((self.Nx + 1, self.Ny)) - - # the sigmas must be assembled by the solver - self.sigma_x = np.zeros_like(self.Ex) - self.sigma_y = np.zeros_like(self.Ey) - self.sigma_z = np.zeros_like(self.Ex) - self.sigma_star_x = np.zeros_like(self.Hz) - self.sigma_star_y = np.zeros_like(self.Hz) - self.sigma_star_z = np.zeros_like(self.Hz) - - # we assemble these after the sigmas - self.Ax = np.zeros_like(self.sigma_x) - self.Ay = np.zeros_like(self.sigma_y) - self.Az = np.zeros_like(self.sigma_z) - self.Bx = np.zeros_like(self.sigma_x) - self.By = np.zeros_like(self.sigma_y) - self.Bz = np.zeros_like(self.sigma_z) - self.Cx = np.zeros_like(self.sigma_star_x) - self.Cy = np.zeros_like(self.sigma_star_y) - self.Cz = np.zeros_like(self.sigma_star_z) - self.Dx = np.zeros_like(self.sigma_star_x) - self.Dy = np.zeros_like(self.sigma_star_y) - self.Dz = np.zeros_like(self.sigma_star_z) - - self.lx_block = lx_block - self.ly_block = ly_block - self.rx_block = rx_block - self.ry_block = ry_block - - self.C1 = self.dt / (self.dx * mu_0) - self.C2 = self.dt / (self.dy * mu_0) - self.C4 = self.dt / (self.dy * eps_0) - self.C5 = self.dt / (self.dx * eps_0) - self.C3 = self.dt / eps_0 - self.C6 = self.dt / eps_0 - - def assemble_coeffs(self): - self.Ax = (2 * eps_0 - self.dt * self.sigma_x) / (2 * eps_0 + self.dt * self.sigma_x) - self.Ay = (2 * eps_0 - self.dt * self.sigma_y) / (2 * eps_0 + self.dt * self.sigma_y) - self.Az = (2 * eps_0 - self.dt * self.sigma_z) / (2 * eps_0 + self.dt * self.sigma_z) - self.Bx = 2 * self.dt / (2 * eps_0 + self.dt * self.sigma_x) - self.By = 2 * self.dt / (2 * eps_0 + self.dt * self.sigma_y) - self.Bz = 2 * self.dt / (2 * eps_0 + self.dt * self.sigma_z) - self.Cx = (2 * mu_0 - self.dt * self.sigma_star_x) / (2 * mu_0 + self.dt * self.sigma_star_x) - self.Cy = (2 * mu_0 - self.dt * self.sigma_star_y) / (2 * mu_0 + self.dt * self.sigma_star_y) - self.Cz = (2 * mu_0 - self.dt * self.sigma_star_z) / (2 * mu_0 + self.dt * self.sigma_star_z) - self.Dx = 2 * self.dt / (2 * mu_0 + self.dt * self.sigma_star_x) - self.Dy = 2 * self.dt / (2 * mu_0 + self.dt * self.sigma_star_y) - self.Dz = 2 * self.dt / (2 * mu_0 + self.dt * self.sigma_star_z) - - def advance_h_fdtd(self): - Hz = self.Hz - Ex = self.Ex - Ey = self.Ey - - for ii in range(self.Nx): - for jj in range(self.Ny): - self.Hzx[ii, jj] = self.Cx[ii, jj] * self.Hzx[ii, jj] - self.Dx[ii, jj] / self.dx * (self.Ey[ii + 1, jj] - - self.Ey[ii, jj]) - self.Hzy[ii, jj] = self.Cy[ii, jj] * self.Hzy[ii, jj] + self.Dy[ii, jj] / self.dy * (self.Ex[ii, jj + 1] - - self.Ex[ii, jj]) - - self.Hz = self.Hzx + self.Hzy - - def advance_e_fdtd(self): - Hz = self.Hz - Ex = self.Ex - Ey = self.Ey - - for ii in range(self.Nx): - for jj in range(1, self.Ny): - self.Exy[ii, jj] = self.Ay[ii, jj] * self.Exy[ii, jj] + self.By[ii, jj] / self.dy * ( - self.Hz[ii, jj] - self.Hz[ii, jj - 1]) - self.Exz[ii, jj] = self.Az[ii, jj] * self.Exz[ii, jj] - - for ii in range(1, self.Nx): - for jj in range(self.Ny): - self.Eyz[ii, jj] = self.Az[ii, jj] * self.Eyz[ii, jj] - self.Eyx[ii, jj] = self.Ax[ii, jj] * self.Eyx[ii, jj] - self.Bx[ii, jj] / self.dx * ( - self.Hz[ii, jj] - self.Hz[ii - 1, jj]) - - self.Ex = self.Exy + self.Exz - self.Ey = self.Eyx + self.Eyz - - def update_e_boundary(self): - - if self.lx_block is not None: - for jj in range(self.Ny): - self.Eyz[0, jj] = self.Az[0, jj] * self.Eyz[0, jj] - self.Eyx[0, jj] = self.Ax[0, jj] * self.Eyx[0, jj] - self.Bx[0, jj] / self.dx * ( - self.Hz[0, jj] - self.lx_block.Hz[-1, jj]) - - if self.rx_block is not None: - for jj in range(self.Ny): - self.Eyz[-1, jj] = self.Az[-1, jj] * self.Eyz[-1, jj] - self.Eyx[-1, jj] = self.Ax[-1, jj] * self.Eyx[-1, jj] - self.Bx[-1, jj] / self.dx * ( - self.rx_block.Hz[0, jj] - self.Hz[-1, jj]) - - self.Ey = self.Eyx + self.Eyz - - if self.ly_block is not None: - for ii in range(self.Nx): - self.Exy[ii, 0] = self.Ay[ii, 0] * self.Exy[ii, 0] + self.By[ii, 0] / self.dy * ( - self.Hz[ii, 0] - self.ly_block.Hz[ii, - 1]) - self.Exz[ii, 0] = self.Az[ii, 0] * self.Exz[ii, 0] - - if self.ry_block is not None: - for ii in range(self.Nx): - self.Exy[ii, -1] = self.Ay[ii, -1] * self.Exy[ii, -1] + self.By[ii, -1] / self.dy * ( - self.ry_block.Hz[ii, 0] - self.Hz[ii, - 1]) - self.Exz[ii, -1] = self.Az[ii, -1] * self.Exz[ii, -1] - - self.Ex = self.Exy + self.Exz From 5d6afecb6122b606bca0dee84583cd8183bc8dcf Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:54:29 +0100 Subject: [PATCH 09/39] Delete build/lib/wakis/pmlBlock3D.py --- build/lib/wakis/pmlBlock3D.py | 265 ---------------------------------- 1 file changed, 265 deletions(-) delete mode 100644 build/lib/wakis/pmlBlock3D.py diff --git a/build/lib/wakis/pmlBlock3D.py b/build/lib/wakis/pmlBlock3D.py deleted file mode 100644 index ae98d0b..0000000 --- a/build/lib/wakis/pmlBlock3D.py +++ /dev/null @@ -1,265 +0,0 @@ -import numpy as np -from scipy.constants import c as c_light, epsilon_0 as eps_0, mu_0 as mu_0 - - -class PmlBlock3D: - def __init__(self, Nx, Ny, Nz, dt, dx, dy, dz, i_block, j_block, k_block, lx_block=None, ly_block=None, lz_block=None, rx_block=None, - ry_block=None, rz_block=None): - self.Nx = Nx - self.Ny = Ny - self.Nz = Nz - self.dt = dt - self.dx = dx - self.dy = dy - self.dz = dz - - self.i_block = i_block - self.j_block = j_block - self.k_block = k_block - - # Solver3D constructor takes care of populating this matrix (should we ensure this somehow?) - self.blocks_mat = np.full((3, 3, 3), None) - - self.Ex = np.zeros((Nx, Ny + 1, Nz + 1)) - self.Ey = np.zeros((Nx + 1, Ny, Nz + 1)) - self.Ez = np.zeros((Nx + 1, Ny + 1, Nz)) - self.Exy = np.zeros((Nx, Ny + 1, Nz + 1)) - self.Exz = np.zeros((Nx, Ny + 1, Nz + 1)) - self.Eyx = np.zeros((Nx + 1, Ny, Nz + 1)) - self.Eyz = np.zeros((Nx + 1, Ny, Nz + 1)) - self.Ezx = np.zeros((Nx + 1, Ny + 1, Nz)) - self.Ezy = np.zeros((Nx + 1, Ny + 1, Nz)) - self.Hx = np.zeros((Nx + 1, Ny, Nz)) - self.Hy = np.zeros((Nx, Ny + 1, Nz)) - self.Hz = np.zeros((Nx, Ny, Nz + 1)) - self.Hzx = np.zeros((Nx, Ny, Nz + 1)) - self.Hzy = np.zeros((Nx, Ny, Nz + 1)) - self.Hxy = np.zeros((Nx + 1, Ny, Nz)) - self.Hxz = np.zeros((Nx + 1, Ny, Nz)) - self.Hyx = np.zeros((Nx, Ny + 1, Nz)) - self.Hyz = np.zeros((Nx, Ny + 1, Nz)) - self.Jx = np.zeros((self.Nx, self.Ny + 1, self.Nz + 1)) - self.Jy = np.zeros((self.Nx + 1, self.Ny, self.Nz + 1)) - self.Jz = np.zeros((self.Nx + 1, self.Ny + 1, self.Nz)) - - # the sigmas must be assembled by the solver - self.sigma_x = np.zeros((Nx + 1, Ny + 1, Nz + 1)) - self.sigma_y = np.zeros((Nx + 1, Ny + 1, Nz + 1)) - self.sigma_z = np.zeros((Nx + 1, Ny + 1, Nz + 1)) - self.sigma_star_x = np.zeros((Nx, Ny + 1, Nz + 1)) - self.sigma_star_y = np.zeros((Nx + 1, Ny, Nz + 1)) - self.sigma_star_z = np.zeros((Nx + 1, Ny + 1, Nz)) - - # we assemble these after the sigmas - self.Ax = np.zeros_like(self.sigma_x) - self.Ay = np.zeros_like(self.sigma_y) - self.Az = np.zeros_like(self.sigma_z) - self.Bx = np.zeros_like(self.sigma_x) - self.By = np.zeros_like(self.sigma_y) - self.Bz = np.zeros_like(self.sigma_z) - self.Cx = np.zeros_like(self.sigma_star_x) - self.Cy = np.zeros_like(self.sigma_star_y) - self.Cz = np.zeros_like(self.sigma_star_z) - self.Dx = np.zeros_like(self.sigma_star_x) - self.Dy = np.zeros_like(self.sigma_star_y) - self.Dz = np.zeros_like(self.sigma_star_z) - - self.lx_block = lx_block - self.ly_block = ly_block - self.lz_block = lz_block - self.rx_block = rx_block - self.ry_block = ry_block - self.rz_block = rz_block - - self.C1 = self.dt / (self.dx * mu_0) - self.C2 = self.dt / (self.dy * mu_0) - self.C4 = self.dt / (self.dy * eps_0) - self.C5 = self.dt / (self.dx * eps_0) - self.C3 = self.dt / eps_0 - self.C6 = self.dt / eps_0 - - def assemble_coeffs(self): - self.Ax = (2 * eps_0 - self.dt * self.sigma_x) / (2 * eps_0 + self.dt * self.sigma_x) - self.Ay = (2 * eps_0 - self.dt * self.sigma_y) / (2 * eps_0 + self.dt * self.sigma_y) - self.Az = (2 * eps_0 - self.dt * self.sigma_z) / (2 * eps_0 + self.dt * self.sigma_z) - self.Bx = 2 * self.dt / (2 * eps_0 + self.dt * self.sigma_x) - self.By = 2 * self.dt / (2 * eps_0 + self.dt * self.sigma_y) - self.Bz = 2 * self.dt / (2 * eps_0 + self.dt * self.sigma_z) - self.Cx = (2 * mu_0 - self.dt * self.sigma_star_x) / (2 * mu_0 + self.dt * self.sigma_star_x) - self.Cy = (2 * mu_0 - self.dt * self.sigma_star_y) / (2 * mu_0 + self.dt * self.sigma_star_y) - self.Cz = (2 * mu_0 - self.dt * self.sigma_star_z) / (2 * mu_0 + self.dt * self.sigma_star_z) - self.Dx = 2 * self.dt / (2 * mu_0 + self.dt * self.sigma_star_x) - self.Dy = 2 * self.dt / (2 * mu_0 + self.dt * self.sigma_star_y) - self.Dz = 2 * self.dt / (2 * mu_0 + self.dt * self.sigma_star_z) - - def advance_h_fdtd(self): - Hz = self.Hz - Ex = self.Ex - Ey = self.Ey - - for ii in range(self.Nx): - for jj in range(self.Ny): - for kk in range(self.Nz + 1): - self.Hzx[ii, jj, kk] = self.Cx[ii, jj, kk] * self.Hzx[ii, jj, kk] - self.Dx[ii, jj, kk] / self.dx * (self.Ey[ii + 1, jj, kk] - - self.Ey[ii, jj, kk]) - self.Hzy[ii, jj, kk] = self.Cy[ii, jj, kk] * self.Hzy[ii, jj, kk] + self.Dy[ii, jj, kk] / self.dy * (self.Ex[ii, jj + 1, kk] - - self.Ex[ii, jj, kk]) - - for ii in range(self.Nx + 1): - for jj in range(self.Ny): - for kk in range(self.Nz): - self.Hxy[ii, jj, kk] = self.Cy[ii, jj, kk] * self.Hxy[ii, jj, kk] - self.Dy[ii, jj, kk] / self.dy * (self.Ez[ii, jj + 1, kk] - - self.Ez[ii, jj, kk]) - self.Hxz[ii, jj, kk] = self.Cz[ii, jj, kk] * self.Hxz[ii, jj, kk] + self.Dz[ii, jj, kk] / self.dz * (self.Ey[ii, jj, kk + 1] - - self.Ey[ii, jj, kk]) - - for ii in range(self.Nx): - for jj in range(self.Ny + 1): - for kk in range(self.Nz): - self.Hyx[ii, jj, kk] = self.Cx[ii, jj, kk] * self.Hyx[ii, jj, kk] + self.Dx[ii, jj, kk] / self.dx * (self.Ez[ii + 1, jj, kk] - - self.Ez[ii, jj, kk]) - self.Hyz[ii, jj, kk] = self.Cz[ii, jj, kk] * self.Hyz[ii, jj, kk] - self.Dz[ii, jj, kk] / self.dz * (self.Ex[ii, jj, kk + 1] - - self.Ex[ii, jj, kk]) - - def advance_e_fdtd(self): - Hz = self.Hz - Ex = self.Ex - Ey = self.Ey - - for ii in range(self.Nx): - for jj in range(1, self.Ny): - for kk in range(self.Nz + 1): - self.Exy[ii, jj, kk] = self.Ay[ii, jj, kk] * self.Exy[ii, jj, kk] + self.By[ii, jj, kk] / self.dy * ( - self.Hz[ii, jj, kk] - self.Hz[ii, jj - 1, kk]) - - for ii in range(self.Nx): - for jj in range(self.Ny + 1): - for kk in range(1, self.Nz): - self.Exz[ii, jj, kk] = self.Az[ii, jj, kk] * self.Exz[ii, jj, kk] - self.Bz[ii, jj, kk] / self.dz * ( - self.Hy[ii, jj, kk] - self.Hy[ii, jj, kk - 1]) - - for ii in range(1, self.Nx): - for jj in range(self.Ny): - for kk in range(self.Nz + 1): - self.Eyx[ii, jj, kk] = self.Ax[ii, jj, kk] * self.Eyx[ii, jj, kk] - self.Bx[ii, jj, kk] / self.dx * ( - self.Hz[ii, jj, kk] - self.Hz[ii - 1, jj, kk]) - - for ii in range(self.Nx + 1): - for jj in range(self.Ny): - for kk in range(1, self.Nz): - self.Eyz[ii, jj, kk] = self.Az[ii, jj, kk] * self.Eyz[ii, jj, kk] + self.Bz[ii, jj, kk] / self.dz * ( - self.Hx[ii, jj, kk] - self.Hx[ii, jj, kk - 1]) - - for ii in range(1, self.Nx): - for jj in range(self.Ny + 1): - for kk in range(self.Nz): - self.Ezx[ii, jj, kk] = self.Ax[ii, jj, kk]*self.Ezx[ii, jj, kk] + self.Bx[ii, jj, kk] / self.dx * ( - self.Hy[ii, jj, kk] - self.Hy[ii - 1, jj, kk]) - - for ii in range(self.Nx + 1): - for jj in range(1, self.Ny): - for kk in range(self.Nz): - self.Ezy[ii, jj, kk] = self.Ay[ii, jj, kk]*self.Ezy[ii, jj, kk] - self.By[ii, jj, kk] / self.dy * ( - self.Hx[ii, jj, kk] - self.Hx[ii, jj - 1, kk]) - - def sum_e_fields(self): - self.Ex = self.Exy + self.Exz - self.Ey = self.Eyx + self.Eyz - self.Ez = self.Ezx + self.Ezy - - def sum_h_fields(self): - self.Hx = self.Hxy + self.Hxz - self.Hy = self.Hyx + self.Hyz - self.Hz = self.Hzx + self.Hzy - - def update_e_boundary(self): - - i_block = self.i_block - j_block = self.j_block - k_block = self.k_block - blocks_mat = self.blocks_mat - Nx = self.Nx - Ny = self.Ny - Nz = self.Nz - Ax = self.Ax - Ay = self.Ay - Az = self.Az - Bx = self.Bx - By = self.By - Bz = self.Bz - Exy = self.Exy - Exz = self.Exz - Eyx = self.Eyx - Eyz = self.Eyz - Ezx = self.Ezx - Ezy = self.Ezy - Hx = self.Hx - Hy = self.Hy - Hz = self.Hz - dx = self.dx - dy = self.dy - dz = self.dz - - # Separate update on edges doesn't seem to be needed as derivatives are taken in one direction per time - # Update E on "lower" faces - if i_block > 0 and blocks_mat[i_block - 1, j_block, k_block] is not None: - for jj in range(Ny): - for kk in range(Nz + 1): - Eyx[0, jj, kk] = Ax[0, jj, kk] * Eyx[0, jj, kk] - Bx[0, jj, kk] / dx * (Hz[0, jj, kk] - - blocks_mat[i_block - 1, j_block, k_block].Hz[-1, jj, kk]) - for jj in range(Ny + 1): - for kk in range(Nz): - Ezx[0, jj, kk] = Ax[0, jj, kk] * Ezx[0, jj, kk] + Bx[0, jj, kk] / dx * (Hy[0, jj, kk] - - blocks_mat[i_block - 1, j_block, k_block].Hy[-1, jj, kk]) - - if j_block > 0 and blocks_mat[i_block, j_block - 1, k_block] is not None: - for ii in range(Nx): - for kk in range(Nz + 1): - Exy[ii, 0, kk] = Ay[ii, 0, kk] * Exy[ii, 0, kk] + By[ii, 0, kk] / dy * (Hz[ii, 0, kk] - - blocks_mat[i_block, j_block - 1, k_block].Hz[ii, -1, kk]) - for ii in range(Nx + 1): - for kk in range(Nz): - Ezy[ii, 0, kk] = Ay[ii, 0, kk] * Ezy[ii, 0, kk] - By[ii, 0, kk] / dy * (Hx[ii, 0, kk] - - blocks_mat[i_block, j_block - 1, k_block].Hx[ii, -1, kk]) - - if k_block > 0 and blocks_mat[i_block, j_block, k_block - 1] is not None: - for ii in range(Nx + 1): - for jj in range(Ny): - Eyz[ii, jj, 0] = Az[ii, jj, 0] * Eyz[ii, jj, 0] + Bz[ii, jj, 0] / dz * (Hx[ii, jj, 0] - - blocks_mat[i_block, j_block, k_block - 1].Hx[ii, jj, - 1]) - for ii in range(Nx): - for jj in range(Ny + 1): - Exz[ii, jj, 0] = Az[ii, jj, 0] * Exz[ii, jj, 0] - Bz[ii, jj, 0] / dz * (Hy[ii, jj, 0] - - blocks_mat[i_block, j_block, k_block - 1].Hy[ii, jj, - 1]) - - # Update E on "upper" faces - if i_block < 2 and blocks_mat[i_block + 1, j_block, k_block] is not None: - for jj in range(Ny): - for kk in range(Nz + 1): - Eyx[Nx, jj, kk] = Ax[Nx, jj, kk] * Eyx[Nx, jj, kk] - Bx[Nx, jj, kk] / dx * ( - blocks_mat[i_block + 1, j_block, k_block].Hz[0, jj, kk] - Hz[-1, jj, kk]) - for jj in range(Ny + 1): - for kk in range(Nz): - Ezx[Nx, jj, kk] = Ax[Nx, jj, kk] * Ezx[Nx, jj, kk] + Bx[Nx, jj, kk] / dx * ( - blocks_mat[i_block + 1, j_block, k_block].Hy[0, jj, kk] - Hy[-1, jj, kk]) - - if j_block < 2 and blocks_mat[i_block, j_block + 1, k_block] is not None: - for ii in range(Nx): - for kk in range(Nz + 1): - Exy[ii, Ny, kk] = Ay[ii, Ny, kk] * Exy[ii, Ny, kk] + By[ii, Ny, kk] / dy * ( - blocks_mat[i_block, j_block + 1, k_block].Hz[ii, 0, kk] - Hz[ii, -1, kk]) - for ii in range(Nx + 1): - for kk in range(Nz): - Ezy[ii, Ny, kk] = Ay[ii, Ny, kk] * Ezy[ii, Ny, kk] - By[ii, Ny, kk] / dy * ( - blocks_mat[i_block, j_block + 1, k_block].Hx[ii, 0, kk] - Hx[ii, -1, kk]) - - if k_block < 2 and blocks_mat[i_block, j_block, k_block + 1] is not None: - for ii in range(Nx + 1): - for jj in range(Ny): - Eyz[ii, jj, Nz] = Az[ii, jj, Nz] * Eyz[ii, jj, Nz] + Bz[ii, jj, Nz] / dz * ( - blocks_mat[i_block, j_block, k_block + 1].Hx[ii, jj, 0] - Hx[ii, jj, -1]) - for ii in range(Nx): - for jj in range(Ny + 1): - Exz[ii, jj, Nz] = Az[ii, jj, Nz] * Exz[ii, jj, Nz] - Bz[ii, jj, Nz] / dz * ( - blocks_mat[i_block, j_block, k_block + 1].Hy[ii, jj, 0] - Hy[ii, jj, -1]) - From 44df15c2a969f2972f8a8a64ddba664deae4d88a Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:54:36 +0100 Subject: [PATCH 10/39] Delete build/lib/wakis/routines.py --- build/lib/wakis/routines.py | 336 ------------------------------------ 1 file changed, 336 deletions(-) delete mode 100644 build/lib/wakis/routines.py diff --git a/build/lib/wakis/routines.py b/build/lib/wakis/routines.py deleted file mode 100644 index 46da44f..0000000 --- a/build/lib/wakis/routines.py +++ /dev/null @@ -1,336 +0,0 @@ -# copyright ################################# # -# This file is part of the wakis Package. # -# Copyright (c) CERN, 2024. # -# ########################################### # - -import numpy as np -import h5py -from tqdm import tqdm -from scipy.constants import c as c_light - -from wakis.sources import Beam - -class RoutinesMixin(): - - def emsolve(self, Nt, source=None, callback=None, - save=False, fields=['E'], components=['Abs'], save_every=1, subdomain=None, - plot=False, plot_every=1, use_etd=False, - plot3d=False, **kwargs): - ''' - Run the simulation and save the selected field components in HDF5 files - for every timestep. Each field will be saved in a separate HDF5 file 'Xy.h5' - where X is the field and y the component. - - Parameters: - ---------- - Nt: int - Number of timesteps to run - source: source object - source object from `sources.py` defining the time-dependednt source. - It should have an update function `source.update(solver, t)` - save: bool - Flag to enable saving the field in HDF5 format - fields: list, default ['E'] - 3D field magnitude ('E', 'H', or 'J') to save - 'Ex', 'Hy', etc., is also accepted and will override - the `components` parameter. - components: list, default ['z'] - Field compoonent ('x', 'y', 'z', 'Abs') to save. It will be overriden - if a component is specified in the`field` parameter - save_every: int, default 1 - Number of timesteps between saves - subdomain: list, default None - Slice [x,y,z] of the domain to be saved - plot: bool, default False - Flag to enable 2D plotting - plot3d: bool, default False - Flag to enable 3D plotting - plot_every: int - Number of timesteps between consecutive plots - **kwargs: - Keyword arguments to be passed to the Plot2D function. - * Default kwargs used for 2D plotting: - {'field':'E', 'component':'z', - 'plane':'ZY', 'pos':0.5, 'title':'Ez', - 'cmap':'rainbow', 'patch_reverse':True, - 'off_screen': True, 'interpolation':'spline36'} - * Default kwargs used for 3D plotting: - {'field':'E', 'component':'z', - 'add_stl':None, 'stl_opacity':0.0, 'stl_colors':'white', - 'title':'Ez', 'cmap':'jet', 'clip_volume':False, 'clip_normal':'-y', - 'field_on_stl':True, 'field_opacity':1.0, - 'off_screen':True, 'zoom':1.0, 'nan_opacity':1.0} - - Raises: - ------- - ImportError: - If the hdf5 dependency cannot be imported - - Dependencies: - ------------- - h5py - ''' - self.Nt = Nt - if source is not None: self.source = source - - if save: - - hfs = {} - for field in fields: - if len(field) == 1: - for component in components: - hfs[field+component] = h5py.File(field+component+'.h5', 'w') - else: - hfs[field] = h5py.File(field+'.h5', 'w') - - for hf in hfs: - hf['x'], hf['y'], hf['z'] = self.x, self.y, self.z - hf['t'] = np.arange(0, Nt*self.dt, save_every*self.dt) - - if subdomain is not None: - xx, yy, zz = subdomain - else: - xx, yy, zz = slice(0,self.Nx), slice(0,self.Ny), slice(0,self.Nz) - - if plot: - plotkw = {'field':'E', 'component':'z', - 'plane':'ZY', 'pos':0.5, 'cmap':'rainbow', - 'patch_reverse':True, 'title':'Ez', - 'off_screen': True, 'interpolation':'spline36'} - plotkw.update(kwargs) - - if plot3d: - plotkw = {'field':'E', 'component':'z', - 'add_stl':None, 'stl_opacity':0.0, 'stl_colors':'white', - 'title':'Ez', 'cmap':'jet', 'clip_volume':False, 'clip_normal':'-y', - 'field_on_stl':True, 'field_opacity':1.0, - 'off_screen':True, 'zoom':1.0, 'nan_opacity':1.0} - - plotkw.update(kwargs) - - # get ABC values - if self.activate_abc: - E_abc_2, H_abc_2 = self.get_abc() - E_abc_1, H_abc_1 = self.get_abc() - - # Time loop - for n in tqdm(range(Nt)): - - if source is not None: - source.update(self, n*self.dt) - - if save and n%save_every == 0: - for field in hfs.keys(): - try: - d = getattr(self, field[0])[xx,yy,zz,field[1:]] - except: - raise(f'Component {field} not valid. Input must have a \ - field ["E", "H", "J"] and a component ["x", "y", "z", "Abs"]') - - # Save timestep in HDF5 - hfs[field]['#'+str(n).zfill(5)] = d - - # Advance - self.one_step() - - # Plot - if plot and n%plot_every == 0: - self.plot2D(n=n, **plotkw) - - if plot3d and n%plot_every == 0: - self.plot3D(n=n, **plotkw) - - # ABC BCs - if self.activate_abc: - self.update_abc(E_abc_2, H_abc_2) # n-2 - E_abc_2, H_abc_2 = E_abc_1, H_abc_1 # n-1 - E_abc_1, H_abc_1 = self.get_abc() # n - - # Callback func(solver, t) - if callback is not None: - callback(self, n*self.dt) - - # End - if save: - for hf in hfs: - hf.close() - - def wakesolve(self, wakelength, - wake=None, - callback=None, - compute_plane='both', - plot=False, plot_from=None, plot_every=1, plot_until=None, - save_J=False, - add_space=None, #for legacy - use_edt=None, #deprecated - **kwargs): - ''' - Run the EM simulation and compute the longitudinal (z) and transverse (x,y) - wake potential WP(s) and impedance Z(s). - - The `Ez` field is saved every timestep in a subdomain (xtest, ytest, z) around - the beam trajectory in HDF5 format file `Ez.h5`. - - The computed results are available as Solver class attributes: - - wake potential: WP (longitudinal), WPx, WPy (transverse) [V/pC] - - impedance: Z (longitudinal), Zx, Zy (transverse) [Ohm] - - beam charge distribution: lambdas (distance) [C/m] lambdaf (spectrum) [C] - - Parameters: - ----------- - wakelength: float - Desired length of the wake in [m] to be computed - - Maximum simulation time in [s] can be computed from the wakelength parameter as: - .. math:: t_{max} = t_{inj} + (wakelength + (z_{max}-z_{min}))/c - wake: Wake obj, default None - `Wake()` object containing the information needed to run - the wake solver calculation. See Wake() docstring for more information. - Can be passed at `Solver()` instantiation as parameter too. - save_J: bool, default False - Flag to enable saving the current J in a diferent HDF5 file 'Jz.h5' - plot: bool, default False - Flag to enable 2D plotting - plot_every: int - Number of timesteps between consecutive plots - **kwargs: - Keyword arguments to be passed to the Plot2D function. - Default kwargs used: - {'plane':'ZY', 'pos':0.5, 'title':'Ez', - 'cmap':'rainbow', 'patch_reverse':True, - 'off_screen': True, 'interpolation':'spline36'} - - Raises: - ------- - AttributeError: - If the Wake object is not provided - ImportError: - If the hdf5 dependency cannot be imported - - Dependencies: - ------------- - h5py - ''' - - if wake is not None: self.wake = wake - if self.wake is None: - raise('Wake solver information not passed to the solver instantiation') - - if add_space is not None: #legacy support - self.wake.skip_cells = add_space - - # plot params defaults - if plot: - plotkw = {'plane':'ZY', 'pos':0.5, 'title':'Ez', - 'cmap':'rainbow', 'patch_reverse':True, - 'off_screen': True, 'interpolation':'spline36'} - plotkw.update(kwargs) - - # integration path (test position) - self.xtest, self.ytest = self.wake.xtest, self.wake.ytest - self.ixt, self.iyt = np.abs(self.x-self.xtest).argmin(), np.abs(self.y-self.ytest).argmin() - if compute_plane.lower() == 'longitudinal': - xx, yy = self.ixt, self.iyt - else: - xx, yy = slice(self.ixt-1, self.ixt+2), slice(self.iyt-1, self.iyt+2) - - # Compute simulation time - self.wake.wakelength = wakelength - self.ti = self.wake.ti - self.v = self.wake.v - if self.use_mpi: #E- should it be zmin, zmax instead? - z = self.Z # use global coords - zz = slice(0, self.NZ) - else: - z = self.z - zz = slice(0, self.Nz) - - tmax = (wakelength + self.ti*self.v + (z.max()-z.min()))/self.v #[s] - Nt = int(tmax/self.dt) - self.tmax, self.Nt = tmax, Nt - - # Add beam source - beam = Beam(q=self.wake.q, - sigmaz=self.wake.sigmaz, - beta=self.wake.beta, - xsource=self.wake.xsource, ysource=self.wake.ysource, - ) - - # hdf5 - self.Ez_file = self.wake.Ez_file - hf = None # needed for MPI - if self.use_mpi: - if self.rank==0: - hf = h5py.File(self.Ez_file, 'w') - hf['x'], hf['y'], hf['z'] = self.x[xx], self.y[yy], z[zz] - hf['t'] = np.arange(0, Nt*self.dt, self.dt) - - if save_J: - hfJ = h5py.File('Jz.h5', 'w') - hfJ['x'], hfJ['y'], hfJ['z'] = self.x[xx], self.y[yy], z[zz] - hfJ['t'] = np.arange(0, Nt*self.dt, self.dt) - else: - hf = h5py.File(self.Ez_file, 'w') - hf['x'], hf['y'], hf['z'] = self.x[xx], self.y[yy], z[zz] - hf['t'] = np.arange(0, Nt*self.dt, self.dt) - - if save_J: - hfJ = h5py.File('Jz.h5', 'w') - hfJ['x'], hfJ['y'], hfJ['z'] = self.x[xx], self.y[yy], z[zz] - hfJ['t'] = np.arange(0, Nt*self.dt, self.dt) - - def save_to_h5(self, hf, field, x, y, z, comp, n): - if self.use_mpi: - _field = self.mpi_gather(field, x, y, z, comp) - if self.rank == 0: - hf['#'+str(n).zfill(5)] = _field - else: - hf['#'+str(n).zfill(5)] = getattr(self, field)[x, y, z, comp] - - if plot_until is None: plot_until = Nt - if plot_from is None: plot_from = int(self.ti/self.dt) - - print('Running electromagnetic time-domain simulation...') - for n in tqdm(range(Nt)): - - # Initial condition - beam.update(self, n*self.dt) - - # Save - save_to_h5(self, hf, 'E', xx, yy, zz, 'z', n) - if save_J: - save_to_h5(self, hfJ, 'J', xx, yy, zz, 'z', n) - - # Advance - self.one_step() - - # Plot - if plot: - if n%plot_every == 0 and nplot_from: - self.plot2D(field='E', component='z', n=n, **plotkw) - else: - pass - - # Callback func(solver, t) - if callback is not None: - callback(self, n*self.dt) - - # End of time loop - if self.use_mpi: - if self.rank==0: - hf.close() - if save_J: - hfJ.close() - - # Compute wakefield magnitudes is done inside WakeSolver - self.wake.solve(compute_plane=compute_plane) - else: - hf.close() - if save_J: - hfJ.close() - - # Compute wakefield magnitudes is done inside WakeSolver - self.wake.solve(compute_plane=compute_plane) - - - \ No newline at end of file From 2b2dbef4b80eff7e4b589531898b9e817431d4f9 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:54:45 +0100 Subject: [PATCH 11/39] Delete build/lib/wakis/plotting.py --- build/lib/wakis/plotting.py | 820 ------------------------------------ 1 file changed, 820 deletions(-) delete mode 100644 build/lib/wakis/plotting.py diff --git a/build/lib/wakis/plotting.py b/build/lib/wakis/plotting.py deleted file mode 100644 index cdf2ef6..0000000 --- a/build/lib/wakis/plotting.py +++ /dev/null @@ -1,820 +0,0 @@ -# copyright ################################# # -# This file is part of the wakis Package. # -# Copyright (c) CERN, 2024. # -# ########################################### # - -import numpy as np -import matplotlib.pyplot as plt - -class PlotMixin: - - def plot3D(self, field='E', component='z', clim=None, hide_solids=None, - show_solids=None, add_stl=None, stl_opacity=0.1, stl_colors='white', - title=None, cmap='jet', - clip_interactive=False, clip_normal='-y', - clip_box=False, clip_bounds=None, - off_screen=False, zoom=0.5, camera_position=None, - nan_opacity=1.0, n=None): - ''' - Built-in 3D plotting using PyVista - - Parameters: - ----------- - field: str, default 'E' - 3D field magnitude ('E', 'H', or 'J') to plot - To plot a component 'Ex', 'Hy' is also accepted - component: str, default 'z' - 3D field compoonent ('x', 'y', 'z', 'Abs') to plot. It will be overriden - if a component is defined in field - clim: list, optional - Colorbar limits for the field plot [min, max] - hide_solids: bool, optional - Mask the values inside solid to np.nan. NaNs will be shown in gray, - since there is a bug with the nan_opacity parameter - show_solids: bool, optional - Mask the values outside solid to np.nan. - add_stl: str or list, optional - List or str of stl solids to add to the plot by `pv.add_mesh` - stl_opacity: float, default 0.1 - Opacity of the stl surfaces (0 - Transparent, 1 - Opaque) - stl_colors: str or list of str, default 'white' - Color of the stl surfaces - title: str, optional - Title used to save the screenshot of the 3D plot (Path+Name) if off_screen=True - cmap: str, default 'jet' - Colormap name to use in the field display - clip_interactive: bool, default False - Enable an interactive widget to clip out part of the domain, plane normal is defined by - `clip_normal` parameter - clip_normal: str, default '-y' - Normal direction of the clip_volume interactive plane - clip_box: bool, default False - Enable a static box clipping of the domain. The box bounds are defined by `clip_bounds` parameter - clip_bounds: Default None - List of bounds [xmin, xmax, ymin, ymax, zmin, zmax] of the box to clip if clip_box is active. - field_on_stl : bool, default False - Samples the field on the stl file specified in `add_stl`. - field_opacity : optional, default 1.0 - Sets de opacity of the `field_on_stl` plot - off_screen: bool, default False - Enable plot rendering off screen, for gif frames generation. - Plot will not be rendered if set to True. - n: int, optional - Timestep number to be added to the plot title and figsave title. - ''' - if self.use_mpi: - print('*** plot3D is not supported when `use_mpi=True`') - return - - import pyvista as pv - - if len(field) == 2: #support for e.g. field='Ex' - component = field[1] - field = field[0] - - if title is None: - title = field + component +'3d' - - if self.plotter_active and not off_screen: - self.plotter_active = False - - if not self.plotter_active: - - pl = pv.Plotter(off_screen=off_screen) - - # Plot stl surface(s) - if add_stl is not None: - if type(add_stl) is str: - key = add_stl - surf = self.grid.read_stl(key) - pl.add_mesh(surf, color=stl_colors, opacity=stl_opacity, smooth_shading=True) - - elif type(add_stl) is list: - for i, key in enumerate(add_stl): - surf = self.grid.read_stl(key) - if type(stl_colors) is list: - pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, smooth_shading=True) - else: - pl.add_mesh(surf, color=stl_colors, opacity=stl_opacity, smooth_shading=True) - else: - key = self.grid.stl_solids.keys()[0] - surf = self.grid.read_stl(key) - pl.add_mesh(surf, color=stl_colors, opacity=stl_opacity, smooth_shading=True) - - if camera_position is None: - pl.camera_position = 'zx' - pl.camera.azimuth += 30 - pl.camera.elevation += 30 - else: - pl.camera_position = camera_position - - pl.set_background('mistyrose', top='white') - try: pl.add_logo_widget('docs/img/wakis-logo-pink.png') - except: pass - pl.camera.zoom(zoom) - pl.add_axes() - pl.enable_3_lights() - - if off_screen: - self.plotter_active = True - else: - pl = self.pl - - - # Plot field - if field == 'E': - self.grid.grid.cell_data[field+component] = np.reshape(self.E[:, :, :, component], self.N) - elif field == 'H': - self.grid.grid.cell_data[field+component] = np.reshape(self.H[:, :, :, component], self.N) - elif field == 'J': - self.grid.grid.cell_data[field+component] = np.reshape(self.J[:, :, :, component], self.N) - else: - print("`field` value not valid") - - points = self.grid.grid.cell_data_to_point_data() #interpolate - - # Mask the values inside solid to np.nan - if hide_solids is not None: - tol = np.min([self.dx, self.dy, self.dz])*1e-3 - if type(hide_solids) is str: - surf = self.grid.read_stl(hide_solids) - select = self.grid.grid.select_enclosed_points(surf, tolerance=tol) - mask = select['SelectedPoints'] > 0 - - elif type(hide_solids) is list: - for i, solid in enumerate(hide_solids): - surf = self.grid.read_stl(solid) - select = self.grid.grid.select_enclosed_points(surf, tolerance=tol) - if i == 0: - mask = select['SelectedPoints'] > 0 - else: - mask += select['SelectedPoints'] > 0 - - points[field+component][mask] = np.nan - - # Mask the values outside solid to np.nan - if show_solids is not None: - tol = np.min([self.dx, self.dy, self.dz])*1e-3 - if type(show_solids) is str: - surf = self.grid.read_stl(show_solids) - select = self.grid.grid.select_enclosed_points(surf, tolerance=tol) - mask = select['SelectedPoints'] > 0 - - elif type(show_solids) is list: - for solid in show_solids: - surf = self.grid.read_stl(solid) - select = self.grid.grid.select_enclosed_points(surf, tolerance=tol) - if i == 0: - mask = select['SelectedPoints'] > 0 - else: - mask += select['SelectedPoints'] > 0 - - points[field+component][np.logical_not(mask)] = np.nan - - # Clip a rectangle of the domain - if clip_box: - if clip_bounds is None: - Lx, Ly = (self.grid.xmax-self.grid.xmin), (self.grid.ymax-self.grid.ymin) - clip_bounds = [self.grid.xmax-Lx/2, self.grid.xmax, - self.grid.ymax-Ly/2, self.grid.ymax, - self.grid.zmin, self.grid.zmax] - - ac1 = pl.add_mesh(points.clip_box(bounds=clip_bounds), opacity=nan_opacity, - scalars=field+component, cmap=cmap, clim=clim) - - # Enable an interactive widget to clip out part of the domain with a plane, with clip_normal - elif clip_interactive: - ac1 = pl.add_mesh_clip_plane(points, normal=clip_normal, opacity=1.0, - scalars=field+component, cmap=cmap, clim=clim, - normal_rotation=False, nan_opacity=nan_opacity) - else: - print('Plotting option inconsistent') - - # Save - if n is not None: - pl.add_title(field+component+f' field, timestep={n}', font='times', font_size=12) - title += '_'+str(n).zfill(6) - if off_screen: - pl.screenshot(title+'.png') - pl.remove_actor(ac1) - self.pl = pl - else: - pl.show(full_screen=False) - - def plot3DonSTL(self, field='E', component='z', clim=None, cmap='jet', log_scale=False, - stl_with_field=None, field_opacity=1.0, tolerance=None, - stl_transparent=None, stl_opacity=0.1, stl_colors='white', - clip_plane = False, clip_interactive=False, - clip_normal='-x', clip_origin=[0,0,0], - clip_box=False, clip_bounds=None, - title=None, off_screen=False, n=None, - zoom=0.5, camera_position=None, **kwargs): - ''' - Built-in 3D plotting using PyVista - - Parameters: - ----------- - field: str, default 'E' - 3D field magnitude ('E', 'H', or 'J') to plot - To plot a component 'Ex', 'Hy' is also accepted - component: str, default 'z' - 3D field compoonent ('x', 'y', 'z', 'Abs') to plot. It will be overriden - if a component is defined in field - clim: list, optional - Colorbar limits for the field plot [min, max] - cmap: str or cmap obj, default 'jet' - Colormap to use for the field plot - log_scale: bool, default False - Turns on logarithmic scale colorbar - stl_with_field : list or str - STL str name or list of names to samples the selected field on - field_opacity : optional, default 1.0 - Sets de opacity of the `field_on_stl` plot - tolerance : float, default None - Tolerance to apply to PyVista's sampling algorithm - stl_transparent: list or str, default None - STL name or list of names to add to the scene with the selected transparency and color - stl_opacity: float, default 0.1 - Opacity of the STL solids without field - stl_colors: list or str, default 'white' - str or list of colors to use for each STL solid - clip_interactive: bool, default False - Enable an interactive widget to clip out part of the domain, plane normal is defined by - `clip_normal` parameter - clip_plane: bool, default False - Clip stl_with_field surface with a plane and show field on such plane - clip_normal: str, default '-y' - Normal direction of the clip_volume interactive plane and the clip_plane - clip_origin: list, default [0,0,0] - Origin of the clipping plane for the clip_plane option - clip_box: bool, default False - Enable a static box clipping of the domain. The box bounds are defined by `clip_bounds` parameter - clip_bounds: Default None - List of bounds [xmin, xmax, ymin, ymax, zmin, zmax] of the box to clip if clip_box is active. - off_screen: bool, default False - Enable plot rendering off screen, for gif frames generation. - Plot will not be rendered if set to True. - title: str - Name to use in the .png savefile with off_screen is True - n: int, optional - Timestep number to be added to the plot title and figsave title. - **kwargs: optional - PyVista's add_mesh optional arguments: - https://docs.pyvista.org/api/plotting/_autosummary/pyvista.plotter.add_mesh - ''' - if self.use_mpi: - print('*** plot3D is not supported when `use_mpi=True`') - return - - import pyvista as pv - - if len(field) == 2: #support for e.g. field='Ex' - component = field[1] - field = field[0] - - if title is None: - title = field + component +'3d' - - if self.plotter_active and not off_screen: - self.plotter_active = False - - if not self.plotter_active: - pl = pv.Plotter(off_screen=off_screen, lighting='none') - light = pv.Light(light_type='headlight') - pl.add_light(light) - - # Plot stl surface(s) - if stl_transparent is not None: - if type(stl_transparent) is str: - key = stl_transparent - surf = self.grid.read_stl(key) - pl.add_mesh(surf, color=stl_colors, opacity=stl_opacity, smooth_shading=True) - - elif type(stl_transparent) is list: - for i, key in enumerate(stl_transparent): - surf = self.grid.read_stl(key) - if type(stl_colors) is list: - pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, smooth_shading=True) - else: - pl.add_mesh(surf, color=stl_colors, opacity=stl_opacity, smooth_shading=True) - - if off_screen: - self.plotter_active = True - else: - pl = self.pl - - # Plot field - if field == 'E': - self.grid.grid.cell_data[field+component] = np.reshape(self.E[:, :, :, component], self.N) - - elif field == 'H': - self.grid.grid.cell_data[field+component] = np.reshape(self.H[:, :, :, component], self.N) - - elif field == 'J': - self.grid.grid.cell_data[field+component] = np.reshape(self.J[:, :, :, component], self.N) - else: - print("`field` value not valid") - - points = self.grid.grid.cell_data_to_point_data() #interpolate - - # Interpolate fields on stl - if stl_with_field is not None: - if type(stl_with_field) is str: - key = stl_with_field - surf = self.grid.read_stl(key) - if clip_plane: - try: surf = surf.clip_closed_surface(normal=clip_normal, origin=clip_origin).subdivide_adaptive(max_edge_len=3*self.dz) - except: print("Surface non-manifold, clip with plane skipped") - - fieldonsurf = surf.sample(points, tolerance) - - if clip_interactive: # interactive plotting with a plane - ac1 = pl.add_mesh_clip_plane(fieldonsurf, normal=clip_normal, normal_rotation=False, - scalars=field+component, opacity=field_opacity, - cmap=cmap, clim=clim, log_scale=log_scale, - **kwargs) - - elif clip_box: # Clip a rectangle of the domain - if clip_bounds is None: - Lx, Ly = (self.grid.xmax-self.grid.xmin), (self.grid.ymax-self.grid.ymin) - clip_bounds = [self.grid.xmax-Lx/2, self.grid.xmax, - self.grid.ymax-Ly/2, self.grid.ymax, - self.grid.zmin, self.grid.zmax] - - ac1 = pl.add_mesh(fieldonsurf.clip_box(bounds=clip_bounds), cmap=cmap, clim=clim, - scalars=field+component, opacity=field_opacity, - log_scale=log_scale, - **kwargs) - - else: - ac1 = pl.add_mesh(fieldonsurf, cmap=cmap, clim=clim, - scalars=field+component, opacity=field_opacity, - log_scale=log_scale, - **kwargs) - - elif type(stl_with_field) is list: - for i, key in enumerate(stl_with_field): - surf = self.grid.read_stl(key) - if clip_plane: - try: surf = surf.clip_closed_surface(normal=clip_normal, origin=clip_origin) - except: print("Surface non-manifold, clip with plane skipped") - - fieldonsurf = surf.sample(points) - - if clip_interactive: # interactive plotting with a plane - ac1 = pl.add_mesh_clip_plane(fieldonsurf, normal=clip_normal, normal_rotation=False, - scalars=field+component, opacity=field_opacity, - cmap=cmap, clim=clim, log_scale=log_scale, - **kwargs) - elif clip_box: # Clip a rectangle of the domain - if clip_bounds is None: - Lx, Ly = (self.grid.xmax-self.grid.xmin), (self.grid.ymax-self.grid.ymin) - clip_bounds = [self.grid.xmax-Lx/2, self.grid.xmax, - self.grid.ymax-Ly/2, self.grid.ymax, - self.grid.zmin, self.grid.zmax] - ac1 = pl.add_mesh(fieldonsurf.clip_box(bounds=clip_bounds), cmap=cmap, clim=clim, - scalars=field+component, opacity=field_opacity, - log_scale=log_scale, **kwargs) - else: - ac1 = pl.add_mesh(fieldonsurf, cmap=cmap, clim=clim, - scalars=field+component, opacity=field_opacity, - log_scale=log_scale, - **kwargs) - if camera_position is None: - pl.camera_position = 'zx' - pl.camera.azimuth += 30 - pl.camera.elevation += 30 - else: - pl.camera_position = camera_position - - pl.set_background('mistyrose', top='white') - try: pl.add_logo_widget('docs/img/wakis-logo-pink.png') - except: pass - pl.camera.zoom(zoom) - pl.add_axes() - pl.enable_anti_aliasing() - #pl.enable_3_lights() - - if n is not None: - pl.add_title(field+component+f' field, timestep={n}', font='times', font_size=12) - title += '_'+str(n).zfill(6) - - # Save - if off_screen: - pl.screenshot(title+'.png') - pl.remove_actor(ac1) - self.pl = pl - else: - pl.show(full_screen=False) - - def plot2D(self, field='E', component='z', plane='ZY', pos=0.5, norm=None, - vmin=None, vmax=None, figsize=[8,4], cmap='jet', patch_alpha=0.1, - patch_reverse=False, add_patch=False, title=None, off_screen=False, - n=None, interpolation='antialiased', dpi=100, return_handles=False): - ''' - Built-in 2D plotting of a field slice using matplotlib - - Parameters: - ---------- - field: str, default 'E' - Field magnitude ('E', 'H', or 'J') to plot - To plot a component 'Ex', 'Hy' is also accepted - component: str, default 'z' - Field compoonent ('x', 'y', 'z', 'Abs') to plot. It will be overriden - if a component is defined in field - plane: arr or str, default 'XZ' - Plane where to plot the 2d field cut: array of 2 slices() and 1 int [x,y,z] - or a str 'XY', 'ZY' or 'ZX' - pos: float, default 0.5 - Position of the cutting plane, as a franction of the plane's normal dimension - e.g. plane 'XZ' wil be sitting at y=pos*(ymax-ymin) - norm: str, default None - Plotting scale to pass to matplotlib imshow: 'linear', 'log', 'symlog' - ** Only for matplotlib version >= 3.8 - vmin: list, optional - Colorbar min limit for the field plot - vmax: list, optional - Colorbar max limit for the field plot - figsize: list, default [8,4] - Figure size to pass to the plot initialization - add_patch: str or list, optional - List or str of stl solids to add to the plot by `pv.add_mesh` - patch_alpha: float, default 0.1 - Value for the transparency of the patch if `add_patch = True` - title: str, optional - Title used to save the screenshot of the 3D plot (Path+Name) if off_screen=True. - If n is provided, 'str(n).zfill(6)' will be added to the title. - cmap: str, default 'jet' - Colormap name to use in the field display - off_screen: bool, default False - Enable plot rendering off screen, for gif frames generation. - Plot will not be rendered if set to True. - n: int, optional - Timestep number to be added to the plot title and figsave title. - interpolation: str, default 'antialiased' - Interpolation method to pass to matplotlib imshow e.g., 'none', - 'antialiased', 'nearest', 'bilinear', 'bicubic', 'spline16', 'spline36', - ''' - from mpl_toolkits.axes_grid1 import make_axes_locatable - - Nx, Ny, Nz = self.Nx, self.Ny, self.Nz - xmin, xmax = self.grid.xmin, self.grid.xmax - ymin, ymax = self.grid.ymin, self.grid.ymax - zmin, zmax = self.grid.zmin, self.grid.zmax - _z = self.z - - if self.use_mpi: - zmin, zmax = self.grid.ZMIN, self.grid.ZMAX - Nz = self.grid.NZ - _z = self.grid.Z - - if type(field) is str: - if len(field) == 2: #support for e.g. field='Ex' - component = field[1] - field = field[0] - elif len(field) == 4: #support for e.g. field='EAbs' - component = field[1:] - field = field[0] - - if title is None: - title = field + component +'2d' - - if type(plane) is not str and len(plane) == 3: - x, y, z = plane[0], plane[1], plane[2] - - if type(plane[2]) is int: - cut = f'(x,y,a) a={round(self.z[z],3)}' - xax, yax = 'y', 'x' - extent = [self.y[y].min(), self.y[y].max(), - self.x[x].min(), self.x[x].max()] - - if type(plane[0]) is int: - cut = f'(a,y,z) a={round(self.x[x],3)}' - xax, yax = 'z', 'y' - extent = [_z[z].min(), _z[z].max(), - self.y[y].min(), self.y[y].max()] - - if type(plane[1]) is int: - cut = f'(x,a,z) a={round(self.y[y],3)}' - xax, yax = 'z', 'x' - extent = [_z[z].min(), _z[z].max(), - self.x[x].min(), self.x[x].max()] - - elif plane == 'XY': - x, y, z = slice(0,Nx), slice(0,Ny), int(Nz*pos) #plane XY - cut = f'(x,y,a) a={round(pos*(zmax-zmin)+zmin,3)}' - xax, yax = 'y', 'x' - extent = [ymin, ymax, xmin, xmax] - - elif plane == 'ZY' or plane == 'YZ': - x, y, z = int(Nx*pos), slice(0,Ny), slice(0,Nz) #plane ZY - cut = f'(a,y,z) a={round(pos*(xmax-xmin)+xmin,3)}' - xax, yax = 'z', 'y' - extent = [zmin, zmax, ymin, ymax] - - elif plane == 'ZX' or plane == 'XZ': - x, y, z = slice(0,Nx), int(Ny*pos), slice(0,Nz) #plane XZ - cut = f'(x,a,z) a={round(pos*(ymax-ymin)+ymin,3)}' - xax, yax = 'z', 'x' - extent = [zmin, zmax, xmin, xmax] - - else: - print("Plane needs to be an array of slices [x,y,z] or a str 'XY', 'ZY', 'ZX'") - - if self.use_mpi: # only in rank=0 - _field = self.mpi_gather(field, x=x, y=y, z=z, component=component) - - if self.rank == 0: - fig, ax = plt.subplots(1,1, figsize=figsize, dpi=dpi) - im = ax.imshow(_field, cmap=cmap, norm=norm, - extent=extent, origin='lower', vmin=vmin, vmax=vmax, - interpolation=interpolation) - - fig.colorbar(im, cax=make_axes_locatable(ax).append_axes('right', size='5%', pad=0.05)) - ax.set_title(f'Wakis {field}{component}{cut}') - ax.set_xlabel(xax) - ax.set_ylabel(yax) - - if n is not None: - fig.suptitle('$'+str(field)+'_{'+str(component)+'}$ field, timestep='+str(n)) - title += '_'+str(n).zfill(6) - - fig.tight_layout() - - if off_screen: - fig.savefig(title+'.png') - plt.clf() - plt.close(fig) - elif return_handles: - return fig, ax - else: - plt.show(block=False) - else: - fig, ax = plt.subplots(1,1, figsize=figsize, dpi=dpi) - if field == 'E': - _field = self.E[x, y, z, component] - if field == 'H': - _field = self.H[x, y, z, component] - if field == 'J': - _field = self.J[x, y, z, component] - - im = ax.imshow(_field, cmap=cmap, norm=norm, - extent=extent, origin='lower', vmin=vmin, vmax=vmax, - interpolation=interpolation) - - fig.colorbar(im, cax=make_axes_locatable(ax).append_axes('right', size='5%', pad=0.05)) - ax.set_title(f'Wakis {field}{component}{cut}') - ax.set_xlabel(xax) - ax.set_ylabel(yax) - - # Patch stl - not supported when running MPI - if add_patch is not None: - if type(add_patch) is str: - mask = np.reshape(self.grid.grid[add_patch], (Nx, Ny, Nz)) - patch = np.ones((Nx, Ny, Nz)) - if patch_reverse: - patch[mask] = np.nan - else: - patch[np.logical_not(mask)] = np.nan - ax.imshow(patch[x,y,z], cmap='Greys', extent=extent, origin='lower', alpha=patch_alpha) - - elif type(add_patch) is list: - for solid in add_patch: - mask = np.reshape(self.grid.grid[solid], (Nx, Ny, Nz)) - patch = np.ones((Nx, Ny, Nz)) - if patch_reverse: - patch[mask] = np.nan - else: - patch[np.logical_not(mask)] = np.nan - ax.imshow(patch[x,y,z], cmap='Greys', extent=extent, origin='lower', alpha=patch_alpha) - - if n is not None: - fig.suptitle('$'+str(field)+'_{'+str(component)+'}$ field, timestep='+str(n)) - title += '_'+str(n).zfill(6) - - fig.tight_layout() - - if off_screen: - fig.savefig(title+'.png') - plt.clf() - plt.close(fig) - elif return_handles: - return fig, ax - else: - plt.show(block=False) - - def plot1D(self, field='E', component='z', line='z', pos=[0.5], - xscale='linear', yscale='linear', xlim=None, ylim=None, - figsize=[8,4], title=None, off_screen=False, n=None, - colors=None, **kwargs): - ''' - Built-in 1D plotting of a field line using matplotlib - - Parameters: - ---------- - field: str, default 'E' - Field magnitude ('E', 'H', or 'J') to plot - To plot a component 'Ex', 'Hy' is also accepted - component: str, default 'z' - Field compoonent ('x', 'y', 'z', 'Abs') to plot. It will be overriden - if a component is defined in field - line: str or list, default 'z' - line of indexes to plot. E.g. line=[0, slice(10,Ny-10), 0] - pos: float or list, default 0.5 - Float or list of floats betwwen 0-1 indicating the cut position. - Only used if line is str. - xlim, ylim: tupple - limits for x and y axis (see matplotlib.ax.set_xlim for more) - xscale, yscale: str - scale to use in x and y axes (see matplotlib.ax.set_xscale for more) - figsize: list, default [8,4] - Figure size to pass to the plot initialization - title: str, optional - Title used to save the screenshot of the 3D plot (Path+Name) if off_screen=True. - If n is provided, 'str(n).zfill(6)' will be added to the title. - cmap: str, default 'jet' - Colormap name to use in the field display - off_screen: bool, default False - Enable plot rendering off screen, for gif frames generation. - Plot will not be rendered if set to True. - n: int, optional - Timestep number to be added to the plot title and figsave title. - colors: list, optional - List of matplotlib-compatible colors. len(colors) >= len(pos) - **kwargs: - Keyword arguments to be passed to the `matplotlib.plot` function. - Default kwargs used: - kwargs = {'color':'g', 'lw':1.2, 'ls':'-'} - ''' - Nx, Ny, Nz = self.Nx, self.Ny, self.Nz - xmin, xmax = self.grid.xmin, self.grid.xmax - ymin, ymax = self.grid.ymin, self.grid.ymax - zmin, zmax = self.grid.zmin, self.grid.zmax - _z = self.z - - if self.use_mpi: - zmin, zmax = self.grid.ZMIN, self.grid.ZMAX - Nz = self.grid.NZ - _z = self.grid.Z - if self.rank == 0: - fig, ax = plt.subplots(1,1, figsize=figsize) - else: - fig, ax = plt.subplots(1,1, figsize=figsize) - - - plotkw = {'lw':1.2, 'ls':'-'} - - if colors is None: - colors = ['k', 'tab:red', 'tab:blue', 'tab:green', - 'tab:orange', 'tab:purple', 'tab:pink'] - plotkw.update(kwargs) - - if len(field) == 2: #support for e.g. field='Ex' - component = field[1] - field = field[0] - - if title is None: - title = field + component +'1d' - - if type(pos) is not list: #support for a list of cut positions - pos_arr = [pos] - else: - pos_arr = pos - - for i, pos in enumerate(pos_arr): - - if type(line) is not str and len(line) == 3: - x, y, z = line[0], line[1], line[2] - - #z-axis - if type(line[2]) is slice: - cut = f'(a,b,z) a={round(self.x[x],3)}, b={round(self.y[y],3)}' - xax = 'z' - xx = _z[z] - xlims = (_z[z].min(), _z[z].max()) - - #x-axis - elif type(line[0]) is slice: - cut = f'(x,a,b) a={round(self.y[y],3)}, b={round(_z[z],3)}' - xax = 'x' - xx = self.x[x] - xlims = (self.x[x].min(), self.x[x].max()) - - #y-axis - elif type(line[1]) is slice: - cut = f'(a,y,b) a={round(self.x[x],3)}, b={round(_z[z],3)}' - xax = 'y' - xx = self.y[y] - xlims = (self.y[y].min(), self.y[y].max()) - - elif line.lower() == 'x': - x, y, z = slice(0,Nx), int(Ny*pos), int(Nz*pos) #x-axis - cut = f'(x,a,b) a={round(self.y[y],3)}, b={round(_z[z],3)}' - xax = 'x' - xx = self.x[x] - xlims = (xmin, xmax) - - elif line.lower() == 'y': - x, y, z = int(Nx*pos), slice(0,Ny), int(Nz*pos) #y-axis - cut = f'(a,y,b) a={round(self.x[x],3)}, b={round(_z[z],3)}' - xax = 'y' - xx = self.y[y] - xlims = (ymin, ymax) - - elif line.lower() == 'z': - x, y, z = int(Nx*pos), int(Ny*pos), slice(0,Nz) #z-axis - cut = f'(a,b,z) a={round(self.x[x],3)}, b={round(self.y[y],3)}' - xax = 'z' - xx = _z[z] - xlims = (zmin, zmax) - - else: - print("line needs to be an array of slices [x,y,z] or a str 'x', 'y', 'z'") - - if i == 0: # first one on top - zorder = 10 - else: zorder = i - - if self.use_mpi: # only in rank=0 - _field = self.mpi_gather(field, x=x, y=y, z=z, component=component) - - if self.rank == 0: - ax.plot(xx, _field, color=colors[i], zorder=zorder, - label=f'{field}{component}{cut}', **plotkw) - else: - if field == 'E': - _field = self.E[x, y, z, component] - if field == 'H': - _field = self.H[x, y, z, component] - if field == 'J': - _field = self.J[x, y, z, component] - - ax.plot(xx, _field, color=colors[i], zorder=zorder, - label=f'{field}{component}{cut}', **plotkw) - - if not self.use_mpi: - yax = f'{field}{component} amplitude' - - ax.set_title(f'Wakis {field}{component}'+(len(pos_arr)==1)*f'{cut}') - ax.set_xlabel(xax) - ax.set_ylabel(yax, color=colors[0]) - ax.set_xlim(xlims) - - ax.set_xscale(xscale) - ax.set_yscale(yscale) - - if len(pos_arr) > 1: - ax.legend(loc=1) - - if xlim is not None: - ax.set_xlim(xlim) - if ylim is not None: - ax.set_ylim(ylim) - - if n is not None: - fig.suptitle('$'+field+'_{'+component+'}$ field, timestep='+str(n)) - title += '_'+str(n).zfill(6) - - fig.tight_layout() - - if off_screen: - fig.savefig(title+'.png') - plt.clf() - plt.close(fig) - - else: - plt.show() - - elif self.use_mpi and self.rank == 0: - yax = f'{field}{component} amplitude' - - ax.set_title(f'Wakis {field}{component}'+(len(pos_arr)==1)*f'{cut}') - ax.set_xlabel(xax) - ax.set_ylabel(yax, color=colors[0]) - ax.set_xlim(xlims) - - ax.set_xscale(xscale) - ax.set_yscale(yscale) - - if len(pos_arr) > 1: - ax.legend(loc=1) - - if xlim is not None: - ax.set_xlim(xlim) - if ylim is not None: - ax.set_ylim(ylim) - - if n is not None: - fig.suptitle('$'+field+'_{'+component+'}$ field, timestep='+str(n)) - title += '_'+str(n).zfill(6) - - fig.tight_layout() - - if off_screen: - fig.savefig(title+'.png') - plt.clf() - plt.close(fig) - - else: - plt.show() - From 8d1de64f264a2c9d891a428d252f013c1a57fa69 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:54:54 +0100 Subject: [PATCH 12/39] Delete build/lib/wakis/geometry.py --- build/lib/wakis/geometry.py | 254 ------------------------------------ 1 file changed, 254 deletions(-) delete mode 100644 build/lib/wakis/geometry.py diff --git a/build/lib/wakis/geometry.py b/build/lib/wakis/geometry.py deleted file mode 100644 index 3882458..0000000 --- a/build/lib/wakis/geometry.py +++ /dev/null @@ -1,254 +0,0 @@ -# copyright ################################# # -# This file is part of the wakis Package. # -# Copyright (c) CERN, 2024. # -# ########################################### # - -import numpy as np -import re - -def extract_colors_from_stp(stp_file): - """ - Extracts a mapping from solid names to RGB color values from a STEP (.stp) file. - - Args: - stp_file (str): Path to the STEP file. - - Returns: - dict[str, list[float]]: A dictionary mapping solid names to [R, G, B] colors. - """ - solids, _ = extract_names_from_stp(stp_file) - - colors = [] - stl_colors = {} - - color_pattern = re.compile(r"#\d+=COLOUR_RGB\('?',([\d.]+),([\d.]+),([\d.]+)\);") - - # Extract colors - with open(stp_file, 'r', encoding='utf-8', errors='ignore') as f: - for line in f: - color_match = color_pattern.search(line) - if color_match: - r = float(color_match.group(1)) - g = float(color_match.group(2)) - b = float(color_match.group(3)) - colors.append([r, g, b]) - - # Map solids to colors by order of appearance (colors >=solids) - for i in range(len(list(solids.keys()))): - solid = solids[list(solids.keys())[i]] - solid_re = re.sub(r'[^a-zA-Z0-9_-]', '-', solid) - stl_colors[f'{str(i).zfill(3)}_{solid_re}'] = colors[i] - - return stl_colors - -def extract_materials_from_stp(stp_file): - """ - Extracts a mapping from solid names to materials from a STEP (.stp) file. - - Args: - stp_file (str): Path to the STEP file. - - Returns: - dict[str, str]: A dictionary mapping solid names to material names. - """ - - solids, materials = extract_names_from_stp(stp_file) - stl_materials = {} - for i in range(len(list(solids.keys()))): - solid = solids[list(solids.keys())[i]] - try: - mat = materials[list(solids.keys())[i]].lower() - except: - print(f'Solid #{list(solids.keys())[i]} has no assigned material') - mat = 'None' - - # Remove problematic characters - solid_re = re.sub(r'[^a-zA-Z0-9_-]', '-', solid) - mat_re = re.sub(r'[^a-zA-Z0-9_-]', '-', mat) - stl_materials[f'{str(i).zfill(3)}_{solid_re}'] = mat_re - - return stl_materials - -def extract_solids_from_stp(stp_file, path=''): - if path and not path.endswith('/'): - path += '/' - solids, materials = extract_names_from_stp(stp_file) - stl_solids = {} - for i in range(len(list(solids.keys()))): - solid = solids[list(solids.keys())[i]] - try: - mat = materials[list(solids.keys())[i]] - except: - print(f'Solid #{list(solids.keys())[i]} has no assigned material') - mat = 'None' - - # Remove problematic characters - solid_re = re.sub(r'[^a-zA-Z0-9_-]', '-', solid) - mat_re = re.sub(r'[^a-zA-Z0-9_-]', '-', mat) - name = f'{str(i).zfill(3)}_{solid_re}_{mat_re}' - stl_solids[f'{str(i).zfill(3)}_{solid_re}'] = path + name + '.stl' - - return stl_solids - -def extract_names_from_stp(stp_file): - """ - Extracts solid names and their corresponding materials from a STEP (.stp) file. - - This function parses a given STEP file to identify solid objects and their assigned materials. - The solid names are extracted from `MANIFOLD_SOLID_BREP` statements, while the materials are - linked via `PRESENTATION_LAYER_ASSIGNMENT` statements. - - Args: - stp_file (str): Path to the STEP (.stp) file. - - Returns: - tuple[dict[int, str], dict[int, str]]: - - A dictionary mapping solid IDs to their names. - - A dictionary mapping solid IDs to their corresponding material names. - - Example: - >>> solids, materials = extract_names_from_stp("example.stp") - >>> print(solids) - {37: "Vacuum|Half_cell_dx", 39: "Be window left"} - >>> print(materials) - {37: "Vacuum", 39: "Berillium"} - """ - solid_dict = {} - material_dict = {} - - # Compile regex patterns - #solid_pattern = re.compile(r"#(\d+)=MANIFOLD_SOLID_BREP\('([^']+)'.*;") - solid_pattern = re.compile(r"#(\d+)=ADVANCED_BREP_SHAPE_REPRESENTATION\('([^']*)',\(([^)]*)\),#\d+\);") - material_pattern = re.compile(r"#(\d+)=PRESENTATION_LAYER_ASSIGNMENT\('([^']+)','[^']+',\(#([\d#,]+)\)\);") - - # First pass: extract solids - with open(stp_file, 'r', encoding='utf-8', errors='ignore') as f: - for line in f: - solid_match = solid_pattern.search(line) - if solid_match: - #solid_number = int(solid_match.group(1)) #if MANIFOLD - solid_number = int(solid_match.group(3).split(',')[0].strip().lstrip('#')) - solid_name = solid_match.group(2) - solid_dict[solid_number] = solid_name - - # Second pass: extract materials - with open(stp_file, 'r', encoding='utf-8', errors='ignore') as f: - for line in f: - material_match = material_pattern.search(line) - if material_match: - material_name = material_match.group(2) - solid_numbers = [int(num.strip("#")) for num in material_match.group(3).split(',')] - for solid_number in solid_numbers: - if solid_number in solid_dict: - material_dict[solid_number] = material_name - - return solid_dict, material_dict - -def get_stp_unit_scale(stp_file): - """ - Reads the unit definition from a STEP (.stp or .step) file and determines the - scale factor required to convert the geometry to meters. - - This function: - - Opens and scans the header section of the STEP file. - - Detects the SI base unit definition (e.g., millimeter, centimeter, meter). - - Returns a scale factor to convert the geometry to meters. - - Handles missing or unreadable unit information gracefully. - - Args: - stp_file (str): Path to the STEP (.stp or .step) file. - - Returns: - float: Scale factor to convert STEP geometry to meters. - Defaults to 1.0 if no valid unit information is found. - """ - - unit_map = { - ".MILLI.": 1e-3, - ".CENTI.": 1e-2, - ".DECI.": 1e-1, - ".KILO.": 1e3, - "$": 1.0, # '$' indicates no prefix, i.e. plain meters - } - - try: - with open(stp_file, "r", encoding="utf-8", errors="ignore") as f: - header = f.read(10000) # read only the beginning of the file - - match = re.search( - r"SI_UNIT\s*\(\s*(\.\w+\.)?\s*,\s*\.METRE\.\s*\)", - header, - re.IGNORECASE, - ) - - if match: - prefix = match.group(1).upper() if match.group(1) else "$" - scale = unit_map.get(prefix, 1.0) - print(f"Detected STEP unit: {prefix} → scale to meters: {scale}") - return scale - else: - print("No unit found, files remain in original unit.") - return 1.0 - - except Exception as exc: - print(f"Error reading unit from STEP file: {exc}") - print("Files remain in original unit.") - - return 1.0 - -def generate_stl_solids_from_stp(stp_file, results_path=''): - """ - Extracts solid objects from a STEP (.stp) file and exports them as STL files. - - This function: - - Imports the STEP file using `cadquery`. - - Extracts solid names and their materials using `extract_names_from_stp()`. - - Sanitizes solid and material names by replacing problematic characters. - - Scales the solid to meter using `get_stp_unit_scale()`. - - Saves each solid as an STL file in the current folder (default) or the given path. - - Args: - stp_file (str): Path to the STEP (.stp) file. - results_path (str) (optional): default: '', path to save the STL (.stl) files - - Raises: - Exception: If `cadquery` is not installed, it prompts the user to install it. - - Example: - >>> extract_stl_solids_from_stp("example.stp") - 000_Vacuum-Half_cell_dx_Vacuum.stl - 001_Be_window_left_Berillium.stl - """ - - try: - import cadquery as cq - except: - raise Exception('''This function needs the open-source package `cadquery` - To install it in a conda environment do: - - `pip install cadquery` - - [!] We recommend having a dedicated conda environment to avoid version issues - ''') - - stp = cq.importers.importStep(stp_file) - - scale_factor = get_stp_unit_scale(stp_file) - if scale_factor != 1.0: - print(f"Scaling geometry to meters (factor={scale_factor}).") - scaled_solids = [solid.scale(scale_factor) for solid in stp.objects[0]] - stp.objects = [scaled_solids] - - solid_dict = extract_solids_from_stp(stp_file) - - if not results_path.endswith('/'): - results_path += '/' - - print(f"Generating stl from file: {stp_file}... ") - for i, obj in enumerate(stp.objects[0]): - name = solid_dict[list(solid_dict.keys())[i]] - print(name) - obj.exportStl(results_path+name) - - return solid_dict - From 5f6b3840c0542510f4a34f771d8da403fa278a66 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:55:03 +0100 Subject: [PATCH 13/39] Delete build/lib/wakis/grid2D.py --- build/lib/wakis/grid2D.py | 595 -------------------------------------- 1 file changed, 595 deletions(-) delete mode 100644 build/lib/wakis/grid2D.py diff --git a/build/lib/wakis/grid2D.py b/build/lib/wakis/grid2D.py deleted file mode 100644 index 6d8e8e8..0000000 --- a/build/lib/wakis/grid2D.py +++ /dev/null @@ -1,595 +0,0 @@ -import numpy as np -from numba import jit -import numba - -def seg_length(x_1, y_1, x_2, y_2): - return np.linalg.norm(np.array([x_1 - x_2, y_1 - y_2])) - -@jit('b1(f8, f8)', nopython=True) -def eq(a, b): - tol = 1e-8 - return abs(a - b) < tol - -@jit('b1(f8, f8)', nopython=True) -def neq(a, b): - tol = 1e-8 - return not eq(a, b) - - -class Grid2D: - - """ - Class holding the grid info and the routines for cell extensions. - Constructor arguments: - - xmin,xmax,ymin,ymax: extent of the domain. - - nx, ny: number of cells per direction - - conductors: conductor object - - sol_type: type of solver. 'FDTD' for staircased FDTD, 'DM' for Conformal Dey-Mittra FDTD, - 'ECT' for Extended Cell Technique conformal FDTD - """ - - def __init__(self, xmin, xmax, ymin, ymax, nx, ny, conductors, sol_type): - self.xmin = xmin - self.xmax = xmax - self.ymin = ymin - self.ymax = ymax - self.nx = nx - self.ny = ny - self.dx = (xmax - xmin) / nx - self.dy = (ymax - ymin) / ny - self.conductors = conductors - - self.l_x = np.zeros((nx, ny + 1)) - self.l_y = np.zeros((nx + 1, ny)) - self.S = np.zeros((nx, ny)) - self.S_stab = np.zeros((nx, ny)) - self.S_enl = np.zeros((nx, ny)) - self.S_red = np.zeros((nx, ny)) - self.flag_unst_cell = np.zeros((nx, ny), dtype=bool) - self.flag_int_cell = np.zeros((nx, ny), dtype=bool) - self.flag_bound_cell = np.zeros((nx, ny), dtype=bool) - self.flag_avail_cell = np.zeros((nx, ny), dtype=bool) - self.flag_ext_cell = np.zeros((nx, ny), dtype=bool) - self.flag_intr_cell = np.zeros((nx, ny), dtype=bool) - self.broken = np.zeros((self.nx, self.ny), dtype=bool) - - if (sol_type is not 'FDTD') and (sol_type is not 'DM') and (sol_type is not 'ECT'): - raise ValueError("sol_type must be:\n" + - "\t'FDTD' for standard staircased FDTD\n" + - "\t'DM' for Dey-Mittra conformal FDTD\n" + - "\t'ECT' for Enlarged Cell Technique conformal FDTD") - - if sol_type is 'DM' or sol_type is 'FDTD': - self.compute_edges(in_conductor=self.conductors.in_conductor, - intersec_x=self.conductors.intersec_x, - intersec_y=self.conductors.intersec_y, l_x=self.l_x, l_y=self.l_y, - dx=self.dx, dy=self.dy, nx=self.nx, ny=self.ny, xmin=self.xmin, - ymin=self.ymin) - compute_areas(l_x=self.l_x, l_y=self.l_y, S=self.S, S_red=self.S_red, nx=self.nx, - ny=self.ny, dx=self.dx, dy=self.dy) - mark_cells(l_x=self.l_x, l_y=self.l_y, nx=self.nx, ny=self.ny, dx=self.dx, - dy=self.dy, S=self.S, flag_int_cell=self.flag_int_cell, - S_stab=self.S_stab, flag_unst_cell=self.flag_unst_cell, - flag_bound_cell=self.flag_bound_cell, - flag_avail_cell=self.flag_avail_cell) - elif sol_type is 'ECT': - self.compute_edges(in_conductor=self.conductors.in_conductor, - intersec_x=self.conductors.intersec_x, - intersec_y=self.conductors.intersec_y, l_x=self.l_x, l_y=self.l_y, - dx=self.dx, dy=self.dy, nx=self.nx, ny=self.ny, xmin=self.xmin, - ymin=self.ymin) - compute_areas(l_x=self.l_x, l_y=self.l_y, S=self.S, S_red=self.S_red, nx=self.nx, - ny=self.ny, dx=self.dx, dy=self.dy) - mark_cells(l_x=self.l_x, l_y=self.l_y, nx=self.nx, ny=self.ny, dx=self.dx, - dy=self.dy, S=self.S, flag_int_cell=self.flag_int_cell, - S_stab=self.S_stab, flag_unst_cell=self.flag_unst_cell, - flag_bound_cell=self.flag_bound_cell, - flag_avail_cell=self.flag_avail_cell) - # info about intruded cells (i,j,[(i_borrowing,j_borrowing,area_borrowing, )]) - self.borrowing = np.empty((nx, ny), dtype=object) - - for i in range(nx): - for j in range(ny): - self.borrowing[i, j] = [] - self.flag_ext_cell = self.flag_unst_cell.copy() - self.compute_extensions(nx=nx, ny=ny, S=self.S, flag_int_cell=self.flag_int_cell, - S_stab=self.S_stab, S_enl=self.S_enl, S_red=self.S_red, - flag_unst_cell=self.flag_unst_cell, - flag_avail_cell=self.flag_avail_cell, - flag_ext_cell=self.flag_ext_cell, - flag_intr_cell=self.flag_intr_cell, - borrowing=self.borrowing) - - """ - Function to compute the length of the edges of the conformal grid. - Inputs: - - tol: an edge shorter than tol will be considered as of zero length - """ - - @staticmethod - def compute_edges(tol=1e-8, in_conductor=None, intersec_x=None, intersec_y=None, l_x=None, - l_y=None, dx=None, dy=None, nx=None, ny=None, xmin=None, ymin=None): - """ - Notation: - - (x_3, y_3)----- l_x[i, j + 1] --------- - | | - | | - l_y[i, j] (i, j) l_y[i + 1, j] - | | - | | - (x_1, y_1)------- l_x[i, j] -------(x_2, y_2) - - """ - for ii in range(nx): - for jj in range(ny + 1): - x_1 = ii * dx + xmin - y_1 = jj * dy + ymin - x_2 = (ii + 1) * dx + xmin - y_2 = jj * dy + ymin - # if point 1 is in conductor - if in_conductor(x_1, y_1): - # if point 2 is in conductor, length of the l_x[i, j] is zero - if in_conductor(x_2, y_2): - l_x[ii, jj] = 0 - # if point 2 is not in conductor, length of l_x[i, j] is the fractional length - else: - l_x[ii, jj] = seg_length(intersec_x(x_2, y_2), y_1, x_2, y_2) - # if point 1 is not in conductor - else: - # if point 2 is in conductor, length of l_x[i, j] is the fractional length - if in_conductor(x_2, y_2): - l_x[ii, jj] = seg_length(x_1, y_1, intersec_x(x_1, y_1), y_2) - # if point 2 is not in conductor, length of l_x[i, j] is dx - else: - l_x[ii, jj] = dx - - for ii in range(nx + 1): - for jj in range(ny): - x_1 = ii * dx + xmin - y_1 = jj * dy + ymin - x_3 = ii * dx + xmin - y_3 = (jj + 1) * dy + ymin - # if point 1 is in conductor - if in_conductor(x_1, y_1): - # if point 3 to the right is in conductor, length of the l_y[i, j] is zero - if in_conductor(x_3, y_3): - l_y[ii, jj] = 0 - # if point 3 is not in conductor, length of l_y[i, j] is the fractional length - else: - l_y[ii, jj] = seg_length(x_1, intersec_y(x_3, y_3), x_3, y_3) - # if point 1 is not in conductor - else: - # if point 3 is in conductor, length of the l_y[i, j] is the fractional length - if in_conductor(x_3, y_3): - l_y[ii, jj] = seg_length(x_1, y_1, x_3, intersec_y(x_1, y_1)) - # if point 3 is not in conductor, length of l_y[i, j] is dy - else: - l_y[ii, jj] = dy - - # set to zero the length of very small cells - if tol > 0.: - l_x /= dx - l_y /= dy - - low_values_flags = abs(l_x) < tol - high_values_flags = abs(l_x - 1.) < tol - l_x[low_values_flags] = 0 - l_x[high_values_flags] = 1. - - low_values_flags = abs(l_y) < tol - high_values_flags = abs(l_y - 1.) < tol - l_y[low_values_flags] = 0 - l_y[high_values_flags] = 1. - - l_x *= dx - l_y *= dy - - """ - Function to compute the extension of the unstable cells - """ - - @staticmethod - def compute_extensions(nx=None, ny=None, S=None, flag_int_cell=None, - S_stab=None, S_enl=None, S_red=None, - flag_unst_cell=None, - flag_avail_cell=None, - flag_ext_cell=None, - flag_intr_cell = None, - borrowing=None, l_verbose=True, - kk=0): - - N = np.sum(flag_ext_cell) - - if l_verbose: - print(kk) - print('ext cells: %d' % N) - # Do the simple one-cell extension - Grid2D._compute_extensions_one_cell(nx=nx, ny=ny, S=S, S_stab=S_stab, S_enl=S_enl, - S_red=S_red, flag_avail_cell=flag_avail_cell, - flag_ext_cell=flag_ext_cell, - flag_intr_cell=flag_intr_cell, - borrowing=borrowing) - - N_one_cell = (N - np.sum(flag_ext_cell)) - if l_verbose: - print('one cell exts: %d' % N_one_cell) - # If any cell could not be extended do the four-cell extension - ''' - if np.sum(flag_ext_cell) > 0: - N = np.sum(flag_ext_cell) - Grid2D._compute_extensions_four_cells(nx=nx, ny=ny, S=S, flag_int_cell=flag_int_cell, - S_stab=S_stab, S_enl=S_enl, S_red=S_red, - flag_unst_cell=flag_unst_cell, - flag_avail_cell=flag_avail_cell, - flag_ext_cell=flag_ext_cell, - flag_intr_cell = flag_intr_cell, - borrowing=borrowing) - N_four_cells = (N - np.sum(flag_ext_cell)) - if 1: #l_verbose: - print('four cell exts: %d' % N_four_cells) - ''' - # If any cell could not be extended do the eight-cell extension - if np.sum(flag_ext_cell) > 0: - N = np.sum(flag_ext_cell) - Grid2D._compute_extensions_eight_cells(nx=nx, ny=ny, S=S, flag_int_cell=flag_int_cell, - S_stab=S_stab, S_enl=S_enl, S_red=S_red, - flag_unst_cell=flag_unst_cell, - flag_avail_cell=flag_avail_cell, - flag_ext_cell=flag_ext_cell, - flag_intr_cell=flag_intr_cell, - borrowing=borrowing) - N_eight_cells = (N - np.sum(flag_ext_cell)) - if 1: #l_verbose: - print('eight cell exts: %d' % N_eight_cells) - # If any cell could not be extended the algorithm failed - if np.sum(flag_ext_cell) > 0: - N = (np.sum(flag_ext_cell)) - raise RuntimeError(str(N) + ' cells could not be extended.\n' + - 'Please refine the mesh') - - """ - Function to compute the one-cell extension of the unstable cells - """ - - @staticmethod - def _compute_extensions_one_cell(nx=None, ny=None, S=None, - S_stab=None, S_enl=None, S_red=None, flag_avail_cell=None, - flag_ext_cell=None, flag_intr_cell=None, borrowing=None): - - for ii in range(0, nx): - for jj in range(0, ny): - if flag_ext_cell[ii, jj]: - S_ext = S_stab[ii, jj] - S[ii, jj] - if (S[ii - 1, jj] > S_ext and flag_avail_cell[ii - 1, jj]): - denom = S[ii - 1, jj] - patch = S_ext * S[ii - 1, jj] / denom - if S_red[ii - 1, jj] - patch > 0: - S_red[ii - 1, jj] -= patch - borrowing[ii, jj].append([ii - 1, jj, patch, None]) - flag_intr_cell[ii-1, jj] = True - S_enl[ii, jj] = S[ii, jj] + patch - flag_ext_cell[ii, jj] = False - if (S[ii, jj - 1] > S_ext and flag_avail_cell[ii, jj - 1] and flag_ext_cell[ii, jj]): - denom = S[ii, jj - 1] - patch = S_ext * S[ii, jj - 1] / denom - if S_red[ii, jj - 1] - patch > 0: - S_red[ii, jj - 1] -= patch - borrowing[ii, jj].append([ii, jj - 1, patch, None]) - flag_intr_cell[ii, jj-1] = True - S_enl[ii, jj] = S[ii, jj] + patch - flag_ext_cell[ii, jj] = False - if (S[ii, jj + 1] > S_ext and flag_avail_cell[ii, jj + 1] - and flag_ext_cell[ii, jj]): - denom = S[ii, jj + 1] - patch = S_ext * S[ii, jj + 1] / denom - if S_red[ii, jj + 1] - patch > 0: - S_red[ii, jj + 1] -= patch - borrowing[ii, jj].append([ii, jj + 1, patch, None]) - flag_intr_cell[ii, jj+1] = True - S_enl[ii, jj] = S[ii, jj] + patch - flag_ext_cell[ii, jj] = False - if (S[ii + 1, jj] > S_ext and flag_avail_cell[ii + 1, jj] - and flag_ext_cell[ii, jj]): - denom = S[ii + 1, jj] - patch = S_ext * S[ii + 1, jj] / denom - if S_red[ii + 1, jj] - patch > 0: - S_red[ii + 1, jj] -= patch - borrowing[ii, jj].append([ii + 1, jj, patch, None]) - flag_intr_cell[ii+1, jj] = True - S_enl[ii, jj] = S[ii, jj] + patch - flag_ext_cell[ii, jj] = False - - - """ - Function to compute the four-cell extension of the unstable cells - """ - - @staticmethod - def _compute_extensions_four_cells(nx=None, ny=None, S=None, flag_int_cell=None, - S_stab=None, S_enl=None, S_red=None, flag_unst_cell=None, - flag_avail_cell=None, flag_ext_cell=None, - flag_intr_cell=None, borrowing=None): - for ii in range(0, nx): - for jj in range(0, ny): - local_avail = flag_avail_cell.copy() - if (flag_unst_cell[ii, jj] and flag_int_cell[ii, jj] - and flag_ext_cell[ii, jj]): - denom = ((flag_avail_cell[ii - 1, jj]) * S[ii - 1, jj] + ( - flag_avail_cell[ii + 1, jj]) * S[ii + 1, jj] + - (flag_avail_cell[ii, jj - 1]) * S[ii, jj - 1] + ( - flag_avail_cell[ii, jj + 1]) * S[ii, jj + 1]) - S_ext = S_stab[ii, jj] - S[ii, jj] - neg_cell = True - # idea: if any cell would reach negative area it is locally not available. - # then denom has to be recomputed from scratch - - while denom >= S_ext and neg_cell: - neg_cell = False - if local_avail[ii - 1, jj]: - patch = S_ext * S[ii - 1, jj] / denom - if S_red[ii - 1, jj] - patch <= 0: - neg_cell = True - local_avail[ii - 1, jj] = False - if local_avail[ii + 1, jj]: - patch = S_ext * S[ii + 1, jj] / denom - if S_red[ii + 1, jj] - patch <= 0: - neg_cell = True - local_avail[ii + 1, jj] = False - if local_avail[ii, jj - 1]: - patch = S_ext * S[ii, jj - 1] / denom - if S_red[ii, jj - 1] - patch <= 0: - neg_cell = True - local_avail[ii, jj - 1] = False - if local_avail[ii, jj + 1]: - patch = S_ext * S[ii, jj + 1] / denom - if S_red[ii, jj + 1] - patch <= 0: - neg_cell = True - local_avail[ii, jj + 1] = False - denom = ((local_avail[ii - 1, jj]) * S[ii - 1, jj] + - (local_avail[ii + 1, jj]) * S[ii + 1, jj] + - (local_avail[ii, jj - 1]) * S[ii, jj - 1] + - (local_avail[ii, jj + 1]) * S[ii, jj + 1]) - - # If possible, do 4-cell extension - if denom >= S_ext: - S_enl[ii, jj] = S[ii, jj] - if local_avail[ii - 1, jj]: - patch = S_ext * S[ii - 1, jj] / denom - borrowing[ii, jj].append([ii - 1, jj, patch, None]) - flag_intr_cell[ii-1, jj] = True - S_enl[ii, jj] += patch - S_red[ii - 1, jj] -= patch - if local_avail[ii + 1, jj]: - patch = S_ext * S[ii + 1, jj] / denom - borrowing[ii, jj].append([ii + 1, jj, patch, None]) - flag_intr_cell[ii+1, jj] = True - S_enl[ii, jj] += patch - S_red[ii + 1, jj] -= patch - if local_avail[ii, jj - 1]: - patch = S_ext * S[ii, jj - 1] / denom - borrowing[ii, jj].append([ii, jj - 1, patch, None]) - flag_intr_cell[ii, jj-1] = True - S_enl[ii, jj] += patch - S_red[ii, jj - 1] -= patch - if local_avail[ii, jj + 1]: - patch = S_ext * S[ii, jj + 1] / denom - borrowing[ii, jj].append([ii, jj + 1, patch, None]) - flag_intr_cell[ii, jj+1] = True - S_enl[ii, jj] += patch - S_red[ii, jj + 1] -= patch - - flag_ext_cell[ii, jj] = False - - """ - Function to compute the eight-cell extension of the unstable cells - """ - - @staticmethod - def _compute_extensions_eight_cells(nx=None, ny=None, S=None, flag_int_cell=None, - S_stab=None, S_enl=None, S_red=None, flag_unst_cell=None, - flag_avail_cell=None, flag_ext_cell=None, - flag_intr_cell=None, borrowing=None): - for ii in range(0, nx): - for jj in range(0, ny): - local_avail = flag_avail_cell.copy() - if (flag_unst_cell[ii, jj] and flag_int_cell[ii, jj] - and flag_ext_cell[ii, jj]): - S_enl[ii, jj] = S[ii, jj] - S_ext = S_stab[ii, jj] - S[ii, jj] - - denom = ((flag_avail_cell[ii - 1, jj]) * S[ii - 1, jj] + - (flag_avail_cell[ii + 1, jj]) * S[ii + 1, jj] + - (flag_avail_cell[ii, jj - 1]) * S[ii, jj - 1] + - (flag_avail_cell[ii, jj + 1]) * S[ii, jj + 1] + - (flag_avail_cell[ii - 1, jj - 1]) * S[ii - 1, jj - 1] + - (flag_avail_cell[ii + 1, jj - 1]) * S[ii + 1, jj - 1] + - (flag_avail_cell[ii - 1, jj + 1]) * S[ii - 1, jj + 1] + - (flag_avail_cell[ii + 1, jj + 1]) * S[ii + 1, jj + 1]) - - neg_cell = True - while denom >= S_ext and neg_cell: - neg_cell = False - if local_avail[ii - 1, jj]: - patch = S_ext * S[ii - 1, jj] / denom - if S_red[ii - 1, jj] - patch <= 0: - neg_cell = True - local_avail[ii - 1, jj] = False - if local_avail[ii + 1, jj]: - patch = S_ext * S[ii + 1, jj] / denom - if S_red[ii + 1, jj] - patch <= 0: - neg_cell = True - local_avail[ii + 1, jj] = False - if local_avail[ii, jj - 1]: - patch = S_ext * S[ii, jj - 1] / denom - if S_red[ii, jj - 1] - patch <= 0: - neg_cell = True - local_avail[ii, jj - 1] = False - if local_avail[ii, jj + 1]: - patch = S_ext * S[ii, jj + 1] / denom - if S_red[ii, jj + 1] - patch <= 0: - neg_cell = True - local_avail[ii, jj + 1] = False - if local_avail[ii - 1, jj - 1]: - patch = S_ext * S[ii - 1, jj - 1] / denom - if S_red[ii - 1, jj - 1] - patch <= 0: - neg_cell = True - local_avail[ii - 1, jj - 1] = False - if local_avail[ii + 1, jj - 1]: - patch = S_ext * S[ii + 1, jj - 1] / denom - if S_red[ii + 1, jj - 1] - patch <= 0: - neg_cell = True - local_avail[ii + 1, jj - 1] = False - if local_avail[ii - 1, jj + 1]: - patch = S_ext * S[ii - 1, jj + 1] / denom - if S_red[ii - 1, jj + 1] - patch <= 0: - neg_cell = True - local_avail[ii - 1, jj + 1] = False - if local_avail[ii + 1, jj + 1]: - patch = S_ext * S[ii + 1, jj + 1] / denom - if S_red[ii + 1, jj + 1] - patch <= 0: - neg_cell = True - local_avail[ii + 1, jj + 1] = False - - denom = ((local_avail[ii - 1, jj]) * S[ii - 1, jj] + - (local_avail[ii + 1, jj]) * S[ii + 1, jj] + - (local_avail[ii, jj - 1]) * S[ii, jj - 1] + - (local_avail[ii, jj + 1]) * S[ii, jj + 1] + - (local_avail[ii - 1, jj - 1]) * S[ii - 1, jj - 1] + - (local_avail[ii + 1, jj - 1]) * S[ii + 1, jj - 1] + - (local_avail[ii - 1, jj + 1]) * S[ii - 1, jj + 1] + - (local_avail[ii + 1, jj + 1]) * S[ii + 1, jj + 1]) - - if denom >= S_ext: - S_enl[ii, jj] = S[ii, jj] - if local_avail[ii - 1, jj]: - patch = S_ext * S[ii - 1, jj] / denom - borrowing[ii, jj].append([ii - 1, jj, patch, None]) - flag_intr_cell[ii-1, jj] = True - S_enl[ii, jj] += patch - S_red[ii - 1, jj] -= patch - if local_avail[ii + 1, jj]: - patch = S_ext * S[ii + 1, jj] / denom - borrowing[ii, jj].append([ii + 1, jj, patch, None]) - flag_intr_cell[ii+1, jj] = True - S_enl[ii, jj] += patch - S_red[ii + 1, jj] -= patch - if local_avail[ii, jj - 1]: - patch = S_ext * S[ii, jj - 1] / denom - borrowing[ii, jj].append([ii, jj - 1, patch, None]) - flag_intr_cell[ii, jj-1] = True - S_enl[ii, jj] += patch - S_red[ii, jj - 1] -= patch - if local_avail[ii, jj + 1]: - patch = S_ext * S[ii, jj + 1] / denom - borrowing[ii, jj].append([ii, jj + 1, patch, None]) - flag_intr_cell[ii, jj+1] = True - S_enl[ii, jj] += patch - S_red[ii, jj + 1] -= patch - if local_avail[ii - 1, jj - 1]: - patch = S_ext * S[ii - 1, jj - 1] / denom - borrowing[ii, jj].append([ii - 1, jj - 1, patch, None]) - flag_intr_cell[ii-1, jj-1] = True - S_enl[ii, jj] += patch - S_red[ii - 1, jj - 1] -= patch - if local_avail[ii + 1, jj - 1]: - patch = S_ext * S[ii + 1, jj - 1] / denom - borrowing[ii, jj].append([ii + 1, jj - 1, patch, None]) - flag_intr_cell[ii+1, jj-1] = True - S_enl[ii, jj] += patch - S_red[ii + 1, jj - 1] -= patch - if local_avail[ii - 1, jj + 1]: - patch = S_ext * S[ii - 1, jj + 1] / denom - borrowing[ii, jj].append([ii - 1, jj + 1, patch, None]) - flag_intr_cell[ii-1, jj+1] = True - S_enl[ii, jj] += patch - S_red[ii - 1, jj + 1] -= patch - if local_avail[ii + 1, jj + 1]: - patch = S_ext * S[ii + 1, jj + 1] / denom - borrowing[ii, jj].append([ii + 1, jj + 1, patch, None]) - flag_intr_cell[ii+1, jj+1] = True - S_enl[ii, jj] += patch - S_red[ii + 1, jj + 1] -= patch - - flag_ext_cell[ii, jj] = False - -""" -Function to compute the area of the cells of the conformal grid. -""" -@jit('(f8[:,:], f8[:,:], f8[:,:], f8[:,:], i4, i4, f8, f8)', nopython=True) -def compute_areas(l_x, l_y, S, S_red, nx, ny, dx, dy): - # Normalize the grid lengths for robustness - l_x /= dx - l_y /= dy - - # Loop over the cells - for ii in range(nx): - for jj in range(ny): - # If at least 3 edges have full length consider the area as full - # (this also takes care of the case in which an edge lies exactly on the boundary) - #count_cane = eq(l_x[ii, jj], 1.0) + eq(l_y[ii, jj], 1.0) + eq(l_x[ii, jj + 1], 1.0) + eq(l_y[ii + 1, jj], 1.0) - if (np.sum(np.array([eq(l_x[ii, jj], 1.0), eq(l_y[ii, jj], 1.0), - eq(l_x[ii, jj + 1], 1.0), eq(l_y[ii + 1, jj], 1.0)])) >= 3): - #if count_cane >=3: - S[ii, jj] = 1.0 - elif (eq(l_x[ii, jj], 0.) and neq(l_y[ii, jj], 0.) - and neq(l_x[ii, jj + 1], 0.) and neq(l_y[ii + 1, jj], 0)): - S[ii, jj] = 0.5 * (l_y[ii, jj] + l_y[ii + 1, jj]) * l_x[ii, jj + 1] - elif (eq(l_x[ii, jj + 1], 0.) and neq(l_y[ii, jj], 0.) - and neq(l_x[ii, jj], 0.) and neq(l_y[ii + 1, jj], 0.)): - S[ii, jj] = 0.5 * (l_y[ii, jj] + l_y[ii + 1, jj]) * l_x[ii, jj] - elif (eq(l_y[ii, jj], 0.) and neq(l_x[ii, jj], 0.) - and neq(l_x[ii, jj + 1], 0.) and neq(l_y[ii + 1, jj], 0.)): - S[ii, jj] = 0.5 * (l_x[ii, jj] + l_x[ii, jj + 1]) * l_y[ii + 1, jj] - elif (eq(l_y[ii + 1, jj], 0.) and neq(l_x[ii, jj], 0.) and neq(l_y[ii, jj], 0.) - and neq(l_x[ii, jj + 1], 0.)): - S[ii, jj] = 0.5 * (l_x[ii, jj] + l_x[ii, jj + 1]) * l_y[ii, jj] - elif (eq(l_x[ii, jj], 0.) and eq(l_y[ii, jj], 0.) - and neq(l_x[ii, jj + 1], 0.) and neq(l_y[ii + 1, jj], 0.)): - S[ii, jj] = 0.5 * l_x[ii, jj + 1] * l_y[ii + 1, jj] - elif (eq(l_x[ii, jj], 0.) and eq(l_y[ii + 1, jj], 0.) - and neq(l_x[ii, jj + 1], 0.) and neq(l_y[ii, jj], 0.)): - S[ii, jj] = 0.5 * l_x[ii, jj + 1] * l_y[ii, jj] - elif (eq(l_x[ii, jj + 1], 0.) and eq(l_y[ii + 1, jj], 0.) - and neq(l_x[ii, jj], 0.) and neq(l_y[ii, jj], 0.)): - S[ii, jj] = 0.5 * l_x[ii, jj] * l_y[ii, jj] - elif (eq(l_x[ii, jj + 1], 0.) and eq(l_y[ii, jj], 0.) - and neq(l_x[ii, jj], 0.) and neq(l_y[ii + 1, jj], 0.)): - S[ii, jj] = 0.5 * l_x[ii, jj] * l_y[ii + 1, jj] - elif (0. < l_x[ii, jj] <= 1. and 0. < l_y[ii, jj] <= 1. - and eq(l_x[ii, jj + 1], 1.) and eq(l_y[ii + 1, jj], 1.)): - S[ii, jj] = 1. - 0.5 * (1. - l_x[ii, jj]) * (1. - l_y[ii, jj]) - elif (0. < l_x[ii, jj] <= 1. and 0. < l_y[ii + 1, jj] <= 1. - and eq(l_x[ii, jj + 1], 1.) and eq(l_y[ii, jj], 1.)): - S[ii, jj] = 1. - 0.5 * (1. - l_x[ii, jj]) * (1. - l_y[ii + 1, jj]) - elif (0. < l_x[ii, jj + 1] <= 1. and 0. < l_y[ii + 1, jj] <= 1. - and eq(l_x[ii, jj], 1.) and eq(l_y[ii, jj], 1.)): - S[ii, jj] = 1. - 0.5 * (1. - l_x[ii, jj + 1]) * (1. - l_y[ii + 1, jj]) - elif (0. < l_x[ii, jj + 1] <= 1. and 0. < l_y[ii, jj] <= 1. - and eq(l_x[ii, jj], 1.) and eq(l_y[ii + 1, jj], 1.)): - S[ii, jj] = 1. - 0.5 * (1. - l_x[ii, jj + 1]) * (1. - l_y[ii, jj]) - - l_x *= dx - l_y *= dy - -""" -Function to mark wich cells are interior (int), require extension (unst), -are on the boundary(bound), are available for intrusion (avail) -""" -@jit('(f8[:,:], f8[:,:], i4, i4, f8, f8, f8[:,:], b1[:,:], f8[:,:], b1[:,:], b1[:,:], b1[:,:])', nopython=True) -def mark_cells(l_x, l_y, nx, ny, dx, dy, S, flag_int_cell, S_stab, - flag_unst_cell, flag_bound_cell, flag_avail_cell): - for ii in range(nx): - for jj in range(ny): - flag_int_cell[ii, jj] = S[ii, jj] > 0 - if flag_int_cell[ii, jj]: - S_stab[ii, jj] = 0.5 * np.max( - np.array([l_x[ii, jj] * dy, l_x[ii, jj + 1] * dy, - l_y[ii, jj] * dx, - l_y[ii + 1, jj] * dx])) - flag_unst_cell[ii, jj] = S[ii, jj] < S_stab[ii, jj] - flag_bound_cell[ii, jj] = ((0 < l_x[ii, jj] < dx) or - (0 < l_y[ii, jj] < dy) or - (0 < l_x[ii, jj + 1] < dx) or - (0 < l_y[ii + 1, jj] < dy)) - flag_avail_cell[ii, jj] = flag_int_cell[ii, jj] and (not flag_unst_cell[ii, jj]) From c8776ea6919e381a441af9b1b4b8ea250c6865e2 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:55:14 +0100 Subject: [PATCH 14/39] Delete build/lib/wakis/grid3D.py --- build/lib/wakis/grid3D.py | 415 -------------------------------------- 1 file changed, 415 deletions(-) delete mode 100644 build/lib/wakis/grid3D.py diff --git a/build/lib/wakis/grid3D.py b/build/lib/wakis/grid3D.py deleted file mode 100644 index f08aad4..0000000 --- a/build/lib/wakis/grid3D.py +++ /dev/null @@ -1,415 +0,0 @@ -import numpy as np -from grid2D import Grid2D -from grid2D import compute_areas as compute_areas_2D, mark_cells as mark_cells_2D -from numba import jit -from field import Field - -def seg_length(x_1, y_1, z_1, x_2, y_2, z_2): - return np.linalg.norm(np.array([x_1 - x_2, y_1 - y_2, z_1 - z_2])) - - -def eq(a, b, tol=1e-8): - return abs(a - b) < tol - - -def neq(a, b, tol=1e-8): - return not eq(a, b, tol) - - # Undo the normalization - S *= dx * dy - l_x *= dx - l_y *= dy - S_red[:] = S.copy()[:] - -class Grid3D: - """ - Class holding the grid info and the routines for cell extensions. - Constructor arguments: - - xmin, xmax, ymin, ymax, zmin, zmax: extent of the domain. - - nx, ny, nz: number of cells per direction - - conductors: conductor object - - sol_type: type of solver. 'FDTD' for staircased FDTD, 'DM' for Conformal Dey-Mittra FDTD, - 'ECT' for Enlarged Cell Technique conformal FDTD - """ - - def __init__(self, xmin, xmax, ymin, ymax, zmin, zmax, nx, ny, nz, conductors, sol_type): - - self.xmin = xmin - self.xmax = xmax - self.ymin = ymin - self.ymax = ymax - self.zmin = zmin - self.zmax = zmax - self.nx = nx - self.ny = ny - self.nz = nz - self.dx = (xmax - xmin) / nx - self.dy = (ymax - ymin) / ny - self.dz = (ymax - ymin) / nz - self.conductors = conductors - - self.l_x = np.zeros((nx, ny + 1, nz + 1)) - self.l_y = np.zeros((nx + 1, ny, nz + 1)) - self.l_z = np.zeros((nx + 1, ny + 1, nz)) - self.Sxy = np.zeros((nx, ny, nz + 1)) - self.Syz = np.zeros((nx + 1, ny, nz)) - self.Szx = np.zeros((nx, ny + 1, nz)) - self.Sxy_stab = np.zeros_like(self.Sxy) - self.Sxy_enl = np.zeros_like(self.Sxy) - self.Sxy_red = np.zeros_like(self.Sxy) - self.Syz_stab = np.zeros_like(self.Syz) - self.Syz_enl = np.zeros_like(self.Syz) - self.Syz_red = np.zeros_like(self.Syz) - self.Szx_stab = np.zeros_like(self.Szx) - self.Szx_enl = np.zeros_like(self.Szx) - self.Szx_red = np.zeros_like(self.Szx) - self.flag_unst_cell_xy = np.zeros_like(self.Sxy, dtype=bool) - self.flag_intr_cell_xy = np.zeros_like(self.Sxy, dtype=bool) - self.flag_int_cell_xy = np.zeros_like(self.Sxy, dtype=bool) - self.flag_bound_cell_xy = np.zeros_like(self.Sxy, dtype=bool) - self.flag_avail_cell_xy = np.zeros_like(self.Sxy, dtype=bool) - self.flag_ext_cell_xy = np.zeros_like(self.Sxy, dtype=bool) - self.flag_unst_cell_yz = np.zeros_like(self.Syz, dtype=bool) - self.flag_intr_cell_yz = np.zeros_like(self.Syz, dtype=bool) - self.flag_int_cell_yz = np.zeros_like(self.Syz, dtype=bool) - self.flag_bound_cell_yz = np.zeros_like(self.Syz, dtype=bool) - self.flag_avail_cell_yz = np.zeros_like(self.Syz, dtype=bool) - self.flag_ext_cell_yz = np.zeros_like(self.Syz, dtype=bool) - self.flag_unst_cell_zx = np.zeros_like(self.Szx, dtype=bool) - self.flag_intr_cell_zx = np.zeros_like(self.Szx, dtype=bool) - self.flag_int_cell_zx = np.zeros_like(self.Szx, dtype=bool) - self.flag_bound_cell_zx = np.zeros_like(self.Szx, dtype=bool) - self.flag_avail_cell_zx = np.zeros_like(self.Szx, dtype=bool) - self.flag_ext_cell_zx = np.zeros_like(self.Szx, dtype=bool) - self.broken_xy = np.zeros_like(self.Sxy, dtype=bool) - self.broken_yz = np.zeros_like(self.Syz, dtype=bool) - self.broken_zx = np.zeros_like(self.Szx, dtype=bool) - - if (sol_type is not 'FDTD') and (sol_type is not 'DM') and (sol_type is not 'ECT') and (sol_type is not 'FIT'): - raise ValueError("sol_type must be:\n" + - "\t'FDTD' for standard staircased FDTD\n" + - "\t'DM' for Dey-Mittra conformal FDTD\n" + - "\t'ECT' for Enlarged Cell Technique conformal FDTD") - - self.compute_edges() - if sol_type is 'DM' or sol_type is 'FDTD': #or sol_type is 'FIT': - self.compute_areas(self.l_x, self.l_y, self.l_z, self.Sxy, self.Syz, self.Szx, - self.Sxy_red, self.Syz_red, self.Szx_red, - self.nx, self.ny, self.nz, self.dx, self.dy, self.dz) - self.mark_cells(self.l_x, self.l_y, self.l_z, self.nx, self.ny, self.nz, self.dx, self.dy, self.dz, - self.Sxy, self.Syz, self.Szx, self.flag_int_cell_xy, self.flag_int_cell_yz, self.flag_int_cell_zx, - self.Sxy_stab, self.Syz_stab, self.Szx_stab, self.flag_unst_cell_xy, self.flag_unst_cell_yz, - self.flag_unst_cell_zx, self.flag_bound_cell_xy, self.flag_bound_cell_yz, self.flag_bound_cell_zx, - self.flag_avail_cell_xy, self.flag_avail_cell_yz, self.flag_avail_cell_zx) - elif sol_type is 'ECT': - self.compute_areas(self.l_x, self.l_y, self.l_z, self.Sxy, self.Syz, self.Szx, - self.Sxy_red, self.Syz_red, self.Szx_red, - self.nx, self.ny, self.nz, self.dx, self.dy, self.dz) - self.mark_cells(self.l_x, self.l_y, self.l_z, self.nx, self.ny, self.nz, self.dx, self.dy, self.dz, - self.Sxy, self.Syz, self.Szx, self.flag_int_cell_xy, self.flag_int_cell_yz, self.flag_int_cell_zx, - self.Sxy_stab, self.Syz_stab, self.Szx_stab, self.flag_unst_cell_xy, self.flag_unst_cell_yz, - self.flag_unst_cell_zx, self.flag_bound_cell_xy, self.flag_bound_cell_yz, self.flag_bound_cell_zx, - self.flag_avail_cell_xy, self.flag_avail_cell_yz, self.flag_avail_cell_zx) - - - # info about intruded cells (i,j,[(i_borrowing,j_borrowing,area_borrowing, )]) - self.borrowing_xy = np.empty((nx, ny, nz + 1), dtype=object) - self.borrowing_yz = np.empty((nx + 1, ny, nz), dtype=object) - self.borrowing_zx = np.empty((nx, ny + 1, nz), dtype=object) - - for ii in range(nx): - for jj in range(ny): - for kk in range(nz + 1): - self.borrowing_xy[ii, jj, kk] = [] - for ii in range(nx + 1): - for jj in range(ny): - for kk in range(nz): - self.borrowing_yz[ii, jj, kk] = [] - for ii in range(nx): - for jj in range(ny + 1): - for kk in range(nz): - self.borrowing_zx[ii, jj, kk] = [] - - self.flag_ext_cell_xy = self.flag_unst_cell_xy.copy() - self.flag_ext_cell_yz = self.flag_unst_cell_yz.copy() - self.flag_ext_cell_zx = self.flag_unst_cell_zx.copy() - self.mark_cells(self.l_x, self.l_y, self.l_z, self.nx, self.ny, self.nz, self.dx, self.dy, self.dz, - self.Sxy, self.Syz, self.Szx, self.flag_int_cell_xy, self.flag_int_cell_yz, self.flag_int_cell_zx, - self.Sxy_stab, self.Syz_stab, self.Szx_stab, self.flag_unst_cell_xy, self.flag_unst_cell_yz, - self.flag_unst_cell_zx, self.flag_bound_cell_xy, self.flag_bound_cell_yz, self.flag_bound_cell_zx, - self.flag_avail_cell_xy, self.flag_avail_cell_yz, self.flag_avail_cell_zx) - self.compute_extensions() - - - elif sol_type is 'FIT': - - # primal Grid G - self.x = np.linspace(self.xmin, self.xmax, self.nx+1) - self.y = np.linspace(self.ymin, self.ymax, self.ny+1) - self.z = np.linspace(self.zmin, self.zmax, self.nz+1) - - Y, X, Z = np.meshgrid(self.y, self.x, self.z) - - self.L = Field(self.nx, self.ny, self.nz) - self.L.field_x = X[1:, 1:, 1:] - X[:-1, :-1, :-1] - self.L.field_y = Y[1:, 1:, 1:] - Y[:-1, :-1, :-1] - self.L.field_z = Z[1:, 1:, 1:] - Z[:-1, :-1, :-1] - - self.iA = Field(self.nx, self.ny, self.nz) - self.iA.field_x = np.divide(1.0, self.L.field_y * self.L.field_z) - self.iA.field_y = np.divide(1.0, self.L.field_x * self.L.field_z) - self.iA.field_z = np.divide(1.0, self.L.field_x * self.L.field_y) - - # tilde grid ~G - #self.itA = self.iA - #self.tL = self.L - - self.tx = (self.x[1:]+self.x[:-1])/2 - self.ty = (self.y[1:]+self.y[:-1])/2 - self.tz = (self.z[1:]+self.z[:-1])/2 - - self.tx = np.append(self.tx, self.tx[-1]) - self.ty = np.append(self.ty, self.ty[-1]) - self.tz = np.append(self.tz, self.tz[-1]) - - tY, tX, tZ = np.meshgrid(self.ty, self.tx, self.tz) - - self.tL = Field(self.nx, self.ny, self.nz) - self.tL.field_x = tX[1:, 1:, 1:] - tX[:-1, :-1, :-1] - self.tL.field_y = tY[1:, 1:, 1:] - tY[:-1, :-1, :-1] - self.tL.field_z = tZ[1:, 1:, 1:] - tZ[:-1, :-1, :-1] - - self.itA = Field(self.nx, self.ny, self.nz) - aux = self.tL.field_y * self.tL.field_z - self.itA.field_x = np.divide(1.0, aux, out=np.zeros_like(aux), where=aux!=0) - aux = self.tL.field_x * self.tL.field_z - self.itA.field_y = np.divide(1.0, aux, out=np.zeros_like(aux), where=aux!=0) - aux = self.tL.field_x * self.tL.field_y - self.itA.field_z = np.divide(1.0, aux, out=np.zeros_like(aux), where=aux!=0) - del aux - - self.compute_areas(self.l_x, self.l_y, self.l_z, self.Sxy, self.Syz, self.Szx, - self.Sxy_red, self.Syz_red, self.Szx_red, - self.nx, self.ny, self.nz, self.dx, self.dy, self.dz) - self.mark_cells(self.l_x, self.l_y, self.l_z, self.nx, self.ny, self.nz, self.dx, self.dy, self.dz, - self.Sxy, self.Syz, self.Szx, self.flag_int_cell_xy, self.flag_int_cell_yz, self.flag_int_cell_zx, - self.Sxy_stab, self.Syz_stab, self.Szx_stab, self.flag_unst_cell_xy, self.flag_unst_cell_yz, - self.flag_unst_cell_zx, self.flag_bound_cell_xy, self.flag_bound_cell_yz, self.flag_bound_cell_zx, - self.flag_avail_cell_xy, self.flag_avail_cell_yz, self.flag_avail_cell_zx) - - """ - Function to compute the length of the edges of the conformal grid. - Inputs: - - tol: an edge shorter than tol will be considered as of zero length - """ - - def compute_edges(self, tol=1e-8): - # edges xy - for kk in range(self.nz + 1): - for ii in range(self.nx): - for jj in range(self.ny + 1): - x_1 = ii * self.dx + self.xmin - y_1 = jj * self.dy + self.ymin - x_2 = (ii + 1) * self.dx + self.xmin - y_2 = jj * self.dy + self.ymin - z = kk * self.dz + self.zmin - # if point 1 is in conductor - if self.conductors.in_conductor(x_1, y_1, z): - # if point 2 is in conductor, length of the l_x[i, j] is zero - if self.conductors.in_conductor(x_2, y_2, z): - self.l_x[ii, jj, kk] = 0 - # if point 2 is not in conductor, length of l_x[i, j] - # is the fractional length - else: - self.l_x[ii, jj, kk] = seg_length(self.conductors.intersec_x(x_2, - y_2, z), - y_2, z, x_2, y_2, z) - # if point 1 is not in conductor - else: - # if point 2 is in conductor, length of l_x[i, j] is the fractional length - if self.conductors.in_conductor(x_2, y_2, z): - self.l_x[ii, jj, kk] = seg_length(x_1, y_1, z, - self.conductors.intersec_x(x_1, y_1, - z), - y_1, z) - # if point 2 is not in conductor, length of l_x[i, j] is dx - else: - self.l_x[ii, jj, kk] = self.dx - - for kk in range(self.nz + 1): - for ii in range(self.nx + 1): - for jj in range(self.ny): - x_1 = ii * self.dx + self.xmin - y_1 = jj * self.dy + self.ymin - x_3 = ii * self.dx + self.xmin - y_3 = (jj + 1) * self.dy + self.ymin - z = kk * self.dz + self.zmin - # if point 1 is in conductor - if self.conductors.in_conductor(x_1, y_1, z): - # if point 3 to the right is in conductor, length of the l_y[i, j] is zero - if self.conductors.in_conductor(x_3, y_3, z): - self.l_y[ii, jj, kk] = 0 - # if point 3 is not in conductor, length of l_y[i, j] - # is the fractional length - else: - self.l_y[ii, jj, kk] = seg_length(x_3, - self.conductors.intersec_y(x_3, y_3, - z), - z, x_3, y_3, z) - # if point 1 is not in conductor - else: - # if point 3 is in conductor, length of the l_y[i, j] - # is the fractional length - if self.conductors.in_conductor(x_3, y_3, z): - self.l_y[ii, jj, kk] = seg_length(x_1, y_1, z, x_1, - self.conductors.intersec_y(x_1, y_1, - z), - z) - # if point 3 is not in conductor, length of l_y[i, j] is dy - else: - self.l_y[ii, jj, kk] = self.dy - - for jj in range(self.ny + 1): - for ii in range(self.nx + 1): - for kk in range(self.nz): - y = jj*self.dy + self.ymin - x_1 = ii * self.dx + self.xmin - z_1 = kk * self.dz + self.zmin - x_3 = ii * self.dx + self.xmin - z_3 = (kk + 1) * self.dz + self.zmin - # if point 1 is in conductor - if self.conductors.in_conductor(x_1, y, z_1): - # if point 3 to the right is in conductor, length of the l_y[i, j] is zero - if self.conductors.in_conductor(x_3, y, z_3): - self.l_z[ii, jj, kk] = 0 - # if point 3 is not in conductor, length of l_y[i, j] - # is the fractional length - else: - self.l_z[ii, jj, kk] = seg_length(x_3, y, - self.conductors.intersec_z(x_3, - y, z_3), - x_3, y, z_3) - # if point 1 is not in conductor - else: - # if point 3 is in conductor, length of the l_y[i, j] - # is the fractional length - if self.conductors.in_conductor(x_3, y, z_3): - self.l_z[ii, jj, kk] = seg_length(x_1, y, z_1, x_1, y, - self.conductors.intersec_z(x_1, y, - z_1)) - # if point 3 is not in conductor, length of l_y[i, j] is dy - else: - self.l_z[ii, jj, kk] = self.dz - - # set to zero the length of very small cells - if tol > 0.: - self.l_x /= self.dx - self.l_y /= self.dy - self.l_z /= self.dz - - low_values_flags = abs(self.l_x) < tol - high_values_flags = abs(self.l_x - 1.) < tol - self.l_x[low_values_flags] = 0 - self.l_x[high_values_flags] = 1. - - low_values_flags = abs(self.l_y) < tol - high_values_flags = abs(self.l_y - 1.) < tol - self.l_y[low_values_flags] = 0 - self.l_y[high_values_flags] = 1. - - low_values_flags = abs(self.l_z) < tol - high_values_flags = abs(self.l_z - 1.) < tol - self.l_z[low_values_flags] = 0 - self.l_z[high_values_flags] = 1. - - self.l_x *= self.dx - self.l_y *= self.dy - self.l_z *= self.dz - - - """ - Function to compute the area of the cells of the conformal grid. - """ - @staticmethod - @jit('(f8[:,:,:], f8[:,:,:], f8[:,:,:], f8[:,:,:], f8[:,:,:], f8[:,:,:], f8[:,:,:], f8[:,:,:], f8[:,:,:], i4, i4, i4, f8, f8, f8)', nopython=True) - def compute_areas(l_x, l_y, l_z, Sxy, Syz, Szx, Sxy_red, Syz_red, Szx_red, nx, ny, nz, dx, dy, dz): - - for kk in range(nz + 1): - compute_areas_2D(l_x[:, :, kk], l_y[:, :, kk], Sxy[:, :, kk], Sxy_red[:, :, kk], - nx, ny, dx, dy) - - for ii in range(nx + 1): - compute_areas_2D(l_y[ii, :, :], l_z[ii, :, :], Syz[ii, :, :], Syz_red[ii, :, :], - ny, nz, dy, dz) - - for jj in range(ny + 1): - compute_areas_2D(l_x[:, jj, :], l_z[:, jj, :], Szx[:, jj, :], Szx_red[:, jj, :], - nx, nz, dx, dz) - - """ - Function to mark wich cells are interior (int), require extension (unst), - are on the boundary(bound), are available for intrusion (avail) - """ - @staticmethod - @jit('(f8[:,:,:], f8[:,:,:], f8[:,:,:], i4, i4, i4, f8, f8, f8, f8[:,:,:], f8[:,:,:], f8[:,:,:], b1[:,:,:], b1[:,:,:], b1[:,:,:], f8[:,:,:], f8[:,:,:], f8[:,:,:], b1[:,:,:], b1[:,:,:], b1[:,:,:], b1[:,:,:], b1[:,:,:], b1[:,:,:], b1[:,:,:], b1[:,:,:], b1[:,:,:])', nopython=True) - def mark_cells(l_x, l_y, l_z, nx, ny, nz, dx, dy, dz, Sxy, Syz, Szx, flag_int_cell_xy, flag_int_cell_yz, flag_int_cell_zx, - Sxy_stab, Syz_stab, Szx_stab, flag_unst_cell_xy, flag_unst_cell_yz, flag_unst_cell_zx, flag_bound_cell_xy, - flag_bound_cell_yz, flag_bound_cell_zx, flag_avail_cell_xy, flag_avail_cell_yz, flag_avail_cell_zx): - - for kk in range(nz + 1): - mark_cells_2D(l_x[:, :, kk], l_y[:, :, kk], nx, ny, dx, dy, Sxy[:, :, kk], - flag_int_cell_xy[:, :, kk], Sxy_stab[:, :, kk], flag_unst_cell_xy[:, :, kk], - flag_bound_cell_xy[:, :, kk], flag_avail_cell_xy[:, :, kk]) - - for ii in range(nx + 1): - mark_cells_2D(l_y[ii, :, :], l_z[ii, :, :], ny, nz, dy, dz, Syz[ii, :, :], - flag_int_cell_yz[ii, :, :], Syz_stab[ii, :, :], flag_unst_cell_yz[ii, :, :], - flag_bound_cell_yz[ii, :, :], flag_avail_cell_yz[ii, :, :]) - - for jj in range(ny + 1): - mark_cells_2D(l_x[:, jj, :], l_z[:, jj, :], nx, nz, dx, dz, Szx[:, jj, :], - flag_int_cell_zx[:, jj, :], Szx_stab[:, jj, :], flag_unst_cell_zx[:, jj, :], - flag_bound_cell_zx[:, jj, :], flag_avail_cell_zx[:, jj, :]) - - """ - Function to compute the extension of the unstable cells - """ - - def compute_extensions(self): - #breakpoint() - for kk in range(self.nz + 1): - #breakpoint() - Grid2D.compute_extensions(nx=self.nx, ny=self.ny, S=self.Sxy[:, :, kk], - flag_int_cell=self.flag_int_cell_xy[:, :, kk], - S_stab=self.Sxy_stab[:, :, kk], S_enl=self.Sxy_enl[:, :, kk], - S_red=self.Sxy_red[:, :, kk], - flag_unst_cell=self.flag_unst_cell_xy[:, :, kk], - flag_avail_cell=self.flag_avail_cell_xy[:, :, kk], - flag_ext_cell=self.flag_ext_cell_xy[:, :, kk], - flag_intr_cell=self.flag_intr_cell_xy[:,:,kk], - borrowing=self.borrowing_xy[:, :, kk], - kk=kk, l_verbose=False) - - for ii in range(self.nx + 1): - Grid2D.compute_extensions(nx=self.ny, ny=self.nz, S=self.Syz[ii, :, :], - flag_int_cell=self.flag_int_cell_yz[ii, :, :], - S_stab=self.Syz_stab[ii, :, :], S_enl=self.Syz_enl[ii, :, :], - S_red=self.Syz_red[ii, :, :], - flag_unst_cell=self.flag_unst_cell_yz[ii, :, :], - flag_avail_cell=self.flag_avail_cell_yz[ii, :, :], - flag_ext_cell=self.flag_ext_cell_yz[ii, :, :], - borrowing=self.borrowing_yz[ii, :, :], - flag_intr_cell=self.flag_intr_cell_yz[ii,:,:], - kk = ii, l_verbose=False) - - for jj in range(self.ny + 1): - Grid2D.compute_extensions(nx=self.nx, ny=self.nz, S=self.Szx[:, jj, :], - flag_int_cell=self.flag_int_cell_zx[:, jj, :], - S_stab=self.Szx_stab[:, jj, :], S_enl=self.Szx_enl[:, jj, :], - S_red=self.Szx_red[:, jj, :], - flag_unst_cell=self.flag_unst_cell_zx[:, jj, :], - flag_avail_cell=self.flag_avail_cell_zx[:, jj, :], - flag_ext_cell=self.flag_ext_cell_zx[:, jj, :], - flag_intr_cell=self.flag_intr_cell_zx[:,jj,:], - borrowing=self.borrowing_zx[:, jj, :], - kk = jj, l_verbose=False) From 9865cd57ac2ab1c1c0cc9b715ea33e710bc4780e Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:55:22 +0100 Subject: [PATCH 15/39] Delete build/lib/wakis/gridFIT3D.py --- build/lib/wakis/gridFIT3D.py | 841 ----------------------------------- 1 file changed, 841 deletions(-) delete mode 100644 build/lib/wakis/gridFIT3D.py diff --git a/build/lib/wakis/gridFIT3D.py b/build/lib/wakis/gridFIT3D.py deleted file mode 100644 index d6e447b..0000000 --- a/build/lib/wakis/gridFIT3D.py +++ /dev/null @@ -1,841 +0,0 @@ -# copyright ################################# # -# This file is part of the wakis Package. # -# Copyright (c) CERN, 2024. # -# ########################################### # - -import numpy as np -import pyvista as pv -from functools import partial -from scipy.optimize import least_squares - -from .field import Field - -try: - from mpi4py import MPI - imported_mpi = True -except ImportError: - imported_mpi = False - -class GridFIT3D: - """ - Class holding the grid information and - stl importing handling using PyVista - - Parameters - ---------- - xmin, xmax, ymin, ymax, zmin, zmax: float - extent of the domain. - Nx, Ny, Nz: int - number of cells per direction - stl_solids: dict, optional - stl files to import in the domain. - {'Solid 1': stl_1, 'Solid 2': stl_2, ...} - If stl files are not in the same folder, - add the path to the file name. - stl_materials: dict, optional - Material properties associated with stl - {'Solid 1': [eps1, mu1], - 'Solid 2': [eps1, mu1], - ...} - stl_rotate: list or dict, optional - Angle of rotation to apply to the stl models: [rot_x, rot_y, rot_z] - - if list, it will be applied to all stls in `stl_solids` - - if dict, it must contain the same keys as `stl_solids`, - indicating the rotation angle per stl - stl_scale: float or dict, optional - Scaling value to apply to the stl model to convert to [m] - - if float, it will be applied to all stl in `stl_solids` - - if dict, it must contain the same keys as `stl_solids` - tol: float, default 1e-3 - Tolerance factor for stl import, used in grid.select_enclosed_points. - Importing tolerance is computed by: tol*min(dx,dy,dz). - verbose: int or bool, default 1 - Enable verbose ouput on the terminal - """ - - def __init__(self, xmin, xmax, ymin, ymax, zmin, zmax, - Nx, Ny, Nz, - x=None, y=None, z=None, - use_mpi=False, - use_mesh_refinement=False, refinement_method='insert', refinement_tol=1e-8, - snap_points=None, snap_tol=1e-5, snap_solids=None, - stl_solids=None, stl_materials=None, - stl_rotate=[0., 0., 0.], stl_translate=[0., 0., 0.], stl_scale=1.0, - stl_colors=None, verbose=1, stl_tol=1e-3): - - if verbose: print('Generating grid...') - self.verbose = verbose - self.use_mpi = use_mpi - self.use_mesh_refinement = use_mesh_refinement - - # domain limits - self.xmin = xmin - self.xmax = xmax - self.ymin = ymin - self.ymax = ymax - self.zmin = zmin - self.zmax = zmax - self.Nx = Nx - self.Ny = Ny - self.Nz = Nz - self.dx = (xmax - xmin) / Nx - self.dy = (ymax - ymin) / Ny - self.dz = (zmax - zmin) / Nz - - # stl info - self.stl_solids = stl_solids - self.stl_materials = stl_materials - self.stl_rotate = stl_rotate - self.stl_translate = stl_translate - self.stl_scale = stl_scale - self.stl_colors = stl_colors - if stl_solids is not None: - self._prepare_stl_dicts() - - # MPI subdivide domain - if self.use_mpi: - self.ZMIN = None - self.ZMAX = None - self.NZ = None - self.Z = None - if imported_mpi: - self.mpi_initialize() - if self.verbose: print(f"MPI initialized for {self.rank} of {self.size}") - else: - raise ImportError("*** mpi4py is required when use_mpi=True but was not found") - - # primal Grid G base axis x, y, z - self.x = x - self.y = y - self.z = z - self.refinement_method = refinement_method - self.snap_points = snap_points - self.snap_tol = snap_tol - self.snap_solids = snap_solids # if None, use all stl_solids - - if self.x is not None and self.y is not None and self.z is not None: - # allow user to set the grid axis manually - self.Nx = len(self.x) - 1 - self.Ny = len(self.y) - 1 - self.Nz = len(self.z) - 1 - self.dx = np.min(np.diff(self.x)) - self.dy = np.min(np.diff(self.y)) - self.dz = np.min(np.diff(self.z)) - - elif self.use_mesh_refinement: - if verbose: print('Applying mesh refinement...') - if self.snap_points is None and stl_solids is not None: - self.compute_snap_points(snap_solids=snap_solids, snap_tol=snap_tol) - self.refine_xyz_axis(method=refinement_method, tol=refinement_tol) # obtain self.x, self.y, self.z - else: - self.x = np.linspace(self.xmin, self.xmax, self.Nx+1) - self.y = np.linspace(self.ymin, self.ymax, self.Ny+1) - self.z = np.linspace(self.zmin, self.zmax, self.Nz+1) - - # grid G and tilde grid ~G, lengths and inverse areas - self.compute_grid() - - # tolerance for stl import tol*min(dx,dy,dz) - if verbose: print('Importing STL solids...') - self.tol = stl_tol - if stl_solids is not None: - self.mark_cells_in_stl() - if stl_colors is None: - self.assign_colors() - - def compute_grid(self): - X, Y, Z = np.meshgrid(self.x, self.y, self.z, indexing='ij') - self.grid = pv.StructuredGrid(X.transpose(), Y.transpose(), Z.transpose()) - - self.L = Field(self.Nx, self.Ny, self.Nz) - self.L.field_x = X[1:, 1:, 1:] - X[:-1, :-1, :-1] - self.L.field_y = Y[1:, 1:, 1:] - Y[:-1, :-1, :-1] - self.L.field_z = Z[1:, 1:, 1:] - Z[:-1, :-1, :-1] - - self.iA = Field(self.Nx, self.Ny, self.Nz) - self.iA.field_x = np.divide(1.0, self.L.field_y * self.L.field_z) - self.iA.field_y = np.divide(1.0, self.L.field_x * self.L.field_z) - self.iA.field_z = np.divide(1.0, self.L.field_x * self.L.field_y) - - # tilde grid ~G - self.tx = (self.x[1:]+self.x[:-1])/2 - self.ty = (self.y[1:]+self.y[:-1])/2 - self.tz = (self.z[1:]+self.z[:-1])/2 - - self.tx = np.append(self.tx, self.tx[-1]) - self.ty = np.append(self.ty, self.ty[-1]) - self.tz = np.append(self.tz, self.tz[-1]) - - tX, tY, tZ = np.meshgrid(self.tx, self.ty, self.tz, indexing='ij') - - self.tL = Field(self.Nx, self.Ny, self.Nz) - self.tL.field_x = tX[1:, 1:, 1:] - tX[:-1, :-1, :-1] - self.tL.field_y = tY[1:, 1:, 1:] - tY[:-1, :-1, :-1] - self.tL.field_z = tZ[1:, 1:, 1:] - tZ[:-1, :-1, :-1] - - self.itA = Field(self.Nx, self.Ny, self.Nz) - aux = self.tL.field_y * self.tL.field_z - self.itA.field_x = np.divide(1.0, aux, out=np.zeros_like(aux), where=aux!=0) - aux = self.tL.field_x * self.tL.field_z - self.itA.field_y = np.divide(1.0, aux, out=np.zeros_like(aux), where=aux!=0) - aux = self.tL.field_x * self.tL.field_y - self.itA.field_z = np.divide(1.0, aux, out=np.zeros_like(aux), where=aux!=0) - del aux - - def mpi_initialize(self): - comm = MPI.COMM_WORLD # Get MPI communicator - self.comm = comm - self.rank = self.comm.Get_rank() - self.size = self.comm.Get_size() - - # Error handling for Nz < size - if self.Nz < self.size: - raise ValueError(f"Nz ({self.Nz}) must be greater than or equal to the number of MPI processes ({self.size}).") - - # global z quantities [ALLCAPS] - self.ZMIN = self.zmin - self.ZMAX = self.zmax - self.NZ = self.Nz - self.Nz%(self.size) # ensure multiple of MPI size - self.Z = np.linspace(self.ZMIN, self.ZMAX, self.NZ+1)[:-1] + self.dz/2 - - if self.verbose and self.rank==0: - print(f"Global grid ZMIN={self.ZMIN}, ZMAX={self.ZMAX}, NZ={self.NZ}") - - # MPI subdomain quantities - self.Nz = self.NZ // (self.size) - self.dz = (self.ZMAX - self.ZMIN) / self.NZ - self.zmin = self.rank * self.Nz * self.dz + self.ZMIN - self.zmax = (self.rank+1) * self.Nz * self.dz + self.ZMIN - - if self.verbose: print(f"MPI rank {self.rank} of {self.size} initialized with \ - zmin={self.zmin}, zmax={self.zmax}, Nz={self.Nz}") - # Add ghost cells - self.n_ghosts = 1 - if self.rank > 0: - self.zmin += - self.n_ghosts * self.dz - self.Nz += self.n_ghosts - if self.rank < (self.size-1): - self.zmax += self.n_ghosts * self.dz - self.Nz += self.n_ghosts - - # Support for single core - if self.rank == 0 and self.size == 1: - self.zmax += self.n_ghosts * self.dz - self.Nz += self.n_ghosts - - def mpi_gather_asGrid(self): - _grid = None - if self.rank == 0: - print(f"Generating global grid from {self.ZMIN} to {self.ZMAX}") - _grid = GridFIT3D(self.xmin, self.xmax, - self.ymin, self.ymax, - self.ZMIN, self.ZMAX, - self.Nx, self.Ny, self.NZ, - use_mpi=False, - stl_solids=self.stl_solids, - stl_materials=self.stl_materials, - stl_scale=self.stl_scale, - stl_rotate=self.stl_rotate, - stl_translate=self.stl_translate, - stl_colors=self.stl_colors, - verbose=self.verbose, - tol=self.tol, - ) - return _grid - - def _prepare_stl_dicts(self): - if type(self.stl_solids) is not dict: - if type(self.stl_solids) is str: - self.stl_solids = {'Solid 1' : self.stl_solids} - else: - raise Exception('Attribute `stl_solids` must contain a string or a dictionary') - - if type(self.stl_rotate) is not dict: - # if not a dict, the same values will be applied to all solids - stl_rotate = {} - for key in self.stl_solids.keys(): - stl_rotate[key] = self.stl_rotate - self.stl_rotate = stl_rotate - - if type(self.stl_scale) is not dict: - # if not a dict, the same values will be applied to all solids - stl_scale = {} - for key in self.stl_solids.keys(): - stl_scale[key] = self.stl_scale - self.stl_scale = stl_scale - - if type(self.stl_translate) is not dict: - # if not a dict, the same values will be applied to all solids - stl_translate = {} - for key in self.stl_solids.keys(): - stl_translate[key] = self.stl_translate - self.stl_translate = stl_translate - - def mark_cells_in_stl(self): - # Obtain masks with grid cells inside each stl solid - tol = np.min([self.dx, self.dy, self.dz])*self.tol - for key in self.stl_solids.keys(): - - surf = self.read_stl(key) - - # mark cells in stl [True == in stl, False == out stl] - try: - select = self.grid.select_enclosed_points(surf, tolerance=tol) - except: - select = self.grid.select_enclosed_points(surf, tolerance=tol, check_surface=False) - self.grid[key] = select.point_data_to_cell_data()['SelectedPoints'] > tol - - def read_stl(self, key): - # import stl - surf = pv.read(self.stl_solids[key]) - - # rotate - surf = surf.rotate_x(self.stl_rotate[key][0]) - surf = surf.rotate_y(self.stl_rotate[key][1]) - surf = surf.rotate_z(self.stl_rotate[key][2]) - - # translate - surf = surf.translate(self.stl_translate[key]) - - # scale - surf = surf.scale(self.stl_scale[key]) - - return surf - - def compute_snap_points(self, snap_solids=None, snap_tol=1e-8): - if self.verbose: print('* Calculating snappy points...') - # Support for user-defined stl_keys as list - if snap_solids is None: - snap_solids = self.stl_solids.keys() - - # Union of all the surfaces - # [TODO]: should use | for union instead or +? - model = None - for key in snap_solids: - solid = self.read_stl(key) - if model is None: - model = solid - else: - model = model + solid - - edges = model.extract_feature_edges(boundary_edges=True, manifold_edges=False) - - # Extract points lying in the X-Z plane (Y ≈ 0) - xz_plane_points = edges.points[np.abs(edges.points[:, 1]) < snap_tol] - # Extract points lying in the Y-Z plane (X ≈ 0) - yz_plane_points = edges.points[np.abs(edges.points[:, 0]) < snap_tol] - # Extract points lying in the X-Y plane (Z ≈ 0) - xy_plane_points = edges.points[np.abs(edges.points[:, 2]) < 1e-5] - - self.snap_points = np.r_[xz_plane_points, yz_plane_points, xy_plane_points] - - # get the unique x, y, z coordinates - x_snaps = np.unique(np.round(self.snap_points[:, 0], 5)) - y_snaps = np.unique(np.round(self.snap_points[:, 1], 5)) - z_snaps = np.unique(np.round(self.snap_points[:, 2], 5)) - - # Include simulation domain bounds - self.x_snaps = np.unique(np.concatenate(([self.xmin], x_snaps, [self.xmax]))) - self.y_snaps = np.unique(np.concatenate(([self.ymin], y_snaps, [self.ymax]))) - self.z_snaps = np.unique(np.concatenate(([self.zmin], z_snaps, [self.zmax]))) - - def plot_snap_points(self, snap_solids=None, snap_tol=1e-8): - # TODO - # Support for user-defined stl_keys as list - if snap_solids is None: - snap_solids = self.stl_solids.keys() - - # Union of all the surfaces - # [TODO]: should use | for union instead or +? - model = None - for key in snap_solids: - solid = self.read_stl(key) - if model is None: - model = solid - else: - model = model + solid - - edges = model.extract_feature_edges(boundary_edges=True, manifold_edges=False) - - # Extract points lying in the X-Z plane (Y ≈ 0) - xz_plane_points = edges.points[np.abs(edges.points[:, 1]) < snap_tol] - # Extract points lying in the Y-Z plane (X ≈ 0) - yz_plane_points = edges.points[np.abs(edges.points[:, 0]) < snap_tol] - # Extract points lying in the X-Y plane (Z ≈ 0) - xy_plane_points = edges.points[np.abs(edges.points[:, 2]) < 1e-5] - - xz_cloud = pv.PolyData(xz_plane_points) - yz_cloud = pv.PolyData(yz_plane_points) - xy_cloud = pv.PolyData(xy_plane_points) - - pl = pv.Plotter() - pl.add_mesh(model, color='white', opacity=0.5, label='base STL') - pl.add_mesh(edges, color='black', line_width=5, opacity=0.8,) - pl.add_mesh(xz_cloud, color='green', point_size=20, render_points_as_spheres=True, label='XZ plane points') - pl.add_mesh(yz_cloud, color='orange', point_size=20, render_points_as_spheres=True, label='YZ plane points') - pl.add_mesh(xy_cloud, color='magenta', point_size=20, render_points_as_spheres=True, label='XY plane points') - pl.add_legend() - pl.show() - - def refine_axis(self, xmin, xmax, Nx, x_snaps, - method='insert', tol=1e-12): - - # Loss function to minimize cell size spread - def loss_function(x, x0, is_snap): - # avoid moving snap points - penalty_snap = np.sum((x[is_snap] - x0[is_snap])**2) * 1000 - # avoid gaps < uniform gap - dx = np.diff(x) - threshold = 1/(len(x)-1) # or a hardcoded `min_spacing` - penalty_small_gaps = np.sum((threshold - dx[dx < threshold])**2) - # avoid large spread in gap length - dx = np.diff(x) - penalty_variance = np.std(dx) * 10 - #return penalty_snap + penalty_small_gaps + penalty_variance - return np.hstack([penalty_snap, penalty_small_gaps, penalty_variance]) - - # Uniformly distributed points as initial guess - x_snaps = (x_snaps-xmin)/(xmax-xmin) # normalize to [0,1] - - if method == 'insert': - x0 = np.unique(np.append(x_snaps, np.linspace(0, 1, Nx - len(x_snaps)))) - - elif method == 'neighbor': - x = np.linspace(0, 1, Nx) - dx = np.diff(x)[0] - mask = np.zeros_like(x, dtype=bool) - i=0 - for s in x_snaps: - m = np.isclose(x, s, rtol=0.0, atol=dx/2) - if np.sum(m)>0: - x[np.argmax(m)] = s - x0 = x.copy() - - elif method == 'subdivision': - # x = snaps - while len(x) < Nx: - #idx of segments sorted min -> max - idx_max_diffs = np.argsort(np.diff(x))[-1] # take bigger - - print(f"Bigger segment starts at {x[idx_max_diffs]}") - # compute new point in the middle of the segment - val = x[idx_max_diffs] + (x[idx_max_diffs + 1] - x[idx_max_diffs]) / 2 - - # insert the new point - x = np.insert(x, idx_max_diffs+1, val) - x = np.unique(x) - print(f"Inserted point {val} at index {idx_max_diffs}") - x0 = x.copy() - else: - raise ValueError(f"Method {method} not supported. Use 'insert', 'neighbor' or 'subdivision'.") - - # minimize segment length spread for the test points - is_snap = np.isin(x0, x_snaps) - result = least_squares(loss_function, - x0=x0.copy(), - bounds=(0,1),#(zmin, zmax), - jac='3-point', - method='dogbox', - loss='arctan', - gtol=tol, - ftol=tol, - xtol=tol, - verbose=1, - args=(x0.copy(), is_snap.copy()), - ) - # transform back to [xmin, xmax] - return result.x*(xmax-xmin)+xmin - - def refine_xyz_axis(self, method='insert', tol=1e-6): - '''Refine the grid in the x, y, z axis - using the snap points extracted from the stl solids. - The snap points are used to refine the grid ''' - - if self.verbose: print(f'* Refining x axis with {len(self.x_snaps)} snaps...') - self.x = self.refine_axis(self.xmin, self.xmax, self.Nx+1, self.x_snaps, - method=method, tol=tol) - - if self.verbose: print(f'* Refining y axis with {len(self.y_snaps)} snaps...') - self.y = self.refine_axis(self.ymin, self.ymax, self.Ny+1, self.y_snaps, - method=method, tol=tol) - - if self.verbose: print(f'* Refining z axis with {len(self.z_snaps)} snaps...') - self.z = self.refine_axis(self.zmin, self.zmax, self.Nz+1, self.z_snaps, - method=method, tol=tol) - - self.Nx = len(self.x) - 1 - self.Ny = len(self.y) - 1 - self.Nz = len(self.z) - 1 - self.dx = np.min(np.diff(self.x)) #TODO: should this be an array? - self.dy = np.min(np.diff(self.y)) - self.dz = np.min(np.diff(self.z)) - - print(f"Refined grid: Nx = {len(self.x)}, Ny ={len(self.y)}, Nz = {len(self.z)}") - - def assign_colors(self): - '''Classify colors assigned to each solid - based on the categories in `material_colors` dict - inside `materials.py` - ''' - self.stl_colors = {} - - for key in self.stl_solids: - mat = self.stl_materials[key] - if type(mat) is str: - self.stl_colors[key] = mat - elif len(mat) == 2: - if mat[0] is np.inf: #eps_r - self.stl_colors[key] = 'pec' - elif mat[0] > 1.0: #eps_r - self.stl_colors[key] = 'dielectric' - else: - self.stl_colors[key] = 'vacuum' - elif len(mat) == 3: - self.stl_colors[key] = 'lossy metal' - else: - self.stl_colors[key] = 'other' - - def plot_solids(self, bounding_box=False, show_grid=False, anti_aliasing=None, - opacity=1.0, specular=0.5, offscreen=False, **kwargs): - """ - Generates a 3D visualization of the imported STL geometries using PyVista. - - Parameters: - ----------- - bounding_box : bool, optional - If True, adds a bounding box around the plotted geometry (default: False). - - show_grid : bool, optional - If True, adds the grid's mesh wireframe to the display (default: False). - - anti_aliasing : str or None, optional - Enables anti-aliasing if provided. Valid values depend on PyVista settings (default: None). - - opacity : float, optional - Controls the transparency of the plotted solids. A value of 1.0 is fully opaque, - while 0.0 is fully transparent (default: 1.0). - - specular : float, optional - Adjusts the specular lighting effect on the surface. Higher values increase shininess (default: 0.5). - - **kwargs : dict - Additional keyword arguments passed to `pyvista.add_mesh()`, allowing customization of the mesh rendering. - - Notes: - ------ - - Colors are determined by the `GridFIT3D.stl_colors` attribute dictionary if not None - - Solids labeled as 'vacuum' are rendered with a default opacity of 0.3 for visibility.e. - - The camera is positioned at an angle to provide better depth perception. - - If `bounding_box=True`, a bounding box is drawn around the model. - - If `anti_aliasing` is specified, it is enabled to improve rendering quality. - - """ - - from .materials import material_colors - - pl = pv.Plotter() - pl.add_mesh(self.grid, opacity=0., name='grid', show_scalar_bar=False) - for key in self.stl_solids: - try: - color = material_colors[self.stl_colors[key]] # match library e.g. 'vacuum' - except: - color = self.stl_colors[key] # specifies color e.g. 'tab:red' - - if self.stl_colors[key] == 'vacuum' or self.stl_materials[key] == 'vacuum': - _opacity = 0.3 - else: - _opacity = opacity - pl.add_mesh(self.read_stl(key), color=color, - opacity=_opacity, specular=specular, smooth_shading=True, - **kwargs) - - pl.set_background('mistyrose', top='white') - try: pl.add_logo_widget('../docs/img/wakis-logo-pink.png') - except: pass - pl.camera_position = 'zx' - pl.camera.azimuth += 30 - pl.camera.elevation += 30 - pl.add_axes() - - if anti_aliasing is not None: - pl.enable_anti_aliasing(anti_aliasing) - - if bounding_box: - pl.add_bounding_box() - - if show_grid: - pl.add_mesh(self.grid, style='wireframe', color='grey', opacity=0.3, name='grid') - - if offscreen: - pl.export_html('grid_plot_solids.html') - else: - pl.show() - - def plot_stl_mask(self, stl_solid, cmap='viridis', bounding_box=True, show_grid=True, - add_stl='all', stl_opacity=0., stl_colors=None, - xmax=None, ymax=None, zmax=None, - anti_aliasing='ssaa', offscreen=False): - - """ - Interactive 3D visualization of the structured grid mask and imported STL geometries. - - This routine uses PyVista to display the grid scalar field corresponding to a - chosen STL mask. It provides interactive slider widgets to clip the domain - along the X, Y, and Z directions. At each slider position, the clipped - scalar field is shown with a colormap while the grid structure is shown - as a 2D slice in wireframe. Optionally, one or more STL geometries can - be added to the scene, along with a bounding box of the simulation domain. - - Parameters - ---------- - stl_solid : str - Key name of the `stl_solids` dictionary to retrieve the mask for - visualization (used as the scalar field). - cmap : str, default 'viridis' - Colormap used to visualize the clipped scalar values. - bounding_box : bool, default True - If True, add a static wireframe bounding box of the simulation domain. - show_grid : bool, default True - If True, adds the computational grid overlay on the clipped slice - add_stl : {'all', str, list[str]}, default 'all' - STL geometries to add: - * 'all' → add all STL solids found in `self.stl_solids` - * str → add a single STL solid by key - * list → add a list of STL solids by key - If None, no STL surfaces are added. - stl_opacity : float, default 0.0 - Opacity of the STL surfaces (0 = fully transparent, 1 = fully opaque). - stl_colors : str, list[str], dict, or None, default None - Color(s) of the STL surfaces: - * str → single color for all STL surfaces - * list → per-solid colors, in order - * dict → mapping from STL key to color (using `material_colors`) - * None → use default colors from `self.stl_colors` - xmax, ymax, zmax : float, optional - Initial clipping positions along each axis. If None, use the - maximum domain extent. - anti_aliasing : {'ssaa', 'fxaa', None}, default 'ssaa' - Anti-aliasing mode passed to `pl.enable_anti_aliasing`. - offscreen : bool, default False - If True, render offscreen and export the scene to - ``grid_stl_mask_.html``. If False, open an interactive window. - - Notes - ----- - * Three sliders (X, Y, Z) control clipping of the scalar field by a box - along the respective axis. The clipped scalar field is shown with the - given colormap. A simultaneous 2D slice of the grid is displayed in - wireframe at the clip location. - * STL solids can be visualized in transparent mode to show the relation - between the structured grid and the geometry. - * A static domain bounding box can be added for reference. - * Camera, background, and lighting are pre-configured for clarity - """ - from .materials import material_colors - if stl_colors is None: - stl_colors = self.stl_colors - - if xmax is None: xmax = self.xmax - if ymax is None: ymax = self.ymax - if zmax is None: zmax = self.zmax - - pv.global_theme.allow_empty_mesh = True - pl = pv.Plotter() - vals = {'x':xmax, 'y':ymax, 'z':zmax} - - # --- Update function --- - def update_clip(val, axis="x"): - vals[axis] = val - # define bounds dynamically - if axis == "x": - slice_obj = self.grid.slice(normal="x", origin=(val, 0, 0)) - elif axis == "y": - slice_obj = self.grid.slice(normal="y", origin=(0, val, 0)) - else: # z - slice_obj = self.grid.slice(normal="z", origin=(0, 0, val)) - - # add clipped volume (scalars) - pl.add_mesh( - self.grid.clip_box(bounds=(self.xmin, vals['x'], - self.ymin, vals['y'], - self.zmin, vals['z']), invert=False), - scalars=stl_solid, - cmap=cmap, - name="clip", - ) - - # add slice wireframe (grid structure) - if show_grid: - pl.add_mesh(slice_obj, style="wireframe", color="grey", name="slice") - - # Plot stl surface(s) - if add_stl is not None: - if type(add_stl) is str: #add all stl solids - if add_stl.lower() == 'all': - for i, key in enumerate(self.stl_solids): - surf = self.read_stl(key) - if type(stl_colors) is dict: - if type(stl_colors[key]) is str: - pl.add_mesh(surf, color=material_colors[stl_colors[key]], opacity=stl_opacity, silhouette=dict(color=material_colors[stl_colors[key]]), name=key) - else: - pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=dict(color=stl_colors[key]), name=key) - elif type(stl_colors) is list: - pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, silhouette=dict(color=stl_colors[i]), name=key) - else: - pl.add_mesh(surf, color='white', opacity=stl_opacity, silhouette=True, name=key) - else: #add 1 selected stl solid - key = add_stl - surf = self.read_stl(key) - if type(stl_colors[key]) is str: - pl.add_mesh(surf, color=material_colors[stl_colors[key]], opacity=stl_opacity, silhouette=dict(color=material_colors[stl_colors[key]]), name=key) - else: - pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=dict(color=stl_colors[key]), name=key) - - elif type(add_stl) is list: #add selected list of stl solids - for i, key in enumerate(add_stl): - surf = self.read_stl(key) - if type(stl_colors[key]) is dict: - if type(stl_colors) is str: - pl.add_mesh(surf, color=material_colors[stl_colors[key]], opacity=stl_opacity, silhouette=dict(color=material_colors[stl_colors[key]]), name=key) - else: - pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=dict(color=stl_colors[key]), name=key) - elif type(stl_colors) is list: - pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, silhouette=dict(color=stl_colors[i]), name=key) - else: - pl.add_mesh(surf, color='white', opacity=stl_opacity, silhouette=True, name=key) - - # --- Sliders (placed side-by-side vertically) --- - pl.add_slider_widget( - lambda val: update_clip(val, "x"), - [self.xmin, self.xmax], - value=xmax, title="X Clip", - pointa=(0.8, 0.8), pointb=(0.95, 0.8), # top-right - style='modern', - ) - - pl.add_slider_widget( - lambda val: update_clip(val, "y"), - [self.ymin, self.ymax], - value=ymax, title="Y Clip", - pointa=(0.8, 0.6), pointb=(0.95, 0.6), # middle-right - style='modern', - ) - - pl.add_slider_widget( - lambda val: update_clip(val, "z"), - [self.zmin, self.zmax], - value=zmax, title="Z Clip", - pointa=(0.8, 0.4), pointb=(0.95, 0.4), # lower-right - style='modern', - ) - - # Camera orientation - pl.camera_position = 'zx' - pl.camera.azimuth += 30 - pl.camera.elevation += 30 - pl.set_background('mistyrose', top='white') - try: pl.add_logo_widget('../docs/img/wakis-logo-pink.png') - except: pass - pl.add_axes() - pl.enable_3_lights() - pl.enable_anti_aliasing(anti_aliasing) - - if bounding_box: - pl.add_mesh(pv.Box(bounds=(self.xmin, self.xmax, self.ymin, self.ymax, self.zmin, self.zmax)), - style="wireframe", color="black", line_width=2, name="domain_box") - - if offscreen: - pl.export_html(f'grid_stl_mask_{stl_solid}.html') - else: - pl.show() - - def inspect(self, add_stl=None, stl_opacity=0.5, stl_colors=None, - anti_aliasing='ssaa', offscreen=False): - - '''3D plot using pyvista to visualize - the structured grid and - the imported stl geometries - - Parameters - --- - add_stl: str or list, optional - List or str of stl solids to add to the plot by `pv.add_mesh` - stl_opacity: float, default 0.1 - Opacity of the stl surfaces (0 - Transparent, 1 - Opaque) - stl_colors: str or list of str, default 'white' - Color of the stl surfaces - ''' - from .materials import material_colors - if stl_colors is None: - stl_colors = self.stl_colors - - pv.global_theme.allow_empty_mesh = True - pl = pv.Plotter() - pl.add_mesh(self.grid, show_edges=True, cmap=['white', 'white'], name='grid') - def clip(widget): - # Plot structured grid - b = widget.bounds - x = self.x[np.logical_and(self.x>=b[0], self.x<=b[1])] - y = self.y[np.logical_and(self.y>=b[2], self.y<=b[3])] - z = self.z[np.logical_and(self.z>=b[4], self.z<=b[5])] - X, Y, Z = np.meshgrid(x, y, z, indexing='ij') - grid = pv.StructuredGrid(X.transpose(), Y.transpose(), Z.transpose()) - - pl.add_mesh(grid, show_edges=True, cmap=['white', 'white'], name='grid') - # Plot stl surface(s) - if add_stl is not None: #add 1 selected stl solid - if type(add_stl) is str: - key = add_stl - surf = self.read_stl(key) - surf = surf.clip_box(widget.bounds, invert=False) - if type(stl_colors[key]) is str: - pl.add_mesh(surf, color=material_colors[self.stl_colors[key]], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) - else: - pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) - - elif type(add_stl) is list: #add selected list of stl solids - for i, key in enumerate(add_stl): - surf = self.read_stl(key) - surf = surf.clip_box(widget.bounds, invert=False) - if type(stl_colors) is dict: - if type(stl_colors[key]) is str: - pl.add_mesh(surf, color=material_colors[stl_colors[key]], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) - else: - pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) - elif type(stl_colors) is list: - pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) - else: - pl.add_mesh(surf, color='white', opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) - - else: #add all stl solids - for i, key in enumerate(self.stl_solids): - surf = self.read_stl(key) - surf = surf.clip_box(widget.bounds, invert=False) - if type(stl_colors) is dict: - if type(stl_colors[key]) is str: - pl.add_mesh(surf, color=material_colors[stl_colors[key]], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) - else: - pl.add_mesh(surf, color=stl_colors[key], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) - elif type(stl_colors) is list: - pl.add_mesh(surf, color=stl_colors[i], opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) - else: - pl.add_mesh(surf, color='white', opacity=stl_opacity, silhouette=True, smooth_shading=True, name=key) - - _ = pl.add_box_widget(callback=clip, rotation_enabled=False) - - # Camera orientation - pl.camera_position = 'zx' - pl.camera.azimuth += 30 - pl.camera.elevation += 30 - pl.set_background('mistyrose', top='white') - try: pl.add_logo_widget('../docs/img/wakis-logo-pink.png') - except: pass - #pl.camera.zoom(zoom) - pl.add_axes() - pl.enable_3_lights() - pl.enable_anti_aliasing(anti_aliasing) - - if offscreen: - pl.export_html('grid_inspect.html') - else: - pl.show() \ No newline at end of file From 8ea836ff6a8380d6d2bd3aa66e0b3bf8cf72d5e8 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:55:31 +0100 Subject: [PATCH 16/39] Delete build/lib/wakis/materials.py --- build/lib/wakis/materials.py | 50 ------------------------------------ 1 file changed, 50 deletions(-) delete mode 100644 build/lib/wakis/materials.py diff --git a/build/lib/wakis/materials.py b/build/lib/wakis/materials.py deleted file mode 100644 index 155d8e2..0000000 --- a/build/lib/wakis/materials.py +++ /dev/null @@ -1,50 +0,0 @@ -# copyright ################################# # -# This file is part of the wakis Package. # -# Copyright (c) CERN, 2024. # -# ########################################### # - -''' -Material library dictionary - -Format (non-conductive): -{ - 'material key' : [eps_r, mu_r], -} - -Format (conductive): -{ - 'material key' : [eps_r, mu_r, sigma[S/m]], -} - -! Note: -* 'material key' in lower case only -* eps = eps_r*eps_0 and mu = mu_r*mu_0 -''' - -import numpy as np - -material_lib = { - 'pec' : [np.inf, 1.], - - 'vacuum' : [1.0, 1.0], - - 'dielectric' : [10., 1.0], - - 'lossy metal' : [10, 1.0, 10], - - 'copper' : [5.8e+07, 1.0, 5.8e+07], - - 'berillium' : [2.5e+07, 1.0, 2.5e+07], -} - -material_colors = { - 'pec' : 'silver', - - 'vacuum' : 'tab:blue', - - 'dielectric' : 'tab:green', - - 'lossy metal' : 'tab:orange', - - 'other' : 'white', -} \ No newline at end of file From 7dfbf304b4219d3d41df2ba13b5ed402f51b153f Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:55:39 +0100 Subject: [PATCH 17/39] Delete build/lib/wakis/solver2D.py --- build/lib/wakis/solver2D.py | 484 ------------------------------------ 1 file changed, 484 deletions(-) delete mode 100644 build/lib/wakis/solver2D.py diff --git a/build/lib/wakis/solver2D.py b/build/lib/wakis/solver2D.py deleted file mode 100644 index 2b89458..0000000 --- a/build/lib/wakis/solver2D.py +++ /dev/null @@ -1,484 +0,0 @@ -import numpy as np -from scipy.constants import c as c_light, epsilon_0 as eps_0, mu_0 as mu_0 -from pmlBlock2D import PmlBlock2D - - -def eq(a, b, tol=1e-8): - return abs(a - b) < tol - - -def neq(a, b, tol=1e-8): - return not eq(a, b, tol) - - -class EMSolver2D: - def __init__(self, grid, sol_type, cfln, i_s, j_s, bc_low, bc_high, - N_pml_low=None, N_pml_high=None): - self.grid = grid - self.type = type - self.cfln = cfln - - self.dt = cfln / (c_light * np.sqrt(1 / self.grid.dx ** 2 + 1 / self.grid.dy ** 2)) - self.dx = self.grid.dx - self.dy = self.grid.dy - self.Nx = self.grid.nx - self.Ny = self.grid.ny - self.sol_type = sol_type - - self.N_pml_low = np.zeros(2, dtype=int) - self.N_pml_high = np.zeros(2, dtype=int) - self.bc_low = bc_low - self.bc_high = bc_high - - if bc_low[0] == 'pml': - self.N_pml_low[0] = 10 if N_pml_low is None else N_pml_low[0] - if bc_low[1] == 'pml': - self.N_pml_low[1] = 10 if N_pml_low is None else N_pml_low[1] - if bc_high[0] == 'pml': - self.N_pml_high[0] = 10 if N_pml_high is None else N_pml_high[0] - if bc_high[1] == 'pml': - self.N_pml_high[1] = 10 if N_pml_high is None else N_pml_high[1] - - self.blocks = [] - self.pml_ly = None - self.pml_lx = None - self.pml_rx = None - self.pml_ry = None - self.pml_lxly = None - self.pml_rxly = None - self.pml_lxry = None - self.pml_rxry = None - - if bc_low[0] is 'pml': - self.pml_lx = PmlBlock2D(self.N_pml_low[0], self.Ny, self.dt, self.dx, self.dy) - self.blocks.append(self.pml_lx) - if bc_low[1] is 'pml': - self.pml_lxly = PmlBlock2D(self.N_pml_low[0], self.N_pml_low[1], self.dt, self.dx, self.dy) - self.blocks.append(self.pml_lxly) - if bc_high[1] is 'pml': - self.pml_lxry = PmlBlock2D(self.N_pml_low[0], self.N_pml_high[1], self.dt, self.dx, self.dy) - self.blocks.append(self.pml_lxry) - - if bc_high[0] is 'pml': - self.pml_rx = PmlBlock2D(self.N_pml_high[0], self.Ny, self.dt, self.dx, self.dy) - self.blocks.append(self.pml_rx) - if bc_low[1] is 'pml': - self.pml_rxry = PmlBlock2D(self.N_pml_high[0], self.N_pml_high[1], self.dt, self.dx, self.dy) - self.blocks.append(self.pml_rxry) - if bc_high[1] is 'pml': - self.pml_rxly = PmlBlock2D(self.N_pml_high[0], self.N_pml_low[1], self.dt, self.dx, self.dy) - self.blocks.append(self.pml_rxly) - - if bc_low[1] is 'pml': - self.pml_ly = PmlBlock2D(self.Nx, self.N_pml_low[1], self.dt, self.dx, self.dy) - self.blocks.append(self.pml_ly) - - if bc_high[1] is 'pml': - self.pml_ry = PmlBlock2D(self.Nx, self.N_pml_high[1], self.dt, self.dx, self.dy) - self.blocks.append(self.pml_ry) - - self.organize_pmls() - self.alpha_pml = 3 - self.R0_pml = 0.001 - - self.assemble_conductivities_pmls() - - self.assemble_coeffs_pmls() - - self.N_tot_x = self.Nx + self.N_pml_low[0] + self.N_pml_high[0] - self.N_tot_y = self.Ny + self.N_pml_low[0] + self.N_pml_high[0] - self.sol_type = sol_type - - self.Ex = np.zeros((self.Nx, self.Ny + 1)) - self.Ey = np.zeros((self.Nx + 1, self.Ny)) - self.Hz = np.zeros((self.Nx, self.Ny)) - self.V_new = np.zeros((self.Nx, self.Ny)) - self.Jx = np.zeros((self.Nx, self.Ny + 1)) - self.Jy = np.zeros((self.Nx + 1, self.Ny)) - - self.rho = np.zeros((self.Nx, self.Ny)) - - if (sol_type is not 'FDTD') and (sol_type is not 'DM') and (sol_type is not 'ECT'): - raise ValueError("sol_type must be:\n" + - "\t'FDTD' for standard staircased FDTD\n" + - "\t'DM' for Dey-Mittra conformal FDTD\n" + - "\t'ECT' for Enlarged Cell Technique conformal FDTD") - - if sol_type is 'DM' or sol_type is 'ECT': - self.Vxy = np.zeros((self.Nx, self.Ny)) - if sol_type is 'ECT': - self.V_enl = np.zeros((self.Nx, self.Ny)) - - if sol_type is 'ECT' or sol_type is 'DM': - self.C1 = self.dt / mu_0 - self.C4 = self.dt / (eps_0 * self.dy) - self.C5 = self.dt / (eps_0 * self.dx) - self.C3 = self.dt / eps_0 - self.C6 = self.dt / eps_0 - - if sol_type is 'FDTD': - Z_0 = np.sqrt(mu_0 / eps_0) - - self.C1 = self.dt / (self.dx * mu_0) - self.C2 = self.dt / (self.dy * mu_0) - self.C4 = self.dt / (self.dy * eps_0) - self.C5 = self.dt / (self.dx * eps_0) - self.C3 = self.dt / eps_0 - self.C6 = self.dt / eps_0 - - # indices for the source - self.i_s = i_s - self.j_s = j_s - - self.time = 0 - - def organize_pmls(self): - if self.bc_low[0] == 'pml': - self.pml_lx.rx_block = self - if self.bc_low[1] is 'pml': - self.pml_lx.ly_block = self.pml_lxly - self.pml_lxly.ry_block = self.pml_lx - self.pml_lxly.rx_block = self.pml_ly - if self.bc_high[1] is 'pml': - self.pml_lx.ry_block = self.pml_lxry - self.pml_lxry.ly_block = self.pml_lx - self.pml_lxry.rx_block = self.pml_ry - - if self.bc_high[0] is 'pml': - self.pml_rx.lx_block = self - if self.bc_high[1] is 'pml': - self.pml_rx.ry_block = self.pml_rxry - self.pml_rxly.lx_block = self.pml_ly - self.pml_rxly.ry_block = self.pml_rx - if self.bc_low[1] is 'pml': - self.pml_rx.ly_block = self.pml_rxly - self.pml_rxry.lx_block = self.pml_ry - self.pml_rxry.ly_block = self.pml_rx - - - if self.bc_low[1] is 'pml': - self.pml_ly.ry_block = self - if self.bc_low[0] is 'pml': - self.pml_ly.lx_block = self.pml_lxly - if self.bc_low[0] is 'pml': - self.pml_ly.rx_block = self.pml_rxly - - if self.bc_high[1] is 'pml': - self.pml_ry.ly_block = self - if self.bc_low[0] is 'pml': - self.pml_ry.lx_block = self.pml_lxry - if self.bc_high[0] is 'pml': - self.pml_ry.rx_block = self.pml_rxry - - def assemble_conductivities_pmls(self): - sigma_m_low_x = 0 - sigma_m_high_x = 0 - sigma_m_low_y = 0 - sigma_m_high_y = 0 - if self.bc_low[0] is 'pml': - sigma_m_low_x = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_low[0]-1)*self.dx) * np.log(self.R0_pml) - if self.bc_low[1] is 'pml': - sigma_m_low_y = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_low[1]-1)*self.dy) * np.log(self.R0_pml) - if self.bc_high[0] is 'pml': - sigma_m_high_x = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_high[0]-1)*self.dy) * np.log(self.R0_pml) - if self.bc_high[1] is 'pml': - sigma_m_high_y = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_high[1]-1)*self.dy) * np.log(self.R0_pml) - - if self.bc_low[0] is 'pml': - for n in range(self.N_pml_low[0]): - self.pml_lx.sigma_x[-(n+1), :] = sigma_m_low_x * (n / (self.N_pml_low[0])) ** self.alpha_pml - if self.bc_low[1] is 'pml': - for n in range((self.N_pml_low[1])): - self.pml_lxly.sigma_y[:, -(n+1)] = sigma_m_low_y * (n / (self.N_pml_low[1])) ** self.alpha_pml - for n in range((self.N_pml_low[0])): - self.pml_lxly.sigma_x[-(n+1), :] = sigma_m_low_x * (n / (self.N_pml_low[0])) ** self.alpha_pml - if self.bc_high[1] is 'pml': - for n in range(self.N_pml_high[1]): - self.pml_lxry.sigma_y[:, n] = sigma_m_high_y * (n / (self.N_pml_high[1])) ** self.alpha_pml - for n in range(self.N_pml_low[0]): - self.pml_lxry.sigma_x[-(n+1), :] = sigma_m_low_x * (n / (self.N_pml_low[0])) ** self.alpha_pml - - if self.bc_high[0] is 'pml': - for n in range(self.N_pml_high[0]): - self.pml_rx.sigma_x[n, :] = sigma_m_high_x * (n / (self.N_pml_high[0])) ** self.alpha_pml - if self.bc_high[1] is 'pml': - for n in range(self.N_pml_high[0]): - self.pml_rxry.sigma_x[n, :] = sigma_m_high_x * (n / (self.N_pml_high[0])) ** self.alpha_pml - for n in range(self.N_pml_high[1]): - self.pml_rxry.sigma_y[:, n] = sigma_m_high_y * (n / (self.N_pml_high[1])) ** self.alpha_pml - if self.bc_low[1] == 'pml': - for n in range(self.N_pml_low[1]): - self.pml_rxly.sigma_y[:, -(n+1)] = sigma_m_low_y * (n / (self.N_pml_low[1])) ** self.alpha_pml - for n in range(self.N_pml_high[0]): - self.pml_rxly.sigma_x[n, :] = sigma_m_high_x * (n / (self.N_pml_high[0])) ** self.alpha_pml - - if self.bc_low[1] is 'pml': - for n in range(self.N_pml_low[1]): - self.pml_ly.sigma_y[:, -(n+1)] = sigma_m_low_y * (n / (self.N_pml_low[1])) ** self.alpha_pml - - if self.bc_high[1] is 'pml': - for n in range(self.N_pml_high[1]): - self.pml_ry.sigma_y[:, n] = sigma_m_high_y * (n / (self.N_pml_high[1])) ** self.alpha_pml - - if self.pml_lx is not None: - self.pml_lx.sigma_star_x = self.pml_lx.sigma_x * mu_0 / eps_0 - self.pml_lx.sigma_star_y = self.pml_lx.sigma_y * mu_0 / eps_0 - if self.pml_ly is not None: - self.pml_ly.sigma_star_x = self.pml_ly.sigma_x * mu_0 / eps_0 - self.pml_ly.sigma_star_y = self.pml_ly.sigma_y * mu_0 / eps_0 - if self.pml_rx is not None: - self.pml_rx.sigma_star_x = self.pml_rx.sigma_x * mu_0 / eps_0 - self.pml_rx.sigma_star_y = self.pml_rx.sigma_y * mu_0 / eps_0 - if self.pml_ry is not None: - self.pml_ry.sigma_star_x = self.pml_ry.sigma_x * mu_0 / eps_0 - self.pml_ry.sigma_star_y = self.pml_ry.sigma_y * mu_0 / eps_0 - if self.pml_lxly is not None: - self.pml_lxly.sigma_star_x = self.pml_lxly.sigma_x * mu_0 / eps_0 - self.pml_lxly.sigma_star_y = self.pml_lxly.sigma_y * mu_0 / eps_0 - if self.pml_lxry is not None: - self.pml_lxry.sigma_star_x = self.pml_lxry.sigma_x * mu_0 / eps_0 - self.pml_lxry.sigma_star_y = self.pml_lxry.sigma_y * mu_0 / eps_0 - if self.pml_rxry is not None: - self.pml_rxry.sigma_star_x = self.pml_rxry.sigma_x * mu_0 / eps_0 - self.pml_rxry.sigma_star_y = self.pml_rxry.sigma_y * mu_0 / eps_0 - if self.pml_rxly is not None: - self.pml_rxly.sigma_star_x = self.pml_rxly.sigma_x * mu_0 / eps_0 - self.pml_rxly.sigma_star_y = self.pml_rxly.sigma_y * mu_0 / eps_0 - - def assemble_coeffs_pmls(self): - if self.bc_low[0] is 'pml': - self.pml_lx.assemble_coeffs() - if self.bc_low[1] is 'pml': - self.pml_lxly.assemble_coeffs() - if self.bc_high[1] is 'pml': - self.pml_lxry.assemble_coeffs() - - if self.bc_high[0] is 'pml': - self.pml_rx.assemble_coeffs() - if self.bc_low[0] is 'pml': - self.pml_rxry.assemble_coeffs() - if self.bc_high[1] is 'pml': - self.pml_rxly.assemble_coeffs() - - if self.bc_low[1] is 'pml': - self.pml_ly.assemble_coeffs() - - if self.bc_high[1] is 'pml': - self.pml_ry.assemble_coeffs() - - def update_e_boundary(self): - Ex = self.Ex - Ey = self.Ey - Hz = self.Hz - if self.pml_lx is not None: - for jj in range(self.Ny): - Ey[0, jj] = Ey[0, jj] - self.C3 * self.Jy[0, jj] - self.C5 * (Hz[0, jj] - self.pml_lx.Hz[-1, jj]) - - if self.pml_rx is not None: - for jj in range(self.Ny): - Ey[-1, jj] = Ey[-1, jj] - self.C3 * self.Jy[-1, jj] - self.C5 * (self.pml_rx.Hz[0, jj] - Hz[-1, jj]) - - if self.pml_ly is not None: - for ii in range(self.Nx): - Ex[ii, 0] = Ex[ii, 0] - self.C3 * self.Jx[ii, 0] + self.C4 * (Hz[ii, 0] - self.pml_ly.Hz[ii, -1]) - - if self.pml_ry is not None: - for ii in range(self.Nx): - Ex[ii, -1] = Ex[ii, -1] - self.C3 * self.Jx[ii, -1] + self.C4 * (self.pml_ry.Hz[ii, 0] - Hz[ii, -1]) - - for block in self.blocks: - block.update_e_boundary() - - def gauss(self, t): - tau = 10 * self.dt - if t < 6 * tau: - return 100 * np.exp(-(t - 3 * tau) ** 2 / tau ** 2) - else: - return 0. - - def one_step(self): - if self.sol_type == 'ECT': - self.compute_v_and_rho() - self.one_step_ect(Nx=self.Nx, Ny=self.Ny, V_enl=self.V_enl, - rho=self.rho, Hz=self.Hz, C1=self.C1, - flag_int_cell=self.grid.flag_int_cell, - flag_unst_cell=self.grid.flag_unst_cell, - flag_intr_cell=self.grid.flag_intr_cell, - S=self.grid.S, - borrowing=self.grid.borrowing, S_enl=self.grid.S_enl, - S_red=self.grid.S_red, V_new = self.V_new, dt=self.dt) - for block in self.blocks: - block.advance_h_fdtd() - - self.advance_e_dm() - - for block in self.blocks: - block.advance_e_fdtd() - - self.update_e_boundary() - - self.time += self.dt - - if self.sol_type == 'FDTD': - self.one_step_fdtd() - if self.sol_type == 'DM': - self.one_step_dm() - - def one_step_fdtd(self): - Z_0 = np.sqrt(mu_0 / eps_0) - Ex = self.Ex - Ey = self.Ey - Hz = self.Hz - - self.advance_h_fdtd() - for block in self.blocks: - block.advance_h_fdtd() - self.advance_e_fdtd() - for block in self.blocks: - block.advance_e_fdtd() - self.update_e_boundary() - - self.time += self.dt - - def one_step_dm(self): - self.compute_v_and_rho() - - for i in range(self.Nx): - for j in range(self.Ny): - if self.grid.flag_int_cell[i, j]: - self.Hz[i, j] = self.Hz[i, j] - self.dt / (mu_0 * self.grid.S[i, j]) * self.Vxy[i, j] - - for block in self.blocks: - block.advance_h_fdtd() - - self.advance_e_dm() - - for block in self.blocks: - block.advance_e_fdtd() - - self.update_e_boundary_dm() - - self.time += self.dt - - def advance_h_fdtd(self): - Ex = self.Ex - Ey = self.Ey - Hz = self.Hz - for ii in range(self.Nx): - for jj in range(self.Ny): - if self.grid.flag_int_cell[ii, jj]: - Hz[ii, jj] = (Hz[ii, jj] - self.C1 * (Ey[ii + 1, jj] - Ey[ii, jj]) + - self.C2 * (Ex[ii, jj + 1]- Ex[ii, jj])) - - def advance_e_fdtd(self): - Z_0 = np.sqrt(mu_0 / eps_0) - Ex = self.Ex - Ey = self.Ey - Hz = self.Hz - for ii in range(self.Nx): - for jj in range(1, self.Ny): - if self.grid.flag_int_cell[ii, jj]: - if self.grid.l_x[ii, jj] > 0: - Ex[ii, jj] = Ex[ii, jj] - self.C3 * self.Jx[ii, jj] + self.C4 * ( - Hz[ii, jj] - Hz[ii, jj - 1]) - - for ii in range(1, self.Nx): - for jj in range(self.Ny): - if self.grid.flag_int_cell[ii, jj]: - if self.grid.l_y[ii, jj] > 0: - Ey[ii, jj] = Ey[ii, jj] - self.C3 * self.Jy[ii, jj] - self.C5 * ( - Hz[ii, jj] - Hz[ii - 1, jj]) - - @staticmethod - def one_step_ect(Nx=None, Ny=None, V_enl=None, rho=None, Hz=None, C1=None, flag_int_cell=None, - flag_unst_cell=None, flag_intr_cell=None, S=None, borrowing=None, S_enl=None, - S_red=None, V_new=None, dt = None, comp=None, kk=None): - - #if dt==None: dt = self.dt - - V_enl = np.zeros((Nx, Ny)) - - # take care of unstable cells - for ii in range(Nx): - for jj in range(Ny): - if flag_int_cell[ii, jj] and flag_unst_cell[ii, jj]: - - V_enl[ii, jj] = rho[ii, jj] * S[ii, jj] - - if len(borrowing[ii, jj]) == 0: - print('error in one_step_ect') - for (ip, jp, patch, _) in borrowing[ii, jj]: - - V_enl[ii, jj] += rho[ip, jp] * patch - - rho_enl = V_enl[ii, jj] / S_enl[ii, jj] - - # communicate to the intruded cell the intruding rho - for (ip, jp, patch, _) in borrowing[ii, jj]: - V_enl[ip, jp] += rho_enl * patch - - Hz[ii, jj] = Hz[ii, jj] - dt/mu_0 * rho_enl - - - # take care of stable cells - for ii in range(Nx): - for jj in range(Ny): - if flag_int_cell[ii, jj] and not flag_unst_cell[ii, jj]: - # stable cell which hasn't been intruded - if not flag_intr_cell[ii, jj]: - Hz[ii, jj] = Hz[ii, jj] - dt/mu_0 * rho[ii, jj] - # stable cell which has been intruded - else: - V_enl[ii, jj] += rho[ii, jj] * S_red[ii, jj] - Hz[ii, jj] = Hz[ii, jj] - dt/mu_0 * V_enl[ii, jj] / S[ii, jj] - - - def compute_v_and_rho(self): - l_y = self.grid.l_y - l_x = self.grid.l_x - for ii in range(self.Nx): - for jj in range(self.Ny): - if self.grid.flag_int_cell[ii, jj]: - self.Vxy[ii, jj] = ( - self.Ey[ii + 1, jj] * l_y[ii + 1, jj] - self.Ey[ii, jj] * l_y[ - ii, jj] - - self.Ex[ii, jj + 1] * l_x[ii, jj + 1] + self.Ex[ii, jj] * l_x[ - ii, jj]) - if self.sol_type != 'DM': - self.rho[ii, jj] = self.Vxy[ii, jj] / self.grid.S[ii, jj] - - def advance_e_dm(self): - for ii in range(self.Nx): - for jj in range(1, self.Ny): - if self.grid.l_x[ii, jj] > 0: - self.Ex[ii, jj] = self.Ex[ii, jj] + self.dt / (eps_0 * self.dy) * ( - self.Hz[ii, jj] - self.Hz[ii, jj - 1]) - self.C3 * self.Jx[ii, jj] - - for ii in range(1, self.Nx): - for jj in range(self.Ny): - if self.grid.l_y[ii, jj] > 0: - self.Ey[ii, jj] = self.Ey[ii, jj] - self.dt / (eps_0 * self.dx) * ( - self.Hz[ii, jj] - self.Hz[ii - 1, jj]) - self.C3 * self.Jy[ii, jj] - - def update_e_boundary_dm(self): - Ex = self.Ex - Ey = self.Ey - Hz = self.Hz - if self.pml_lx is not None: - for jj in range(self.Ny): - Ey[0, jj] = Ey[0, jj] - self.C3 * self.Jy[0, jj] - self.dt / (eps_0 * self.dy) * (Hz[0, jj] - self.pml_lx.Hz[-1, jj]) - - if self.pml_rx is not None: - for jj in range(self.Ny): - Ey[-1, jj] = Ey[-1, jj] - self.C3 * self.Jy[-1, jj] - self.dt / (eps_0 * self.dy) * (self.pml_rx.Hz[0, jj] - Hz[-1, jj]) - - if self.pml_ly is not None: - for ii in range(self.Nx): - Ex[ii, 0] = Ex[ii, 0] - self.C3 * self.Jx[ii, 0] + self.dt / (eps_0 * self.dy) * (Hz[ii, 0] - self.pml_ly.Hz[ii, -1]) - - if self.pml_ry is not None: - for ii in range(self.Nx): - Ex[ii, -1] = Ex[ii, -1] - self.C3 * self.Jx[ii, -1] + self.dt / (eps_0 * self.dy) * (self.pml_ry.Hz[ii, 0] - Hz[ii, -1]) - - for block in self.blocks: - block.update_e_boundary() From 9807faf868ff3c30b5f0bc43232222eb132bf1c8 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:55:47 +0100 Subject: [PATCH 18/39] Delete build/lib/wakis/solver3D.py --- build/lib/wakis/solver3D.py | 1126 ----------------------------------- 1 file changed, 1126 deletions(-) delete mode 100644 build/lib/wakis/solver3D.py diff --git a/build/lib/wakis/solver3D.py b/build/lib/wakis/solver3D.py deleted file mode 100644 index 1d9e42b..0000000 --- a/build/lib/wakis/solver3D.py +++ /dev/null @@ -1,1126 +0,0 @@ -import numpy as np -from scipy.constants import c as c_light, epsilon_0 as eps_0, mu_0 as mu_0 -from solver2D import EMSolver2D -from pmlBlock3D import PmlBlock3D - -from numba import jit - -def eq(a, b, tol=1e-8): - return abs(a - b) < tol - - -def neq(a, b, tol=1e-8): - return not eq(a, b, tol) - - -class EMSolver3D: - def __init__(self, grid, sol_type, cfln=0.5, - bc_low=['Dirichlet', 'Dirichlet', 'Dirichlet'], - bc_high=['Dirichlet', 'Dirichlet', 'Dirichlet'], - i_s=0, j_s=0, k_s=0, N_pml_low=None, N_pml_high=None): - - self.grid = grid - self.type = type - self.cfln = cfln - - self.dt = cfln / (c_light * np.sqrt(1 / self.grid.dx ** 2 + 1 / self.grid.dy ** 2 + - 1 / self.grid.dz ** 2)) - self.dx = self.grid.dx - self.dy = self.grid.dy - self.dz = self.grid.dz - self.Nx = self.grid.nx - self.Ny = self.grid.ny - self.Nz = self.grid.nz - self.sol_type = sol_type - - self.N_pml_low = np.zeros(3, dtype=int) - self.N_pml_high = np.zeros(3, dtype=int) - self.bc_low = bc_low - self.bc_high = bc_high - - Nx = self.Nx - Ny = self.Ny - Nz = self.Nz - - self.sigma_x = np.zeros((Nx + 1, Ny + 1, Nz + 1)) - self.sigma_y = np.zeros((Nx + 1, Ny + 1, Nz + 1)) - self.sigma_z = np.zeros((Nx + 1, Ny + 1, Nz + 1)) - self.sigma_star_x = np.zeros((Nx, Ny + 1, Nz + 1)) - self.sigma_star_y = np.zeros((Nx + 1, Ny, Nz + 1)) - self.sigma_star_z = np.zeros((Nx + 1, Ny + 1, Nz)) - - if bc_low[0] == 'pml': - self.N_pml_low[0] = 10 if N_pml_low is None else N_pml_low[0] - if bc_low[1] == 'pml': - self.N_pml_low[1] = 10 if N_pml_low is None else N_pml_low[1] - if bc_low[2] == 'pml': - self.N_pml_low[2] = 10 if N_pml_low is None else N_pml_low[2] - if bc_high[0] == 'pml': - self.N_pml_high[0] = 10 if N_pml_high is None else N_pml_high[0] - if bc_high[1] == 'pml': - self.N_pml_high[1] = 10 if N_pml_high is None else N_pml_high[1] - if bc_high[2] == 'pml': - self.N_pml_high[2] = 10 if N_pml_high is None else N_pml_high[2] - - self.blocks = [] - - self.blocks_mat = np.full((3, 3, 3), None) - self.blocks_mat[1, 1, 1] = self - - self.connect_pmls() - - self.alpha_pml = 3 - self.R0_pml = 0.001 - - self.assemble_conductivities_pmls() - self.assemble_coeffs_pmls() - - self.Ex = np.zeros((self.Nx, self.Ny + 1, self.Nz + 1)) - self.Ey = np.zeros((self.Nx + 1, self.Ny, self.Nz + 1)) - self.Ez = np.zeros((self.Nx + 1, self.Ny + 1, self.Nz)) - self.Hx = np.zeros((self.Nx + 1, self.Ny, self.Nz)) - self.Hy = np.zeros((self.Nx, self.Ny + 1, self.Nz)) - self.Hz = np.zeros((self.Nx, self.Ny, self.Nz + 1)) - self.Jx = np.zeros((self.Nx, self.Ny + 1, self.Nz + 1)) - self.Jy = np.zeros((self.Nx + 1, self.Ny, self.Nz + 1)) - self.Jz = np.zeros((self.Nx + 1, self.Ny + 1, self.Nz)) - self.rho_xy = np.zeros((self.Nx, self.Ny, self.Nz + 1)) - self.rho_yz = np.zeros((self.Nx + 1, self.Ny, self.Nz)) - self.rho_zx = np.zeros((self.Nx, self.Ny + 1, self.Nz)) - - if (sol_type is not 'FDTD') and (sol_type is not 'DM') and (sol_type is not 'ECT'): - raise ValueError("sol_type must be:\n" + - "\t'FDTD' for standard staircased FDTD\n" + - "\t'DM' for Dey-Mittra conformal FDTD\n" + - "\t'ECT' for Enlarged Cell Technique conformal FDTD") - - if sol_type is 'DM' or sol_type is 'ECT': - self.Vxy = np.zeros((self.Nx, self.Ny, self.Nz + 1)) - self.Vyz = np.zeros((self.Nx + 1, self.Ny, self.Nz)) - self.Vzx = np.zeros((self.Nx, self.Ny + 1, self.Nz)) - if sol_type is 'ECT': - self.Vxy_enl = np.zeros((self.Nx, self.Ny, self.Nz + 1)) - self.Vyz_enl = np.zeros((self.Nx + 1, self.Ny, self.Nz)) - self.Vzx_enl = np.zeros((self.Nx, self.Ny + 1, self.Nz)) - - self.C1 = self.dt / (self.dx * mu_0) - self.C2 = self.dt / (self.dy * mu_0) - self.C7 = self.dt / (self.dz * mu_0) - self.C4 = self.dt / (self.dy * eps_0) - self.C5 = self.dt / (self.dx * eps_0) - self.C8 = self.dt / (self.dz * eps_0) - self.C3 = self.dt / eps_0 - self.C6 = self.dt / eps_0 - - self.CN = self.dt / mu_0 - - # indices for the source - self.i_s = i_s - self.j_s = j_s - self.k_s = k_s - - self.time = 0 - - def connect_pmls(self): - bc_low = self.bc_low - bc_high = self.bc_high - if bc_low[0] is 'pml': - i_block = 0 - j_block = 1 - k_block = 1 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.Ny, self.Nz, self.dt, - self.dx, - self.dy, self.dz, i_block, j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_low[1] is 'pml': - i_block = 0 - j_block = 0 - k_block = 1 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.N_pml_low[1], self.Nz, - self.dt, self.dx, self.dy, self.dz, i_block, - j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_low[2] is 'pml': - i_block = 0 - j_block = 0 - k_block = 0 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.N_pml_low[1], - self.N_pml_low[2], self.dt, self.dx, - self.dy, - self.dz, i_block, j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_high[2] is 'pml': - i_block = 0 - j_block = 0 - k_block = 2 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.N_pml_low[1], - self.N_pml_high[2], self.dt, self.dx, - self.dy, - self.dz, i_block, j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_high[1] is 'pml': - i_block = 0 - j_block = 2 - k_block = 1 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.N_pml_high[1], self.Nz, - self.dt, self.dx, self.dy, self.dz, i_block, - j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_low[2] is 'pml': - i_block = 0 - j_block = 2 - k_block = 0 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.N_pml_high[1], - self.N_pml_low[2], self.dt, self.dx, - self.dy, - self.dz, i_block, j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_high[2] is 'pml': - i_block = 0 - j_block = 2 - k_block = 2 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.N_pml_high[1], - self.N_pml_high[2], self.dt, self.dx, - self.dy, - self.dz, i_block, j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_low[2] is 'pml': - i_block = 0 - j_block = 1 - k_block = 0 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.Ny, self.N_pml_low[2], - self.dt, self.dx, self.dy, self.dz, i_block, - j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_high[2] is 'pml': - i_block = 0 - j_block = 1 - k_block = 2 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_low[0], self.Ny, self.N_pml_high[2], - self.dt, self.dx, self.dy, self.dz, i_block, - j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - - if bc_high[0] is 'pml': - i_block = 2 - j_block = 1 - k_block = 1 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.Ny, self.Nz, self.dt, - self.dx, self.dy, self.dz, i_block, j_block, - k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_low[1] is 'pml': - i_block = 2 - j_block = 0 - k_block = 1 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.N_pml_low[1], self.Nz, - self.dt, self.dx, self.dy, self.dz, i_block, - j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_low[2] is 'pml': - i_block = 2 - j_block = 0 - k_block = 0 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.N_pml_low[1], - self.N_pml_low[2], self.dt, self.dx, - self.dy, - self.dz, i_block, j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_high[2] is 'pml': - i_block = 2 - j_block = 0 - k_block = 2 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.N_pml_low[1], - self.N_pml_high[2], self.dt, self.dx, - self.dy, - self.dz, i_block, j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_high[1] is 'pml': - i_block = 2 - j_block = 2 - k_block = 1 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.N_pml_high[1], self.Nz, - self.dt, self.dx, self.dy, self.dz, i_block, - j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_low[2] is 'pml': - i_block = 2 - j_block = 2 - k_block = 0 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.N_pml_high[1], - self.N_pml_low[2], self.dt, self.dx, - self.dy, - self.dz, i_block, j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_high[2] is 'pml': - i_block = 2 - j_block = 2 - k_block = 2 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.N_pml_high[1], - self.N_pml_high[2], self.dt, self.dx, - self.dy, - self.dz, i_block, j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_low[2] is 'pml': - i_block = 2 - j_block = 1 - k_block = 0 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.Ny, self.N_pml_low[2], - self.dt, self.dx, self.dy, self.dz, i_block, - j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_high[2] is 'pml': - i_block = 2 - j_block = 1 - k_block = 2 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.N_pml_high[0], self.Ny, self.N_pml_high[2], - self.dt, self.dx, self.dy, self.dz, i_block, - j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - - if bc_low[1] is 'pml': - i_block = 1 - j_block = 0 - k_block = 1 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.Nx, self.N_pml_low[1], self.Nz, self.dt, - self.dx, - self.dy, self.dz, i_block, j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_low[2] is 'pml': - i_block = 1 - j_block = 0 - k_block = 0 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.Nx, self.N_pml_low[1], self.N_pml_low[2], - self.dt, self.dx, self.dy, self.dz, i_block, - j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_high[2] is 'pml': - i_block = 1 - j_block = 0 - k_block = 2 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.Nx, self.N_pml_low[1], self.N_pml_high[2], - self.dt, self.dx, self.dy, self.dz, i_block, - j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - - if bc_high[1] is 'pml': - i_block = 1 - j_block = 2 - k_block = 1 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.Nx, self.N_pml_high[1], self.Nz, self.dt, - self.dx, self.dy, self.dz, i_block, j_block, - k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_low[2] is 'pml': - i_block = 1 - j_block = 2 - k_block = 0 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.Nx, self.N_pml_low[1], self.N_pml_low[2], - self.dt, self.dx, self.dy, self.dz, i_block, - j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - if bc_high[2] is 'pml': - i_block = 1 - j_block = 2 - k_block = 2 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.Nx, self.N_pml_low[1], self.N_pml_high[2], - self.dt, self.dx, self.dy, self.dz, i_block, - j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - - if bc_low[2] is 'pml': - i_block = 1 - j_block = 1 - k_block = 0 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.Nx, self.Ny, self.N_pml_low[2], self.dt, - self.dx, - self.dy, self.dz, i_block, j_block, k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - - if bc_high[2] is 'pml': - i_block = 1 - j_block = 1 - k_block = 2 - self.blocks_mat[i_block, j_block, k_block] = PmlBlock3D(self.Nx, self.Ny, self.N_pml_high[2], self.dt, - self.dx, self.dy, self.dz, i_block, j_block, - k_block) - self.blocks.append(self.blocks_mat[i_block, j_block, k_block]) - - for block in self.blocks: - block.blocks_mat = self.blocks_mat - - def assemble_conductivities_pmls(self): - sigma_m_low_x = 0 - sigma_m_high_x = 0 - sigma_m_low_y = 0 - sigma_m_high_y = 0 - sigma_m_low_z = 0 - sigma_m_high_z = 0 - Z_0_2 = mu_0 / eps_0 - - if self.bc_low[0] is 'pml': - sigma_m_low_x = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_low[0] - 1) * self.dx) * np.log(self.R0_pml) - if self.bc_low[1] is 'pml': - sigma_m_low_y = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_low[1] - 1) * self.dy) * np.log(self.R0_pml) - if self.bc_low[2] is 'pml': - sigma_m_low_z = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_low[2] - 1) * self.dy) * np.log(self.R0_pml) - if self.bc_high[0] is 'pml': - sigma_m_high_x = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_high[0] - 1) * self.dy) * np.log(self.R0_pml) - if self.bc_high[1] is 'pml': - sigma_m_high_y = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_high[1] - 1) * self.dy) * np.log(self.R0_pml) - if self.bc_high[2] is 'pml': - sigma_m_high_z = -(self.alpha_pml + 1) * eps_0 * c_light / (2 * (self.N_pml_high[2] - 1) * self.dy) * np.log(self.R0_pml) - - - if self.bc_low[0] is 'pml': - (i_block, j_block, k_block) = (0, 1, 1) - for n in range(self.N_pml_low[0]): - self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 - if self.bc_low[1] is 'pml': - (i_block, j_block, k_block) = (0, 0, 1) - for n in range((self.N_pml_low[0])): - self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_low[1])): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 - if self.bc_low[2] is 'pml': - (i_block, j_block, k_block) = (0, 0, 0) - for n in range((self.N_pml_low[0])): - self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_low[1])): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_low[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 - if self.bc_low[2] is 'pml': - (i_block, j_block, k_block) = (0, 0, 2) - for n in range((self.N_pml_low[0])): - self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_low[1])): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_high[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 - if self.bc_high[1] is 'pml': - (i_block, j_block, k_block) = (0, 2, 1) - for n in range(self.N_pml_low[0]): - self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 - for n in range(self.N_pml_high[1]): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 - - if self.bc_low[2] is 'pml': - (i_block, j_block, k_block) = (0, 2, 0) - for n in range((self.N_pml_low[0])): - self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_high[1])): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_low[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 - if self.bc_low[2] is 'pml': - (i_block, j_block, k_block) = (0, 2, 2) - for n in range((self.N_pml_low[0])): - self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_high[1])): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_high[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 - if self.bc_low[2] is 'pml': - (i_block, j_block, k_block) = (0, 1, 0) - for n in range((self.N_pml_low[0])): - self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_low[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 - if self.bc_high[2] is 'pml': - (i_block, j_block, k_block) = (0, 1, 2) - for n in range((self.N_pml_low[0])): - self.blocks_mat[i_block, j_block, k_block].sigma_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[-(n + 1), :, :] = sigma_m_low_x * ( - n / (self.N_pml_low[0])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_high[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 - - if self.bc_high[0] is 'pml': - (i_block, j_block, k_block) = (2, 1, 1) - for n in range(self.N_pml_high[0]): - self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 - if self.bc_low[1] is 'pml': - (i_block, j_block, k_block) = (2, 0, 1) - for n in range(self.N_pml_high[0]): - self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_low[1])): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 - if self.bc_low[2] is 'pml': - (i_block, j_block, k_block) = (2, 0, 0) - for n in range(self.N_pml_high[0]): - self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_low[1])): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_low[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 - if self.bc_low[2] is 'pml': - (i_block, j_block, k_block) = (2, 0, 2) - for n in range(self.N_pml_high[0]): - self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_low[1])): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_high[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 - if self.bc_high[1] is 'pml': - (i_block, j_block, k_block) = (2, 2, 1) - for n in range(self.N_pml_high[0]): - self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 - for n in range(self.N_pml_high[1]): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 - if self.bc_low[2] is 'pml': - (i_block, j_block, k_block) = (2, 2, 0) - for n in range(self.N_pml_high[0]): - self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_high[1])): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_low[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 - if self.bc_low[2] is 'pml': - (i_block, j_block, k_block) = (2, 2, 2) - for n in range(self.N_pml_high[0]): - self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_high[1])): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_high[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 - if self.bc_low[2] is 'pml': - (i_block, j_block, k_block) = (2, 1, 0) - for n in range(self.N_pml_high[0]): - self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_low[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 - if self.bc_high[2] is 'pml': - (i_block, j_block, k_block) = (2, 1, 2) - for n in range(self.N_pml_high[0]): - self.blocks_mat[i_block, j_block, k_block].sigma_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_x[n, :, :] = sigma_m_high_x * ( - n / (self.N_pml_high[0])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_high[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 - - if self.bc_low[1] is 'pml': - (i_block, j_block, k_block) = (1, 0, 1) - for n in range(self.N_pml_low[1]): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 - if self.bc_low[2] is 'pml': - (i_block, j_block, k_block) = (1, 0, 0) - for n in range(self.N_pml_low[1]): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_low[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 - if self.bc_high[2] is 'pml': - (i_block, j_block, k_block) = (1, 0, 2) - for n in range(self.N_pml_low[1]): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, -(n + 1), :] = sigma_m_low_y * ( - n / (self.N_pml_low[1])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_high[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 - - if self.bc_high[1] is 'pml': - (i_block, j_block, k_block) = (1, 2, 1) - for n in range(self.N_pml_high[1]): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 - if self.bc_low[2] is 'pml': - (i_block, j_block, k_block) = (1, 2, 0) - for n in range(self.N_pml_high[1]): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_low[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 - if self.bc_high[2] is 'pml': - (i_block, j_block, k_block) = (1, 2, 2) - for n in range(self.N_pml_high[1]): - self.blocks_mat[i_block, j_block, k_block].sigma_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_y[:, n, :] = sigma_m_high_y * ( - n / (self.N_pml_high[1])) ** self.alpha_pml*Z_0_2 - for n in range((self.N_pml_high[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 - - if self.bc_low[2] is 'pml': - (i_block, j_block, k_block) = (1, 1, 0) - for n in range((self.N_pml_low[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, -(n + 1)] = sigma_m_low_z * ( - n / (self.N_pml_low[2])) ** self.alpha_pml*Z_0_2 - - if self.bc_high[2] is 'pml': - (i_block, j_block, k_block) = (1, 1, 2) - for n in range((self.N_pml_high[2])): - self.blocks_mat[i_block, j_block, k_block].sigma_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml - self.blocks_mat[i_block, j_block, k_block].sigma_star_z[:, :, n] = sigma_m_high_z * ( - n / (self.N_pml_high[2])) ** self.alpha_pml*Z_0_2 - - #for i_block in range(3): - # for j_block in range(3): - # for k_block in range(3): - # block = self.blocks_mat[i_block, j_block, k_block] - # if block is not None: - # if not (i_block == 1 and j_block == 1 and k_block == 1): - # block.sigma_star_x = block.sigma_x * mu_0 / eps_0 - # block.sigma_star_y = block.sigma_y * mu_0 / eps_0 - # block.sigma_star_z = block.sigma_z * mu_0 / eps_0 - - def assemble_coeffs_pmls(self): - for i_block in range(3): - for j_block in range(3): - for k_block in range(3): - if self.blocks_mat[i_block, j_block, k_block] is not None: - if not (i_block == 1 and j_block == 1 and k_block == 1): - self.blocks_mat[i_block, j_block, k_block].assemble_coeffs() - - def update_e_boundary(self): - Ex = self.Ex - Ey = self.Ey - Ez = self.Ez - Hx = self.Hx - Hy = self.Hy - Hz = self.Hz - - Nx = self.Nx - Ny = self.Ny - Nz = self.Nz - - # Update E on "lower" faces - if self.blocks_mat[0, 1, 1] is not None: - for jj in range(self.Ny): - for kk in range(1, self.Nz): - Ey[0, jj, kk] = (Ey[0, jj, kk] - self.C3 * self.Jy[0, jj, kk] + - self.C8 * (Hx[0, jj, kk] - Hx[0, jj, kk - 1]) - - self.C5 * (Hz[0, jj, kk] - self.blocks_mat[0, 1, 1].Hz[-1, jj, kk])) - - for jj in range(1, self.Ny): - for kk in range(self.Nz): - Ez[0, jj, kk] = (Ez[0, jj, kk] - self.C3 * self.Jz[0, jj, kk] + - self.C5 * (Hy[0, jj, kk] - self.blocks_mat[0, 1, 1].Hy[-1, jj, kk]) - - self.C4 * (Hx[0, jj, kk] - Hx[0, jj - 1, kk])) - - if self.blocks_mat[1, 0, 1] is not None: - for ii in range(self.Nx): - for kk in range(1, self.Nz): - Ex[ii, 0, kk] = (Ex[ii, 0, kk] - self.C3 * self.Jx[ii, 0, kk] + - self.C4 * (Hz[ii, 0, kk] - self.blocks_mat[1, 0, 1].Hz[ii, -1, kk]) - - self.C8 * (Hy[ii, 0, kk] - Hy[ii, 0, kk - 1])) - for ii in range(1, self.Nx): - for kk in range(self.Nz): - Ez[ii, 0, kk] = (Ez[ii, 0, kk] - self.C3 * self.Jz[ii, 0, kk] + - self.C5 * (Hy[ii, 0, kk] - Hy[ii - 1, 0, kk]) - - self.C4 * (Hx[ii, 0, kk] - self.blocks_mat[1, 0, 1].Hx[ii, -1, kk])) - - if self.blocks_mat[1, 1, 0] is not None: - for ii in range(self.Nx): - for jj in range(1, self.Ny): - Ex[ii, jj, 0] = (Ex[ii, jj, 0] - self.C3 * self.Jx[ii, jj, 0] + - self.C4 * (Hz[ii, jj, 0] - Hz[ii, jj - 1, 0]) - - self.C8 * (Hy[ii, jj, 0] - self.blocks_mat[1, 1, 0].Hy[ii, jj, -1])) - for ii in range(1, self.Nx): - for jj in range(self.Ny): - Ey[ii, jj, 0] = (Ey[ii, jj, 0] - self.C3 * self.Jy[ii, jj, 0] + - self.C8 * (Hx[ii, jj, 0] - self.blocks_mat[1, 1, 0].Hx[ii, jj, -1]) - - self.C5 * (Hz[ii, jj, 0] - Hz[ii - 1, jj, 0])) - - # Update E on "upper" faces - if self.blocks_mat[2, 1, 1] is not None: - for jj in range(self.Ny): - for kk in range(1, self.Nz): - Ey[Nx, jj, kk] = (Ey[Nx, jj, kk] - self.C3 * self.Jy[Nx, jj, kk] + - self.C8 * (Hx[Nx, jj, kk] - Hx[Nx, jj, kk - 1]) - - self.C5 * (self.blocks_mat[2, 1, 1].Hz[0, jj, kk] - Hz[Nx - 1, jj, kk])) - for jj in range(1, self.Ny): - for kk in range(self.Nz): - Ez[Nx, jj, kk] = (Ez[Nx, jj, kk] - self.C3 * self.Jz[Nx, jj, kk] + - self.C5 * (self.blocks_mat[2, 1, 1].Hy[0, jj, kk] - Hy[Nx - 1, jj, kk]) - - self.C4 * (Hx[Nx, jj, kk] - Hx[Nx, jj - 1, kk])) - - if self.blocks_mat[1, 2, 1] is not None: - for ii in range(self.Nx): - for kk in range(1, self.Nz): - Ex[ii, Ny, kk] = (Ex[ii, Ny, kk] - self.C3 * self.Jx[ii, Ny, kk] + - self.C4 * (self.blocks_mat[1, 2, 1].Hz[ii, 0, kk] - Hz[ii, Ny - 1, kk]) - - self.C8 * (Hy[ii, Ny, kk] - Hy[ii, Ny, kk - 1])) - for ii in range(1, self.Nx): - for kk in range(self.Nz): - Ez[ii, Ny, kk] = (Ez[ii, Ny, kk] - self.C3 * self.Jz[ii, Ny, kk] + - self.C5 * (Hy[ii, Ny, kk] - Hy[ii - 1, Ny, kk]) - - self.C4 * (self.blocks_mat[1, 2, 1].Hx[ii, 0, kk] - Hx[ii, Ny - 1, kk])) - - if self.blocks_mat[1, 1, 2] is not None: - for ii in range(self.Nx): - for jj in range(1, self.Ny): - Ex[ii, jj, Nz] = (Ex[ii, jj, Nz] - self.C3 * self.Jx[ii, jj, Nz] + - self.C4 * (Hz[ii, jj, Nz] - Hz[ii, jj - 1, Nz]) - - self.C8 * (self.blocks_mat[1, 1, 2].Hy[ii, jj, 0] - Hy[ii, jj, Nz - 1])) - for ii in range(1, self.Nx): - for jj in range(self.Ny): - Ey[ii, jj, Nz] = (Ey[ii, jj, Nz] - self.C3 * self.Jy[ii, jj, Nz] + - self.C8 * (self.blocks_mat[1, 1, 2].Hx[ii, jj, 0] - Hx[ii, jj, Nz - 1]) - - self.C5 * (Hz[ii, jj, Nz] - Hz[ii - 1, jj, Nz])) - # Update Ez on edges (xy) - if self.blocks_mat[0, 1, 1] is not None and self.blocks_mat[1, 0, 1] is not None: - for kk in range(Nz): - Ez[0, 0, kk] = (Ez[0, 0, kk] - self.C3 * self.Jz[0, 0, kk] + - self.C5 * (Hy[0, 0, kk] - self.blocks_mat[0, 1, 1].Hy[-1, 0, kk]) - - self.C4 * (Hx[0, 0, kk] - self.blocks_mat[1, 0, 1].Hx[0, -1, kk])) - if self.blocks_mat[0, 1, 1] is not None and self.blocks_mat[1, 2, 1] is not None: - for kk in range(Nz): - Ez[0, Ny, kk] = (Ez[0, Ny, kk] - self.C3 * self.Jz[0, Ny, kk] + - self.C5 * (Hy[0, Ny, kk] - self.blocks_mat[0, 1, 1].Hy[-1, Ny, kk]) - - self.C4 * (self.blocks_mat[1, 2, 1].Hx[0, 0, kk] - Hx[0, Ny - 1, kk])) - if self.blocks_mat[2, 1, 1] is not None and self.blocks_mat[1, 0, 1] is not None: - for kk in range(Nz): - Ez[Nx, 0, kk] = (Ez[Nx, 0, kk] - self.C3 * self.Jz[Nx, 0, kk] + - self.C5 * (self.blocks_mat[2, 1, 1].Hy[0, 0, kk] - Hy[Nx - 1, 0, kk]) - - self.C4 * (Hx[Nx, 0, kk] - self.blocks_mat[1, 0, 1].Hx[Nx, -1, kk])) - if self.blocks_mat[2, 1, 1] is not None and self.blocks_mat[1, 2, 1] is not None: - for kk in range(Nz): - Ez[Nx, Ny, kk] = (Ez[Nx, Ny, kk] - self.C3 * self.Jz[Nx, Ny, kk] + - self.C5 * (self.blocks_mat[2, 1, 1].Hy[0, Ny, kk] - Hy[Nx - 1, Ny, kk]) - - self.C4 * (self.blocks_mat[1, 2, 1].Hx[Nx, 0, kk] - Hx[Nx, Ny - 1, kk])) - # Update Ex on edges (yz) - if self.blocks_mat[1, 0, 1] is not None and self.blocks_mat[1, 1, 0] is not None: - for ii in range(Nx): - Ex[ii, 0, 0] = (Ex[ii, 0, 0] - self.C3 * self.Jx[ii, 0, 0] + - self.C4 * (Hz[ii, 0, 0] - self.blocks_mat[1, 0, 1].Hz[ii, -1, 0]) - - self.C8 * (Hy[ii, 0, 0] - self.blocks_mat[1, 1, 0].Hy[ii, 0, -1])) - if self.blocks_mat[1, 0, 1] is not None and self.blocks_mat[1, 1, 2] is not None: - for ii in range(Nx): - Ex[ii, 0, Nz] = (Ex[ii, 0, Nz] - self.C3 * self.Jx[ii, 0, Nz] + - self.C4 * (Hz[ii, 0, Nz] - self.blocks_mat[1, 0, 1].Hz[ii, -1, Nz]) - - self.C8 * (self.blocks_mat[1, 1, 2].Hy[ii, 0, 0] - Hy[ii, 0, Nz - 1])) - if self.blocks_mat[1, 2, 1] is not None and self.blocks_mat[1, 1, 0] is not None: - for ii in range(Nx): - Ex[ii, Ny, 0] = (Ex[ii, Ny, 0] - self.C3 * self.Jx[ii, Ny, 0] + - self.C4 * (self.blocks_mat[1, 2, 1].Hz[ii, 0, 0] - Hz[ii, Ny - 1, 0]) - - self.C8 * (Hy[ii, Ny, 0] - self.blocks_mat[1, 1, 0].Hy[ii, Ny, -1])) - if self.blocks_mat[1, 2, 1] is not None and self.blocks_mat[1, 1, 2] is not None: - for ii in range(Nx): - Ex[ii, Ny, Nz] = (Ex[ii, Ny, Nz] - self.C3 * self.Jx[ii, Ny, Nz] + - self.C4 * (self.blocks_mat[1, 2, 1].Hz[ii, 0, Nz] - Hz[ii, Ny - 1, Nz]) - - self.C8 * (self.blocks_mat[1, 1, 2].Hy[ii, Ny, 0] - Hy[ii, Ny, Nz - 1])) - - # Update Ey on edges (xz) - if self.blocks_mat[0, 1, 1] is not None and self.blocks_mat[1, 1, 0] is not None: - for jj in range(Ny): - Ey[0, jj, 0] = (Ey[0, jj, 0] - self.C3 * self.Jy[0, jj, 0] + - self.C8 * (Hx[0, jj, 0] - self.blocks_mat[1, 1, 0].Hx[0, jj, -1]) - - self.C5 * (Hz[0, jj, 0] - self.blocks_mat[0, 1, 1].Hz[-1, jj, 0])) - if self.blocks_mat[0, 1, 1] is not None and self.blocks_mat[1, 1, 2] is not None: - for jj in range(Ny): - Ey[0, jj, Nz] = (Ey[0, jj, Nz] - self.C3 * self.Jy[0, jj, Nz] + - self.C8 * (self.blocks_mat[1, 1, 2].Hx[0, jj, 0] - Hx[0, jj, Nz - 1]) - - self.C5 * (Hz[0, jj, Nz] - self.blocks_mat[0, 1, 1].Hz[-1, jj, Nz])) - if self.blocks_mat[2, 1, 1] is not None and self.blocks_mat[1, 1, 0] is not None: - for jj in range(Ny): - Ey[Nx, jj, 0] = (Ey[Nx, jj, 0] - self.C3 * self.Jy[Nx, jj, 0] + - self.C8 * (Hx[Nx, jj, 0] - self.blocks_mat[1, 1, 0].Hx[Nx, jj, -1]) - - self.C5 * (self.blocks_mat[2, 1, 1].Hz[0, jj, 0] - Hz[Nx - 1, jj, 0])) - if self.blocks_mat[2, 1, 1] is not None and self.blocks_mat[1, 1, 2] is not None: - for jj in range(Ny): - Ey[Nx, jj, Nz] = (Ey[Nx, jj, Nz] - self.C3 * self.Jy[Nx, jj, Nz] + - self.C8 * (self.blocks_mat[1, 1, 2].Hx[Nx, jj, 0] - Hx[Nx, jj, Nz - 1]) - - self.C5 * (self.blocks_mat[2, 1, 1].Hz[0, jj, Nz] - Hz[Nx - 1, jj, Nz])) - - def gauss(self, t): - tau = 10 * self.dt - if t < 6 * tau: - return 100 * np.exp(-(t - 3 * tau) ** 2 / tau ** 2) - else: - return 0. - - def one_step(self): - if self.sol_type == 'ECT': - self.one_step_ect() - if self.sol_type == 'FDTD': - self.one_step_fdtd() - if self.sol_type == 'DM': - self.one_step_dm() - - self.time += self.dt - - def one_step_ect(self): - self.compute_v_and_rho() - self.advance_h_ect() - for block in self.blocks: - block.advance_h_fdtd() - block.sum_h_fields() - self.advance_e_dm() - self.update_e_boundary() - for block in self.blocks: - block.advance_e_fdtd() - block.update_e_boundary() - block.sum_e_fields() - - def one_step_fdtd(self): - self.advance_h_fdtd(self.grid.Sxy, self.grid.Syz, self.grid.Szx, self.Ex, self.Ey, - self.Ez, self.Hx, self.Hy, self.Hz, self.Nx, self.Ny, self.Nz, - self.C1, self.C2, self.C7) - - for block in self.blocks: - block.advance_h_fdtd() - - block.sum_h_fields() - self.advance_e_fdtd(self.grid.l_x, self.grid.l_y, self.grid.l_z, self.Ex, self.Ey, self.Ez, - self.Hx, self.Hy, self.Hz, self.Jx, self.Jy, self.Jz, self.Nx, self.Ny, - self.Nz, self.C3, self.C4, self.C5, self.C8) - self.update_e_boundary() - for block in self.blocks: - block.advance_e_fdtd() - block.update_e_boundary() - block.sum_e_fields() - - @staticmethod - @jit(nopython=True) - def advance_h_fdtd(Sxy, Syz, Szx, Ex, Ey, Ez, Hx, Hy, Hz, Nx, Ny, Nz, C1, C2, C7): - - # Compute cell voltages - for ii in range(Nx + 1): - for jj in range(Ny): - for kk in range(Nz): - if Syz[ii, jj, kk] > 0: - Hx[ii, jj, kk] = (Hx[ii, jj, kk] - - C2 * (Ez[ii, jj + 1, kk] - Ez[ii, jj, kk]) + - C7 * (Ey[ii, jj, kk + 1] - Ey[ii, jj, kk])) - - for ii in range(Nx): - for jj in range(Ny + 1): - for kk in range(Nz): - if Szx[ii, jj, kk] > 0: - Hy[ii, jj, kk] = (Hy[ii, jj, kk] - - C7 * (Ex[ii, jj, kk + 1] - Ex[ii, jj, kk]) + - C1 * (Ez[ii + 1, jj, kk] - Ez[ii, jj, kk])) - - for ii in range(Nx): - for jj in range(Ny): - for kk in range(Nz + 1): - if Sxy[ii, jj, kk] > 0: - Hz[ii, jj, kk] = (Hz[ii, jj, kk] - - C1 * (Ey[ii + 1, jj, kk] - Ey[ii, jj, kk]) + - C2 * (Ex[ii, jj + 1, kk] - Ex[ii, jj, kk])) - - @staticmethod - @jit(nopython=True) - def advance_e_fdtd(l_x, l_y, l_z, Ex, Ey, Ez, Hx, Hy, Hz, Jx, Jy, Jz, Nx, Ny, Nz, C3, C4, C5, C8): - - for ii in range(Nx): - for jj in range(1, Ny): - for kk in range(1, Nz): - if l_x[ii, jj, kk] > 0: - Ex[ii, jj, kk] = (Ex[ii, jj, kk] - C3 * Jx[ii, jj, kk] + - C4 * (Hz[ii, jj, kk] - Hz[ii, jj - 1, kk]) - - C8 * (Hy[ii, jj, kk] - Hy[ii, jj, kk - 1])) - - for ii in range(1, Nx): - for jj in range(Ny): - for kk in range(1, Nz): - if l_y[ii, jj, kk] > 0: - Ey[ii, jj, kk] = (Ey[ii, jj, kk] - C3 * Jy[ii, jj, kk] + - C8 * (Hx[ii, jj, kk] - Hx[ii, jj, kk - 1]) - - C5 * (Hz[ii, jj, kk] - Hz[ii - 1, jj, kk])) - - for ii in range(1, Nx): - for jj in range(1, Ny): - for kk in range(Nz): - if l_z[ii, jj, kk] > 0: - Ez[ii, jj, kk] = (Ez[ii, jj, kk] - C3 * Jz[ii, jj, kk] + - C5 * (Hy[ii, jj, kk] - Hy[ii - 1, jj, kk]) - - C4 * (Hx[ii, jj, kk] - Hx[ii, jj - 1, kk])) - - def one_step_dm(self): - self.compute_v_and_rho() - - for i in range(self.Nx): - for j in range(self.Ny): - for k in range(self.Nz + 1): - if self.grid.flag_int_cell_xy[i, j, k]: - self.Hz[i, j, k] = (self.Hz[i, j, k] - - self.dt / (mu_0 * self.grid.Sxy[i, j, k]) * - self.Vxy[i, j, k]) - - for i in range(self.Nx + 1): - for j in range(self.Ny): - for k in range(self.Nz): - if self.grid.flag_int_cell_yz[i, j, k]: - self.Hx[i, j, k] = (self.Hx[i, j, k] - - self.dt / (mu_0 * self.grid.Syz[i, j, k]) * - self.Vyz[i, j, k]) - - for i in range(self.Nx): - for j in range(self.Ny + 1): - for k in range(self.Nz): - if self.grid.flag_int_cell_zx[i, j, k]: - self.Hy[i, j, k] = (self.Hy[i, j, k] - - self.dt / (mu_0 * self.grid.Szx[i, j, k]) * - self.Vzx[i, j, k]) - - self.advance_e_dm() - - def advance_h_ect(self, dt=None): - - for ii in range(self.Nx + 1): - EMSolver2D.advance_h_ect(Nx=self.Ny, Ny=self.Nz, V_enl=self.Vyz_enl[ii, :, :], - rho=self.rho_yz[ii, :, :], Hz=self.Hx[ii, :, :], C1=self.CN, - flag_int_cell=self.grid.flag_int_cell_yz[ii, :, :], - flag_unst_cell=self.grid.flag_unst_cell_yz[ii, :, :], - flag_intr_cell=self.grid.flag_intr_cell_yz[ii,:,:], - S=self.grid.Syz[ii, :, :], - borrowing=self.grid.borrowing_yz[ii, :, :], - S_enl=self.grid.Syz_enl[ii, :, :], - S_red=self.grid.Syz_red[ii, :, :], dt=dt, comp='x', kk=ii) - - for jj in range(self.Ny + 1): - EMSolver2D.advance_h_ect(Nx=self.Nx, Ny=self.Nz, V_enl=self.Vzx_enl[:, jj, :], - rho=self.rho_zx[:, jj, :], Hz=self.Hy[:, jj, :], C1=self.CN, - flag_int_cell=self.grid.flag_int_cell_zx[:, jj, :], - flag_unst_cell=self.grid.flag_unst_cell_zx[:, jj, :], - flag_intr_cell=self.grid.flag_intr_cell_zx[:,jj,:], - S=self.grid.Szx[:, jj, :], - borrowing=self.grid.borrowing_zx[:, jj, :], - S_enl=self.grid.Szx_enl[:, jj, :], - S_red=self.grid.Szx_red[:, jj, :], dt=dt, comp='y', kk=jj) - - for kk in range(self.Nz + 1): - EMSolver2D.advance_h_ect(Nx=self.Nx, Ny=self.Ny, V_enl=self.Vxy_enl[:, :, kk], - rho=self.rho_xy[:, :, kk], Hz=self.Hz[:, :, kk], C1=self.CN, - flag_int_cell=self.grid.flag_int_cell_xy[:, :, kk], - flag_unst_cell=self.grid.flag_unst_cell_xy[:, :, kk], - flag_intr_cell=self.grid.flag_intr_cell_xy[:,:,kk], - S=self.grid.Sxy[:, :, kk], - borrowing=self.grid.borrowing_xy[:, :, kk], - S_enl=self.grid.Sxy_enl[:, :, kk], - S_red=self.grid.Sxy_red[:, :, kk], dt=dt, comp='z', kk=kk) - - def compute_v_and_rho(self): - l_x = self.grid.l_x - l_y = self.grid.l_y - l_z = self.grid.l_z - - for ii in range(self.Nx): - for jj in range(self.Ny): - for kk in range(self.Nz + 1): - if self.grid.flag_int_cell_xy[ii, jj, kk]: - self.Vxy[ii, jj, kk] = (self.Ex[ii, jj, kk] * l_x[ii, jj, kk] - - self.Ex[ii, jj + 1, kk] * l_x[ii, jj + 1, kk] + - self.Ey[ii + 1, jj, kk] * l_y[ii + 1, jj, kk] - - self.Ey[ii, jj, kk] * l_y[ii, jj, kk]) - - if self.sol_type != 'DM': - self.rho_xy[ii, jj, kk] = (self.Vxy[ii, jj, kk] / - self.grid.Sxy[ii, jj, kk]) - - - for ii in range(self.Nx + 1): - for jj in range(self.Ny): - for kk in range(self.Nz): - if self.grid.flag_int_cell_yz[ii, jj, kk]: - self.Vyz[ii, jj, kk] = (self.Ey[ii, jj, kk] * l_y[ii, jj, kk] - - self.Ey[ii, jj, kk + 1] * l_y[ii, jj, kk + 1] + - self.Ez[ii, jj + 1, kk] * l_z[ii, jj + 1, kk] - - self.Ez[ii, jj, kk] * l_z[ii, jj, kk]) - - if self.sol_type != 'DM': - self.rho_yz[ii, jj, kk] = (self.Vyz[ii, jj, kk] / - self.grid.Syz[ii, jj, kk]) - - - for ii in range(self.Nx): - for jj in range(self.Ny + 1): - for kk in range(self.Nz): - if self.grid.flag_int_cell_zx[ii, jj, kk]: - self.Vzx[ii, jj, kk] = (self.Ez[ii, jj, kk] * l_z[ii, jj, kk] - - self.Ez[ii + 1, jj, kk] * l_z[ii + 1, jj, kk] + - self.Ex[ii, jj, kk + 1] * l_x[ii, jj, kk + 1] - - self.Ex[ii, jj, kk] * l_x[ii, jj, kk]) - - if self.sol_type != 'DM': - self.rho_zx[ii, jj, kk] = (self.Vzx[ii, jj, kk] / - self.grid.Szx[ii, jj, kk]) - - def advance_e_dm(self, dt=None): - - if dt==None: dt = self.dt - - Ex = self.Ex - Ey = self.Ey - Ez = self.Ez - Hx = self.Hx - Hy = self.Hy - Hz = self.Hz - - C4 = dt / (self.dy * eps_0) - C5 = dt / (self.dx * eps_0) - C8 = dt / (self.dz * eps_0) - C3 = dt / eps_0 - - for ii in range(self.Nx): - for jj in range(1, self.Ny): - for kk in range(1, self.Nz): - if self.grid.l_x[ii, jj, kk] > 0: - Ex[ii, jj, kk] = (Ex[ii, jj, kk] - self.C3 * self.Jx[ii, jj, kk] + - self.C4 * (Hz[ii, jj, kk] - Hz[ii, jj - 1, kk]) - - self.C8 * (Hy[ii, jj, kk] - Hy[ii, jj, kk - 1])) - - for ii in range(1, self.Nx): - for jj in range(self.Ny): - for kk in range(1, self.Nz): - if self.grid.l_y[ii, jj, kk] > 0: - Ey[ii, jj, kk] = (Ey[ii, jj, kk] - self.C3 * self.Jy[ii, jj, kk] + - self.C8 * (Hx[ii, jj, kk] - Hx[ii, jj, kk - 1]) - - self.C5 * (Hz[ii, jj, kk] - Hz[ii - 1, jj, kk])) - - for ii in range(1, self.Nx): - for jj in range(1, self.Ny): - for kk in range(self.Nz): - if self.grid.l_z[ii, jj, kk] > 0: - Ez[ii, jj, kk] = (Ez[ii, jj, kk] - self.C3 * self.Jz[ii, jj, kk] + - self.C5 * (Hy[ii, jj, kk] - Hy[ii - 1, jj, kk]) - - self.C4 * (Hx[ii, jj, kk] - Hx[ii, jj - 1, kk])) From 17df67e035f0340023c1ac059b64ecf40b79ac06 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:55:57 +0100 Subject: [PATCH 19/39] Delete build/lib/wakis/solverFIT3D.py --- build/lib/wakis/solverFIT3D.py | 1110 -------------------------------- 1 file changed, 1110 deletions(-) delete mode 100644 build/lib/wakis/solverFIT3D.py diff --git a/build/lib/wakis/solverFIT3D.py b/build/lib/wakis/solverFIT3D.py deleted file mode 100644 index c142264..0000000 --- a/build/lib/wakis/solverFIT3D.py +++ /dev/null @@ -1,1110 +0,0 @@ -# copyright ################################# # -# This file is part of the wakis Package. # -# Copyright (c) CERN, 2024. # -# ########################################### # - -from tqdm import tqdm - -import numpy as np -import time -import h5py - -from scipy.constants import c as c_light, epsilon_0 as eps_0, mu_0 as mu_0 -from scipy.sparse import csc_matrix as sparse_mat -from scipy.sparse import diags, hstack, vstack - -from .field import Field -from .materials import material_lib -from .plotting import PlotMixin -from .routines import RoutinesMixin - -try: - from cupyx.scipy.sparse import csc_matrix as gpu_sparse_mat - imported_cupyx = True -except ImportError: - imported_cupyx = False - -try: - from sparse_dot_mkl import csr_matrix as mkl_sparse_mat, dot_product_mkl - imported_mkl = True -except ImportError: - imported_mkl = False - - -class SolverFIT3D(PlotMixin, RoutinesMixin): - - def __init__(self, grid, wake=None, cfln=0.5, dt=None, - bc_low=['Periodic', 'Periodic', 'Periodic'], - bc_high=['Periodic', 'Periodic', 'Periodic'], - use_stl=False, use_conductors=False, - use_gpu=False, use_mpi=False, dtype=np.float64, - n_pml=10, bg=[1.0, 1.0], verbose=1): - ''' - Class holding the 3D time-domain electromagnetic solver - algorithm based on the Finite Integration Technique (FIT) - - Parameters: - ----------- - grid: GridFIT3D object - Instance of GridFIT3D class containing the simulation mesh and the - imported geometry - wake: WakeSolver object, optional - Instance of WakeSolver class containing the beam parameters. Needed to - run a wakefield simulation to compute wake potential and impedance - cfln: float, default 0.5 - Convergence condition by Courant–Friedrichs–Lewy, used to compute the - simulation timestep - dt: float, optional - Simulation timestep. If not None, it overrides the cfln-based timestep - bc_low: list, default ['Periodic', 'Periodic', 'Periodic'] - Domain box boundary conditions for X-, Y-, Z- - bc_high: list, default ['Periodic', 'Periodic', 'Periodic'] - Domain box boundary conditions for X+, Y+, Z+ - use_conductors: bool, default False - If true, enables geometry import based on elements from `conductors.py` - use_stl: bool, default False - If true, activates all the solids and materials passed to the `grid` object - use_gpu: bool, default False, - Using cupyx, enables GPU accelerated computation of every timestep - bg: list, default [1.0, 1.0] - Background material for the simulation box [eps_r, mu_r, sigma]. Default is vacuum. - It supports any material from the material library in `materials.py`, of a - custom list of floats. If conductivity (sigma) is passed, - it enables flag: use_conductivity - verbose: int or bool, default True - Enable verbose ouput on the terminal if 1 or True - - Attributes - ---------- - E: Field object - Object to access the Electric field data in [V/m]. - E.g.: solver.E[:,:,n,'z'] gives a 2D numpy.ndarray fieldmap of Ez component, located at the n-th cell in z - H: Field object - Object to access the Magnetic field data in [A/m]. - E.g.: solver.H[i,j,k,'x'] gives a point value of Hx component, located at the i,j,k cell - J: Field object - Object to access the Current density field data in [A/m^2]. - ieps: Field object - Object to access the ε^-1 tensor containing 1/permittivity values in the 3 dimensions. - imu: Field object - Object to access the μ^-1 tensor containing 1/permeability values in the 3 dimensions. - sigma: Field object - Object to access the condutcity σ tensor in the 3 dimensions. - ''' - - self.verbose = verbose - if verbose: t0 = time.time() - - # Flags - self.step_0 = True - self.nstep = int(0) - self.plotter_active = False - self.use_conductors = use_conductors - self.use_stl = use_stl - self.use_gpu = use_gpu - self.use_mpi = use_mpi - self.activate_abc = False # Will turn true if abc BCs are chosen - self.activate_pml = False # Will turn true if pml BCs are chosen - self.use_conductivity = False # Will turn true if conductive material or pml is added - self.imported_mkl = imported_mkl # Use MKL backend when available - self.one_step = self._one_step - if use_stl: - self.use_conductors = False - - # Grid - self.grid = grid - - self.Nx = self.grid.Nx - self.Ny = self.grid.Ny - self.Nz = self.grid.Nz - self.N = self.Nx*self.Ny*self.Nz - - self.dx = self.grid.dx - self.dy = self.grid.dy - self.dz = self.grid.dz - - self.x = self.grid.x[:-1]+self.dx/2 - self.y = self.grid.y[:-1]+self.dy/2 - self.z = self.grid.z[:-1]+self.dz/2 - - self.L = self.grid.L - self.iA = self.grid.iA - self.tL = self.grid.tL - self.itA = self.grid.itA - - # Wake computation - self.wake = wake - - # Fields - self.dtype = dtype - self.E = Field(self.Nx, self.Ny, self.Nz, use_gpu=self.use_gpu, dtype=self.dtype) - self.H = Field(self.Nx, self.Ny, self.Nz, use_gpu=self.use_gpu, dtype=self.dtype) - self.J = Field(self.Nx, self.Ny, self.Nz, use_gpu=self.use_gpu, dtype=self.dtype) - - # MPI init - if self.use_mpi: - if self.grid.use_mpi: - self.mpi_initialize() - self.one_step = self.mpi_one_step - else: - print('*** Grid not subdivided for MPI, set `use_mpi`=True also in `GridFIT3D` to enable MPI') - - # Matrices - if verbose: print('Assembling operator matrices...') - N = self.N - self.Px = diags([-1, 1], [0, 1], shape=(N, N), dtype=np.int8) - self.Py = diags([-1, 1], [0, self.Nx], shape=(N, N), dtype=np.int8) - self.Pz = diags([-1, 1], [0, self.Nx*self.Ny], shape=(N, N), dtype=np.int8) - - # original grid - self.Ds = diags(self.L.toarray(), shape=(3*N, 3*N), dtype=self.dtype) - self.iDa = diags(self.iA.toarray(), shape=(3*N, 3*N), dtype=self.dtype) - - # tilde grid - self.tDs = diags(self.tL.toarray(), shape=(3*N, 3*N), dtype=self.dtype) - self.itDa = diags(self.itA.toarray(), shape=(3*N, 3*N), dtype=self.dtype) - - # Curl matrix - self.C = vstack([ - hstack([sparse_mat((N,N)), -self.Pz, self.Py]), - hstack([self.Pz, sparse_mat((N,N)), -self.Px]), - hstack([-self.Py, self.Px, sparse_mat((N,N))]) - ], dtype=np.int8) - - # Boundaries - if verbose: print('Applying boundary conditions...') - self.bc_low = bc_low - self.bc_high = bc_high - self.apply_bc_to_C() - - # Materials - if verbose: print('Adding material tensors...') - if type(bg) is str: - bg = material_lib[bg.lower()] - - if len(bg) == 3: - self.eps_bg, self.mu_bg, self.sigma_bg = bg[0]*eps_0, bg[1]*mu_0, bg[2] - self.use_conductivity = True - else: - self.eps_bg, self.mu_bg, self.sigma_bg = bg[0]*eps_0, bg[1]*mu_0, 0.0 - - self.ieps = Field(self.Nx, self.Ny, self.Nz, use_ones=True, dtype=self.dtype)*(1./self.eps_bg) - self.imu = Field(self.Nx, self.Ny, self.Nz, use_ones=True, dtype=self.dtype)*(1./self.mu_bg) - self.sigma = Field(self.Nx, self.Ny, self.Nz, use_ones=True, dtype=self.dtype)*self.sigma_bg - - if self.use_stl: - self.apply_stl() - - # Fill PML BCs - if self.activate_pml: - if verbose: print('Filling PML sigmas...') - self.n_pml = n_pml - self.pml_lo = 5e-3 - self.pml_hi = 1.e-1 - self.pml_func = np.geomspace - self.fill_pml_sigmas() - - # Timestep calculation - if verbose: print('Calculating maximal stable timestep...') - self.cfln = cfln - if dt is None: - self.dt = cfln / (c_light * np.sqrt(1 / self.grid.dx ** 2 + 1 / self.grid.dy ** 2 + 1 / self.grid.dz ** 2)) - else: - self.dt = dt - self.dt = dtype(self.dt) - - if self.use_conductivity: # relaxation time criterion tau - - mask = np.logical_and(self.sigma.toarray()!=0, #for non-conductive - self.ieps.toarray()!=0) #for PEC eps=inf - - self.tau = (1/self.ieps.toarray()[mask]) / \ - self.sigma.toarray()[mask] - - if self.dt > self.tau.min(): - self.dt = self.tau.min() - - # Pre-computing - if verbose: print('Pre-computing...') - self.iDeps = diags(self.ieps.toarray(), shape=(3*N, 3*N), dtype=self.dtype) - self.iDmu = diags(self.imu.toarray(), shape=(3*N, 3*N), dtype=self.dtype) - self.Dsigma = diags(self.sigma.toarray(), shape=(3*N, 3*N), dtype=self.dtype) - - self.tDsiDmuiDaC = self.tDs * self.iDmu * self.iDa * self.C - self.itDaiDepsDstC = self.itDa * self.iDeps * self.Ds * self.C.transpose() - - if imported_mkl and not self.use_gpu: # MKL backend for CPU - print('Using MKL backend for time-stepping...') - self.tDsiDmuiDaC = mkl_sparse_mat(self.tDsiDmuiDaC) - self.itDaiDepsDstC = mkl_sparse_mat(self.itDaiDepsDstC) - self.one_step = self.one_step_mkl - - # Move to GPU - if use_gpu: - if verbose: print('Moving to GPU...') - if imported_cupyx: - self.tDsiDmuiDaC = gpu_sparse_mat(self.tDsiDmuiDaC) - self.itDaiDepsDstC = gpu_sparse_mat(self.itDaiDepsDstC) - self.ieps.to_gpu() - self.sigma.to_gpu() - else: - raise ImportError('*** cupyx could not be imported, please check CUDA installation') - - if verbose: print(f'Total initialization time: {time.time() - t0} s') - - def update_tensors(self, tensor='all'): - '''Update tensor matrices after - Field ieps, imu or sigma have been modified - and pre-compute the time-stepping matrices - - Parameters: - ----------- - tensor : str, default 'all' - Name of the tensor to update: 'ieps', 'imu', 'sigma' - for permitivity, permeability and conductivity, respectively. - If left to default 'all', all thre tensors will be recomputed. - ''' - if self.verbose: print(f'Re-computing tensor "{tensor}"...') - - if tensor == 'ieps': - self.iDeps = diags(self.ieps.toarray(), shape=(3*self.N, 3*self.N), dtype=self.dtype) - elif tensor =='imu': - self.iDmu = diags(self.imu.toarray(), shape=(3*self.N, 3*self.N), dtype=self.dtype) - elif tensor == 'sigma': - self.Dsigma = diags(self.sigma.toarray(), shape=(3*self.N, 3*self.N), dtype=self.dtype) - elif tensor == 'all': - self.iDeps = diags(self.ieps.toarray(), shape=(3*self.N, 3*self.N), dtype=self.dtype) - self.iDmu = diags(self.imu.toarray(), shape=(3*self.N, 3*self.N), dtype=self.dtype) - self.Dsigma = diags(self.sigma.toarray(), shape=(3*self.N, 3*self.N), dtype=self.dtype) - - if self.verbose: print('Re-Pre-computing ...') - self.tDsiDmuiDaC = self.tDs * self.iDmu * self.iDa * self.C - self.itDaiDepsDstC = self.itDa * self.iDeps * self.Ds * self.C.transpose() - self.step_0 = False - - def _one_step(self): - if self.step_0: - self.set_ghosts_to_0() - self.step_0 = False - self.attrcleanup() - - self.H.fromarray(self.H.toarray() - - self.dt*self.tDsiDmuiDaC*self.E.toarray() - ) - - self.E.fromarray(self.E.toarray() + - self.dt*(self.itDaiDepsDstC*self.H.toarray() - - self.ieps.toarray()*self.J.toarray() - ) - ) - - #include current computation - if self.use_conductivity: - self.J.fromarray(self.sigma.toarray()*self.E.toarray()) - - def one_step_mkl(self): - if self.step_0: - self.set_ghosts_to_0() - self.step_0 = False - self.attrcleanup() - - self.H.fromarray(self.H.toarray() - - self.dt*dot_product_mkl(self.tDsiDmuiDaC,self.E.toarray()) - ) - - self.E.fromarray(self.E.toarray() + - self.dt*(dot_product_mkl(self.itDaiDepsDstC,self.H.toarray()) - - self.ieps.toarray()*self.J.toarray() - ) - ) - - #include current computation - if self.use_conductivity: - self.J.fromarray(self.sigma.toarray()*self.E.toarray()) - - def mpi_initialize(self): - self.comm = self.grid.comm - self.rank = self.grid.rank - self.size = self.grid.size - - self.NZ = self.grid.NZ - self.ZMIN = self.grid.ZMIN - self.ZMAX = self.grid.ZMAX - self.Z = self.grid.Z - - def mpi_one_step(self): - if self.step_0: - self.set_ghosts_to_0() - self.step_0 = False - self.attrcleanup() - - self.H.fromarray(self.H.toarray() - - self.dt*self.tDsiDmuiDaC*self.E.toarray() - ) - - self.mpi_communicate(self.H) - self.mpi_communicate(self.J) - self.E.fromarray(self.E.toarray() + - self.dt*(self.itDaiDepsDstC * self.H.toarray() - - self.ieps.toarray()*self.J.toarray() - ) - ) - - self.mpi_communicate(self.E) - # include current computation - if self.use_conductivity: - self.J.fromarray(self.sigma.toarray()*self.E.toarray()) - - def mpi_communicate(self, field): - if self.use_gpu: - field.from_gpu() - - # ghosts lo - if self.rank > 0: - for d in ['x','y','z']: - self.comm.Sendrecv(field[:, :, 1, d], - recvbuf=field[:, :, 0, d], - dest=self.rank-1, sendtag=0, - source=self.rank-1, recvtag=1) - # ghosts hi - if self.rank < self.size - 1: - for d in ['x','y','z']: - self.comm.Sendrecv(field[:, :, -2, d], - recvbuf=field[:, :, -1, d], - dest=self.rank+1, sendtag=1, - source=self.rank+1, recvtag=0) - - if self.use_gpu: - field.to_gpu() - - def mpi_gather(self, field, x=None, y=None, z=None, component=None): - ''' - Gather a specific component or slice of a distributed field from all MPI ranks. - - This function collects a selected component of a field (E, H, J, or custom) - from all MPI processes along the z-axis and reconstructs the global field data - on the root rank (rank 0). The user can specify slices or single indices - along x, y, and z to control the subset of data gathered. - - Parameters - ---------- - field : str or Field obj - The field to gather. If a string, it must begin with one of: - - `'E'`, `'H'`, or `'J'` followed optionally by a component label - (e.g., `'Ex'`, `'Hz'`, `'JAbs'`). - If no component is specified, defaults to `'z'`. - - x : int or slice, optional - Range of x-indices to gather. If None, defaults to the full x-range. - - y : int or slice, optional - Range of y-indices to gather. If None, defaults to the full y-range. - - z : int or slice, optional - Range of z-indices to gather. If None, defaults to the full z-range. - - component : str or slice, optional - Component of the field to gather ('x', 'y', 'z', or a slice). - If None and not inferred from `field`, defaults to `'z'`. - - Returns - ------- - numpy.ndarray or None - The gathered field values assembled on rank 0 with shape depending on - the selected slices along (x, y, z). Returns `None` on non-root ranks. - - Notes - ----- - - Assumes field data is distributed along the z-dimension. - - Automatically handles removal of ghost cells in reconstruction. - - Field components are inferred from the input `field` string or the - `component` argument. - - This method supports full 3D subvolume extraction or 1D/2D slices - for performance diagnostics and visualization. - - Examples - -------- - >>> # Gather Ex component on full domain on rank 0 - >>> global_Ex = solver.mpi_gather('Ex') - - >>> # Gather a 2D yz-slice at x=10 of the J field - >>> yz_J = solver.mpi_gather('J', x=10) - ''' - - if x is None: - x = slice(0, self.Nx) - if y is None: - y = slice(0, self.Ny) - if z is None: - z = slice(0, self.NZ) - - if type(field) is str: - if len(field) == 2: #support for e.g. field='Ex' - component = field[1] - field = field[0] - elif len(field) == 4: #support for Abs - component = field[1:] - field = field[0] - elif component is None: - component = 'z' - print("[!] `component` not specified, using default component='z'") - - if field == 'E': - local = self.E[x, y, :, component].ravel() - elif field == 'H': - local = self.H[x, y, :, component].ravel() - elif field == 'J': - local = self.J[x, y, :, component].ravel() - else: - if component is None: - component = 'z' - print("[!] `component` not specified, using default component='z'") - local = field[x, y, :, component].ravel() - - buffer = self.comm.gather(local, root=0) - _field = None - - if self.rank == 0: - if type(x) is int and type(y) is int: # 1d array at x=a, y=b - nz = self.NZ//self.size - _field = np.zeros((self.NZ)) - for r in range(self.size): - zz = np.s_[r*nz:(r+1)*nz] - if r == 0: - _field[zz] = np.reshape(buffer[r], (nz+self.grid.n_ghosts))[:-1] - elif r == (self.size-1): - _field[zz] = np.reshape(buffer[r], (nz+self.grid.n_ghosts))[1:] - else: - _field[zz] = np.reshape(buffer[r], (nz+2*self.grid.n_ghosts))[1:-1] - _field = _field[z] - - elif type(x) is int: # 2d slice at x=a - ny = y.stop-y.start - nz = self.NZ//self.size - _field = np.zeros((ny, self.NZ)) - for r in range(self.size): - zz = np.s_[r*nz:(r+1)*nz] - if r == 0: - _field[:, zz] = np.reshape(buffer[r], (ny, nz+self.grid.n_ghosts))[:, :-1] - elif r == (self.size-1): - _field[:, zz] = np.reshape(buffer[r], (ny, nz+self.grid.n_ghosts))[:, 1:] - else: - _field[:, zz] = np.reshape(buffer[r], (ny, nz+2*self.grid.n_ghosts))[:, 1:-1] - _field = _field[:, z] - - elif type(y) is int: # 2d slice at y=a - nx = x.stop-x.start - nz = self.NZ//self.size - _field = np.zeros((nx, self.NZ)) - for r in range(self.size): - zz = np.s_[r*nz:(r+1)*nz] - if r == 0: - _field[:, zz] = np.reshape(buffer[r], (nx, nz+self.grid.n_ghosts))[:, :-1] - elif r == (self.size-1): - _field[:, zz] = np.reshape(buffer[r], (nx, nz+self.grid.n_ghosts))[:, 1:] - else: - _field[:, zz] = np.reshape(buffer[r], (nx, nz+2*self.grid.n_ghosts))[:, 1:-1] - _field = _field[:, z] - - else: # both type slice -> 3d array - nx = x.stop-x.start - ny = y.stop-y.start - nz = self.NZ//self.size - _field = np.zeros((nx, ny, self.NZ)) - for r in range(self.size): - zz = np.s_[r*nz:(r+1)*nz] - if r == 0: - _field[:, :, zz] = np.reshape(buffer[r], (nx, ny, nz+self.grid.n_ghosts))[:, :, :-1] - elif r == (self.size-1): - _field[:, :, zz] = np.reshape(buffer[r], (nx, ny, nz+self.grid.n_ghosts))[:, :, 1:] - else: - _field[:, :, zz] = np.reshape(buffer[r], (nx, ny, nz+2*self.grid.n_ghosts))[:, :, 1:-1] - _field = _field[:, :, z] - - return _field - - def mpi_gather_asField(self, field): - ''' - Gather distributed field data from all MPI ranks and return a global Field object. - - This method collects the specified electromagnetic field data (E, H, or J) - from all MPI processes and assembles it into a single global `Field` object - on the root rank (rank 0). The field data can be specified as a string - ('E', 'H', or 'J') or as a `wakis.Field` object. - - Parameters - ---------- - field : str or Field obj - The field to gather. If a string, it must be one of: - - `'E'` for the electric field - - `'H'` for the magnetic field - - `'J'` for the current density - - Passing a `wakis.Field` is also supported (e.g. ieps, imu, sigma) - - Returns - ------- - Field - A `wakis.Field` object containing the gathered global field data - with shape (Nx, Ny, NZ, 3). Only returned on rank 0. On other - ranks, the returned value is undefined and should not be used. - - Notes - ----- - - This method assumes the field is distributed along the `z`-axis. - - Ghost cells are removed appropriately when reassembling the global field. - ''' - - _field = Field(self.Nx, self.Ny, self.NZ) - - for d in ['x','y','z']: - if type(field) is str: - if field == 'E': - local = self.E[:, :, :,d].ravel() - elif field == 'H': - local = self.H[:, :, :,d].ravel() - elif field == 'J': - local = self.J[:, :, :,d].ravel() - else: - local = field[:, :, :, d].ravel() - - buffer = self.comm.gather(local, root=0) - if self.rank == 0: - nz = self.NZ//self.size - for r in range(self.size): - zz = np.s_[r*nz:(r+1)*nz] - if r == 0: - _field[:, :, zz, d] = np.reshape(buffer[r], (self.Nx, self.Ny, nz+self.grid.n_ghosts))[:, :, :-1] - elif r == (self.size-1): - _field[:, :, zz, d] = np.reshape(buffer[r], (self.Nx, self.Ny, nz+self.grid.n_ghosts))[:, :, 1:] - else: - _field[:, :, zz, d] = np.reshape(buffer[r], (self.Nx, self.Ny, nz+2*self.grid.n_ghosts))[:, :, 1:-1] - - return _field - - def apply_bc_to_C(self): - ''' - Modifies rows or columns of C and tDs and itDa matrices - according to bc_low and bc_high - ''' - xlo, ylo, zlo = 1., 1., 1. - xhi, yhi, zhi = 1., 1., 1. - - # Check BCs for internal MPI subdomains - if self.use_mpi and self.grid.use_mpi: - if self.rank > 0: - self.bc_low=['pec', 'pec', 'mpi'] - - if self.rank < self.size - 1: - self.bc_high=['pec', 'pec', 'mpi'] - - # Perodic: out == in - if any(True for x in self.bc_low if x.lower() == 'periodic'): - if self.bc_low[0].lower() == 'periodic' and self.bc_high[0].lower() == 'periodic': - self.tL[-1, :, :, 'x'] = self.L[0, :, :, 'x'] - self.itA[-1, :, :, 'y'] = self.iA[0, :, :, 'y'] - self.itA[-1, :, :, 'z'] = self.iA[0, :, :, 'z'] - - if self.bc_low[1].lower() == 'periodic' and self.bc_high[1].lower() == 'periodic': - self.tL[:, -1, :, 'y'] = self.L[:, 0, :, 'y'] - self.itA[:, -1, :, 'x'] = self.iA[:, 0, :, 'x'] - self.itA[:, -1, :, 'z'] = self.iA[:, 0, :, 'z'] - - if self.bc_low[2].lower() == 'periodic' and self.bc_high[2].lower() == 'periodic': - self.tL[:, :, -1, 'z'] = self.L[:, :, 0, 'z'] - self.itA[:, :, -1, 'x'] = self.iA[:, :, 0, 'x'] - self.itA[:, :, -1, 'y'] = self.iA[:, :, 0, 'y'] - - self.tDs = diags(self.tL.toarray(), shape=(3*self.N, 3*self.N), dtype=float) - self.itDa = diags(self.itA.toarray(), shape=(3*self.N, 3*self.N), dtype=float) - - # Dirichlet PEC: tangential E field = 0 at boundary - if any(True for x in self.bc_low if x.lower() in ('electric','pec','pml')) \ - or any(True for x in self.bc_high if x.lower() in ('electric','pec','pml')): - if self.bc_low[0].lower() in ('electric','pec', 'pml'): - xlo = 0 - if self.bc_low[1].lower() in ('electric','pec', 'pml'): - ylo = 0 - if self.bc_low[2].lower() in ('electric','pec', 'pml'): - zlo = 0 - if self.bc_high[0].lower() in ('electric','pec', 'pml'): - xhi = 0 - if self.bc_high[1].lower() in ('electric','pec', 'pml'): - yhi = 0 - if self.bc_high[2].lower() in ('electric','pec', 'pml'): - zhi = 0 - - # Assemble matrix - self.BC = Field(self.Nx, self.Ny, self.Nz, dtype=np.int8, use_ones=True) - - for d in ['x', 'y', 'z']: #tangential to zero - if d != 'x': - self.BC[0, :, :, d] = xlo - self.BC[-1, :, :, d] = xhi - if d != 'y': - self.BC[:, 0, :, d] = ylo - self.BC[:, -1, :, d] = yhi - if d != 'z': - self.BC[:, :, 0, d] = zlo - self.BC[:, :, -1, d] = zhi - - self.Dbc = diags(self.BC.toarray(), - shape=(3*self.N, 3*self.N), - dtype=np.int8 - ) - - # Update C (columns) - self.C = self.C*self.Dbc - - - # Dirichlet PMC: tangential H field = 0 at boundary - if any(True for x in self.bc_low if x.lower() in ('magnetic','pmc')) \ - or any(True for x in self.bc_high if x.lower() in ('magnetic','pmc')): - if self.bc_low[0].lower() == 'magnetic' or self.bc_low[0] == 'pmc': - xlo = 0 - if self.bc_low[1].lower() == 'magnetic' or self.bc_low[1] == 'pmc': - ylo = 0 - if self.bc_low[2].lower() == 'magnetic' or self.bc_low[2] == 'pmc': - zlo = 0 - if self.bc_high[0].lower() == 'magnetic' or self.bc_high[0] == 'pmc': - xhi = 0 - if self.bc_high[1].lower() == 'magnetic' or self.bc_high[1] == 'pmc': - yhi = 0 - if self.bc_high[2].lower() == 'magnetic' or self.bc_high[2] == 'pmc': - zhi = 0 - - # Assemble matrix - self.BC = Field(self.Nx, self.Ny, self.Nz, dtype=np.int8, use_ones=True) - - for d in ['x', 'y', 'z']: #tangential to zero - if d != 'x': - self.BC[0, :, :, d] = xlo - self.BC[-1, :, :, d] = xhi - if d != 'y': - self.BC[:, 0, :, d] = ylo - self.BC[:, -1, :, d] = yhi - if d != 'z': - self.BC[:, :, 0, d] = zlo - self.BC[:, :, -1, d] = zhi - - self.Dbc = diags(self.BC.toarray(), - shape=(3*self.N, 3*self.N), - dtype=np.int8 - ) - - # Update C (rows) - self.C = self.Dbc*self.C - - # Absorbing boundary conditions ABC - if any(True for x in self.bc_low if x.lower() == 'abc') \ - or any(True for x in self.bc_high if x.lower() == 'abc'): - if self.bc_high[0].lower() == 'abc': - self.tL[-1, :, :, 'x'] = self.L[0, :, :, 'x'] - self.itA[-1, :, :, 'y'] = self.iA[0, :, :, 'y'] - self.itA[-1, :, :, 'z'] = self.iA[0, :, :, 'z'] - - if self.bc_high[1].lower() == 'abc': - self.tL[:, -1, :, 'y'] = self.L[:, 0, :, 'y'] - self.itA[:, -1, :, 'x'] = self.iA[:, 0, :, 'x'] - self.itA[:, -1, :, 'z'] = self.iA[:, 0, :, 'z'] - - if self.bc_high[2].lower() == 'abc': - self.tL[:, :, -1, 'z'] = self.L[:, :, 0, 'z'] - self.itA[:, :, -1, 'x'] = self.iA[:, :, 0, 'x'] - self.itA[:, :, -1, 'y'] = self.iA[:, :, 0, 'y'] - - self.tDs = diags(self.tL.toarray(), shape=(3*self.N, 3*self.N), dtype=float) - self.itDa = diags(self.itA.toarray(), shape=(3*self.N, 3*self.N), dtype=float) - self.activate_abc = True - - # Perfect Matching Layers (PML) - if any(True for x in self.bc_low if x.lower() == 'pml') \ - or any(True for x in self.bc_high if x.lower() == 'pml'): - self.activate_pml = True - self.use_conductivity = True - - def fill_pml_sigmas(self): - ''' - Routine to calculate pml sigmas and apply them - to the conductivity tensor sigma - - [IN-PROGRESS] - ''' - - # Initialize - sx, sy, sz = np.zeros(self.Nx), np.zeros(self.Ny), np.zeros(self.Nz) - pml_exp = 2 - - # Fill - if self.bc_low[0].lower() == 'pml': - #sx[0:self.n_pml] = eps_0/(2*self.dt)*((self.x[self.n_pml] - self.x[:self.n_pml])/(self.n_pml*self.dx))**pml_exp - sx[0:self.n_pml] = np.linspace( self.pml_hi, self.pml_lo, self.n_pml) - for d in ['x', 'y', 'z']: - for i in range(self.n_pml): - self.ieps[i, :, :, d] = 1./eps_0 - self.sigma[i, :, :, d] = sx[i] - #if sx[i] > 0 : self.ieps[i, :, :, d] = 1/(eps_0+sx[i]*(2*self.dt)) - - if self.bc_low[1].lower() == 'pml': - #sy[0:self.n_pml] = 1/(2*self.dt)*((self.y[self.n_pml] - self.y[:self.n_pml])/(self.n_pml*self.dy))**pml_exp - sy[0:self.n_pml] = self.pml_func( self.pml_hi, self.pml_lo, self.n_pml) - for d in ['x', 'y', 'z']: - for j in range(self.n_pml): - self.ieps[:, j, :, d] = 1./eps_0 - self.sigma[:, j, :, d] = sy[j] - #if sy[j] > 0 : self.ieps[:, j, :, d] = 1/(eps_0+sy[j]*(2*self.dt)) - - if self.bc_low[2].lower() == 'pml': - #sz[0:self.n_pml] = eps_0/(2*self.dt)*((self.z[self.n_pml] - self.z[:self.n_pml])/(self.n_pml*self.dz))**pml_exp - sz[0:self.n_pml] = self.pml_func( self.pml_hi, self.pml_lo, self.n_pml) - for d in ['x', 'y', 'z']: - for k in range(self.n_pml): - self.ieps[:, :, k, d] = 1./eps_0 - self.sigma[:, :, k, d] = sz[k] - #if sz[k] > 0. : self.ieps[:, :, k, d] = 1/(np.mean(sz[:self.n_pml])*eps_0) - - if self.bc_high[0].lower() == 'pml': - #sx[-self.n_pml:] = 1/(2*self.dt)*((self.x[-self.n_pml:] - self.x[-self.n_pml])/(self.n_pml*self.dx))**pml_exp - sx[-self.n_pml:] = self.pml_func( self.pml_lo, self.pml_hi, self.n_pml) - for d in ['x', 'y', 'z']: - for i in range(self.n_pml): - i +=1 - self.ieps[-i, :, :, d] = 1./eps_0 - self.sigma[-i, :, :, d] = sx[-i] - #if sx[-i] > 0 : self.ieps[-i, :, :, d] = 1/(eps_0+sx[-i]*(2*self.dt)) - - if self.bc_high[1].lower() == 'pml': - #sy[-self.n_pml:] = 1/(2*self.dt)*((self.y[-self.n_pml:] - self.y[-self.n_pml])/(self.n_pml*self.dy))**pml_exp - sy[-self.n_pml:] = self.pml_func( self.pml_lo, self.pml_hi, self.n_pml) - for d in ['x', 'y', 'z']: - for j in range(self.n_pml): - j +=1 - self.ieps[:, -j, :, d] = 1./eps_0 - self.sigma[:, -j, :, d] = sy[-j] - #if sy[-j] > 0 : self.ieps[:, -j, :, d] = 1/(eps_0+sy[-j]*(2*self.dt)) - - if self.bc_high[2].lower() == 'pml': - #sz[-self.n_pml:] = eps_0/(2*self.dt)*((self.z[-self.n_pml:] - self.z[-self.n_pml])/(self.n_pml*self.dz))**pml_exp - sz[-self.n_pml:] = self.pml_func( self.pml_lo, self.pml_hi, self.n_pml) - for d in ['x', 'y', 'z']: - for k in range(self.n_pml): - k +=1 - self.ieps[:, :, -k, d] = 1./eps_0 - self.sigma[:, :, -k, d] = sz[-k] - #self.ieps[:, :, -k, d] = 1/(np.mean(sz[-self.n_pml:])*eps_0) - - def get_abc(self): - ''' - Save the n-2 timestep to apply ABC - ''' - E_abc, H_abc = {}, {} - - if self.bc_low[0].lower() == 'abc': - E_abc[0] = {} - H_abc[0] = {} - for d in ['x', 'y', 'z']: - E_abc[0][d+'lo'] = self.E[1, :, :, d] - H_abc[0][d+'lo'] = self.H[1, :, :, d] - - if self.bc_low[1].lower() == 'abc': - E_abc[1] = {} - H_abc[1] = {} - for d in ['x', 'y', 'z']: - E_abc[1][d+'lo'] = self.E[:, 1, :, d] - H_abc[1][d+'lo'] = self.H[:, 1, :, d] - - if self.bc_low[2].lower() == 'abc': - E_abc[2] = {} - H_abc[2] = {} - for d in ['x', 'y', 'z']: - E_abc[2][d+'lo'] = self.E[:, :, 1, d] - H_abc[2][d+'lo'] = self.H[:, :, 1, d] - - if self.bc_high[0].lower() == 'abc': - E_abc[0] = {} - H_abc[0] = {} - for d in ['x', 'y', 'z']: - E_abc[0][d+'hi'] = self.E[-1, :, :, d] - H_abc[0][d+'hi'] = self.H[-1, :, :, d] - - if self.bc_high[1].lower() == 'abc': - E_abc[1] = {} - H_abc[1] = {} - for d in ['x', 'y', 'z']: - E_abc[1][d+'hi'] = self.E[:, -1, :, d] - H_abc[1][d+'hi'] = self.H[:, -1, :, d] - - if self.bc_high[2].lower() == 'abc': - E_abc[2] = {} - H_abc[2] = {} - for d in ['x', 'y', 'z']: - E_abc[2][d+'hi'] = self.E[:, :, -1, d] - H_abc[2][d+'hi'] = self.H[:, :, -1, d] - - return E_abc, H_abc - - def update_abc(self, E_abc, H_abc): - ''' - Apply ABC algo to the selected BC, - to be applied after each timestep - ''' - - if self.bc_low[0].lower() == 'abc': - for d in ['x', 'y', 'z']: - self.E[0, :, :, d] = E_abc[0][d+'lo'] - self.H[0, :, :, d] = H_abc[0][d+'lo'] - - if self.bc_low[1].lower() == 'abc': - for d in ['x', 'y', 'z']: - self.E[:, 0, :, d] = E_abc[1][d+'lo'] - self.H[:, 0, :, d] = H_abc[1][d+'lo'] - - if self.bc_low[2].lower() == 'abc': - for d in ['x', 'y', 'z']: - self.E[:, :, 0, d] = E_abc[2][d+'lo'] - self.H[:, :, 0, d] = H_abc[2][d+'lo'] - - if self.bc_high[0].lower() == 'abc': - for d in ['x', 'y', 'z']: - self.E[-1, :, :, d] = E_abc[0][d+'hi'] - self.H[-1, :, :, d] = H_abc[0][d+'hi'] - - if self.bc_high[1].lower() == 'abc': - for d in ['x', 'y', 'z']: - self.E[:, -1, :, d] = E_abc[1][d+'hi'] - self.H[:, -1, :, d] = H_abc[1][d+'hi'] - - if self.bc_high[2].lower() == 'abc': - for d in ['x', 'y', 'z']: - self.E[:, :, -1, d] = E_abc[2][d+'hi'] - self.H[:, :, -1, d] = H_abc[2][d+'hi'] - - def set_ghosts_to_0(self): - ''' - Cleanup for initial conditions if they are - accidentally applied to the ghost cells - ''' - # Set H ghost quantities to 0 - for d in ['x', 'y', 'z']: #tangential to zero - if d != 'x': - self.H[-1, :, :, d] = 0. - if d != 'y': - self.H[:, -1, :, d] = 0. - if d != 'z': - self.H[:, :, -1, d] = 0. - - # Set E ghost quantities to 0 - self.E[-1, :, :, 'x'] = 0. - self.E[:, -1, :, 'y'] = 0. - self.E[:, :, -1, 'z'] = 0. - - def apply_conductors(self): - ''' - Set the 1/epsilon values inside the PEC conductors to zero - ''' - self.flag_in_conductors = self.grid.flag_int_cell_yz[:-1,:,:] \ - + self.grid.flag_int_cell_zx[:,:-1,:] \ - + self.grid.flag_int_cell_xy[:,:,:-1] - - self.ieps *= self.flag_in_conductors - - def set_field_in_conductors_to_0(self): - ''' - Cleanup for initial conditions if they are - accidentally applied to the conductors - ''' - self.flag_cleanup = self.grid.flag_int_cell_yz[:-1,:,:] \ - + self.grid.flag_int_cell_zx[:,:-1,:] \ - + self.grid.flag_int_cell_xy[:,:,:-1] - - self.H *= self.flag_cleanup - self.E *= self.flag_cleanup - - def apply_stl(self): - ''' - Mask the cells inside the stl and assing the material - defined by the user - - * Note: stl material should contain **relative** epsilon and mu - ** Note 2: when assigning the stl material, the default values - 1./eps_0 and 1./mu_0 are substracted - ''' - grid = self.grid.grid - self.stl_solids = self.grid.stl_solids - self.stl_materials = self.grid.stl_materials - - for key in self.stl_solids.keys(): - - mask = np.reshape(grid[key], (self.Nx, self.Ny, self.Nz)).astype(int) - - if type(self.stl_materials[key]) is str: - # Retrieve from material library - mat_key = self.stl_materials[key].lower() - - eps = material_lib[mat_key][0]*eps_0 - mu = material_lib[mat_key][1]*mu_0 - - # Setting to zero - self.ieps += self.ieps * (-1.0*mask) - self.imu += self.imu * (-1.0*mask) - - # Adding new values - self.ieps += mask * 1./eps - self.imu += mask * 1./mu - - # Conductivity - if len(material_lib[mat_key]) == 3: - sigma = material_lib[mat_key][2] - self.sigma += self.sigma * (-1.0*mask) - self.sigma += mask * sigma - self.use_conductivity = True - - elif self.sigma_bg > 0.0: # assumed sigma = 0 - self.sigma += self.sigma * (-1.0*mask) - - else: - # From input - eps = self.stl_materials[key][0]*eps_0 - mu = self.stl_materials[key][1]*mu_0 - - # Setting to zero - self.ieps += self.ieps * (-1.0*mask) - self.imu += self.imu * (-1.0*mask) - - # Adding new values - self.ieps += mask * 1./eps - self.imu += mask * 1./mu - - # Conductivity - if len(self.stl_materials[key]) == 3: - sigma = self.stl_materials[key][2] - self.sigma += self.sigma * (-1.0*mask) - self.sigma += mask * sigma - self.use_conductivity = True - - elif self.sigma_bg > 0.0: # assumed sigma = 0 - self.sigma += self.sigma * (-1.0*mask) - - def attrcleanup(self): - # Fields - del self.L, self.tL, self.iA, self.itA - if hasattr(self, 'BC'): - del self.BC - del self.Dbc - - # Matrices - del self.Px, self.Py, self.Pz - del self.Ds, self.iDa, self.tDs, self.itDa - del self.C - - def save_state(self, filename="solver_state.h5", close=True): - """Save the solver state to an HDF5 file. - - This function saves only the key state variables (`H`, `E`, `J`) that are updated - in `one_step()`, storing them as datasets in an HDF5 file. - - Parameters: - ----------- - filename : str, optional - The name of the HDF5 file where the state will be stored. Default is "solver_state.h5". - close : bool, optional (default=True) - - If True, the HDF5 file is closed after saving, and the function returns nothing. - - If False, the function returns an open HDF5 file object, which must be closed manually. - - Returns: - -------- - h5py.File or None - - If `close=True`, nothing is returned. - - If `close=False`, returns an open `h5py.File` object for further manipulation. - """ - - if self.use_mpi: # MPI savestate - E = self.mpi_gather_asField('E') - H = self.mpi_gather_asField('H') - J = self.mpi_gather_asField('J') - - if self.rank == 0: - state = h5py.File(filename, "w") - state.create_dataset("H", data=self.H.toarray()) - state.create_dataset("E", data=self.E.toarray()) - state.create_dataset("J", data=self.J.toarray()) - # TODO: check for MPI-GPU - - elif self.use_gpu: # GPU savestate - state = h5py.File(filename, "w") - state.create_dataset("H", data=self.H.toarray().get()) - state.create_dataset("E", data=self.E.toarray().get()) - state.create_dataset("J", data=self.J.toarray().get()) - - else: # CPU savestate - state = h5py.File(filename, "w") - state.create_dataset("H", data=self.H.toarray()) - state.create_dataset("E", data=self.E.toarray()) - state.create_dataset("J", data=self.J.toarray()) - - if close: - state.close() - else: - return state # Caller must close this manually - - def load_state(self, filename="solver_state.h5"): - """Load the solver state from an HDF5 file. - - Reads the saved state variables (`H`, `E`, `J`) from the specified HDF5 file - and restores them to the solver. - - Parameters: - ----------- - filename : str, optional - The name of the HDF5 file to load the solver state from. Default is "solver_state.h5". - - Returns: - -------- - None - """ - state = h5py.File(filename, "r") - - self.E.fromarray(state["E"][:]) - self.H.fromarray(state["H"][:]) - self.J.fromarray(state["J"][:]) - - # TODO: support MPI loadstate - - state.close() - - def read_state(self, filename="solver_state.h5"): - """Open an HDF5 file for reading without loading its contents. - - This function returns an open `h5py.File` object, allowing the caller - to manually inspect or extract data as needed. The file must be closed - by the caller after use. - - Parameters: - ----------- - filename : str, optional - The name of the HDF5 file to open. Default is "solver_state.h5". - - Returns: - -------- - h5py.File - An open HDF5 file object in read mode. - """ - return h5py.File(filename, "r") - - def reset_fields(self): - """ - Resets the electromagnetic field components to zero. - - This function clears the electric field (E), magnetic field (H), and - current density (J) by setting all their components to zero in the - simulation domain. It ensures a clean restart for a new simulation. - - Notes - ----- - - This method is useful when reusing an existing simulation object - without reinitializing all attributes. - """ - for d in ['x', 'y', 'z']: - self.E[:, :, :, d] = 0.0 - self.H[:, :, :, d] = 0.0 - self.J[:, :, :, d] = 0.0 \ No newline at end of file From a2dd44e9df565944a87dd7ae8797586bfab2cd64 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:56:04 +0100 Subject: [PATCH 20/39] Delete build/lib/wakis/sources.py --- build/lib/wakis/sources.py | 418 ------------------------------------- 1 file changed, 418 deletions(-) delete mode 100644 build/lib/wakis/sources.py diff --git a/build/lib/wakis/sources.py b/build/lib/wakis/sources.py deleted file mode 100644 index ad2f205..0000000 --- a/build/lib/wakis/sources.py +++ /dev/null @@ -1,418 +0,0 @@ -# copyright ################################# # -# This file is part of the wakis Package. # -# Copyright (c) CERN, 2024. # -# ########################################### # - -''' -The `sources.py` script containts different classes -defining a time-dependent sources to be installed -in the electromagnetic simulation. - -All sources need an update function that will be called -every simulation timestep, e.g.: - def update(self, t, *args, **kwargs)` -''' - -import numpy as np -import matplotlib.pyplot as plt -from scipy.constants import mu_0, c as c_light - -class Beam: - def __init__(self, xsource=0., ysource=0., beta=1.0, - q=1e-9, sigmaz=None, ti=None): - ''' - Updates the current J every timestep - to introduce a gaussian beam - moving in +z direction - - Parameters - --- - xsource, ysource: float, default 0. - Transverse position of the source [m] - beta: float, default 1.0 - Relativistic beta of the beam [0-1.0] - q: float, default 1e-9 - Beam charge [C] - sigmaz: float, default None - Beam longitudinal sigma [m] - ti: float, default 9.55*sigmaz - Injection time [s] - ''' - - self.xsource, self.ysource = xsource, ysource - self.sigmaz = sigmaz - self.q = q - self.beta = beta - self.v = c_light*beta - if ti is not None: - self.ti = ti - else: self.ti = 8.548921333333334*self.sigmaz/self.v - self.is_first_update = True - - def update(self, solver, t): - if self.is_first_update: - self.ixs, self.iys = np.abs(solver.x-self.xsource).argmin(), np.abs(solver.y-self.ysource).argmin() - self.is_first_update = False - if hasattr(solver, 'ZMIN'): # support for MPI - self.zmin = solver.ZMIN + solver.dz/2 - else: - self.zmin = solver.z.min() - # reference shift - s0 = self.zmin - self.v*self.ti - s = solver.z - self.v*t - # gaussian - profile = 1/np.sqrt(2*np.pi*self.sigmaz**2)*np.exp(-(s-s0)**2/(2*self.sigmaz**2)) - # update - solver.J[self.ixs,self.iys,:,'z'] = self.q*self.v*profile/solver.dx/solver.dy - - def plot(self, t): - # reference shift - s0 = - self.v*self.ti - s = - self.v*t - # gaussian - profile = 1/np.sqrt(2*np.pi*self.sigmaz**2)*np.exp(-(s-s0)**2/(2*self.sigmaz**2)) - source = self.q*self.v*profile - - fig, ax = plt.subplots() - ax.plot(t, source, 'darkorange') - ax.set_xlabel('Time [s]') - ax.set_ylabel('Current Jz [Cm/s]', color='darkorange') - ax.set_ylim(0., +np.abs(source).max()*1.3) - - fig.tight_layout() - plt.show() - -class PlaneWave: - def __init__(self, xs=None, ys=None, zs=0, nodes=None, f=None, - amplitude=1.0, beta=1.0, phase=0): - ''' - Updates the fields E and H every timestep - to introduce a planewave excitation at the given - xs, ys slice, moving in z+ direction - - Parameters - --- - xs, ys: slice, default 0:N - Transverse positions of the source (indexes) - zs: int or slice, default 0 - Injection position in z - nodes: float, default 15 - Number of nodes between z.min and z.max - f: float, default nodes/T - Frequency of the plane wave [Hz]. It overrides nodes param. - beta: float, default 1. - Relativistic beta - - # TODO: support different directions - ''' - # Check inputs and update self - self.nodes = nodes - self.beta = beta - self.xs, self.ys = xs, ys - self.zs = zs - self.f = f - self.amplitude = amplitude - self.is_first_update = True - self.phase = phase - - self.vp = self.beta*c_light # wavefront velocity beta*c - self.w = 2*np.pi*self.f # ang. frequency - self.kz = self.w/c_light # wave number - self.tmax = np.inf - - if self.nodes is not None: - self.tmax = self.nodes/self.f - - def update(self, solver, t): - if self.is_first_update: - if self.xs is None: - self.xs = slice(0, solver.Nx) - if self.ys is None: - self.ys = slice(0, solver.Ny) - - self.is_first_update = False - - if t <= self.tmax: - solver.H[self.xs,self.ys,self.zs,'y'] = -self.amplitude * np.cos(self.w*t+self.phase) - solver.E[self.xs,self.ys,self.zs,'x'] = self.amplitude * np.cos(self.w*t+self.phase) /(self.kz/(mu_0*self.vp)) - else: - solver.H[self.xs,self.ys,self.zs,'y'] = 0. - solver.E[self.xs,self.ys,self.zs,'x'] = 0. - - def plot(self, t): - fig, ax = plt.subplots() - - sourceH = -self.amplitude * np.cos(self.w*t+self.phase) - sourceE = self.amplitude * np.cos(self.w*t+self.phase) /(self.kz/(mu_0*self.vp)) - - sourceH[t>self.tmax] = 0. - sourceE[t>self.tmax] = 0. - - ax.plot(t, sourceH, 'b') - ax.set_xlabel('Time [s]') - ax.set_ylabel('Magnetic field Hy [A/m]', color='b') - ax.set_ylim(-np.abs(sourceH).max(), +np.abs(sourceH).max()) - - axx = ax.twinx() - axx.plot(t, sourceE, 'r') - axx.set_ylabel('Electric field Ex [V/m]', color='r') - axx.set_ylim(-np.abs(sourceE).max(), +np.abs(sourceE).max()) - - fig.tight_layout() - plt.show() - -class WavePacket: - def __init__(self, xs=None, ys=None, zs=0, - sigmaz=None, sigmaxy=None, tinj=None, - wavelength=None, f=None, - amplitude=1.0, beta=1.0, phase=0): - ''' - Updates E and H fields every timestep to - introduce a 2d gaussian wave packetat the - given xs, ys slice travelling in z+ direction - - Parameters - ---- - xs, ys: slice, default 0:N - Transverse positions of the source [index] - zs: int, default 0 - Longitudinal position of the source [index] - sigmaz: float, default 10*dz - Longitudinal gaussian sigma [m] - sigmaxy: float, default 5*dx - Longitudinal gaussian sigma [m] - tinj: - Injection time delay [m], default 6*sigmaz - wavelength: float, default 10*dz - Wave packet wavelength [m] f=c/wavelength - f: float, default None - Wave packet frequency [Hz], overrides wavelength - beta: float, default 1. - Relativistic beta - ''' - - self.beta = beta - self.xs, self.ys = xs, ys - self.zs = zs - self.f = f - self.wavelength = wavelength - self.sigmaxy = sigmaxy - self.sigmaz = sigmaz - self.tinj = tinj - self.amplitude = amplitude - self.phase = phase - - if self.f is None: - self.f = c_light/self.wavelength - self.w = 2*np.pi*self.f - - if self.tinj is None: - self.tinj = 6*self.sigmaz - - self.is_first_update = True - - def update(self, solver, t): - if self.is_first_update: - if self.xs is None: - self.xs = slice(0, solver.Nx) - if self.ys is None: - self.ys = slice(0, solver.Ny) - if self.sigmaz is None: - self.sigmaz = 10*solver.dz - if self.sigmaxy is None: - self.sigmaxy = 5*solver.dx - if self.tinj is None: - self.tinj = 6*self.sigmaz - - self.is_first_update = False - - # reference shift - s0 = solver.z.min()-self.tinj - s = solver.z.min()-self.beta*c_light*t - - # 2d gaussian - X, Y = np.meshgrid(solver.x[self.xs], solver.y[self.ys]) - gaussxy = np.exp(-(X**2+Y**2)/(2*self.sigmaxy**2)) - gausst = np.exp(-(s-s0)**2/(2*self.sigmaz**2)) - - # Update - solver.H[self.xs,self.ys,self.zs,'y'] = -self.amplitude*np.cos(self.w*t+self.phase)*gaussxy*gausst - solver.E[self.xs,self.ys,self.zs,'x'] = self.amplitude*mu_0*c_light*np.cos(self.w*t+self.phase)*gaussxy*gausst - - def plot(self, t, zmin=0): - fig, ax = plt.subplots() - - # compute source evolution - s0 = zmin-self.tinj - s = zmin-self.beta*c_light*t - gausst = np.exp(-(s-s0)**2/(2*self.sigmaz**2)) - - sourceH = -self.amplitude*np.cos(self.w*t+self.phase)*gausst - ax.plot(t, sourceH, 'b') - ax.set_xlabel('Time [s]') - ax.set_ylabel('Magnetic field Hy [A/m]', color='b') - ax.set_ylim(-np.abs(sourceH).max(), +np.abs(sourceH).max()) - - sourceE = self.amplitude*mu_0*c_light*np.cos(self.w*t+self.phase)*gausst - axx = ax.twinx() - axx.plot(t, sourceE, 'r') - axx.set_ylabel('Electric field Ex [V/m]', color='r') - axx.set_ylim(-np.abs(sourceE).max(), +np.abs(sourceE).max()) - - fig.tight_layout() - plt.show() - -class Dipole: - def __init__(self, field='E', component='z', - xs=None, ys=None, zs=None, - nodes=10, f=None, amplitude=1.0, - phase=0): - ''' - Updates the given field and component every timestep - to introduce a dipole-like sinusoidal excitation - - Parameters - --- - field: str, default 'E' - Field to add to source to. Supports component e.g. 'Ex' - component: str, default 'z' - If not specified in field, component of the field to add the source to - xs, ys, zs: int or slice, default N/2 - Positions of the source (indexes) - nodes: float, default 15 - Number of nodes between z.min and z.max - f: float, default nodes/T - Frequency of the plane wave [Hz]. It overrides nodes param. - ''' - # Check inputs and update self - self.nodes = nodes - self.xs, self.ys, self.zs = xs, ys, zs - self.f = f - self.field = field - self.component = component - self.amplitude = amplitude - self.phase = phase - - if len(field) == 2: #support for e.g. field='Ex' - self.component = field[1] - self.field = field[0] - - self.is_first_update = True - - - def update(self, solver, t): - if self.is_first_update: - if self.xs is None: - self.xs = int(solver.Nx/2) - if self.ys is None: - self.ys = int(solver.Ny/2) - if self.zs is None: - self.zs = int(solver.Nz/2) - if self.f is None: - T = (solver.z.max()-solver.z.min())/c_light - self.f = self.nodes/T - - self.w = 2*np.pi*self.f - self.is_first_update = False - - if self.field == 'E': - solver.E[self.xs,self.ys,self.zs,self.component] = self.amplitude*np.sin(self.w*t+self.phase) - elif self.field == 'H': - solver.H[self.xs,self.ys,self.zs,self.component] = self.amplitude*np.sin(self.w*t+self.phase) - elif self.field == 'J': - solver.J[self.xs,self.ys,self.zs,self.component] = self.amplitude*np.sin(self.w*t+self.phase) - else: - print(f'Field "{self.field}" not valid, should be "E", "H" or "J"]') - -class Pulse: - def __init__(self, field='E', component='z', - xs=None, ys=None, zs=None, - shape='Harris', L=None, amplitude=1.0, - delay=0.): - ''' - Injects an electromagnetic pulse at the given - source point (xs, ys, zs), with the selected - shape, length and amplitude. - - Parameters - ---- - field: str, default 'E' - Field to add to source to. Supports component e.g. 'Ex' - component: str, default 'z' - If not specified in field, component of the field to add the source to - xs, ys, zs: int or slice, default N/2 - Positions of the source (indexes) - shape: str, default 'Harris' - Profile of the pulse in time: ['Harris', 'Gaussian'] - L: float, default 50*dt - width of the pulse (~10*sigma) - - Note: injection time for the gaussian pulse t0=5*L to ensure smooth derivative. - ''' - # Check inputs and update self - - self.xs, self.ys, self.zs = xs, ys, zs - self.field = field - self.component = component - self.amplitude = amplitude - self.shape = shape - self.L = L - self.delay = delay - - if len(field) == 2: #support for e.g. field='Ex' - self.component = field[1] - self.field = field[0] - - if shape.lower() == 'harris': - self.tprofile = self.harris_pulse - elif shape.lower() == 'gaussian': - self.tprofile = self.gaussian_pulse - elif shape.lower() == 'rectangular': - self.tprofile = self.rectangular_pulse - else: - print('** shape does not, match available types: "Harris", "Gaussian", "Rectangular"') - - self.is_first_update = True - - def harris_pulse(self, t): - t = t*c_light - self.delay - try: - if t0.: - return 1.0 - else: - return 0.0 - - def update(self, solver, t): - if self.is_first_update: - if self.xs is None: - self.xs = int(solver.Nx/2) - if self.ys is None: - self.ys = int(solver.Ny/2) - if self.zs is None: - self.zs = int(solver.Nz/2) - if self.L is None: - self.L = 50*solver.dt - - self.is_first_update = False - - if self.field == 'E': - solver.E[self.xs,self.ys,self.zs,self.component] = self.amplitude*self.tprofile(t) - elif self.field == 'H': - solver.H[self.xs,self.ys,self.zs,self.component] = self.amplitude*self.tprofile(t) - elif self.field == 'J': - solver.J[self.xs,self.ys,self.zs,self.component] = self.amplitude*self.tprofile(t) - else: - print(f'Field "{self.field}" not valid, should be "E", "H" or "J"]') \ No newline at end of file From c4ee763aa9e17fdb3a6c4211d38f2c70225c8fd2 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 09:56:12 +0100 Subject: [PATCH 21/39] Delete build/lib/wakis/wakeSolver.py --- build/lib/wakis/wakeSolver.py | 1288 --------------------------------- 1 file changed, 1288 deletions(-) delete mode 100644 build/lib/wakis/wakeSolver.py diff --git a/build/lib/wakis/wakeSolver.py b/build/lib/wakis/wakeSolver.py deleted file mode 100644 index 328dedb..0000000 --- a/build/lib/wakis/wakeSolver.py +++ /dev/null @@ -1,1288 +0,0 @@ -# copyright ################################# # -# This file is part of the wakis Package. # -# Copyright (c) CERN, 2024. # -# ########################################### # - -import time -import os -import glob -import shutil -import h5py -import numpy as np -from tqdm import tqdm -from scipy.constants import c as c_light - -class WakeSolver(): - ''' Class for wake potential and impedance - calculation from 3D time domain E fields - ''' - - def __init__(self, wakelength=None, q=1e-9, sigmaz=1e-3, beta=1.0, - xsource=0., ysource=0., xtest=0., ytest=0., - chargedist=None, ti=None, - compute_plane='both', skip_cells=0, add_space=None, - Ez_file='Ez.h5', save=True, results_folder='results/', - verbose=0, logfile=False): - ''' - Parameters - ---------- - wakelength : float, optional - Wakelength to be simulated. If not provided, it will be calculated from the Ez field data. - q : float - Beam total charge in [C] - sigmaz : float - Beam sigma in the longitudinal direction [m] - beta : float, deafult 1.0 - Ratio of beam's velocity to the speed of light c [a.u.] - xsource : float, default 0. - Beam center in the transverse plane, x-dir [m] - ysource : float, default 0. - Beam center in the transverse plane, y-dir [m] - xtest : float, default 0. - Integration path center in the transverse plane, x-dir [m] - ytest : float, default 0. - Integration path center in the transverse plane, y-dir [m] - ti : float, optional - Injection time, when beam enters domain [s]. If not provided, - the default value ti=8.53*sigmaz will be used - chargedist : dict or str, default None - If not provided, an analytic gaussian with sigmaz and q will be used. - When str, specifies the filename containing the charge distribution data - When dict, should contain the charge distribution data in keys (e.g.): {'X','Y'} - 'X' : longitudinal coordinate [m] - 'Y' : charge distribution in [C/m] - Ez_file : str, default 'Ez.h5' - hdf5 file containing Ez(x,y,z) data for every timestep - save: bool, default True - Flag to enable saving the wake potential, impedance and charge distribution - results in `.txt` files. - - Longitudinal: WP.txt, Z.txt. - - Transverse: WPx.txt, WPy.txt, Zx.txt, Zy.txt - - Charge distribution: lambda.txt, spectrum.txt - verbose: bool, default 0 - Controls the level of verbose in the terminal output - logfile: bool, default False - Creates a `wake.log` file with the summary of the input parameters - and calculations performed - - Attributes - ---------- - Ezt : ndarray - Matrix (nz x nt) containing Ez(x_test, y_test, z, t) - where nz = len(z), nt = len(t) - s : ndarray - Wakelegth vector s=c_light*t-z [m] representing the distance between - the source and the integration point. Goes from -ti*c_light - to the simulated wakelength where ti is the beam injection time. - WP : ndarray - Longitudinal wake potential WP(s) [V/pC] - WP_3d : ndarray - Longitudinal wake potential in 3d WP(x,y,s). Shape = (2*n+1, 2*n+1, len(s)) - where n = n_transverse_cells and s the wakelength array [V/pC] - Z : ndarray - Longitudinal impedance [Ohm] computed by the fourier-transformation of the - longitudinal component of the wake potential, which is divided by the - fourier-transformed charge distribution line function lambda(s) using a - single-sided DFT with 1000 samples. - WPx : ndarray - Trasnverse wake potential in x direction WPx(s) [V/pC] - WPy : ndarray - Transverse wake potential in y direction WP(s) [V/pC] - Zx : ndarray - Trasnverse impedance in x-dir Zx(f) [Ohm] - Zy : ndarray - Transverse impedance in y-dir Zy(f) [Ohm] - lambdas : ndarray - Linear charge distribution of the passing beam λ(s) [C/m] - lambdaf : ndarray - Charge distribution spectrum λ(f) [C] - dx : float - Ez field mesh step in transverse plane, x-dir [m] - dy : float - Ez field mesh step in transverse plane, y-dir [m] - x : ndarray - vector containing x-coordinates for field monitor [m] - y : ndarray - vector containing y-coordinates for field monitor [m] - n_transverse_cells : int, default 1 - Number of transverse cells used for the 3d calculation: 2*n+1 - This determines de size of the 3d wake potential - - ''' - - #beam - self.q = q - self.sigmaz = sigmaz - self.beta = beta - self.v = self.beta*c_light - self.xsource, self.ysource = xsource, ysource - self.xtest, self.ytest = xtest, ytest - self.chargedist = chargedist - self.ti = ti - self.skip_cells = skip_cells - self.compute_plane = compute_plane - self.DE_model = None - - if add_space is not None: #legacy support for add_space - self.skip_cells = add_space - - # Injection time - if ti is not None: - self.ti = ti - else: - #ti = 8.548921333333334*self.sigmaz/self.v #injection time as in CST for beta = 1 - ti = 8.548921333333334*self.sigmaz/(np.sqrt(self.beta)*self.v) #injection time as in CST for beta <=1 - self.ti = ti - - #field - self.Ez_file = Ez_file - self.Ez_hf = None - self.Ezt = None #Ez(x_t, y_t, z, t) - self.t = None - self.xf, self.yf, self.zf = None, None, None #field subdomain - self.x, self.y, self.z = None, None, None #full simulation domain - - #solver init - self.wakelength = wakelength - self.s = None - self.lambdas = None - self.WP = None - self.WP_3d = None - self.n_transverse_cells = 1 - self.WPx, self.WPy = None, None - self.f = None - self.Z = None - self.Zx, self.Zy = None, None - self.lambdaf = None - - #user - self.verbose = verbose - self.save = save - self.logfile = logfile - self.folder = results_folder - - if self.save: - if not os.path.exists(self.folder): - os.mkdir(self.folder) - - # create log - if self.log: - self.params_to_log() - - def solve(self, compute_plane=None, **kwargs): - ''' - Perform the wake potential and impedance for - longitudinal and transverse plane - ''' - if compute_plane is None: - compute_plane = self.compute_plane - - for key, val in kwargs.items(): - setattr(self, key, val) - - t0 = time.time() - - if compute_plane.lower() == 'both' or 'transverse': - # Obtain longitudinal Wake potential - self.calc_long_WP_3d() - - #Obtain transverse Wake potential - self.calc_trans_WP() - - #Obtain the longitudinal impedance - self.calc_long_Z() - - #Obtain transverse impedance - self.calc_trans_Z() - - elif compute_plane == 'longitudinal': - # Obtain longitudinal Wake potential - self.calc_long_WP() - - #Obtain the longitudinal impedance - self.calc_long_Z() - - #Elapsed time - t1 = time.time() - totalt = t1-t0 - self.log('Calculation terminated in %ds' %totalt) - - def update_long_WP(self, t): - '''WIP - calculation of wake potential on the fly - TODO: simplify logic, add transverse WP - ''' - - it = int(t/self.dt) - if it == 0: - # --- setup once before time loop --- - # self.s already computed as in your calc_long_WP - self.s = np.arange(-self.ti*self.v, self.wake.wakelength, self.dt*self.v) # 1D array of s-values (m) - self.WP = np.zeros_like(self.s, dtype=np.float64) # accumulator - - # --- inside your time loop, after Ez is updated for current timestep `it` --- - # get Ez at the probe (shape nz,) - # if you store Ez as Ezt[k, it] style, do: - Ez_curr = self.E[self.Nx//2, self.Ny//2, :, 'z'] # or extract from self.E at (xmid, ymid, :) if not using Ezt - - # compute s values for each z where current Ez contributes - s_vals = self.v * t - self.v * self.ti + np.min(self.z) - self.z - # convert to fractional index in s-array - idxf = (s_vals - self.s[0]) / (self.v*self.dt) # float indices - - # mask in-range contributions - mask = (idxf >= 0.0) & (idxf <= (len(self.s) - 1)) - if not np.any(mask): - # no contributions fall into s range this timestep - pass - else: - idxf_m = idxf[mask] - Ez_m = Ez_curr[mask] - - # integer floor indices - i0 = np.floor(idxf_m).astype(int) - frac = idxf_m - i0 - - # handle points that land exactly on last bin: we add all to last bin - last_bin_mask = (i0 >= len(self.s) - 1) - if np.any(last_bin_mask): - # assign entirely to last bin (no i0+1 available) - self.WP[-1] += np.sum(Ez_m[last_bin_mask]) * self.dz - # drop those from other accumulation - keep = ~last_bin_mask - i0 = i0[keep] - frac = frac[keep] - Ez_m = Ez_m[keep] - - # accumulate with linear interpolation weights - if i0.size: - # faster: use bincount - left = np.bincount(i0, weights=Ez_m * (1.0 - frac) * self.dz, minlength=len(self.s)) - right = np.bincount(i0 + 1, weights=Ez_m * frac * self.dz, minlength=len(self.s)) - self.WP += left + right - - if it == self.Nt-1: - WP = WP/(self.q*1e12) - - def calc_long_WP(self, Ezt=None,**kwargs): - ''' - Obtains the wake potential from the pre-computed longitudinal - Ez(z,t) field from the specified solver. - Parameters can be passed as **kwargs. - - Parameters - ---------- - t : ndarray - vector containing time values [s] - z : ndarray - vector containint z-coordinates [m] - sigmaz : float - Beam longitudinal sigma, to calculate injection time [m] - q : float - Beam charge, to normalize wake potential - ti : float, default 8.53*sigmaz/c - Injection time needed to set the negative part of s vector - and wakelength - Ezt : ndarray, default None - Matrix (nz x nt) containing Ez(x_test, y_test, z, t) - where nz = len(z), nt = len(t) - Ez_file : str, default 'Ez.h5' - HDF5 file containing the Ez(x, y, z) field data - for every timestep. Needed only if Ezt is not provided. - ''' - for key, val in kwargs.items(): - setattr(self, key, val) - - # Read h5 - if Ezt is not None: - self.Ezt = Ezt - - elif self.Ez_hf is None: - self.read_Ez() - - # time variables - nt = len(self.t) - dt = self.t[2]-self.t[1] - ti = self.ti - - # longitudinal variables - if self.zf is None: self.zf = self.z - dz = self.zf[2]-self.zf[1] - zmax = np.max(self.zf) #should it be domain's edge instead? - zmin = np.min(self.zf) - - if self.skip_cells !=0: - zz = slice(self.skip_cells, -self.skip_cells) - else: - zz = np.s_[:] - z = self.zf[zz] - nz = len(z) - - # Set Wake length and s - if self.wakelength is not None: - wakelength = self.wakelength - else: - wakelength = nt*dt*self.v - (zmax-zmin) - ti*self.v - self.wakelength = wakelength - - s = np.arange(-self.ti*self.v, wakelength, dt*self.v) - - self.log('Max simulated time = '+str(round(self.t[-1]*1.0e9,4))+' ns') - self.log('Wakelength = '+str(round(wakelength,3))+'m') - - # Initialize - WP = np.zeros_like(s) - keys = list(self.Ez_hf.keys()) - - # check for rounding errors - if nt > len(keys)-4: - nt = len(keys)-4 - self.log('*** rounding error in number of timesteps') - - # Assembly Ez field - self.log('Assembling Ez field...') - Ezt = np.zeros((nz,nt)) #Assembly Ez field - Ez = self.Ez_hf[keys[0]] - - if len(Ez.shape) == 3: - for n in range(nt): - Ez = self.Ez_hf[keys[n]] - Ezt[:, n] = Ez[Ez.shape[0]//2+1,Ez.shape[1]//2+1,zz] - - elif len(Ez.shape) == 1: - for n in range(nt): - Ezt[:, n] = self.Ez_hf[keys[n]] - self.Ezt = Ezt - - # integral of (Ez(xtest, ytest, z, t=(s+z)/c))dz - print('Calculating longitudinal wake potential WP(s)...') - with tqdm(total=len(s)*len(z)) as pbar: - for n in range(len(s)): - for k in range(nz): - ts = (z[k]+s[n])/self.v-zmin/self.v-self.t[0]+ti - it = int(ts/dt) #find index for t - if it < nt: - WP[n] = WP[n]+(Ezt[k, it])*dz #compute integral - pbar.update(1) - - WP = WP/(self.q*1e12) # [V/pC] - - self.s = s - self.WP = WP - - if self.save: - np.savetxt(self.folder+'WP.txt', np.c_[self.s,self.WP], header=' s [m]'+' '*20+'WP [V/pC]'+'\n'+'-'*48) - - def calc_long_WP_3d(self, **kwargs): - ''' - Obtains the 3d wake potential from the pre-computed Ez(x,y,z) - field from the specified solver. The calculation - Parameters can be passed as **kwargs. - - Parameters - ---------- - Ez_file : str, default 'Ez.h5' - HDF5 file containing the Ez(x,y,z) field data for every timestep - t : ndarray - vector containing time values [s] - z : ndarray - vector containing z-coordinates [m] - q : float - Total beam charge in [C]. Default is 1e9 C - n_transverse_cells : int, default 1 - Number of transverse cells used for the 3d calculation: 2*n+1 - This determines de size of the 3d wake potential - ''' - self.log('\n') - self.log('Longitudinal wake potential') - self.log('-'*24) - - for key, val in kwargs.items(): - setattr(self, key, val) - - # Read h5 - if self.Ez_hf is None: - self.read_Ez() - - # time variables - nt = len(self.t) - dt = self.t[2]-self.t[1] - ti = self.ti - - # longitudinal varianles - if self.zf is None: self.zf = self.z - dz = self.zf[2]-self.zf[1] - zmax = np.max(self.zf) - zmin = np.min(self.zf) - - if self.skip_cells !=0: - zz = slice(self.skip_cells, -self.skip_cells) - else: - zz = np.s_[:] - z = self.zf[zz] - nz = len(z) - - # Set Wake length and s - if self.wakelength is not None: - wakelength = self.wakelength - else: - wakelength = nt*dt*self.v - (zmax-zmin) - ti*self.v - self.wakelength = wakelength - - s = np.arange(-self.ti*self.v, wakelength, dt*self.v) - - self.log(f'* Max simulated time = {np.max(self.t)} s') - self.log(f'* Wakelength = {wakelength} m') - - #field subvolume in No.cells for x, y - i0, j0 = self.n_transverse_cells, self.n_transverse_cells - WP = np.zeros_like(s) - WP_3d = np.zeros((i0*2+1,j0*2+1,len(s))) - Ezt = np.zeros((nz,nt)) - keys = list(self.Ez_hf.keys()) - - # check for rounding errors - if nt > len(keys)-4: - nt = len(keys)-4 - self.log('*** rounding error in number of timesteps') - - print('Calculating longitudinal wake potential WP(s)') - with tqdm(total=len(s)*(i0*2+1)*(j0*2+1)) as pbar: - for i in range(-i0,i0+1,1): - for j in range(-j0,j0+1,1): - - # Assembly Ez field - for n in range(nt): - Ez = self.Ez_hf[keys[n]] - Ezt[:, n] = Ez[Ez.shape[0]//2+i,Ez.shape[1]//2+j, zz] - - # integral of (Ez(xtest, ytest, z, t=(s+z)/c))dz - for n in range(len(s)): - for k in range(nz): - ts = (z[k]+s[n])/self.v-zmin/self.v-self.t[0]+ti - it = int(ts/dt) #find index for t - if it < nt: - WP[n] = WP[n]+(Ezt[k, it])*dz #compute integral - - pbar.update(1) - - WP = WP/(self.q*1e12) # [V/pC] - WP_3d[i0+i,j0+j,:] = WP - - self.s = s - self.WP = WP_3d[i0,j0,:] - self.WP_3d = WP_3d - - self.log(f'Elapsed time {pbar.format_dict["elapsed"]} s') - - if self.save: - np.savetxt(self.folder+'WP.txt', np.c_[self.s, self.WP], header=' s [m]'+' '*20+'WP [V/pC]'+'\n'+'-'*48) - - def calc_trans_WP(self, **kwargs): - ''' - Obtains the transverse wake potential from the longitudinal - wake potential in 3d using the Panofsky-Wenzel theorem using a - second-order scheme for the gradient calculation - - Parameters - ---------- - WP_3d : ndarray - Longitudinal wake potential in 3d WP(x,y,s). Shape = (2*n+1, 2*n+1, len(s)) - where n = n_transverse_cells and s the wakelength array - s : ndarray - Wakelegth vector s=c*t-z representing the distance between - the source and the integration point. Goes from -8.53*sigmat to WL - where sigmat = sigmaz/c and WL is the Wakelength - dx : float - Ez field mesh step in transverse plane, x-dir [m] - dy : float - Ez field mesh step in transverse plane, y-dir [m] - x : ndarray, optional - vector containing x-coordinates [m] - y : ndarray, optional - vector containing y-coordinates [m] - n_transverse_cells : int, default 1 - Number of transverse cells used for the 3d calculation: 2*n+1 - This determines de size of the 3d wake potential - ''' - - for key, val in kwargs.items(): - setattr(self, key, val) - - self.log('\n') - self.log('Transverse wake potential') - self.log('-'*24) - self.log(f'* No. transverse cells = {self.n_transverse_cells}') - - # Obtain dx, dy, ds - if 'dx' in kwargs.keys() and 'dy' in kwargs.keys(): - dx = kwargs['dx'] - dy = kwargs['dy'] - else: - dx=self.xf[2]-self.xf[1] - dy=self.yf[2]-self.yf[1] - - ds = self.s[2]-self.s[1] - i0, j0 = self.n_transverse_cells, self.n_transverse_cells - - # Initialize variables - WPx = np.zeros_like(self.s) - WPy = np.zeros_like(self.s) - int_WP = np.zeros_like(self.WP_3d) - - print('Calculating transverse wake potential WPx, WPy...') - # Obtain the transverse wake potential - with tqdm(total=len(self.s)*(i0*2+1)*(j0*2+1)) as pbar: - for n in range(len(self.s)): - for i in range(-i0,i0+1,1): - for j in range(-j0,j0+1,1): - # Perform the integral - int_WP[i0+i,j0+j,n]=np.sum(self.WP_3d[i0+i,j0+j,0:n])*ds - pbar.update(1) - - # Perform the gradient (second order scheme) - WPx[n] = - (int_WP[i0+1,j0,n]-int_WP[i0-1,j0,n])/(2*dx) - WPy[n] = - (int_WP[i0,j0+1,n]-int_WP[i0,j0-1,n])/(2*dy) - - self.WPx = WPx - self.WPy = WPy - - self.log(f'Elapsed time {pbar.format_dict["elapsed"]} s') - - if self.save: - np.savetxt(self.folder+'WPx.txt', np.c_[self.s,self.WPx], header=' s [m]'+' '*20+'WP [V/pC]'+'\n'+'-'*48) - np.savetxt(self.folder+'WPy.txt', np.c_[self.s,self.WPx], header=' s [m]'+' '*20+'WP [V/pC]'+'\n'+'-'*48) - - def calc_long_Z(self, samples=1001, fmax=None, **kwargs): - ''' - Obtains the longitudinal impedance from the longitudinal - wake potential and the beam charge distribution using a - single-sided DFT with 1000 samples. - Parameters can be passed as **kwargs - - Parameters - ---------- - WP : ndarray - Longitudinal wake potential WP(s) - s : ndarray - Wakelegth vector s=c*t-z representing the distance between - the source and the integration point. Goes from -8.53*sigmat to WL - where sigmat = sigmaz/c and WL is the Wakelength - lambdas : ndarray - Charge distribution λ(s) interpolated to s axis, normalized by the beam charge - chargedist : ndarray, optional - Charge distribution λ(z). Not needed if lambdas is specified - q : float, optional - Total beam charge in [C]. Not needed if lambdas is specified - z : ndarray - vector containing z-coordinates [m]. Not needed if lambdas is specified - sigmaz : float - Beam sigma in the longitudinal direction [m]. - Used to calculate maximum frequency of interest fmax=c/(3*sigmaz) - ''' - self.log('\n') - self.log('Longitudinal impedance') - self.log('-'*24) - - for key, val in kwargs.items(): - setattr(self, key, val) - - print('Calculating longitudinal impedance Z...') - self.log(f'Single sided DFT with number of samples = {samples}') - - # setup charge distribution in s - if self.lambdas is None and self.chargedist is not None: - self.calc_lambdas() - elif self.lambdas is None and self.chargedist is None: - self.calc_lambdas_analytic() - try: - self.log('! Using analytic charge distribution λ(s) since no data was provided') - except: #ascii encoder error handling - self.log('! Using analytic charge distribution since no data was provided') - - # Set up the DFT computation - ds = np.mean(self.s[1:]-self.s[:-1]) - if fmax is None: - fmax = self.v/self.sigmaz/3 #max frequency of interest - N = int((self.v/ds)//fmax*samples) #to obtain a 1000 sample single-sided DFT - - # Obtain DFTs - is it v or c? - lambdafft = np.fft.fft(self.lambdas*self.v, n=N) - WPfft = np.fft.fft(self.WP*1e12, n=N) - ffft = np.fft.fftfreq(len(WPfft), ds/self.v) - - # Mask invalid frequencies - mask = np.logical_and(ffft >= 0 , ffft < fmax) - WPf = WPfft[mask]*ds - lambdaf = lambdafft[mask]*ds - self.f = ffft[mask] # Positive frequencies - - # Compute the impedance - self.Z = - WPf / lambdaf - self.lambdaf = lambdaf - - if self.save: - np.savetxt(self.folder+'Z.txt', np.c_[self.f, self.Z], header=' f [Hz]'+' '*20+'Z [Ohm]'+'\n'+'-'*48) - np.savetxt(self.folder+'spectrum.txt', np.c_[self.f, self.lambdaf], header=' f [Hz]'+' '*20+'Charge distribution spectrum [C/s]'+'\n'+'-'*48) - - def calc_trans_Z(self, samples=1001): - ''' - Obtains the transverse impedance from the transverse - wake potential and the beam charge distribution using a - single-sided DFT with 1000 samples - Parameters can be passed as **kwargs - ''' - self.log('\n') - self.log('Transverse impedance') - self.log('-'*24) - - print('Calculating transverse impedance Zx, Zy...') - self.log(f'Single sided DFT with number of samples = {samples}') - - # Set up the DFT computation - ds = self.s[2]-self.s[1] - fmax=1*self.v/self.sigmaz/3 - N=int((self.v/ds)//fmax*samples) #to obtain a 1000 sample single-sided DFT - - # Obtain DFTs - - # Normalized charge distribution λ(w) - lambdafft = np.fft.fft(self.lambdas*self.v, n=N) - ffft=np.fft.fftfreq(len(lambdafft), ds/self.v) - mask = np.logical_and(ffft >= 0 , ffft < fmax) - lambdaf = lambdafft[mask]*ds - - # Horizontal impedance Zx⊥(w) - WPxfft = np.fft.fft(self.WPx*1e12, n=N) - WPxf = WPxfft[mask]*ds - - self.Zx = 1j * WPxf / lambdaf - - # Vertical impedance Zy⊥(w) - WPyfft = np.fft.fft(self.WPy*1e12, n=N) - WPyf = WPyfft[mask]*ds - - self.Zy = 1j * WPyf / lambdaf - - self.fx = ffft[mask] - self.fy = ffft[mask] - - if self.save: - np.savetxt(self.folder+'Zx.txt', np.c_[self.fx, self.Zx], header=' f [Hz]'+' '*20+'Zx [Ohm]'+'\n'+'-'*48) - np.savetxt(self.folder+'Zy.txt', np.c_[self.fy, self.Zy], header=' f [Hz]'+' '*20+'Zy [Ohm]'+'\n'+'-'*48) - - def calc_lambdas(self, **kwargs): - '''Obtains normalized charge distribution in terms of s - λ(s) to use in the Impedance calculation - - Parameters - ---------- - s : ndarray - Wakelegth vector s=c*t-z representing the distance between - the source and the integration point. Goes from -8.53*sigmat to WL - where sigmat = sigmaz/c and WL is the Wakelength - chargedist : ndarray, optional - Charge distribution λ(z) - q : float, optional - Total beam charge in [C] - z : ndarray, optional - vector containing z-coordinates of the domain [m] - zf : ndarray, optional - vector containing z-coordinates of the field monitor [m]. N - ''' - for key, val in kwargs.items(): - setattr(self, key, val) - - if type(self.chargedist) is str: - d = self.read_txt(self.chargedist) - keys = list(d.keys()) - z = d[keys[0]] - chargedist = d[keys[1]] - - elif (self.chargedist) is dict: - keys = list(self.chargedist.keys()) - z = self.chargedist[keys[0]] - chargedist = self.chargedist[keys[1]] - - else: - chargedist = self.chargedist - if len(self.z) == len(self.chargedist): - z = self.z - elif len(self.zf) == len(self.chargedist): - z = self.zf - else: - self.log('Dimension error: check input dimensions') - - self.lambdas = np.interp(self.s, z, chargedist/self.q) - - if self.save: - np.savetxt(self.folder+'lambda.txt', np.c_[self.s, self.lambdas], header=' s [Hz]'+' '*20+'Charge distribution [C/m]'+'\n'+'-'*48) - - def calc_lambdas_analytic(self, **kwargs): - '''Obtains normalized charge distribution in s λ(z) - as an analytical gaussian centered in s=0 and std - equal sigmaz - - Parameters - ---------- - s : ndarray - Wakelegth vector s=c*t-z representing the distance between - the source and the integration point. Goes from -8.53*sigmat to WL - where sigmat = sigmaz/c and WL is the Wakelength - sigmaz : float - Beam sigma in the longitudinal direction [m] - ''' - - for key, val in kwargs.items(): - setattr(self, key, val) - - self.lambdas = 1/(self.sigmaz*np.sqrt(2*np.pi))*np.exp(-(self.s**2)/(2*self.sigmaz**2)) - - if self.save: - np.savetxt(self.folder+'lambda.txt', np.c_[self.s, self.lambdas], header=' s [Hz]'+' '*20+'Charge distribution [C/m]'+'\n'+'-'*48) - - def get_SmartBounds(self, freq_data=None, impedance_data=None, - minimum_peak_height=1.0, distance=3, inspect_bounds=True, - Rs_bounds=[0.8, 10], Q_bounds=[0.5, 5], fres_bounds=[-0.01e9, +0.01e9] - ): - import iddefix - - self.log('\nCalculating bounds using the Smart Bound Determination...') - # Smart bounds - # Find the main resonators and estimate the bounds -courtesy of Malthe Raschke! - bounds = iddefix.SmartBoundDetermination(freq_data, np.real(impedance_data), minimum_peak_height=minimum_peak_height, - Rs_bounds=Rs_bounds, Q_bounds=Q_bounds, fres_bounds=fres_bounds) - - bounds.find(minimum_peak_height=minimum_peak_height, distance=distance) - - if inspect_bounds: - bounds.inspect() - bounds.to_table() - return bounds - - bounds.to_table() - return bounds - - - def get_DEmodel_fitting(self, freq_data=None, impedance_data=None, - plane='longitudinal', dim='z', - parameterBounds=None, N_resonators=None, DE_kernel='DE', - maxiter=1e5, cmaes_sigma=0.01, popsize=150, tol=1e-3, - use_minimization=True, minimization_margin=[0.3, 0.2, 0.01], - minimum_peak_height=1.0, distance=3, inspect_bounds=False, - Rs_bounds=[0.8, 10], Q_bounds=[0.5, 5], fres_bounds=[-0.01e9, +0.01e9], - ): - import iddefix - - if freq_data is None or impedance_data is None: - if plane == 'longitudinal' and dim == 'z': - freq_data = self.f - impedance_data = self.Z - elif plane == 'transverse': - if dim == 'x': - freq_data = self.fx - impedance_data = self.Zx - elif dim == 'y': - freq_data = self.fy - impedance_data = self.Zy - else: - raise ValueError('Invalid dimension. Use dim = "x" or "y".') - else: - raise ValueError('Invalid plane or dimension. Use plane = "longitudinal" or "transverse" and choose the dimension dim = "z", "x" or "y".') - - if parameterBounds is None or N_resonators is None: - bounds = self.get_SmartBounds(parameterBounds=parameterBounds, - N_resonators=N_resonators, - minimum_peak_height=minimum_peak_height, - distance=distance, - inspect_bounds=inspect_bounds, - Rs_bounds=Rs_bounds, Q_bounds=Q_bounds, fres_bounds=fres_bounds) - N_resonators = bounds.N_resonators - parameterBounds = bounds.parameterBounds - - # Build the differential evolution model - print('Fitting the impedance using Differential Evolution...') - self.log('\nExtrapolating wake potential using Differential Evolution...') - - objectiveFunction=iddefix.ObjectiveFunctions.sumOfSquaredErrorReal - DE_model = iddefix.EvolutionaryAlgorithm(freq_data, - np.real(impedance_data), - N_resonators=N_resonators, - parameterBounds=parameterBounds, - plane=plane, - fitFunction='impedance', - wake_length=self.wakelength, # in [m] - objectiveFunction=objectiveFunction, - ) - - if DE_kernel == 'DE': - DE_model.run_differential_evolution(maxiter=int(maxiter), - popsize=popsize, - tol=tol, - mutation=(0.3, 0.8), - crossover_rate=0.5) - - elif DE_kernel == 'CMAES': #TODO: fix UnboundLocalError - DE_model.run_cmaes(maxiter=int(maxiter), - popsize=popsize, - sigma=cmaes_sigma,) - - if use_minimization: - self.log('Running minimization algorithm...') - DE_model.run_minimization_algorithm(minimization_margin) - - self.DE_model = DE_model - self.log(DE_model.warning) - if self.verbose: - print(DE_model.warning) - - return DE_model - - def get_extrapolated_wake(self, wakelength, sigma=None, use_minimization=True): - ''' - Get the extrapolated wake potential [V/pC] from the DE model - ''' - if self.DE_model is None: - raise AttributeError('Run get_DEmodel() first to obtain the DE model') - - if sigma is None: - sigma = self.sigmaz/c_light - - # Get the extrapolated wake potential - # TODO: add beta - t = np.arange(self.s[0]/c_light, wakelength/c_light, (self.s[2]-self.s[1])/c_light) - wake_potential = self.DE_model.get_wake_potential(t, sigma=sigma, - use_minimization=use_minimization) - - s = t * c_light # Convert time to distance [m] - return s, -wake_potential*1e-12 # in [V/pC] + CST convention - - def get_extrapolated_wake_function(self, wakelength, use_minimization=True): - ''' - Get the extrapolated wake function (a.k.a. Green function) from the DE model - ''' - if self.DE_model is None: - raise AttributeError('Run get_DEmodel() first to obtain the DE model') - - t = np.arange(self.s[0]/c_light, wakelength/c_light, (self.s[2]-self.s[1])/c_light) - wake_function = self.DE_model.get_wake(t, use_minimization=use_minimization) - return t, wake_function - - def get_extrapolated_impedance(self, f=None, use_minimization=True, - wakelength=None): - ''' - Get the extrapolated impedance [Ohm] from the DE model - ''' - if self.DE_model is None: - raise AttributeError('Run get_DEmodel() first to obtain the DE model') - - if f is None: - f = self.DE_model.frequency_data - - impedance = self.DE_model.get_impedance(frequency_data=f, - use_minimization=use_minimization, - wakelength=wakelength) - return f, impedance - - @staticmethod - def calc_impedance_from_wake(wake, s=None, t=None, fmax=None, - samples=None, verbose=True): - if type(wake) == list: - t = wake[0] - wake = wake[1] - if s is not None: - t = s/c_light - elif s is None and t is None: - raise AttributeError('Provide time data through parameter "t" [s] or "s" [m]') - dt = np.mean(t[1:]-t[:-1]) - - # Maximum frequency: fmax = 1/dt - if fmax is not None: - aux = np.arange(t.min(), t.max(), 1/fmax/2) - wake = np.interp(aux, t, wake) - dt = np.mean(aux[1:]-aux[:-1]); del aux - else: fmax = 1/dt - # Time resolution: fres=(1/len(wake)/dt/2) - - # Obtain DFTs - if samples is not None: - Wfft = np.fft.fft(wake, n=2*samples) - else: - Wfft = np.fft.fft(wake) - - ffft = np.fft.fftfreq(len(Wfft), dt) - - # Mask invalid frequencies - mask = np.logical_and(ffft >= 0 , ffft < fmax) - Z = Wfft[mask]/len(wake)*2 - f = ffft[mask] # Positive frequencies - - if verbose: - print(f'* Number of samples = {len(f)}') - print(f'* Maximum frequency = {f.max()} Hz') - print(f'* Maximum resolution = {np.mean(f[1:]-f[:-1])} Hz') - - return [f, Z] - - @staticmethod - def calc_wake_from_impedance(impedance, f=None, tmax=None, - samples=None, pad=0, verbose=True): - - if len(impedance) == 2: - f = impedance[0] - Z = impedance[1] - elif f is None: - raise AttributeError('Provide frequency data through parameter "f"') - else: - Z = impedance - df = np.mean(f[1:]-f[:-1]) - - # Maximum time: tmax = 1/(f[2]-f[1]) - if tmax is not None: - aux = np.arange(f.min(), f.max(), 1/tmax) - Z = np.interp(aux, f, Z) - df = np.mean(aux[1:]-aux[:-1]); del aux - else: tmax = 1/df - - # Time resolution: tres=(1/len(Z)/(f[2]-f[1])) - # pad = int(1/df/tres - len(Z)) - # wake = np.real(np.fft.ifft(np.pad(Z, pad))) - wake = np.real(-1*np.fft.fft(Z, n=samples)) - wake = np.roll(wake,-1) - # Inverse fourier transform of impedance - t = np.linspace(0, tmax, len(wake)) - - if verbose: - print(f'* Number of samples = {len(t)}') - print(f'* Maximum time = {t.max()} s') - print(f'* Maximum resolution = {np.mean(t[1:]-t[:-1])} s') - - return [t, wake] - - def read_Ez(self, filename=None, return_value=False): - ''' - Read the Ez.h5 file containing the Ez field information - ''' - - if filename is None: - filename = self.Ez_file - - hf = h5py.File(filename, 'r') - print(f'Reading h5 file {filename}' ) - self.log('Size of the h5 file: ' + str(round((os.path.getsize(filename)/10**9),2))+' Gb') - - #Set attributes - self.Ez_hf = hf - self.Ez_file = filename - if 'x' in hf.keys(): - self.xf = np.array(hf['x']) - if 'y' in hf.keys(): - self.yf = np.array(hf['y']) - if 'z' in hf.keys(): - self.zf = np.array(hf['z']) - if 't' in hf.keys(): - self.t = np.array(hf['t']) - - if return_value: - return hf - - def read_txt(self, txt, skiprows=2, delimiter=None, usecols=None): - ''' - Reads txt variables from ascii files and - returns data in a dictionary. Header should - be the first line. - ''' - - try: - load = np.loadtxt(txt, skiprows=skiprows, delimiter=delimiter, usecols=usecols) - except: - load = np.loadtxt(txt, skiprows=skiprows, delimiter=delimiter, - usecols=usecols, dtype=np.complex128) - - try: # keys == header names - with open(txt) as f: - header = f.readline() - - header = header.replace(' ', '') - header = header.replace('#', '') - header = header.replace('\n', '') - header = header.split(']') - - d = {} - for i in range(len(load[0,:])): - d[header[i]+']'] = load[:, i] - - except: #keys == int 0, 1, ... - d = {} - for i in range(len(load[0,:])): - d[i] = load[:, i] - - return d - - def save_txt(self, f_name, x_data=None, y_data=None, x_name='X [-]', y_name='Y [-]'): - """ - Saves x and y data to a text file in a two-column format. - - This function exports the provided `x_data` and `y_data` to a `.txt` file, - formatting the output with a header that includes custom column names. - - Parameters - ---------- - f_name : str - Name of the output file (without the `.txt` extension). - x_data : numpy.ndarray, optional - Array containing x-axis data. If None, the file is not saved. - y_data : numpy.ndarray, optional - Array containing y-axis data. If None, the file is not saved. - x_name : str, optional - Label for the x-axis column in the output file. Default is `"X [-]"`. - y_name : str, optional - Label for the y-axis column in the output file. Default is `"Y [-]"`. - - Notes - ----- - - The data is saved in a two-column format where `x_data` and `y_data` - are combined column-wise. - - If `x_data` or `y_data` is missing, the function prints a warning and does not save a file. - - Examples - -------- - Save two NumPy arrays to `data.txt`: - - >>> x = np.linspace(0, 10, 5) - >>> y = np.sin(x) - >>> save_txt("data", x, y, x_name="Time [s]", y_name="Amplitude") - - The saved file will look like: - - Time [s] Amplitude - -------------------------------- - 0.00 0.00 - 2.50 0.59 - 5.00 -0.99 - 7.50 0.94 - 10.00 -0.54 - """ - if x_data is not None and y_data is not None: - np.savetxt(f_name+'.txt', np.c_[x_data, y_data], header=' '+x_name+' '*20+y_name+'\n'+'-'*48) - else: - print('txt not saved, please provide x_data and y_data') - - def load_results(self, folder): - '''Load all txt from a given folder - - The txt files are generated when - the attribute`save = True` is used - ''' - if not folder.endswith('/'): - folder = folder + '/' - - _, self.lambdas = self.read_txt(folder+'lambda.txt').values() - _, self.WPx = self.read_txt(folder+'WPx.txt').values() - _, self.WPy = self.read_txt(folder+'WPy.txt').values() - self.s, self.WP = self.read_txt(folder+'WP.txt').values() - - _, self.lambdaf = self.read_txt(folder+'spectrum.txt').values() - _, self.Zx = self.read_txt(folder+'Zx.txt').values() - _, self.Zy = self.read_txt(folder+'Zy.txt').values() - self.f, self.Z = self.read_txt(folder+'Z.txt').values() - - self.f = np.abs(self.f) - self.wakelength = self.s[-1] - - def copy(self): - obj = type(self).__new__(self.__class__) - obj.__dict__.update(self.__dict__) - return obj - - def log(self, txt): - - if self.verbose: - print('\x1b[2;37m'+txt+'\x1b[0m') - - if not self.logfile: - return - - title = 'wake' - f = open(title + '.log', "a") - f.write(txt + '\r\n') - f.close() - - def params_to_log(self): - self.log(time.asctime()) - self.log('Wake computation') - self.log('='*24) - self.log(f'* Charge q = {self.q} [C]') - self.log(f'* Beam sigmaz = {self.sigmaz} [m]') - self.log(f'* xsource, ysource = {self.xsource}, {self.ysource} [m]') - self.log(f'* xtest, ytest = {self.xtest}, {self.ytest} [m]') - self.log(f'* Beam injection time ti= {self.ti} [s]') - - if self.chargedist is not None: - if type(self.chargedist) is str: - self.log(f'* Charge distribution file: {self.chargedist}') - else: - self.log(f'* Charge distribution data is provided') - else: - self.log(f'* Charge distribution analytic') - - self.log('\n') - - def read_cst_3d(self, path=None, folder='3d', filename='Ez.h5', units=1e-3): - ''' - Read CST 3d exports folder and store the - Ez field information into a matrix Ez(x,y,z) - for every timestep into a single `.h5` file - compatible with wakis. - - Parameters - ---------- - path: str, default None - Path to the field data - folder: str, default '3d' - Folder containing the CST field data .txt files - filename: str, default 'Ez.h5' - Name of the h5 file that will be generated - ''' - - self.log('Reading 3d CST field exports') - self.log('-'*24) - - if path is None: - path = folder + '/' - - # Rename files with E-02, E-03 - for file in glob.glob(path +'*E-02.txt'): - file=file.split(path) - title=file[1].split('_') - num=title[1].split('E') - num[0]=float(num[0])/100 - - ntitle=title[0]+'_'+str(num[0])+'.txt' - shutil.copy(path+file[1], path+file[1]+'.old') - os.rename(path+file[1], path+ntitle) - - for file in glob.glob(path +'*E-03.txt'): - file=file.split(path) - title=file[1].split('_') - num=title[1].split('E') - num[0]=float(num[0])/1000 - - ntitle=title[0]+'_'+str(num[0])+'.txt' - shutil.copy(path+file[1], path+file[1]+'.old') - os.rename(path+file[1], path+ntitle) - - for file in glob.glob(path +'*_0.txt'): - file=file.split(path) - title=file[1].split('_') - num=title[1].split('.') - num[0]=float(num[0]) - - ntitle=title[0]+'_'+str(num[0])+'.txt' - shutil.copy(path+file[1], path+file[1]+'.old') - os.rename(path+file[1], path+ntitle) - - #sort - try: - def sorter(item): - num=item.split(path)[1].split('_')[1].split('.txt')[0] - return float(num) - fnames = sorted(glob.glob(path+'*.txt'), key=sorter) - except: - fnames = sorted(glob.glob(path+'*.txt')) - - #Get the number of longitudinal and transverse cells used for Ez - i=0 - with open(fnames[0]) as f: - lines=f.readlines() - n_rows = len(lines)-3 #n of rows minus the header - x1=lines[3].split()[0] - - while True: - i+=1 - x2=lines[i+3].split()[0] - if x1==x2: - break - - n_transverse_cells=i - n_longitudinal_cells=int(n_rows/(n_transverse_cells**2)) - - # Create h5 file - if os.path.exists(path+filename): - os.remove(path+filename) - - hf = h5py.File(path+filename, 'w') - - # Initialize variables - Ez=np.zeros((n_transverse_cells, n_transverse_cells, n_longitudinal_cells)) - x=np.zeros((n_transverse_cells)) - y=np.zeros((n_transverse_cells)) - z=np.zeros((n_longitudinal_cells)) - t=[] - - nsteps, i, j, k = 0, 0, 0, 0 - skip=-4 #number of rows to skip - rows=skip - - # Start scan - self.log(f'Scanning files in {path}:') - for file in tqdm(fnames): - #self.log.debug('Scanning file '+ file + '...') - title=file.split(path) - title2=title[1].split('_') - - try: - num=title2[1].split('.txt') - t.append(float(num[0])*1e-9) - except: - t.append(nsteps) - - with open(file) as f: - for line in f: - rows+=1 - columns = line.split() - - if rows>=0 and len(columns)>1: - k=int(rows/n_transverse_cells**2) - j=int(rows/n_transverse_cells-n_transverse_cells*k) - i=int(rows-j*n_transverse_cells-k*n_transverse_cells**2) - - if k>= n_longitudinal_cells: - k = int(n_longitudinal_cells-1) - - Ez[i,j,k]=float(columns[5]) - - x[i]=float(columns[0])*units - y[j]=float(columns[1])*units - z[k]=float(columns[2])*units - - if nsteps == 0: - prefix='0'*5 - hf.create_dataset('Ez_'+prefix+str(nsteps), data=Ez) - else: - prefix='0'*(5-int(np.log10(nsteps))) - hf.create_dataset('Ez_'+prefix+str(nsteps), data=Ez) - - i, j, k = 0, 0, 0 - rows=skip - nsteps+=1 - - #close file - f.close() - - hf['x'] = x - hf['y'] = y - hf['z'] = z - hf['t'] = t - - hf.close() - - #set field info - self.log('Ez field is stored in a matrix with shape '+str(Ez.shape)+' in '+str(int(nsteps))+' datasets') - self.log(f'Finished scanning files - hdf5 file {filename} succesfully generated') - - #Update self - self.xf = x - self.yf = y - self.zf = z - self.t = np.array(t) From 57f29035608a741ec7aaae729aa2092df5a5ae49 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 13:39:57 +0100 Subject: [PATCH 22/39] Refactor log variable names in test file --- tests/test_007_mpi_lossy_cavity.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/test_007_mpi_lossy_cavity.py b/tests/test_007_mpi_lossy_cavity.py index 9769884..9abeb99 100644 --- a/tests/test_007_mpi_lossy_cavity.py +++ b/tests/test_007_mpi_lossy_cavity.py @@ -82,15 +82,15 @@ class TestMPILossyCavity: -6.04105997e+01 ,-3.06532160e+01 ,-1.17749936e+01 ,-3.12574866e+00, -7.35339521e-01 ,-1.13085658e-01 , 7.18247535e-01 , 8.73829036e-02]) - grid_logs = {'Nx': 60, 'Ny': 60, 'Nz': 140, 'dx': 0.008666666348775227, 'dy': 0.008666666348775227, + gridLogs = {'Nx': 60, 'Ny': 60, 'Nz': 140, 'dx': 0.008666666348775227, 'dy': 0.008666666348775227, 'dz': 0.005714285799435207, 'stl_solids': ['tests/stl/007_vacuum_cavity.stl', 'tests/stl/007_lossymetal_shell.stl'], 'stl_materials': ['vacuum', [30, 1.0, 30]], 'gridInitializationTime': 0} - solver_logs = {'use_gpu': False, 'use_mpi': False, 'bc_low': ['pec', 'pec', 'pec'], + solverLogs = {'use_gpu': False, 'use_mpi': False, 'bc_low': ['pec', 'pec', 'pec'], 'bc_high': ['pec', 'pec', 'pec'], 'n_pml': 10, 'bg': 'pec', 'dt': 6.970326728398968e-12, 'solverInitializationTime': 0} - wakeSolver_logs = {'ti': 2.8516132094735135e-09, 'q': 1e-09, 'sigmaz': 0.1, 'beta': 1.0, + wakeSolverLogs = {'ti': 2.8516132094735135e-09, 'q': 1e-09, 'sigmaz': 0.1, 'beta': 1.0, 'xsource': 0.0, 'ysource': 0.0, 'xtest': 0.0, 'ytest': 0.0, 'chargedist': None, 'skip_cells': 10, 'results_folder': 'tests/007_results/', 'wakelength': 10.0, 'simulationTime': 0} @@ -329,11 +329,11 @@ def test_long_impedance(self): def test_log_file(self): global solver - solver.logger.grid_logs["gridInitializationTime"] = 0 #times can vary - solver.logger.solver_logs["solverInitializationTime"] = 0 - solver.logger.wakeSolver_logs["simulationTime"] = 0 - logfile = os.path.join(solver.logger.wakeSolver_logs["results_folder"], "Simulation_Parameters.log") + solver.logger.gridLogs["gridInitializationTime"] = 0 #times can vary + solver.logger.solverLogs["solverInitializationTime"] = 0 + solver.logger.wakeSolverLogs["simulationTime"] = 0 + logfile = os.path.join(solver.logger.wakeSolverLogs["results_folder"], "Simulation_Parameters.log") assert os.path.exists(logfile), "Log file not created" - assert solver.logger.grid_logs == self.grid_logs, "Grid logs do not match expected values" - assert solver.logger.solver_logs == self.solver_logs, "Solver logs do not match expected values" - assert solver.logger.wakeSolver_logs == self.wakeSolver_logs, "WakeSolver logs do not match expected values" \ No newline at end of file + assert solver.logger.grid == self.gridLogs, "Grid logs do not match expected values" + assert solver.logger.solver == self.solverLogs, "Solver logs do not match expected values" + assert solver.logger.wakeSolver == self.wakeSolverLogs, "WakeSolver logs do not match expected values" From baebebee42c282890788055c597b648cfbb9e2de Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 13:41:19 +0100 Subject: [PATCH 23/39] Refactor logger to use 'grid' instead of 'grid_logs' --- wakis/gridFIT3D.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/wakis/gridFIT3D.py b/wakis/gridFIT3D.py index 0213823..0d6c0fd 100644 --- a/wakis/gridFIT3D.py +++ b/wakis/gridFIT3D.py @@ -148,21 +148,21 @@ def __init__(self, xmin, xmax, ymin, ymax, zmin, zmax, # Forward Parameters to logger self.logger = Logger() - self.logger.grid_logs["Nx"] = self.Nx - self.logger.grid_logs["Ny"] = self.Ny - self.logger.grid_logs["Nz"] = self.Nz - self.logger.grid_logs["dx"] = self.dx - self.logger.grid_logs["dy"] = self.dy - self.logger.grid_logs["dz"] = self.dz - self.logger.grid_logs["stl_solids"] = list(self.stl_solids.values()) if self.stl_solids is not None else [] - self.logger.grid_logs["stl_materials"] = list(self.stl_materials.values()) if self.stl_materials is not None else [] + self.logger.grid["Nx"] = self.Nx + self.logger.grid["Ny"] = self.Ny + self.logger.grid["Nz"] = self.Nz + self.logger.grid["dx"] = self.dx + self.logger.grid["dy"] = self.dy + self.logger.grid["dz"] = self.dz + self.logger.grid["stl_solids"] = list(self.stl_solids.values()) if self.stl_solids is not None else [] + self.logger.grid["stl_materials"] = list(self.stl_materials.values()) if self.stl_materials is not None else [] if stl_rotate != [0., 0., 0.]: - self.logger.grid_logs["stl_rotate"] = self.stl_rotate + self.logger.grid["stl_rotate"] = self.stl_rotate if stl_translate != [0., 0., 0.]: - self.logger.grid_logs["stl_translate"] = self.stl_translate + self.logger.grid["stl_translate"] = self.stl_translate if stl_scale != 1.0: - self.logger.grid_logs["stl_scale"] = self.stl_scale - self.logger.grid_logs["gridInitializationTime"] = time.time()-t0 + self.logger.grid["stl_scale"] = self.stl_scale + self.logger.grid["gridInitializationTime"] = time.time()-t0 def compute_grid(self): X, Y, Z = np.meshgrid(self.x, self.y, self.z, indexing='ij') @@ -859,4 +859,4 @@ def clip(widget): if offscreen: pl.export_html('grid_inspect.html') else: - pl.show() \ No newline at end of file + pl.show() From 9a1e2da1092e48040e1031196e4d902cb4138d10 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 13:43:53 +0100 Subject: [PATCH 24/39] Refactor Logger class to stick to camelCase --- wakis/logger.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/wakis/logger.py b/wakis/logger.py index d7fad6f..99f50a3 100644 --- a/wakis/logger.py +++ b/wakis/logger.py @@ -20,28 +20,28 @@ class Logger(): def __init__(self): - self.grid_logs = {} - self.solver_logs = {} - self.wakeSolver_logs = {} + self.grid = {} + self.solver = {} + self.wakeSolver = {} def save_logs(self): """ Save all logs (grid, solver, wakeSolver) into log-file inside the results folder. """ - logfile = os.path.join(self.wakeSolver_logs["results_folder"], "Simulation_Parameters.log") + logfile = os.path.join(self.wakeSolver["results_folder"], "Simulation_Parameters.log") # Write sections - if not os.path.exists(self.wakeSolver_logs["results_folder"]): - os.mkdir(self.wakeSolver_logs["results_folder"]) + if not os.path.exists(self.wakeSolver["results_folder"]): + os.mkdir(self.wakeSolver["results_folder"]) with open(logfile, "w", encoding="utf-8") as fh: fh.write("Simulation Parameters\n") fh.write("""=====================\n\n""") sections = [ - ("WakeSolver Logs", self.wakeSolver_logs), - ("Solver Logs", self.solver_logs), - ("Grid Logs", self.grid_logs), + ("WakeSolver Logs", self.wakeSolver), + ("Solver Logs", self.solver), + ("Grid Logs", self.grid), ] for title, data in sections: @@ -64,4 +64,4 @@ def _convert(obj): clean = _convert(data) fh.write(json.dumps(clean, indent=2, ensure_ascii=False)) - fh.write("\n") \ No newline at end of file + fh.write("\n") From e301c3db627e011f9ba88a815f30ec06f7581d28 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 13:45:26 +0100 Subject: [PATCH 25/39] Fix logger attribute assignment for wakeSolver --- wakis/routines.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/wakis/routines.py b/wakis/routines.py index 5936e66..0c66c8b 100644 --- a/wakis/routines.py +++ b/wakis/routines.py @@ -344,7 +344,7 @@ def save_to_h5(self, hf, field, x, y, z, comp, n): self.wake.solve(compute_plane=compute_plane) # Forward parameters to logger - self.logger.wakeSolver_logs=self.wake.logger.wakeSolver_logs - self.logger.wakeSolver_logs["wakelength"]=wakelength - self.logger.wakeSolver_logs["simulationTime"]=time.time()-t0 - self.logger.save_logs() \ No newline at end of file + self.logger.wakeSolver=self.wake.logger.wakeSolver_logs + self.logger.wakeSolver["wakelength"]=wakelength + self.logger.wakeSolver["simulationTime"]=time.time()-t0 + self.logger.save_logs() From c853b4b22ee64d58cfc52791ba4a811fe7e2cc17 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 13:46:25 +0100 Subject: [PATCH 26/39] Refactor logger attribute assignments in solverFIT3D --- wakis/solverFIT3D.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/wakis/solverFIT3D.py b/wakis/solverFIT3D.py index cec8db8..00f050a 100644 --- a/wakis/solverFIT3D.py +++ b/wakis/solverFIT3D.py @@ -114,7 +114,7 @@ def __init__(self, grid, wake=None, cfln=0.5, dt=None, # Grid self.grid = grid - bg_log = bg + bgLog = bg self.Nx = self.grid.Nx self.Ny = self.grid.Ny self.Nz = self.grid.Nz @@ -255,15 +255,15 @@ def __init__(self, grid, wake=None, cfln=0.5, dt=None, # assign logs self.logger = Logger() - self.logger.grid_logs = self.grid.logger.grid_logs - self.logger.solver_logs["use_gpu"] = use_gpu - self.logger.solver_logs["use_mpi"] = use_mpi - self.logger.solver_logs["bc_low"] = bc_low - self.logger.solver_logs["bc_high"] = bc_high - self.logger.solver_logs["n_pml"] = n_pml - self.logger.solver_logs["bg"] = bg_log - self.logger.solver_logs["dt"] = self.dt - self.logger.solver_logs["solverInitializationTime"] = time.time() - t0 + self.logger.grid = self.grid.logger.grid_logs + self.logger.solver["use_gpu"] = use_gpu + self.logger.solver["use_mpi"] = use_mpi + self.logger.solver["bc_low"] = bc_low + self.logger.solver["bc_high"] = bc_high + self.logger.solver["n_pml"] = n_pml + self.logger.solver["bg"] = bg_log + self.logger.solver["dt"] = self.dt + self.logger.solver["solverInitializationTime"] = time.time() - t0 def update_tensors(self, tensor='all'): '''Update tensor matrices after @@ -1120,4 +1120,4 @@ def reset_fields(self): for d in ['x', 'y', 'z']: self.E[:, :, :, d] = 0.0 self.H[:, :, :, d] = 0.0 - self.J[:, :, :, d] = 0.0 \ No newline at end of file + self.J[:, :, :, d] = 0.0 From b27588b0ff9e8d81ad9ae4115cc3aaf813209040 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 13:47:11 +0100 Subject: [PATCH 27/39] Change logger dictionary from wakeSolver_logs to wakeSolver --- wakis/wakeSolver.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/wakis/wakeSolver.py b/wakis/wakeSolver.py index 0e20525..476ba30 100644 --- a/wakis/wakeSolver.py +++ b/wakis/wakeSolver.py @@ -173,17 +173,17 @@ def __init__(self, wakelength=None, q=1e-9, sigmaz=1e-3, beta=1.0, self.params_to_log() # Forward parameters to logger - self.logger.wakeSolver_logs["ti"]=self.ti - self.logger.wakeSolver_logs["q"]=self.q - self.logger.wakeSolver_logs["sigmaz"]=self.sigmaz - self.logger.wakeSolver_logs["beta"]=self.beta - self.logger.wakeSolver_logs["xsource"]=self.xsource - self.logger.wakeSolver_logs["ysource"]=self.ysource - self.logger.wakeSolver_logs["xtest"]=self.xtest - self.logger.wakeSolver_logs["ytest"]=self.ytest - self.logger.wakeSolver_logs["chargedist"]=self.chargedist - self.logger.wakeSolver_logs["skip_cells"]=self.skip_cells - self.logger.wakeSolver_logs["results_folder"]=self.folder + self.logger.wakeSolver["ti"]=self.ti + self.logger.wakeSolver["q"]=self.q + self.logger.wakeSolver["sigmaz"]=self.sigmaz + self.logger.wakeSolver["beta"]=self.beta + self.logger.wakeSolver["xsource"]=self.xsource + self.logger.wakeSolver["ysource"]=self.ysource + self.logger.wakeSolver["xtest"]=self.xtest + self.logger.wakeSolver["ytest"]=self.ytest + self.logger.wakeSolver["chargedist"]=self.chargedist + self.logger.wakeSolver["skip_cells"]=self.skip_cells + self.logger.wakeSolver["results_folder"]=self.folder def solve(self, compute_plane=None, **kwargs): ''' From 796080a02c3c0f789e0bfe215a04cb6e62df97af Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 13:58:07 +0100 Subject: [PATCH 28/39] Fix logger reference for wakeSolver --- wakis/routines.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wakis/routines.py b/wakis/routines.py index 0c66c8b..0427bbc 100644 --- a/wakis/routines.py +++ b/wakis/routines.py @@ -344,7 +344,7 @@ def save_to_h5(self, hf, field, x, y, z, comp, n): self.wake.solve(compute_plane=compute_plane) # Forward parameters to logger - self.logger.wakeSolver=self.wake.logger.wakeSolver_logs + self.logger.wakeSolver=self.wake.logger.wakeSolver self.logger.wakeSolver["wakelength"]=wakelength self.logger.wakeSolver["simulationTime"]=time.time()-t0 self.logger.save_logs() From c5a0ed5e57adbe2d7d88f336742bb02c7563bb03 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 14:01:07 +0100 Subject: [PATCH 29/39] Fix logger grid assignment in solverFIT3D.py --- wakis/solverFIT3D.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wakis/solverFIT3D.py b/wakis/solverFIT3D.py index 00f050a..574370f 100644 --- a/wakis/solverFIT3D.py +++ b/wakis/solverFIT3D.py @@ -255,7 +255,7 @@ def __init__(self, grid, wake=None, cfln=0.5, dt=None, # assign logs self.logger = Logger() - self.logger.grid = self.grid.logger.grid_logs + self.logger.grid = self.grid.logger.grid self.logger.solver["use_gpu"] = use_gpu self.logger.solver["use_mpi"] = use_mpi self.logger.solver["bc_low"] = bc_low From c54d7f256ec0fc4a4e917ff113c2c8bfc1903d5f Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 14:03:38 +0100 Subject: [PATCH 30/39] Rename bg_log to bgLog in solverFIT3D.py --- wakis/solverFIT3D.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wakis/solverFIT3D.py b/wakis/solverFIT3D.py index 574370f..04c2531 100644 --- a/wakis/solverFIT3D.py +++ b/wakis/solverFIT3D.py @@ -261,7 +261,7 @@ def __init__(self, grid, wake=None, cfln=0.5, dt=None, self.logger.solver["bc_low"] = bc_low self.logger.solver["bc_high"] = bc_high self.logger.solver["n_pml"] = n_pml - self.logger.solver["bg"] = bg_log + self.logger.solver["bg"] = bgLog self.logger.solver["dt"] = self.dt self.logger.solver["solverInitializationTime"] = time.time() - t0 From 6bf3abfbbefc93dbb951583dbb6c0e8f9127109b Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 14:18:55 +0100 Subject: [PATCH 31/39] Update logger attribute access in test case --- tests/test_007_mpi_lossy_cavity.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_007_mpi_lossy_cavity.py b/tests/test_007_mpi_lossy_cavity.py index 9abeb99..3138cd8 100644 --- a/tests/test_007_mpi_lossy_cavity.py +++ b/tests/test_007_mpi_lossy_cavity.py @@ -329,9 +329,9 @@ def test_long_impedance(self): def test_log_file(self): global solver - solver.logger.gridLogs["gridInitializationTime"] = 0 #times can vary - solver.logger.solverLogs["solverInitializationTime"] = 0 - solver.logger.wakeSolverLogs["simulationTime"] = 0 + solver.logger.grid["gridInitializationTime"] = 0 #times can vary + solver.logger.solver["solverInitializationTime"] = 0 + solver.logger.wakeSolver["simulationTime"] = 0 logfile = os.path.join(solver.logger.wakeSolverLogs["results_folder"], "Simulation_Parameters.log") assert os.path.exists(logfile), "Log file not created" assert solver.logger.grid == self.gridLogs, "Grid logs do not match expected values" From 5583963b614a9c805ec502272f8275e847e4acf4 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Tue, 11 Nov 2025 14:27:58 +0100 Subject: [PATCH 32/39] Fix logfile path in MPI lossy cavity test --- tests/test_007_mpi_lossy_cavity.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_007_mpi_lossy_cavity.py b/tests/test_007_mpi_lossy_cavity.py index 3138cd8..9787bb7 100644 --- a/tests/test_007_mpi_lossy_cavity.py +++ b/tests/test_007_mpi_lossy_cavity.py @@ -332,7 +332,7 @@ def test_log_file(self): solver.logger.grid["gridInitializationTime"] = 0 #times can vary solver.logger.solver["solverInitializationTime"] = 0 solver.logger.wakeSolver["simulationTime"] = 0 - logfile = os.path.join(solver.logger.wakeSolverLogs["results_folder"], "Simulation_Parameters.log") + logfile = os.path.join(solver.logger.wakeSolver["results_folder"], "Simulation_Parameters.log") assert os.path.exists(logfile), "Log file not created" assert solver.logger.grid == self.gridLogs, "Grid logs do not match expected values" assert solver.logger.solver == self.solverLogs, "Solver logs do not match expected values" From e198fb98d556dd243b0a526217a547533ccd542a Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Thu, 13 Nov 2025 09:21:12 +0100 Subject: [PATCH 33/39] Exchanged function assign_logger by update_logger --- wakis/gridFIT3D.py | 37 +++++++++++++++------------- wakis/logger.py | 2 +- wakis/solverFIT3D.py | 31 ++++++++++++++---------- wakis/wakeSolver.py | 57 ++++++++++++++------------------------------ 4 files changed, 58 insertions(+), 69 deletions(-) diff --git a/wakis/gridFIT3D.py b/wakis/gridFIT3D.py index 0d6c0fd..c119d18 100644 --- a/wakis/gridFIT3D.py +++ b/wakis/gridFIT3D.py @@ -66,10 +66,12 @@ def __init__(self, xmin, xmax, ymin, ymax, zmin, zmax, stl_colors=None, verbose=1, stl_tol=1e-3): t0 = time.time() + self.logger = Logger() if verbose: print('Generating grid...') self.verbose = verbose self.use_mpi = use_mpi self.use_mesh_refinement = use_mesh_refinement + self.update_logger(['use_mesh_refinement']) # domain limits self.xmin = xmin @@ -84,6 +86,7 @@ def __init__(self, xmin, xmax, ymin, ymax, zmin, zmax, self.dx = (xmax - xmin) / Nx self.dy = (ymax - ymin) / Ny self.dz = (zmax - zmin) / Nz + self.update_logger(['Nx', 'Ny', 'Nz', 'dx', 'dy', 'dz']) # stl info self.stl_solids = stl_solids @@ -92,6 +95,14 @@ def __init__(self, xmin, xmax, ymin, ymax, zmin, zmax, self.stl_translate = stl_translate self.stl_scale = stl_scale self.stl_colors = stl_colors + self.update_logger(['stl_solids', 'stl_materials']) + if stl_rotate != [0., 0., 0.]: + self.update_logger(['stl_rotate']) + if stl_translate != [0., 0., 0.]: + self.update_logger(['stl_translate']) + if stl_scale != 1.0: + self.update_logger(['stl_scale']) + if stl_solids is not None: self._prepare_stl_dicts() @@ -146,23 +157,8 @@ def __init__(self, xmin, xmax, ymin, ymax, zmin, zmax, if stl_colors is None: self.assign_colors() - # Forward Parameters to logger - self.logger = Logger() - self.logger.grid["Nx"] = self.Nx - self.logger.grid["Ny"] = self.Ny - self.logger.grid["Nz"] = self.Nz - self.logger.grid["dx"] = self.dx - self.logger.grid["dy"] = self.dy - self.logger.grid["dz"] = self.dz - self.logger.grid["stl_solids"] = list(self.stl_solids.values()) if self.stl_solids is not None else [] - self.logger.grid["stl_materials"] = list(self.stl_materials.values()) if self.stl_materials is not None else [] - if stl_rotate != [0., 0., 0.]: - self.logger.grid["stl_rotate"] = self.stl_rotate - if stl_translate != [0., 0., 0.]: - self.logger.grid["stl_translate"] = self.stl_translate - if stl_scale != 1.0: - self.logger.grid["stl_scale"] = self.stl_scale - self.logger.grid["gridInitializationTime"] = time.time()-t0 + self.gridInitializationTime = time.time()-t0 + self.update_logger(['gridInitializationTime']) def compute_grid(self): X, Y, Z = np.meshgrid(self.x, self.y, self.z, indexing='ij') @@ -860,3 +856,10 @@ def clip(widget): pl.export_html('grid_inspect.html') else: pl.show() + + def update_logger(self, attrs) + """ + Assigns the parameters handed via attrs to the logger + """ + for atr in attrs: + self.logger.solver[atr] = getattr(self, atr) \ No newline at end of file diff --git a/wakis/logger.py b/wakis/logger.py index 99f50a3..c51b5fe 100644 --- a/wakis/logger.py +++ b/wakis/logger.py @@ -28,7 +28,7 @@ def save_logs(self): """ Save all logs (grid, solver, wakeSolver) into log-file inside the results folder. """ - logfile = os.path.join(self.wakeSolver["results_folder"], "Simulation_Parameters.log") + logfile = os.path.join(self.wakeSolver["results_folder"], "wakis.log") # Write sections if not os.path.exists(self.wakeSolver["results_folder"]): diff --git a/wakis/solverFIT3D.py b/wakis/solverFIT3D.py index 04c2531..04ad372 100644 --- a/wakis/solverFIT3D.py +++ b/wakis/solverFIT3D.py @@ -95,6 +95,7 @@ def __init__(self, grid, wake=None, cfln=0.5, dt=None, self.verbose = verbose t0 = time.time() + self.logger = Logger() # Flags self.step_0 = True @@ -111,10 +112,11 @@ def __init__(self, grid, wake=None, cfln=0.5, dt=None, self.one_step = self._one_step if use_stl: self.use_conductors = False + self.update_logger(['use_gpu, use_mpi']) # Grid self.grid = grid - bgLog = bg + self.background = bg self.Nx = self.grid.Nx self.Ny = self.grid.Ny self.Nz = self.grid.Nz @@ -132,6 +134,7 @@ def __init__(self, grid, wake=None, cfln=0.5, dt=None, self.iA = self.grid.iA self.tL = self.grid.tL self.itA = self.grid.itA + self.update_logger(['grid','background']) # Wake computation self.wake = wake @@ -176,6 +179,7 @@ def __init__(self, grid, wake=None, cfln=0.5, dt=None, if verbose: print('Applying boundary conditions...') self.bc_low = bc_low self.bc_high = bc_high + self.update_logger(['bc_low', 'bc_high']) self.apply_bc_to_C() # Materials @@ -204,6 +208,7 @@ def __init__(self, grid, wake=None, cfln=0.5, dt=None, self.pml_hi = 1.e-1 self.pml_func = np.geomspace self.fill_pml_sigmas() + self.update_logger(['n_pml']) # Timestep calculation if verbose: print('Calculating maximal stable timestep...') @@ -213,6 +218,7 @@ def __init__(self, grid, wake=None, cfln=0.5, dt=None, else: self.dt = dt self.dt = dtype(self.dt) + self.update_logger(['dt']) if self.use_conductivity: # relaxation time criterion tau @@ -253,17 +259,8 @@ def __init__(self, grid, wake=None, cfln=0.5, dt=None, if verbose: print(f'Total initialization time: {time.time() - t0} s') - # assign logs - self.logger = Logger() - self.logger.grid = self.grid.logger.grid - self.logger.solver["use_gpu"] = use_gpu - self.logger.solver["use_mpi"] = use_mpi - self.logger.solver["bc_low"] = bc_low - self.logger.solver["bc_high"] = bc_high - self.logger.solver["n_pml"] = n_pml - self.logger.solver["bg"] = bgLog - self.logger.solver["dt"] = self.dt - self.logger.solver["solverInitializationTime"] = time.time() - t0 + self.solverInitializationTime = time.time() - t0 + self.update_logger(['solverInitializationTime']) def update_tensors(self, tensor='all'): '''Update tensor matrices after @@ -1121,3 +1118,13 @@ def reset_fields(self): self.E[:, :, :, d] = 0.0 self.H[:, :, :, d] = 0.0 self.J[:, :, :, d] = 0.0 + + def update_logger(self, attrs) + """ + Assigns the parameters handed via attrs to the logger + """ + for atr in attrs: + if atr is 'grid': + self.logger.grid = self.grid.logger.grid + else: + self.logger.solver[atr] = getattr(self, atr) \ No newline at end of file diff --git a/wakis/wakeSolver.py b/wakis/wakeSolver.py index 476ba30..3bc0ea9 100644 --- a/wakis/wakeSolver.py +++ b/wakis/wakeSolver.py @@ -168,22 +168,7 @@ def __init__(self, wakelength=None, q=1e-9, sigmaz=1e-3, beta=1.0, if not os.path.exists(self.folder): os.mkdir(self.folder) - # create log - if self.log: - self.params_to_log() - - # Forward parameters to logger - self.logger.wakeSolver["ti"]=self.ti - self.logger.wakeSolver["q"]=self.q - self.logger.wakeSolver["sigmaz"]=self.sigmaz - self.logger.wakeSolver["beta"]=self.beta - self.logger.wakeSolver["xsource"]=self.xsource - self.logger.wakeSolver["ysource"]=self.ysource - self.logger.wakeSolver["xtest"]=self.xtest - self.logger.wakeSolver["ytest"]=self.ytest - self.logger.wakeSolver["chargedist"]=self.chargedist - self.logger.wakeSolver["skip_cells"]=self.skip_cells - self.logger.wakeSolver["results_folder"]=self.folder + self.assign_logs() def solve(self, compute_plane=None, **kwargs): ''' @@ -323,7 +308,7 @@ def calc_long_WP(self, Ezt=None,**kwargs): # longitudinal variables if self.zf is None: self.zf = self.z - dz = self.zf[2]-self.zf[1] + dz = np.diff(self.zf) #self.zf[2]-self.zf[1] zmax = np.max(self.zf) #should it be domain's edge instead? zmin = np.min(self.zf) @@ -378,7 +363,7 @@ def calc_long_WP(self, Ezt=None,**kwargs): ts = (z[k]+s[n])/self.v-zmin/self.v-self.t[0]+ti it = int(ts/dt) #find index for t if it < nt: - WP[n] = WP[n]+(Ezt[k, it])*dz #compute integral + WP[n] = WP[n]+(Ezt[k, it])*dz[k] #compute integral pbar.update(1) WP = WP/(self.q*1e12) # [V/pC] @@ -1126,26 +1111,6 @@ def log(self, txt): if self.verbose: print('\x1b[2;37m'+txt+'\x1b[0m') - - def params_to_log(self): - self.log(time.asctime()) - self.log('Wake computation') - self.log('='*24) - self.log(f'* Charge q = {self.q} [C]') - self.log(f'* Beam sigmaz = {self.sigmaz} [m]') - self.log(f'* xsource, ysource = {self.xsource}, {self.ysource} [m]') - self.log(f'* xtest, ytest = {self.xtest}, {self.ytest} [m]') - self.log(f'* Beam injection time ti= {self.ti} [s]') - - if self.chargedist is not None: - if type(self.chargedist) is str: - self.log(f'* Charge distribution file: {self.chargedist}') - else: - self.log(f'* Charge distribution data is provided') - else: - self.log(f'* Charge distribution analytic') - - self.log('\n') def read_cst_3d(self, path=None, folder='3d', filename='Ez.h5', units=1e-3): ''' @@ -1306,4 +1271,18 @@ def sorter(item): self.zf = z self.t = np.array(t) - + def assign_logs(self) + """ + Assigns the parameters of the wake to the logger + """ + self.logger.wakeSolver["ti"]=self.ti + self.logger.wakeSolver["q"]=self.q + self.logger.wakeSolver["sigmaz"]=self.sigmaz + self.logger.wakeSolver["beta"]=self.beta + self.logger.wakeSolver["xsource"]=self.xsource + self.logger.wakeSolver["ysource"]=self.ysource + self.logger.wakeSolver["xtest"]=self.xtest + self.logger.wakeSolver["ytest"]=self.ytest + self.logger.wakeSolver["chargedist"]=self.chargedist + self.logger.wakeSolver["skip_cells"]=self.skip_cells + self.logger.wakeSolver["results_folder"]=self.folder From f22c2097aa7cf9e0edc6ae023e0fb6d4758282b6 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Thu, 13 Nov 2025 09:35:59 +0100 Subject: [PATCH 34/39] Fixed bugs --- tests/test_007_mpi_lossy_cavity.py | 8 ++++---- wakis/gridFIT3D.py | 2 +- wakis/solverFIT3D.py | 8 ++++---- wakis/wakeSolver.py | 2 +- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/test_007_mpi_lossy_cavity.py b/tests/test_007_mpi_lossy_cavity.py index 9787bb7..5bbc8cc 100644 --- a/tests/test_007_mpi_lossy_cavity.py +++ b/tests/test_007_mpi_lossy_cavity.py @@ -82,12 +82,12 @@ class TestMPILossyCavity: -6.04105997e+01 ,-3.06532160e+01 ,-1.17749936e+01 ,-3.12574866e+00, -7.35339521e-01 ,-1.13085658e-01 , 7.18247535e-01 , 8.73829036e-02]) - gridLogs = {'Nx': 60, 'Ny': 60, 'Nz': 140, 'dx': 0.008666666348775227, 'dy': 0.008666666348775227, + gridLogs = {'use_mesh_refinement': False, 'Nx': 60, 'Ny': 60, 'Nz': 140, 'dx': 0.008666666348775227, 'dy': 0.008666666348775227, 'dz': 0.005714285799435207, 'stl_solids': ['tests/stl/007_vacuum_cavity.stl', 'tests/stl/007_lossymetal_shell.stl'], 'stl_materials': ['vacuum', [30, 1.0, 30]], 'gridInitializationTime': 0} - solverLogs = {'use_gpu': False, 'use_mpi': False, 'bc_low': ['pec', 'pec', 'pec'], - 'bc_high': ['pec', 'pec', 'pec'], 'n_pml': 10, 'bg': 'pec', + solverLogs = {'use_gpu': False, 'use_mpi': False, 'background': 'pec','bc_low': ['pec', 'pec', 'pec'], + 'bc_high': ['pec', 'pec', 'pec'], 'n_pml': 10, 'dt': 6.970326728398968e-12, 'solverInitializationTime': 0} wakeSolverLogs = {'ti': 2.8516132094735135e-09, 'q': 1e-09, 'sigmaz': 0.1, 'beta': 1.0, @@ -332,7 +332,7 @@ def test_log_file(self): solver.logger.grid["gridInitializationTime"] = 0 #times can vary solver.logger.solver["solverInitializationTime"] = 0 solver.logger.wakeSolver["simulationTime"] = 0 - logfile = os.path.join(solver.logger.wakeSolver["results_folder"], "Simulation_Parameters.log") + logfile = os.path.join(solver.logger.wakeSolver["results_folder"], "wakis.log") assert os.path.exists(logfile), "Log file not created" assert solver.logger.grid == self.gridLogs, "Grid logs do not match expected values" assert solver.logger.solver == self.solverLogs, "Solver logs do not match expected values" diff --git a/wakis/gridFIT3D.py b/wakis/gridFIT3D.py index c119d18..5c750f0 100644 --- a/wakis/gridFIT3D.py +++ b/wakis/gridFIT3D.py @@ -857,7 +857,7 @@ def clip(widget): else: pl.show() - def update_logger(self, attrs) + def update_logger(self, attrs): """ Assigns the parameters handed via attrs to the logger """ diff --git a/wakis/solverFIT3D.py b/wakis/solverFIT3D.py index 04ad372..a94fba0 100644 --- a/wakis/solverFIT3D.py +++ b/wakis/solverFIT3D.py @@ -112,7 +112,7 @@ def __init__(self, grid, wake=None, cfln=0.5, dt=None, self.one_step = self._one_step if use_stl: self.use_conductors = False - self.update_logger(['use_gpu, use_mpi']) + self.update_logger(['use_gpu', 'use_mpi']) # Grid self.grid = grid @@ -208,7 +208,7 @@ def __init__(self, grid, wake=None, cfln=0.5, dt=None, self.pml_hi = 1.e-1 self.pml_func = np.geomspace self.fill_pml_sigmas() - self.update_logger(['n_pml']) + self.update_logger(['n_pml']) # Timestep calculation if verbose: print('Calculating maximal stable timestep...') @@ -1119,12 +1119,12 @@ def reset_fields(self): self.H[:, :, :, d] = 0.0 self.J[:, :, :, d] = 0.0 - def update_logger(self, attrs) + def update_logger(self, attrs): """ Assigns the parameters handed via attrs to the logger """ for atr in attrs: - if atr is 'grid': + if atr == 'grid': self.logger.grid = self.grid.logger.grid else: self.logger.solver[atr] = getattr(self, atr) \ No newline at end of file diff --git a/wakis/wakeSolver.py b/wakis/wakeSolver.py index 3bc0ea9..a9c46d5 100644 --- a/wakis/wakeSolver.py +++ b/wakis/wakeSolver.py @@ -1271,7 +1271,7 @@ def sorter(item): self.zf = z self.t = np.array(t) - def assign_logs(self) + def assign_logs(self): """ Assigns the parameters of the wake to the logger """ From d7e3840727129a02c6d7322be96fe9872ae16acd Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Thu, 13 Nov 2025 10:04:04 +0100 Subject: [PATCH 35/39] Fixed bug --- wakis/gridFIT3D.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wakis/gridFIT3D.py b/wakis/gridFIT3D.py index 5c750f0..06ac90f 100644 --- a/wakis/gridFIT3D.py +++ b/wakis/gridFIT3D.py @@ -862,4 +862,4 @@ def update_logger(self, attrs): Assigns the parameters handed via attrs to the logger """ for atr in attrs: - self.logger.solver[atr] = getattr(self, atr) \ No newline at end of file + self.logger.grid[atr] = getattr(self, atr) \ No newline at end of file From a19ab6106441bd615beb914692f36860bead9a3b Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Thu, 13 Nov 2025 10:17:25 +0100 Subject: [PATCH 36/39] Exchange list by dictionary in comparison --- tests/test_007_mpi_lossy_cavity.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_007_mpi_lossy_cavity.py b/tests/test_007_mpi_lossy_cavity.py index 5bbc8cc..35d17a2 100644 --- a/tests/test_007_mpi_lossy_cavity.py +++ b/tests/test_007_mpi_lossy_cavity.py @@ -83,8 +83,8 @@ class TestMPILossyCavity: -7.35339521e-01 ,-1.13085658e-01 , 7.18247535e-01 , 8.73829036e-02]) gridLogs = {'use_mesh_refinement': False, 'Nx': 60, 'Ny': 60, 'Nz': 140, 'dx': 0.008666666348775227, 'dy': 0.008666666348775227, - 'dz': 0.005714285799435207, 'stl_solids': ['tests/stl/007_vacuum_cavity.stl', 'tests/stl/007_lossymetal_shell.stl'], - 'stl_materials': ['vacuum', [30, 1.0, 30]], 'gridInitializationTime': 0} + 'dz': 0.005714285799435207, 'stl_solids': {'cavity': 'tests/stl/007_vacuum_cavity.stl', 'shell': 'tests/stl/007_lossymetal_shell.stl'}, + 'stl_materials': {'cavity': 'vacuum', 'shell': [30, 1.0, 30]}}, 'gridInitializationTime': 0} solverLogs = {'use_gpu': False, 'use_mpi': False, 'background': 'pec','bc_low': ['pec', 'pec', 'pec'], 'bc_high': ['pec', 'pec', 'pec'], 'n_pml': 10, From c1f59e010623fe6a0b4c244dd4b2a8050999b6b8 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Thu, 13 Nov 2025 10:20:49 +0100 Subject: [PATCH 37/39] Fix indentation in gridLogs dictionary --- tests/test_007_mpi_lossy_cavity.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_007_mpi_lossy_cavity.py b/tests/test_007_mpi_lossy_cavity.py index 35d17a2..50756ad 100644 --- a/tests/test_007_mpi_lossy_cavity.py +++ b/tests/test_007_mpi_lossy_cavity.py @@ -84,7 +84,7 @@ class TestMPILossyCavity: gridLogs = {'use_mesh_refinement': False, 'Nx': 60, 'Ny': 60, 'Nz': 140, 'dx': 0.008666666348775227, 'dy': 0.008666666348775227, 'dz': 0.005714285799435207, 'stl_solids': {'cavity': 'tests/stl/007_vacuum_cavity.stl', 'shell': 'tests/stl/007_lossymetal_shell.stl'}, - 'stl_materials': {'cavity': 'vacuum', 'shell': [30, 1.0, 30]}}, 'gridInitializationTime': 0} + 'stl_materials': {'cavity': 'vacuum', 'shell': [30, 1.0, 30]}, 'gridInitializationTime': 0} solverLogs = {'use_gpu': False, 'use_mpi': False, 'background': 'pec','bc_low': ['pec', 'pec', 'pec'], 'bc_high': ['pec', 'pec', 'pec'], 'n_pml': 10, From 9a684901a609119dad70b7750faaed30cad4dc2c Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Thu, 13 Nov 2025 10:33:51 +0100 Subject: [PATCH 38/39] Fix indentation in solverLogs dictionary --- tests/test_007_mpi_lossy_cavity.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_007_mpi_lossy_cavity.py b/tests/test_007_mpi_lossy_cavity.py index 50756ad..968cba2 100644 --- a/tests/test_007_mpi_lossy_cavity.py +++ b/tests/test_007_mpi_lossy_cavity.py @@ -87,7 +87,7 @@ class TestMPILossyCavity: 'stl_materials': {'cavity': 'vacuum', 'shell': [30, 1.0, 30]}, 'gridInitializationTime': 0} solverLogs = {'use_gpu': False, 'use_mpi': False, 'background': 'pec','bc_low': ['pec', 'pec', 'pec'], - 'bc_high': ['pec', 'pec', 'pec'], 'n_pml': 10, + 'bc_high': ['pec', 'pec', 'pec'], 'dt': 6.970326728398968e-12, 'solverInitializationTime': 0} wakeSolverLogs = {'ti': 2.8516132094735135e-09, 'q': 1e-09, 'sigmaz': 0.1, 'beta': 1.0, From 534792c489be5cdbceb102072c2456f9ee89c096 Mon Sep 17 00:00:00 2001 From: Antoniahuber Date: Fri, 14 Nov 2025 17:32:15 +0100 Subject: [PATCH 39/39] Separate changes from another pull request --- wakis/wakeSolver.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/wakis/wakeSolver.py b/wakis/wakeSolver.py index a9c46d5..62392a2 100644 --- a/wakis/wakeSolver.py +++ b/wakis/wakeSolver.py @@ -308,7 +308,7 @@ def calc_long_WP(self, Ezt=None,**kwargs): # longitudinal variables if self.zf is None: self.zf = self.z - dz = np.diff(self.zf) #self.zf[2]-self.zf[1] + dz = self.zf[2]-self.zf[1] zmax = np.max(self.zf) #should it be domain's edge instead? zmin = np.min(self.zf) @@ -363,7 +363,7 @@ def calc_long_WP(self, Ezt=None,**kwargs): ts = (z[k]+s[n])/self.v-zmin/self.v-self.t[0]+ti it = int(ts/dt) #find index for t if it < nt: - WP[n] = WP[n]+(Ezt[k, it])*dz[k] #compute integral + WP[n] = WP[n]+(Ezt[k, it])*dz #compute integral pbar.update(1) WP = WP/(self.q*1e12) # [V/pC]