From bd761b332b1bfd572ae765bd3e20c4d9631cfb34 Mon Sep 17 00:00:00 2001 From: tprebiba Date: Fri, 14 Jun 2024 16:49:22 +0200 Subject: [PATCH] added option to have fixed pic grid --- .../config_tools/spacecharge_config_tools.py | 54 ++++++++++++++----- 1 file changed, 40 insertions(+), 14 deletions(-) diff --git a/xfields/config_tools/spacecharge_config_tools.py b/xfields/config_tools/spacecharge_config_tools.py index 0a327d40..02f678e9 100644 --- a/xfields/config_tools/spacecharge_config_tools.py +++ b/xfields/config_tools/spacecharge_config_tools.py @@ -227,15 +227,22 @@ def __init__(self, self._fftplan = None - def get_pic(self, x_lim, y_lim): - - assert x_lim < self.x_lims[-1] - assert x_lim > self.x_lims[0] - assert y_lim < self.y_lims[-1] - assert y_lim > self.y_lims[0] - - ix = np.argmin(np.abs(x_lim - self.x_lims)) - iy = np.argmin(np.abs(y_lim - self.y_lims)) + def get_pic(self, x_lim, y_lim, + fixed_grid_x=False, fixed_grid_y=False): + + if fixed_grid_x==False: + assert x_lim < self.x_lims[-1] + assert x_lim > self.x_lims[0] + ix = np.argmin(np.abs(x_lim - self.x_lims)) + else: + ix = 0 + + if fixed_grid_y==False: + assert y_lim < self.y_lims[-1] + assert y_lim > self.y_lims[0] + iy = np.argmin(np.abs(y_lim - self.y_lims)) + else: + iy = 0 if (ix, iy) not in self._existing_pics.keys(): print(f'Creating PIC ({ix}, {iy})') @@ -266,6 +273,7 @@ def replace_spacecharge_with_PIC( n_sigmas_range_pic_x, n_sigmas_range_pic_y, nx_grid, ny_grid, nz_grid, n_lims_x, n_lims_y, z_range, solver='FFTSolver2p5D', + grid_extend_in_x=None, grid_extend_in_y=None, _context=None, _buffer=None): @@ -292,6 +300,12 @@ def replace_spacecharge_with_PIC( Number different limits in y for which PIC need to be generated. z_range : float Range of the longitudinal grid. + grid_extend_in_x : float + Extend of the horizontal grid in meters in both positive and negative x-values (fixed grid). + If provided, it will be used instead of n_sigmas_range_pic_x and n_lims_x. + grid_extend_in_y : float + Extend of the vertical grid in meters in both positive and negative y-values (fixed grid). + If provided, it will be used instead of n_sigmas_range_pic_y and n_lims_y. _context : xtrack.Context (optional) Context in which the PIC elements are created. _buffer : xtrack.Buffer (optional) @@ -322,10 +336,22 @@ def replace_spacecharge_with_PIC( all_sigma_x.append(ee.sigma_x) all_sigma_y.append(ee.sigma_y) - x_lim_min = np.min(all_sigma_x) * (n_sigmas_range_pic_x - 0.5) - x_lim_max = np.max(all_sigma_x) * (n_sigmas_range_pic_x + 0.5) - y_lim_min = np.min(all_sigma_y) * (n_sigmas_range_pic_y - 0.5) - y_lim_max = np.max(all_sigma_y) * (n_sigmas_range_pic_y + 0.5) + if grid_extend_in_x is not None: + x_lim_min = grid_extend_in_x + x_lim_max = grid_extend_in_x + fixed_grid_x = True + else: + x_lim_min = np.min(all_sigma_x) * (n_sigmas_range_pic_x - 0.5) + x_lim_max = np.max(all_sigma_x) * (n_sigmas_range_pic_x + 0.5) + fixed_grid_x = False + if grid_extend_in_y is not None: + y_lim_min = grid_extend_in_y + y_lim_max = grid_extend_in_y + fixed_grid_y = True + else: + y_lim_min = np.min(all_sigma_y) * (n_sigmas_range_pic_y - 0.5) + y_lim_max = np.max(all_sigma_y) * (n_sigmas_range_pic_y + 0.5) + fixed_grid_y = False pic_collection = PICCollection( _context=_context, @@ -341,7 +367,7 @@ def replace_spacecharge_with_PIC( for nn, ee in zip(name_sc_elems, all_sc_elems): xlim = n_sigmas_range_pic_x*ee.sigma_x ylim = n_sigmas_range_pic_y*ee.sigma_y - base_sc = pic_collection.get_pic(xlim, ylim) + base_sc = pic_collection.get_pic(xlim, ylim, fixed_grid_x, fixed_grid_y) sc = base_sc.copy(_buffer=base_sc._buffer) sc.length = ee.length line.element_dict[nn] = sc