Merge branch 'xsuite:main' into qgaussian_derivative

xsuite · Jul 15, 2024 · f760b68 · f760b68
2 parents 38aae6b + 25a44ec
commit f760b68
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 17 deletions.
diff --git a/tests/test_slicer.py b/tests/test_slicer.py
@@ -298,48 +298,48 @@ def test_slicer_moments_single_bunch(test_context):
                                atol=1e-12)
             xo.assert_allclose(sl.num_particles[0, :], [0, 0, p.weight.sum()], rtol=0,
                                atol=1e-12)
-            xo.assert_allclose(sl.sum(c1_name), [0, 0, (c1 * p.weight).sum()],
+            xo.assert_allclose(sl.sum(c1_name)[0], [0, 0, (c1 * p.weight).sum()],
                                rtol=0, atol=1e-12)
-            xo.assert_allclose(sl.sum(c2_name), [0, 0, (c2 * p.weight).sum()],
+            xo.assert_allclose(sl.sum(c2_name)[0], [0, 0, (c2 * p.weight).sum()],
                                rtol=0, atol=1e-12)
-            xo.assert_allclose(sl.sum('zeta'),
+            xo.assert_allclose(sl.sum('zeta')[0],
                                [0, 0, (p.zeta * p.weight).sum()],
                                rtol=0, atol=1e-12)
-            xo.assert_allclose(sl.sum(c1_name + c1_name),
+            xo.assert_allclose(sl.sum(c1_name + c1_name)[0],
                                [0, 0, (c1 ** 2 * p.weight).sum()],
                                rtol=0, atol=1e-12)
-            xo.assert_allclose(sl.sum(c2_name + c2_name),
+            xo.assert_allclose(sl.sum(c2_name + c2_name)[0],
                                [0, 0, (c2 ** 2 * p.weight).sum()],
                                rtol=0, atol=1e-12)
-            xo.assert_allclose(sl.sum(c1_name + c2_name),
+            xo.assert_allclose(sl.sum(c1_name + c2_name)[0],
                                [0, 0, (c1 * c2 * p.weight).sum()],
                                rtol=0, atol=1e-12)
-            xo.assert_allclose(sl.sum('zetazeta'),
+            xo.assert_allclose(sl.sum('zetazeta')[0],
                                [0, 0, (p.zeta ** 2 * p.weight).sum()],
                                rtol=0, atol=1e-12)
-            xo.assert_allclose(sl.sum(c1_name + 'zeta'),
+            xo.assert_allclose(sl.sum(c1_name + 'zeta')[0],
                                [0, 0, (c1 * p.zeta * p.weight).sum()],
                                rtol=0, atol=1e-12)
-            xo.assert_allclose(sl.mean(c1_name),
+            xo.assert_allclose(sl.mean(c1_name)[0],
                                [0, 0, (c1 * p.weight).sum() / p.weight.sum()],
                                rtol=0, atol=1e-12)
-            xo.assert_allclose(sl.mean(c2_name),
+            xo.assert_allclose(sl.mean(c2_name)[0],
                                [0, 0, (c2 * p.weight).sum() / p.weight.sum()],
                                rtol=0, atol=1e-12)
-            xo.assert_allclose(sl.mean(c1_name + c1_name),
+            xo.assert_allclose(sl.mean(c1_name + c1_name)[0],
                                [0, 0,
                                 (c1 ** 2 * p.weight).sum() / p.weight.sum()],
                                rtol=0, atol=1e-12)
-            xo.assert_allclose(sl.mean(c2_name + c2_name),
+            xo.assert_allclose(sl.mean(c2_name + c2_name)[0],
                                [0, 0,
                                 (c2 ** 2 * p.weight).sum() / p.weight.sum()],
                                rtol=0, atol=1e-12)
-            xo.assert_allclose(sl.mean(c1_name + c2_name),
+            xo.assert_allclose(sl.mean(c1_name + c2_name)[0],
                                [0, 0,
                                 (c1 * c2 * p.weight).sum() / p.weight.sum()],
                                rtol=0, atol=1e-12)
             xo.assert_allclose(
-                sl.mean(c1_name + 'zeta'),
+                sl.mean(c1_name + 'zeta')[0],
                 [0, 0, (c1 * p.zeta * p.weight).sum() / p.weight.sum()],
                 rtol=0, atol=1e-12)
             xo.assert_allclose(

diff --git a/xfields/_version.py b/xfields/_version.py
@@ -1 +1 @@
-__version__ = "0.20.0"
+__version__ = "0.20.2"
diff --git a/xfields/solvers/fftsolvers.py b/xfields/solvers/fftsolvers.py
@@ -128,6 +128,10 @@ def solve(self, rho):
             _workspace_dev.T[:,:,:] *= (
                         self._gint_rep_transf_dev.T) # phi_rep_hat
         except Exception: # pyopencl does not support array broadcasting (used in 2.5D)
+            # Check F-contiguity
+            assert _workspace_dev.strides[0] == 16
+            assert self._gint_rep_transf_dev.strides[0] == 16
+
             for ii in range(self.nz):
                 _workspace_dev.T[ii,:,:] *= (
                         self._gint_rep_transf_dev.T[0, :, :]) # phi_rep_hat
@@ -192,7 +196,9 @@ def __init__(self, dx, dy, dz, nx, ny, nz, context=None, fftplan=None):
             del(temp_dev)
 
         # Transform the green function
-        gint_rep_transf = np.fft.fftn(gint_rep, axes=(0,1))
+        gint_rep_transf = np.zeros((2*nx, 2*ny),
+                           dtype=np.complex128, order='F')
+        gint_rep_transf[:, :] = np.fft.fftn(gint_rep, axes=(0,1))
 
         # Transfer to GPU (if needed)
         gint_rep_transf_dev = context.nparray_to_context_array(
@@ -247,7 +253,9 @@ def __init__(self, dx, dy, dz, nx, ny, nz, context=None, fftplan=None):
             del(temp_dev)
 
         # Transform the green function
-        gint_rep_transf = np.fft.fftn(gint_rep, axes=(0,1))
+        gint_rep_transf = np.zeros((2*nx, 2*ny),
+                           dtype=np.complex128, order='F')
+        gint_rep_transf[:, :] = np.fft.fftn(gint_rep, axes=(0,1))
 
         # Transfer to GPU (if needed)
         gint_rep_transf_dev = context.nparray_to_context_array(