Feature/integrated wakefield

tests/test_integration_method.py

@@ -0,0 +1,46 @@
+# copyright ############################### #
+# This file is part of the Xwakes Package.  #
+# Copyright (c) CERN, 2025.                 #
+# ######################################### #
+
+import pathlib
+
+import numpy as np
+from scipy.constants import c as clight
+
+import xwakes as xw
+import xtrack as xt
+import xobjects as xo
+import xpart as xp
+
+test_data_folder = pathlib.Path(__file__).parent.joinpath(
+    'test_data').absolute()
+
+def test_xwakes_kick_vs_pyheadtail_table():
+
+    p = xt.Particles(mass0 = xp.PROTON_MASS_EV, gamma0 = 3.14, zeta=np.linspace(-1, 1, 100000))
+    p.x[p.zeta > 0] += 1e-3
+    p.y[p.zeta > 0] += 1e-3
+    p_table = p.copy()
+
+    # Build equivalent WakeFromTable
+    table = xw.read_headtail_file(
+        test_data_folder / 'PS_wall_impedance_Ekin_2.0GeV.wake',
+        wake_file_columns=['time', 'dipolar_x', 'dipolar_y','quadrupolar_x', 'quadrupolar_y'])
+
+    wake_from_table = xw.WakeFromTable(table, columns=['time', 'dipolar_x', 'dipolar_y','quadrupolar_x', 'quadrupolar_y'], 
+                                       method="Integrated")
+    wake_from_table.configure_for_tracking(zeta_range=(-1, 1), num_slices=100)
+
+    # Zotter convention
+    assert table['dipolar_x'].values[1] > 0
+    assert table['dipolar_y'].values[1] > 0
+    assert table['quadrupolar_x'].values[1] < 0
+    assert table['quadrupolar_y'].values[1] < 0
+
+
+    wake_from_table.track(p_table)
+    p_ref = np.load(test_data_folder / 'particles_pyht.npy')
+
+    xo.assert_allclose(p_table.px, p_ref[3], atol=1e-30, rtol=2e-3)
+    xo.assert_allclose(p_table.py, p_ref[4], atol=1e-30, rtol=2e-3)

xwakes/wakefield_from_table.py

@@ -11,7 +11,7 @@
 
 class WakeFromTable(BaseWake):
 
-    def __init__(self, table, columns=None):
+    def __init__(self, table, columns=None, method = "Interpolated"):
 
         if isinstance(table, dict):
             table = pd.DataFrame(table)
@@ -33,7 +33,7 @@ def __init__(self, table, columns=None):
             component = ComponentFromArrays(
                 interpolation_times=table['time'].values,
                 wake_samples=table[cc].values,
-                kind=cc)
+                kind=cc, method = method)
             components.append(component)
 
         self.components = components

xwakes/wit/component.py

@@ -15,6 +15,7 @@
 from numpy.typing import ArrayLike
 from scipy import special as sp
 from scipy.interpolate import interp1d
+from scipy.integrate import quad
 
 if hasattr(np, 'trapezoid'):
     trapz = np.trapezoid # numpy 2.0
@@ -1394,7 +1395,8 @@ def __init__(self,
                  test_exponents: Tuple[int, int] = None,
                  name: str = "Interpolated Component",
                  f_rois: Optional[List[Tuple[float, float]]] = None,
-                 t_rois: Optional[List[Tuple[float, float]]] = None):
+                 t_rois: Optional[List[Tuple[float, float]]] = None,
+                 method: str = "Interpolated"):
         """
         The init of the ComponentFromArrays class
 
@@ -1429,6 +1431,8 @@ def __init__(self,
         self.interpolation_times = interpolation_times
         self.wake_samples = wake_samples
 
+        self.method = method
+
         if self.interpolation_frequencies is not None:
             impedance = interp1d(self.interpolation_frequencies,
                                  self.impedance_samples)
@@ -1458,7 +1462,20 @@ def function_vs_t(self, t, beta0, dt):
         if isscalar:
             t = np.array([t])
 
-        out = self.wake(t)
+        if self.method == "Interpolated":
+            out = self.wake(t)
+        elif self.method == "Integrated":
+            integrated_bins = []
+            if self.interpolation_times is not None:
+                for i, time_bins in enumerate(t[0]):
+                    integrated_bins.append(quad(self.wake, t[0][i], t[0][i] + dt, 
+                                                limit=200)[0]/dt)
+
+            integrated_function = interp1d(t[0], integrated_bins)
+
+            out = integrated_function(t)
+        else:
+            raise NotImplementedError("Only the Interpolated and Integrated methods are supported")
 
         # At the edges of the provided wake table we take half the provided
         # value. This is equivalent to assume the next sample (not provided) is