Add functionality to change spacecharge PIC mesh at runtime

Author: MarkusSQ

xfields/beam_elements/spacecharge.py

@@ -6,7 +6,7 @@
 import numpy as np
 
 from xfields import BiGaussianFieldMap, mean_and_std
-from xfields import TriLinearInterpolatedFieldMap
+from ..fieldmaps.interpolated import TriLinearInterpolatedFieldMap
 from ..longitudinal_profiles import LongitudinalProfileQGaussian
 from ..fieldmaps import BiGaussianFieldMap
 from ..general import _pkg_root
@@ -182,6 +182,48 @@ def track(self, particles):
         # call C tracking kernel
         super().track(particles)
 
+    def set_xy_mesh(self, x_range, y_range, *, zero_fields=True):
+        """
+        Retile the underlying fieldmap in x/y while
+        keeping the grid counts unchanged.
+
+        Args:
+            x_range (tuple[float, float]): New ``(xmin, xmax)`` in
+                meters.
+            y_range (tuple[float, float]): New ``(ymin, ymax)`` in
+                meters.
+            zero_fields (bool): If ``True``, zero stored
+                ``rho``, ``phi`` and their derivatives after retiling.
+        """
+        (xmin, xmax) = map(float, x_range)
+        (ymin, ymax) = map(float, y_range)
+        self.fieldmap.retile_xy(
+            xmin, xmax, ymin, ymax, zero_fields=zero_fields
+        )
+
+    def set_xyz_mesh(self, *, x_range, y_range, z_range,
+                     zero_fields=True):
+        """
+        Retile the underlying fieldmap in x/y/z while
+        keeping the grid counts unchanged.
+
+        Args:
+            x_range (tuple[float, float]): New ``(xmin, xmax)`` in
+                meters.
+            y_range (tuple[float, float]): New ``(ymin, ymax)`` in
+                meters.
+            z_range (tuple[float, float]): New ``(zmin, zmax)`` in
+                meters.
+            zero_fields (bool): If ``True``, zero stored
+                ``rho``, ``phi``, and their derivatives after retiling.
+        """
+        (xmin, xmax) = map(float, x_range)
+        (ymin, ymax) = map(float, y_range)
+        (zmin, zmax) = map(float, z_range)
+        self.fieldmap.retile_xyz(
+            xmin, xmax, ymin, ymax, zmin, zmax, zero_fields=zero_fields
+        )
+
 class SpaceChargeBiGaussian(xt.BeamElement):
 
     _xofields = {

xfields/fieldmaps/interpolated.py

@@ -543,6 +543,116 @@ def generate_solver(self, solver, fftplan):
 
         return solver
 
+    def retile_xy(self, xmin, xmax, ymin, ymax, *, zero_fields=True):
+        """
+        Retile the x/y mesh in-place while keeping the
+        number of cells unchanged.
+
+        Args:
+            xmin (float): New minimum x coordinate in meters.
+            xmax (float): New maximum x coordinate in meters.
+            ymin (float): New minimum y coordinate in meters.
+            ymax (float): New maximum y coordinate in meters.
+            zero_fields (bool): If ``True``, zero stored
+                ``rho``, ``phi`` and their derivatives after retiling.
+        """
+        nx, ny = len(self._x_grid), len(self._y_grid)
+
+        self._x_grid = np.linspace(float(xmin), float(xmax), nx,
+                                   dtype=np.float64)
+        self._y_grid = np.linspace(float(ymin), float(ymax), ny,
+                                   dtype=np.float64)
+
+        if zero_fields:
+            self._rho[...] = 0.0
+            self._phi[...] = 0.0
+            self._dphi_dx[...] = 0.0
+            self._dphi_dy[...] = 0.0
+            self._dphi_dz[...] = 0.0
+
+        self._dx = (xmax - xmin) / (nx - 1)
+        self._dy = (ymax - ymin) / (ny - 1)
+        self._x_min = float(xmin)
+        self._y_min = float(ymin)
+
+        # Update derived volume quantities (dz unchanged)
+        self._cell_volume = self._dx * self._dy * self._dz
+        self._inv_cell_volume = 1.0 / self._cell_volume
+
+        # Refresh the solver geometry
+        scale_dx, scale_dy, scale_dz = self.scale_coordinates_in_solver
+        self.solver.refresh_geometry(
+            self._x_grid * scale_dx,
+            self._y_grid * scale_dy,
+            self._z_grid * scale_dz)
+
+    def retile_xyz(
+            self,
+            xmin,
+            xmax,
+            ymin,
+            ymax,
+            zmin,
+            zmax,
+            *,
+            zero_fields=True
+    ):
+        """
+        Retile the full x/y/z mesh in-place while
+        keeping the number of cells unchanged.
+
+        Args:
+            xmin (float): New minimum x coordinate in meters.
+            xmax (float): New maximum x coordinate in meters.
+            ymin (float): New minimum y coordinate in meters.
+            ymax (float): New maximum y coordinate in meters.
+            zmin (float): New minimum z coordinate in meters.
+            zmax (float): New maximum z coordinate in meters.
+            zero_fields (bool): If ``True``, zero stored
+                ``rho``, ``phi`` and their derivatives after retiling.
+        """
+        nx, ny, nz = self._x_grid.size, self._y_grid.size, self._z_grid.size
+
+        self._x_grid = np.linspace(
+            float(xmin),
+            float(xmax),
+            nx,
+            dtype=np.float64)
+        self._y_grid = np.linspace(
+            float(ymin),
+            float(ymax),
+            ny,
+            dtype=np.float64)
+        self._z_grid = np.linspace(
+            float(zmin),
+            float(zmax),
+            nz,
+            dtype=np.float64)
+
+        if zero_fields:
+            self._rho[...] = 0.0
+            self._phi[...] = 0.0
+            self._dphi_dx[...] = 0.0
+            self._dphi_dy[...] = 0.0
+            self._dphi_dz[...] = 0.0
+
+        self._dx = (xmax - xmin) / (nx - 1)
+        self._dy = (ymax - ymin) / (ny - 1)
+        self._dz = (zmax - zmin) / (nz - 1)
+        self._x_min = float(xmin)
+        self._y_min = float(ymin)
+        self._z_min = float(zmin)
+
+        self._cell_volume = self._dx * self._dy * self._dz
+        self._inv_cell_volume = 1.0 / self._cell_volume
+
+        # Refresh the solver geometry
+        scale_dx, scale_dy, scale_dz = self.scale_coordinates_in_solver
+        self.solver.refresh_geometry(
+            self._x_grid * scale_dx,
+            self._y_grid * scale_dy,
+            self._z_grid * scale_dz)
+
     @property
     def x_grid(self):
         """
@@ -564,6 +674,27 @@ def z_grid(self):
         """
         return self._z_grid
 
+    @property
+    def x_range(self):
+        """
+        Horizontal range.
+        """
+        return (float(self._x_grid[0]), float(self._x_grid[-1]))
+
+    @property
+    def y_range(self):
+        """
+        Vertical range.
+        """
+        return (float(self._y_grid[0]), float(self._y_grid[-1]))
+
+    @property
+    def z_range(self):
+        """
+        Longitudinal range.
+        """
+        return (float(self._z_grid[0]), float(self._z_grid[-1]))
+
     @property
     def nx(self):
         """
@@ -665,5 +796,3 @@ def _configure_grid(vname, v_grid, dv, v_range, nv):
             v_grid = np.linspace(v_range[0], v_range[1], nv)
 
     return v_grid
-
-

xfields/solvers/fftsolvers.py

@@ -101,6 +101,83 @@ def __init__(self, dx, dy, dz, nx, ny, nz, context=None, fftplan=None):
         self._gint_rep_transf_dev = gint_rep_dev
         self.fftplan = fftplan
 
+    def refresh_geometry(self, x_grid, y_grid, z_grid):
+        """
+        Rebuild the 3D Green's function used by the solver.
+
+        The number of cells in each direction must match the original
+        solver.
+
+        Args:
+            x_grid (array-like): Monotonic array of x cell centers in
+                meters.
+            y_grid (array-like): Monotonic array of y cell centers in
+                meters.
+            z_grid (array-like): Monotonic array of z cell centers in
+                meters.
+        """
+        nx, ny, nz = len(x_grid), len(y_grid), len(z_grid)
+        assert (nx, ny, nz) == (self.nx, self.ny, self.nz), (
+            "refresh_geometry requires unchanged (nx, ny, nz)"
+        )
+
+        dx = float(x_grid[1] - x_grid[0])
+        dy = float(y_grid[1] - y_grid[0])
+        dz = float(z_grid[1] - z_grid[0])
+
+        xg_F = np.arange(0, nx + 2) * dx - dx / 2.0
+        yg_F = np.arange(0, ny + 2) * dy - dy / 2.0
+        zg_F = np.arange(0, nz + 2) * dz - dz / 2.0
+        XX_F, YY_F, ZZ_F = np.meshgrid(xg_F, yg_F, zg_F, indexing='ij')
+
+        F_temp = primitive_func_3d(XX_F, YY_F, ZZ_F)
+
+        gint_rep = np.zeros(
+            (2 * nx, 2 * ny, 2 * nz),
+            dtype=np.complex128,
+            order='F')
+
+        gint_rep[:nx + 1, :ny + 1, :nz + 1] = (
+            F_temp[1:, 1:, 1:]
+            - F_temp[:-1, 1:, 1:]
+            - F_temp[1:, :-1, 1:]
+            + F_temp[:-1, :-1, 1:]
+            - F_temp[1:, 1:, :-1]
+            + F_temp[:-1, 1:, :-1]
+            + F_temp[1:, :-1, :-1]
+            - F_temp[:-1, :-1, :-1]
+        )
+
+        gint_rep[nx + 1:, :ny + 1, :nz + 1] = gint_rep[nx - 1:0:-1,
+                                                       :ny + 1,
+                                                       :nz + 1]
+        gint_rep[:nx + 1, ny + 1:, :nz + 1] = gint_rep[:nx + 1,
+                                                       ny - 1:0:-1,
+                                                       :nz + 1]
+        gint_rep[nx + 1:, ny + 1:, :nz + 1] = gint_rep[nx - 1:0:-1,
+                                                       ny - 1:0:-1,
+                                                       :nz + 1]
+        gint_rep[:nx + 1, :ny + 1, nz + 1:] = gint_rep[:nx + 1,
+                                                       :ny + 1,
+                                                       nz - 1:0:-1]
+        gint_rep[nx + 1:, :ny + 1, nz + 1:] = gint_rep[nx - 1:0:-1,
+                                                       :ny + 1,
+                                                       nz - 1:0:-1]
+        gint_rep[:nx + 1, ny + 1:, nz + 1:] = gint_rep[:nx + 1,
+                                                       ny - 1:0:-1,
+                                                       nz - 1:0:-1]
+        gint_rep[nx + 1:, ny + 1:, nz + 1:] = gint_rep[nx - 1:0:-1,
+                                                       ny - 1:0:-1,
+                                                       nz - 1:0:-1]
+
+        gint_rep_dev = self.context.nparray_to_context_array(gint_rep)
+        self.fftplan.transform(gint_rep_dev)
+
+        self.dx = dx
+        self.dy = dy
+        self.dz = dz
+        self._gint_rep_transf_dev = gint_rep_dev
+
     #@profile
     def solve(self, rho):
 
@@ -213,6 +290,57 @@ def __init__(self, dx, dy, dz, nx, ny, nz, context=None, fftplan=None):
         self._gint_rep_transf_dev = gint_rep_transf_dev
         self.fftplan = fftplan
 
+    def refresh_geometry(self, x_grid, y_grid, z_grid):
+        """
+        Rebuild the 2D Green's function used by the solver.
+
+        The number of cells in each direction must match the original
+        solver.
+
+        Args:
+            x_grid (array-like): Monotonic array of x cell centers in
+                meters.
+            y_grid (array-like): Monotonic array of y cell centers in
+                meters.
+            z_grid (array-like): Monotonic array of z cell centers in
+                meters.
+        """
+        nx, ny, nz = len(x_grid), len(y_grid), len(z_grid)
+        assert (nx, ny, nz) == (self.nx, self.ny, self.nz), (
+            "refresh_geometry requires unchanged (nx, ny, nz)"
+        )
+
+        dx = float(x_grid[1] - x_grid[0])
+        dy = float(y_grid[1] - y_grid[0])
+        dz = float(z_grid[1] - z_grid[0])
+
+        xg_F = np.arange(0, nx + 2) * dx - dx / 2.0
+        yg_F = np.arange(0, ny + 2) * dy - dy / 2.0
+        XX_F, YY_F = np.meshgrid(xg_F, yg_F, indexing='ij')
+
+        F_temp = primitive_func_2p5d(XX_F, YY_F)
+
+        gint_rep = np.zeros((2 * nx, 2 * ny), dtype=np.complex128, order='F')
+        gint_rep[:nx + 1, :ny + 1] = (
+            F_temp[1:, 1:]
+            - F_temp[:-1, 1:]
+            - F_temp[1:, :-1]
+            + F_temp[:-1, :-1]
+        )
+        gint_rep[nx + 1:, :ny + 1] = gint_rep[nx - 1:0:-1, :ny + 1]
+        gint_rep[:nx + 1, ny + 1:] = gint_rep[:nx + 1, ny - 1:0:-1]
+        gint_rep[nx + 1:, ny + 1:] = gint_rep[nx - 1:0:-1, ny - 1:0:-1]
+
+        gint_rep_transf = np.fft.fftn(gint_rep, axes=(0, 1))
+        gint_rep_transf_dev = self.context.nparray_to_context_array(
+            np.atleast_3d(gint_rep_transf)
+        )
+
+        self.dx = dx
+        self.dy = dy
+        self.dz = dz
+        self._gint_rep_transf_dev = gint_rep_transf_dev
+
 class FFTSolver2p5DAveraged(Solver):
 
     def __init__(self, dx, dy, dz, nx, ny, nz, context=None, fftplan=None):