From 86cbb42f8cb4306bdc6d5762376b352fd833163e Mon Sep 17 00:00:00 2001 From: Vyacheslav Fedorov Date: Thu, 31 Aug 2023 06:24:18 +0000 Subject: [PATCH] Update solver module --- .gitignore | 3 + docs/conf.py | 10 +- redpic/__init__.py | 4 +- redpic/core/config.py | 12 +- redpic/solver/__init__.py | 260 +++----------------------------------- redpic/solver/base.py | 15 +++ redpic/solver/red.py | 161 +++++++++++++++++++++++ redpic/solver/utils.py | 78 ++++++++++++ setup.py | 12 +- 9 files changed, 294 insertions(+), 261 deletions(-) create mode 100644 redpic/solver/base.py create mode 100644 redpic/solver/red.py create mode 100644 redpic/solver/utils.py diff --git a/.gitignore b/.gitignore index 5838276..d6b98e7 100644 --- a/.gitignore +++ b/.gitignore @@ -124,3 +124,6 @@ dmypy.json # End of https://www.gitignore.io/api/python,jupyternotebooks *.pyc .DS_Store + +# CProfiler data +profiler \ No newline at end of file diff --git a/docs/conf.py b/docs/conf.py index 59c9ced..0ead557 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -20,13 +20,13 @@ # -- Project information ----------------------------------------------------- -project = redpic.config.PROJECT_NAME -copyright = f'{today.year}, {redpic.config.PROJECT_AUTHOR}' -author = redpic.config.PROJECT_AUTHOR +project = redpic.config.REDPIC_NAME +copyright = f'{today.year}, {redpic.config.REDPIC_AUTHOR}' +author = redpic.config.REDPIC_AUTHOR # The short X.Y version. -version = redpic.config.PROJECT_VERSION +version = redpic.config.REDPIC_VERSION # The full version, including alpha/beta/rc tags. -release = redpic.config.PROJECT_VERSION +release = redpic.config.REDPIC_VERSION # -- General configuration --------------------------------------------------- diff --git a/redpic/__init__.py b/redpic/__init__.py index 073975c..2a490c7 100644 --- a/redpic/__init__.py +++ b/redpic/__init__.py @@ -3,8 +3,8 @@ from redpic.core import config from redpic.solver import * -__version__ = config.PROJECT_VERSION +__version__ = config.REDPIC_VERSION -__doc__ = config.PROJECT_DOC +__doc__ = config.REDPIC_DOC __all__ = accelerator.__all__ + beam.__all__ + solver.__all__ diff --git a/redpic/core/config.py b/redpic/core/config.py index 833f80b..d5cc5ef 100644 --- a/redpic/core/config.py +++ b/redpic/core/config.py @@ -4,14 +4,14 @@ logging_config.dictConfig(LOGGING) -PROJECT_NAME = "redpic" +REDPIC_NAME = "redpic" -PROJECT_VERSION = "0.7.50" +REDPIC_VERSION = "0.7.51" -PROJECT_AUTHOR = "Vyacheslav Fedorov" +REDPIC_AUTHOR = "Vyacheslav Fedorov" -PROJECT_AUTHOR_EMAIL = "slava@fuodorov.ru" +REDPIC_AUTHOR_EMAIL = "slava@fuodorov.ru" -PROJECT_DOC = "Relativistic Difference Scheme Particle-in-Cell code (REDPIC)" +REDPIC_DOC = "Relativistic Difference Scheme Particle-in-Cell code (REDPIC)" -PROJECT_URL = "https://github.com/fuodorov/redpic" +REDPIC_URL = "https://github.com/fuodorov/redpic" diff --git a/redpic/solver/__init__.py b/redpic/solver/__init__.py index 1697c31..cd152e4 100644 --- a/redpic/solver/__init__.py +++ b/redpic/solver/__init__.py @@ -1,242 +1,18 @@ -import numpy as np -import pandas as pd -from numba import jit, prange -from scipy import misc - -from redpic import constants as const -from redpic.accelerator import Accelerator -from redpic.beam.base import BaseBeam - -__all__ = ["Simulation"] - - -def get_field_accelerator( - acc: Accelerator, type: str, x: np.array, y: np.array, z: np.array -) -> (np.array, np.array, np.array): - assert type in ("E", "B"), "Check type field! (E or B)" - - dz = acc.dz - offset_x = acc.Dx(z) - offset_xp = acc.Dxp(z) - offset_y = acc.Dy(z) - offset_yp = acc.Dyp(z) - x_corr = x - offset_x - y_corr = y - offset_y - r_corr = np.sqrt(x_corr * x_corr + y_corr * y_corr) - - if type == "E": - Ez = acc.Ez(z) - dEzdz = acc.dEzdz(z) - d2Ezdz2 = misc.derivative(acc.dEzdz, z, dx=dz, n=1) - Ez = Ez - d2Ezdz2 * r_corr * r_corr / 4 # row remainder - Ex = -dEzdz * x_corr / 2 + Ez * offset_xp # row remainder - Ey = -dEzdz * y_corr / 2 + Ez * offset_yp # row remainder - return Ex, Ey, Ez - if type == "B": - Bx = acc.Bx(z) - By = acc.By(z) - Bz = acc.Bz(z) - dBzdz = acc.dBzdz(z) - d2Bzdz2 = misc.derivative(acc.dBzdz, z, dx=dz, n=1) - Gz = acc.Gz(z) - Bz = Bz - d2Bzdz2 * r_corr * r_corr / 4 # row remainder - Bx = Bx + Gz * y_corr - dBzdz * x_corr / 2 + Bz * offset_xp # row remainder - By = By + Gz * x_corr - dBzdz * y_corr / 2 + Bz * offset_yp # row remainder - return Bx, By, Bz - - -@jit(nopython=True, parallel=True, fastmath=True, cache=True, nogil=True) -def sum_field_particles( - x: np.array, y: np.array, z: np.array, z_start: float, z_stop: float -) -> (np.array, np.array, np.array): - n = int(len(x)) - Fx = np.zeros(n) - Fy = np.zeros(n) - Fz = np.zeros(n) - for i in prange(n): # pylint: disable=E1133 - if z_start <= z[i] <= z_stop: - for j in prange(n): # pylint: disable=E1133 - if z_start <= z[j] <= z_stop and i != j: - Fx[i] += (x[i] - x[j]) / ((x[j] - x[i]) ** 2 + (y[j] - y[i]) ** 2 + (z[j] - z[i]) ** 2) ** (3 / 2) - Fy[i] += (y[i] - y[j]) / ((x[j] - x[i]) ** 2 + (y[j] - y[i]) ** 2 + (z[j] - z[i]) ** 2) ** (3 / 2) - Fz[i] += (z[i] - z[j]) / ((x[j] - x[i]) ** 2 + (y[j] - y[i]) ** 2 + (z[j] - z[i]) ** 2) ** (3 / 2) - return Fx, Fy, Fz - - -def get_field_beam( - beam: BaseBeam, acc: Accelerator, type: str, x: np.array, y: np.array, z: np.array -) -> (np.array, np.array, np.array): - assert type in ("E", "B"), "Check type field! (E or B)" - Q = beam.total_charge - q = Q / beam.n - if type == "E": - Ex, Ey, Ez = sum_field_particles(x, y, z, acc.z_start, acc.z_stop) - return ( - const.ke / (4 * np.pi) * q * Ex / 1e6, - const.ke / (4 * np.pi) * q * Ey / 1e6, - const.ke / (4 * np.pi) * q * Ez / 1e6, - ) - if type == "B": - Bx, By, Bz = sum_field_particles(x, y, z, acc.z_start, acc.z_stop) - return const.km / (4 * np.pi) * q * Bx, -const.km / (4 * np.pi) * q * By, 0 * const.km / (4 * np.pi) * q * Bz - - -@jit(nopython=True, parallel=True, fastmath=True, cache=True, nogil=True) -def first_step_red( - x: np.array, - y: np.array, - z: np.array, - px: np.array, - py: np.array, - pz: np.array, - Fx: np.array, - Fy: np.array, - Fz: np.array, - z_start: float, - z_stop: float, -) -> (np.array, np.array, np.array, np.array, np.array, np.array): - n = int(len(x)) - for i in prange(n): # pylint: disable=E1133 - if z_start <= z[i] <= z_stop: - px[i] += 2 * Fx[i] - py[i] += 2 * Fy[i] - pz[i] += 2 * Fz[i] - gamma = (1 + px[i] ** 2 + py[i] ** 2 + pz[i] ** 2) ** (1 / 2) - vx = px[i] / gamma - vy = py[i] / gamma - vz = pz[i] / gamma - x[i] += vx - y[i] += vy - z[i] += vz - else: - gamma = (1 + px[i] ** 2 + py[i] ** 2 + pz[i] ** 2) ** (1 / 2) - vz = pz[i] / gamma - z[i] += vz - return x, y, z, px, py, pz - - -@jit(nopython=True, parallel=True, fastmath=True, cache=True, nogil=True) -def second_step_red( - x: np.array, - y: np.array, - z: np.array, - px: np.array, - py: np.array, - pz: np.array, - Fx: np.array, - Fy: np.array, - Fz: np.array, - z_start: float, - z_stop: float, -) -> (np.array, np.array, np.array, np.array, np.array, np.array): - n = int(len(x)) - for i in prange(n): # pylint: disable=E1133 - if z_start <= z[i] <= z_stop: - gamma = (1 + px[i] ** 2 + py[i] ** 2 + pz[i] ** 2) ** (1 / 2) - vx = px[i] / gamma - vy = py[i] / gamma - vz = pz[i] / gamma - b2 = 1 + Fx[i] ** 2 + Fy[i] ** 2 + Fz[i] ** 2 - b1 = 2 - b2 - b3 = 2 * (vx * Fx[i] + vy * Fy[i] + vz * Fz[i]) - fx = 2 * (vy * Fz[i] - vz * Fy[i]) - fy = 2 * (vz * Fx[i] - vx * Fz[i]) - fz = 2 * (vx * Fy[i] - vy * Fx[i]) - vx = (vx * b1 + fx + Fx[i] * b3) / b2 - vy = (vy * b1 + fy + Fy[i] * b3) / b2 - vz = (vz * b1 + fz + Fz[i] * b3) / b2 - x[i] += vx - y[i] += vy - z[i] += vz - px[i] = vx * gamma - py[i] = vy * gamma - pz[i] = vz * gamma - else: - gamma = (1 + px[i] ** 2 + py[i] ** 2 + pz[i] ** 2) ** (1 / 2) - vz = pz[i] / gamma - z[i] += vz - return x, y, z, px, py, pz - - -class Simulation: - """Simulation of the beam tracking in the accelerator""" - - def __init__(self, beam, accelerator): - self.beam = beam - self.acc = accelerator - self.result = {} - - def track(self, *, n_files: int = 20) -> None: - assert n_files > 0, "The number of files (n_files) must be a positive number!" - - # Init parameterss - Y = self.beam.da - Y[2] = Y[2] + self.acc.z_start - max(Y[2]) # set initial beam position - - z_start = self.acc.z_start - z_stop = self.acc.z_stop - dz = self.acc.dz - dt = dz / const.c - t_max = (z_stop - z_start) / const.c - - m = self.beam.type.mass - q = self.beam.type.charge - - P0 = m * const.c * const.c / (const.e * 1e6) - E0 = m * const.c / (q * dt * 1e6) - B0 = m / (q * dt) - - # RED - for t in np.arange(0, t_max, 2 * dt): - # into the internal system - Y[0], Y[1], Y[2] = Y[0] / dz, Y[1] / dz, Y[2] / dz - Y[3], Y[4], Y[5] = Y[3] / P0, Y[4] / P0, Y[5] / P0 - - # get electric field from accelerator - Ex, Ey, Ez = get_field_accelerator(self.acc, "E", Y[0] * dz, Y[1] * dz, Y[2] * dz) - Ex, Ey, Ez = Ex / E0, Ey / E0, Ez / E0 - - # get electric field from beam - if self.beam.total_charge: - ex, ey, ez = get_field_beam(self.beam, self.acc, "E", Y[0] * dz, Y[1] * dz, Y[2] * dz) - ex, ey, ez = ex / E0, ey / E0, ez / E0 - Ex, Ey, Ez = Ex + ex, Ey + ey, Ez + ez - - # first step RED - Y[0], Y[1], Y[2], Y[3], Y[4], Y[5] = first_step_red( - Y[0], Y[1], Y[2], Y[3], Y[4], Y[5], Ex, Ey, Ez, z_start / dz, z_stop / dz - ) - gamma = np.sqrt(1 + Y[3] * Y[3] + Y[4] * Y[4] + Y[5] * Y[5]) - vz = Y[5] / gamma - - # get magnetic field from accelerator - Bx, By, Bz = get_field_accelerator(self.acc, "B", Y[0] * dz, Y[1] * dz, Y[2] * dz) - Bx, By, Bz = Bx / B0 / gamma, By / B0 / gamma, Bz / B0 / gamma - - # get magnetic field from beam - if self.beam.total_charge: - bx, by, bz = get_field_beam(self.beam, self.acc, "B", Y[0] * dz, Y[1] * dz, Y[2] * dz) - bx, by, bz = bx * vz / B0, by * vz / B0, bz * vz / B0 - Bx, By, Bz = Bx + bx, By + by, Bz + bz - - # second step RED - Y[0], Y[1], Y[2], Y[3], Y[4], Y[5] = second_step_red( - Y[0], Y[1], Y[2], Y[3], Y[4], Y[5], Bx, By, Bz, z_start / dz, z_stop / dz - ) - gamma = np.sqrt(1 + Y[3] * Y[3] + Y[4] * Y[4] + Y[5] * Y[5]) - vz = Y[5] / gamma - - # into the SI - Y[0], Y[1], Y[2] = Y[0] * dz, Y[1] * dz, Y[2] * dz - Y[3], Y[4], Y[5] = Y[3] * P0, Y[4] * P0, Y[5] * P0 - Bx, By, Bz = Bx * gamma * B0, By * gamma * B0, Bz * gamma * B0 - Ex, Ey, Ez = Ex * E0, Ey * E0, Ez * E0 - - progress = t / t_max * 100 - meters = z_start + t * const.c - X = np.row_stack((Y[0], Y[1], Y[2], Y[3], Y[4], Y[5], Bx, By, Bz, Ex, Ey, Ez)) - Xt = np.transpose(X) - df = pd.DataFrame(Xt, columns=["x", "y", "z", "px", "py", "pz", "Bx", "By", "Bz", "Ex", "Ey", "Ez"]) - if progress % (100 // n_files) < 2 * dt / t_max * 100: - self.result.update({round(meters, 3): df}) - print("\rz = %.2f m (%.1f %%) " % (meters, progress), end="") +from redpic.solver.base import BaseSimulation +from redpic.solver.red import REDSimulation +from redpic.solver.utils import ( + get_field_accelerator, + get_field_beam, + sum_field_particles, +) + +Simulation = REDSimulation + +__all__ = [ + "BaseSimulation", + "Simulation", + "REDSimulation", + "get_field_accelerator", + "get_field_beam", + "sum_field_particles", +] diff --git a/redpic/solver/base.py b/redpic/solver/base.py new file mode 100644 index 0000000..f00d32f --- /dev/null +++ b/redpic/solver/base.py @@ -0,0 +1,15 @@ +from abc import ABC, abstractmethod + +from redpic.accelerator import Accelerator +from redpic.beam.base import BaseBeam + + +class BaseSimulation(ABC): + def __init__(self, beam: BaseBeam, accelerator: Accelerator): + self.beam = beam + self.acc = accelerator + self.result = {} + + @abstractmethod + def track(self, *, n_files: int = 20) -> None: + pass diff --git a/redpic/solver/red.py b/redpic/solver/red.py new file mode 100644 index 0000000..9d24fe0 --- /dev/null +++ b/redpic/solver/red.py @@ -0,0 +1,161 @@ +import numpy as np +import pandas as pd +from numba import jit, prange + +from redpic import constants as const +from redpic.solver.base import BaseSimulation +from redpic.solver.utils import get_field_accelerator, get_field_beam + + +class REDSimulation(BaseSimulation): + @staticmethod + @jit(nopython=True, parallel=True, fastmath=True, cache=True, nogil=True) + def _first_step_red( + x: np.array, + y: np.array, + z: np.array, + px: np.array, + py: np.array, + pz: np.array, + Fx: np.array, + Fy: np.array, + Fz: np.array, + z_start: float, + z_stop: float, + ) -> (np.array, np.array, np.array, np.array, np.array, np.array): + n = int(len(x)) + for i in prange(n): # pylint: disable=E1133 + if z_start <= z[i] <= z_stop: + px[i] += 2 * Fx[i] + py[i] += 2 * Fy[i] + pz[i] += 2 * Fz[i] + gamma = (1 + px[i] ** 2 + py[i] ** 2 + pz[i] ** 2) ** (1 / 2) + vx = px[i] / gamma + vy = py[i] / gamma + vz = pz[i] / gamma + x[i] += vx + y[i] += vy + z[i] += vz + else: + gamma = (1 + px[i] ** 2 + py[i] ** 2 + pz[i] ** 2) ** (1 / 2) + vz = pz[i] / gamma + z[i] += vz + return x, y, z, px, py, pz + + @staticmethod + @jit(nopython=True, parallel=True, fastmath=True, cache=True, nogil=True) + def _second_step_red( + x: np.array, + y: np.array, + z: np.array, + px: np.array, + py: np.array, + pz: np.array, + Fx: np.array, + Fy: np.array, + Fz: np.array, + z_start: float, + z_stop: float, + ) -> (np.array, np.array, np.array, np.array, np.array, np.array): + n = int(len(x)) + for i in prange(n): # pylint: disable=E1133 + if z_start <= z[i] <= z_stop: + gamma = (1 + px[i] ** 2 + py[i] ** 2 + pz[i] ** 2) ** (1 / 2) + vx = px[i] / gamma + vy = py[i] / gamma + vz = pz[i] / gamma + b2 = 1 + Fx[i] ** 2 + Fy[i] ** 2 + Fz[i] ** 2 + b1 = 2 - b2 + b3 = 2 * (vx * Fx[i] + vy * Fy[i] + vz * Fz[i]) + fx = 2 * (vy * Fz[i] - vz * Fy[i]) + fy = 2 * (vz * Fx[i] - vx * Fz[i]) + fz = 2 * (vx * Fy[i] - vy * Fx[i]) + vx = (vx * b1 + fx + Fx[i] * b3) / b2 + vy = (vy * b1 + fy + Fy[i] * b3) / b2 + vz = (vz * b1 + fz + Fz[i] * b3) / b2 + x[i] += vx + y[i] += vy + z[i] += vz + px[i] = vx * gamma + py[i] = vy * gamma + pz[i] = vz * gamma + else: + gamma = (1 + px[i] ** 2 + py[i] ** 2 + pz[i] ** 2) ** (1 / 2) + vz = pz[i] / gamma + z[i] += vz + return x, y, z, px, py, pz + + def track(self, *, n_files: int = 20) -> None: + assert n_files > 0, "The number of files (n_files) must be a positive number!" + + # Init parameterss + Y = self.beam.da + Y[2] = Y[2] + self.acc.z_start - max(Y[2]) # set initial beam position + + z_start = self.acc.z_start + z_stop = self.acc.z_stop + dz = self.acc.dz + dt = dz / const.c + t_max = (z_stop - z_start) / const.c + + m = self.beam.type.mass + q = self.beam.type.charge + + P0 = m * const.c * const.c / (const.e * 1e6) + E0 = m * const.c / (q * dt * 1e6) + B0 = m / (q * dt) + + # RED + for t in np.arange(0, t_max, 2 * dt): + # into the internal system + Y[0], Y[1], Y[2] = Y[0] / dz, Y[1] / dz, Y[2] / dz + Y[3], Y[4], Y[5] = Y[3] / P0, Y[4] / P0, Y[5] / P0 + + # get electric field from accelerator + Ex, Ey, Ez = get_field_accelerator(self.acc, "E", Y[0] * dz, Y[1] * dz, Y[2] * dz) + Ex, Ey, Ez = Ex / E0, Ey / E0, Ez / E0 + + # get electric field from beam + if self.beam.total_charge: + ex, ey, ez = get_field_beam(self.beam, self.acc, "E", Y[0] * dz, Y[1] * dz, Y[2] * dz) + ex, ey, ez = ex / E0, ey / E0, ez / E0 + Ex, Ey, Ez = Ex + ex, Ey + ey, Ez + ez + + # first step RED + Y[0], Y[1], Y[2], Y[3], Y[4], Y[5] = self._first_step_red( + Y[0], Y[1], Y[2], Y[3], Y[4], Y[5], Ex, Ey, Ez, z_start / dz, z_stop / dz + ) + gamma = np.sqrt(1 + Y[3] * Y[3] + Y[4] * Y[4] + Y[5] * Y[5]) + vz = Y[5] / gamma + + # get magnetic field from accelerator + Bx, By, Bz = get_field_accelerator(self.acc, "B", Y[0] * dz, Y[1] * dz, Y[2] * dz) + Bx, By, Bz = Bx / B0 / gamma, By / B0 / gamma, Bz / B0 / gamma + + # get magnetic field from beam + if self.beam.total_charge: + bx, by, bz = get_field_beam(self.beam, self.acc, "B", Y[0] * dz, Y[1] * dz, Y[2] * dz) + bx, by, bz = bx * vz / B0, by * vz / B0, bz * vz / B0 + Bx, By, Bz = Bx + bx, By + by, Bz + bz + + # second step RED + Y[0], Y[1], Y[2], Y[3], Y[4], Y[5] = self._second_step_red( + Y[0], Y[1], Y[2], Y[3], Y[4], Y[5], Bx, By, Bz, z_start / dz, z_stop / dz + ) + gamma = np.sqrt(1 + Y[3] * Y[3] + Y[4] * Y[4] + Y[5] * Y[5]) + vz = Y[5] / gamma + + # into the SI + Y[0], Y[1], Y[2] = Y[0] * dz, Y[1] * dz, Y[2] * dz + Y[3], Y[4], Y[5] = Y[3] * P0, Y[4] * P0, Y[5] * P0 + Bx, By, Bz = Bx * gamma * B0, By * gamma * B0, Bz * gamma * B0 + Ex, Ey, Ez = Ex * E0, Ey * E0, Ez * E0 + + progress = t / t_max * 100 + meters = z_start + t * const.c + X = np.row_stack((Y[0], Y[1], Y[2], Y[3], Y[4], Y[5], Bx, By, Bz, Ex, Ey, Ez)) + Xt = np.transpose(X) + df = pd.DataFrame(Xt, columns=["x", "y", "z", "px", "py", "pz", "Bx", "By", "Bz", "Ex", "Ey", "Ez"]) + if progress % (100 // n_files) < 2 * dt / t_max * 100: + self.result.update({round(meters, 3): df}) + print("\rz = %.2f m (%.1f %%) " % (meters, progress), end="") diff --git a/redpic/solver/utils.py b/redpic/solver/utils.py new file mode 100644 index 0000000..49512ea --- /dev/null +++ b/redpic/solver/utils.py @@ -0,0 +1,78 @@ +import numpy as np +from numba import jit, prange +from scipy import misc + +from redpic import constants as const +from redpic.accelerator import Accelerator +from redpic.beam.base import BaseBeam + + +def get_field_accelerator( + acc: Accelerator, type: str, x: np.array, y: np.array, z: np.array +) -> (np.array, np.array, np.array): + assert type in ("E", "B"), "Check type field! (E or B)" + + dz = acc.dz + offset_x = acc.Dx(z) + offset_xp = acc.Dxp(z) + offset_y = acc.Dy(z) + offset_yp = acc.Dyp(z) + x_corr = x - offset_x + y_corr = y - offset_y + r_corr = np.sqrt(x_corr * x_corr + y_corr * y_corr) + + if type == "E": + Ez = acc.Ez(z) + dEzdz = acc.dEzdz(z) + d2Ezdz2 = misc.derivative(acc.dEzdz, z, dx=dz, n=1) + Ez = Ez - d2Ezdz2 * r_corr * r_corr / 4 # row remainder + Ex = -dEzdz * x_corr / 2 + Ez * offset_xp # row remainder + Ey = -dEzdz * y_corr / 2 + Ez * offset_yp # row remainder + return Ex, Ey, Ez + if type == "B": + Bx = acc.Bx(z) + By = acc.By(z) + Bz = acc.Bz(z) + dBzdz = acc.dBzdz(z) + d2Bzdz2 = misc.derivative(acc.dBzdz, z, dx=dz, n=1) + Gz = acc.Gz(z) + Bz = Bz - d2Bzdz2 * r_corr * r_corr / 4 # row remainder + Bx = Bx + Gz * y_corr - dBzdz * x_corr / 2 + Bz * offset_xp # row remainder + By = By + Gz * x_corr - dBzdz * y_corr / 2 + Bz * offset_yp # row remainder + return Bx, By, Bz + + +@jit(nopython=True, parallel=True, fastmath=True, cache=True, nogil=True) +def sum_field_particles( + x: np.array, y: np.array, z: np.array, z_start: float, z_stop: float +) -> (np.array, np.array, np.array): + n = int(len(x)) + Fx = np.zeros(n) + Fy = np.zeros(n) + Fz = np.zeros(n) + for i in prange(n): # pylint: disable=E1133 + if z_start <= z[i] <= z_stop: + for j in prange(n): # pylint: disable=E1133 + if z_start <= z[j] <= z_stop and i != j: + Fx[i] += (x[i] - x[j]) / ((x[j] - x[i]) ** 2 + (y[j] - y[i]) ** 2 + (z[j] - z[i]) ** 2) ** (3 / 2) + Fy[i] += (y[i] - y[j]) / ((x[j] - x[i]) ** 2 + (y[j] - y[i]) ** 2 + (z[j] - z[i]) ** 2) ** (3 / 2) + Fz[i] += (z[i] - z[j]) / ((x[j] - x[i]) ** 2 + (y[j] - y[i]) ** 2 + (z[j] - z[i]) ** 2) ** (3 / 2) + return Fx, Fy, Fz + + +def get_field_beam( + beam: BaseBeam, acc: Accelerator, type: str, x: np.array, y: np.array, z: np.array +) -> (np.array, np.array, np.array): + assert type in ("E", "B"), "Check type field! (E or B)" + Q = beam.total_charge + q = Q / beam.n + if type == "E": + Ex, Ey, Ez = sum_field_particles(x, y, z, acc.z_start, acc.z_stop) + return ( + const.ke / (4 * np.pi) * q * Ex / 1e6, + const.ke / (4 * np.pi) * q * Ey / 1e6, + const.ke / (4 * np.pi) * q * Ez / 1e6, + ) + if type == "B": + Bx, By, Bz = sum_field_particles(x, y, z, acc.z_start, acc.z_stop) + return const.km / (4 * np.pi) * q * Bx, -const.km / (4 * np.pi) * q * By, 0 * const.km / (4 * np.pi) * q * Bz diff --git a/setup.py b/setup.py index c8d555b..359afee 100644 --- a/setup.py +++ b/setup.py @@ -9,14 +9,14 @@ required = fh.read().splitlines() setup( - name=config.PROJECT_NAME, - version=config.PROJECT_VERSION, - author=config.PROJECT_AUTHOR, - author_email=config.PROJECT_AUTHOR_EMAIL, - description=config.PROJECT_DOC, + name=config.REDPIC_NAME, + version=config.REDPIC_VERSION, + author=config.REDPIC_AUTHOR, + author_email=config.REDPIC_AUTHOR_EMAIL, + description=config.REDPIC_DOC, long_description=long_description, long_description_content_type="text/markdown", - url=config.PROJECT_URL, + url=config.REDPIC_URL, packages=find_packages(), install_requires=required, classifiers=[