This should work in all cases

xsuite · Jul 5, 2024 · cb2b8dc · cb2b8dc
1 parent 6bababd
commit cb2b8dc
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/xfields/solvers/fftsolvers.py b/xfields/solvers/fftsolvers.py
@@ -124,11 +124,10 @@ def solve(self, rho):
         _workspace_dev.T[:self.nz, :self.ny, :self.nx] = rho.T
         self.fftplan.transform(_workspace_dev) # rho_rep_hat
 
-        # try:
-        #     _workspace_dev.T[:,:,:] *= (
-        #                 self._gint_rep_transf_dev.T) # phi_rep_hat
-        # except Exception: # pyopencl does not support array broadcasting (used in 2.5D)
-        if True:
+        try:
+            _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