vectorization follow-ups

.github/workflows/tests.yml

@@ -38,6 +38,7 @@ jobs:
       !contains(needs.check_skip_flags.outputs.head-commit-message, '[skip tests]') }}
     env:
       EXP_TAG: "v7.8.5"
+      MACOSX_DEPLOYMENT_TARGET: "15.0"
     strategy:
       fail-fast: true
       matrix:
@@ -156,15 +157,14 @@ jobs:
           echo CC=mpicc >> $GITHUB_ENV
           echo CXX=mpicxx >> $GITHUB_ENV
           echo FC=mpifort >> $GITHUB_ENV
+          echo "$(brew --prefix llvm)/bin" >> $GITHUB_PATH
           echo OMPI_CC="$(brew --prefix llvm)/bin/clang" >> $GITHUB_ENV
           echo OMPI_CXX="$(brew --prefix llvm)/bin/clang++" >> $GITHUB_ENV
           echo OMPI_FC="$(brew --prefix flang)/bin/flang" >> $GITHUB_ENV
 
-          # Fix header search paths for LLVM 21.1.1
-          LLVM_PREFIX=$(brew --prefix llvm)
-          XCODE_SDK_PATH=$(xcrun --show-sdk-path)
-          echo "CXXFLAGS=-isystem ${LLVM_PREFIX}/include/c++/v1 -isystem ${XCODE_SDK_PATH}/usr/include" >> $GITHUB_ENV
-          echo "LDFLAGS=-L$(brew --prefix libomp)/lib -L${LLVM_PREFIX}/lib/c++ -L${LLVM_PREFIX}/lib" >> $GITHUB_ENV
+          # c++ library path is workaround for clang bug https://github.com/llvm/llvm-project/issues/155531
+          sdkroot=$(xcrun --show-sdk-path)
+          echo "CXXFLAGS=-nostdinc++ -isystem $sdkroot/usr/include/c++/v1 -isystem $sdkroot/usr/include" >> $GITHUB_ENV
 
       - name: Build EXP
         if: matrix.gala-exp == '1'

gala/dynamics/mockstream/mockstream.pyx

@@ -129,6 +129,7 @@ cpdef mockstream_dop853(
             atol, rtol, nmax, dt_max,
             nstiff=-1,  # disable stiffness check
             err_if_fail=err_if_fail, log_output=log_output, save_all=1,
+            transposed=0
         )
 
         n = 0

gala/dynamics/nbody/nbody.pyx

@@ -105,7 +105,8 @@ cpdef direct_nbody_dop853(
             atol, rtol, nmax, dt_max,
             nstiff=-1,  # disable stiffness check - TODO: note somewhere
             err_if_fail=err_if_fail, log_output=log_output,
-            save_all=save_all
+            save_all=save_all,
+            transposed=0
         )
         if save_all:
             return np.array(w)

gala/integrate/cyintegrators/dop853.pxd

@@ -28,6 +28,8 @@ cdef extern from "dopri/dop853.h":
 
     void Fwrapper (unsigned ndim, double t, double *w, double *f,
                    CPotential *p, CFrameType *fr, unsigned norbits, unsigned nbody) except +
+    void Fwrapper_T (unsigned ndim, double t, double *w, double *f,
+                   CPotential *p, CFrameType *fr, unsigned norbits, unsigned nbody) except +
     void Fwrapper_direct_nbody(unsigned ndim, double t, double *w, double *f,
                                CPotential *p, CFrameType *fr,
                                unsigned norbits, unsigned nbody, void *args) except + nogil
@@ -67,6 +69,7 @@ cdef dop853_helper(
     int nstiff,
     unsigned err_if_fail,
     unsigned log_output,
+    int transposed,
     unsigned save_all=?
 )
 

gala/integrate/cyintegrators/dop853.pyx

@@ -20,7 +20,7 @@ np.import_array()
 from cpython.exc cimport PyErr_CheckSignals
 from ...potential.potential.cpotential cimport CPotentialWrapper, CPotential
 from ...potential.frame.cframe cimport CFrameWrapper, CFrameType
-from .dop853 cimport dop853, Fwrapper, FcnEqDiff, six_norm, SolTrait, dop853_dense_state_alloc, dop853_dense_state_free, Dop853DenseState
+from .dop853 cimport dop853, Fwrapper_T, FcnEqDiff, six_norm, SolTrait, dop853_dense_state_alloc, dop853_dense_state_free, Dop853DenseState
 
 
 # LEGACY FUNCTION: don't use this (used by lyapunov functionality)
@@ -104,10 +104,12 @@ cdef dop853_helper(
     int nstiff,
     unsigned err_if_fail,
     unsigned log_output,
-    unsigned save_all=1
+    int transposed,  # does F expect transposed input?
+    unsigned save_all=1,
 ):
     cdef:
-        double[:, ::1] w = w0.copy()
+        double[:, ::1] w
+
         int res
         FILE* cfile
 
@@ -120,6 +122,11 @@ cdef dop853_helper(
         Dop853DenseState* state
         double* output_ptr
 
+    if transposed:
+        w = w0.T.copy()
+    else:
+        w = w0.copy()
+
     if save_all:
         output_ptr = &output_w[0, 0]
         state = dense_state.state
@@ -170,16 +177,23 @@ cdef dop853_helper(
     if res < 0 and err_if_fail == 1:
         raise RuntimeError(f"Integration failed with code {res}")
 
-    if save_all:
-        return np.asarray(output_w).reshape((ntimes, norbits, ndim))
+    if transposed:
+        if save_all:
+            return np.asarray(output_w).reshape((ntimes, ndim, norbits)).transpose((0,2,1))
+        else:
+            return np.array(w.T, copy=False).reshape((norbits, ndim))
     else:
-        return np.asarray(w).reshape((norbits, ndim))
+        if save_all:
+            return np.asarray(output_w).reshape((ntimes, norbits, ndim))
+        else:
+            return np.array(w, copy=False).reshape((norbits, ndim))
 
 
 cpdef dop853_integrate_hamiltonian(
     hamiltonian, double[:, ::1] w0, double[::1] t,
     double atol=1E-10, double rtol=1E-10, int nmax=0, double dt_max = 0.,
-    int nstiff=0, int save_all=1, int err_if_fail=1, int log_output=0
+    int nstiff=0, int save_all=1, int err_if_fail=1, int log_output=0,
+    int nbatch=100,
 ):
     """
     CAUTION: Interpretation of axes is different here! We need the
@@ -203,16 +217,32 @@ cpdef dop853_integrate_hamiltonian(
         CPotential* cp = (<CPotentialWrapper>(hamiltonian.potential.c_instance)).cpotential
         CFrameType cf = (<CFrameWrapper>(hamiltonian.frame.c_instance)).cframe
 
-    # 0 below is for nbody - we ignore that in this test particle integration
-    w = dop853_helper(
-        cp, &cf, <FcnEqDiff> Fwrapper,
-        w0, t,
-        ndim, norbits, 0, args, ntimes,
-        atol, rtol, nmax, dt_max,
-        nstiff=nstiff,
-        save_all=save_all, err_if_fail=err_if_fail, log_output=log_output
-    )
     if save_all:
-        return np.asarray(t), np.asarray(w)
+        wres = np.empty((ntimes, norbits, ndim))
+    else:
+        wres = np.empty((norbits, ndim))
+
+    for i in range(0, norbits, nbatch):
+        # do the integration in batches for performance
+        # FUTURE: this batching could probably be done in C directly
+        j = min(i + nbatch, norbits)
+        wbatch = w0[i:j, :]
+        # 0 below is for nbody - we ignore that in this test particle integration
+        wout = dop853_helper(
+            cp, &cf, <FcnEqDiff> Fwrapper_T,
+            wbatch, t,
+            ndim, j - i, 0, NULL, ntimes,
+            atol, rtol, nmax, dt_max,
+            nstiff=nstiff,
+            save_all=save_all, err_if_fail=err_if_fail, log_output=log_output,
+            transposed=1
+        )
+        if save_all:
+            wres[:, i:j, :] = wout
+        else:
+            wres[i:j, :] = wout
+
+    if save_all:
+        return np.asarray(t), np.asarray(wres)
     else:
-        return np.asarray(t[-1:]), np.asarray(w)
+        return np.asarray(t[-1:]), np.asarray(wres)

gala/integrate/cyintegrators/dopri/dop853.cpp

@@ -976,6 +976,17 @@ void Fwrapper(unsigned full_ndim, double t, double *w, double *f, CPotential *p,
   }
 }
 
+void Fwrapper_T(unsigned full_ndim, double t, double *w, double *f, CPotential *p,
+              CFrameType *fr, unsigned norbits, unsigned na, void *args) {
+  /* na can be ignored here - used in nbody wrapper below */
+
+  int i;
+  unsigned ndim = full_ndim / norbits; // phase-space dimensionality
+
+  // call gradient function
+  hamiltonian_gradient_T(p, fr, norbits, t, w, f);
+}
+
 void Fwrapper_direct_nbody(unsigned full_ndim, double t, double *w, double *f,
                            CPotential *p, CFrameType *fr, unsigned norbits,
                            unsigned nbody, void *args) {

gala/integrate/cyintegrators/dopri/dop853.h

@@ -258,6 +258,9 @@ extern long nrejctRead (void);
 extern void Fwrapper (unsigned ndim, double t, double *w, double *f,
                       CPotential *p, CFrameType *fr,
                       unsigned norbits, unsigned nbody, void *args);
+extern void Fwrapper_T (unsigned ndim, double t, double *w, double *f,
+                      CPotential *p, CFrameType *fr,
+                      unsigned norbits, unsigned nbody, void *args);
 extern void Fwrapper_direct_nbody(unsigned ndim, double t, double *w, double *f,
                                   CPotential *p, CFrameType *fr,
                                   unsigned norbits, unsigned nbody,

gala/integrate/cyintegrators/ruth4.pxd

@@ -3,6 +3,6 @@
 
 from ...potential.potential.cpotential cimport CPotential
 
-cdef void c_ruth4_step(CPotential *p, int ndim, double t, double dt,
+cdef void c_ruth4_step(CPotential *p, size_t n, int ndim, double t, double dt,
                        double *cs, double *ds,
                        double *w, double *grad) nogil

gala/integrate/cyintegrators/ruth4.pyx

@@ -21,19 +21,18 @@ from ...potential import NullPotential
 from libc.stdlib cimport malloc, free
 
 
-cdef void c_ruth4_step(CPotential *p, int half_ndim, double t, double dt,
+cdef void c_ruth4_step(CPotential *p, size_t n, int half_ndim, double t, double dt,
                        double *cs, double *ds,
                        double *w, double *grad) nogil:
     cdef:
-        int j, k
+        int j, k, i
 
     for j in range(4):
+        c_gradient(p, n, t, w, grad)
         for k in range(half_ndim):
-             grad[k] = 0.
-        c_gradient(p, 1, t, w, grad)
-        for k in range(half_ndim):
-            w[half_ndim + k] = w[half_ndim + k] - ds[j] * grad[k] * dt
-            w[k] = w[k] + cs[j] * w[half_ndim + k] * dt
+            for i in range(n):
+                w[(half_ndim + k) * n + i] = w[(half_ndim + k) * n + i] - ds[j] * grad[k * n + i] * dt
+                w[k * n + i] = w[k * n + i] + cs[j] * w[(half_ndim + k) * n + i] * dt
 
 cpdef ruth4_integrate_hamiltonian(hamiltonian,
                                   double[:, ::1] w0,
@@ -80,7 +79,7 @@ cpdef ruth4_integrate_hamiltonian(hamiltonian,
         ], dtype='f8')
 
         # temporary array containers
-        double[::1] grad = np.zeros(half_ndim)
+        double[:, ::1] grad = np.zeros((half_ndim, n))
 
         # return arrays
         double[:, :, ::1] all_w
@@ -90,29 +89,29 @@ cpdef ruth4_integrate_hamiltonian(hamiltonian,
         CPotential* cp = (<CPotentialWrapper>(hamiltonian.potential.c_instance)).cpotential
 
     if save_all:
-        all_w = np.zeros((ntimes, n, ndim))
+        all_w = np.zeros((ntimes, ndim, n))
 
         # save initial conditions
-        all_w[0, :, :] = w0.copy()
+        all_w[0, :, :] = w0.T.copy()
 
-    tmp_w = w0.copy()
+    tmp_w = w0.T.copy()
 
     with nogil:
-
         for j in range(1, ntimes, 1):
-            for i in range(n):
-                c_ruth4_step(cp, half_ndim, t[j], dt,
-                             &cs[0], &ds[0],
-                             &tmp_w[i, 0], &grad[0])
+            grad[:] = 0.
+            c_ruth4_step(cp, n, half_ndim, t[j], dt,
+                        &cs[0], &ds[0],
+                        &tmp_w[0, 0], &grad[0, 0])
 
-                if save_all:
-                    for k in range(ndim):
-                        all_w[j, i, k] = tmp_w[i, k]
+            if save_all:
+                for k in range(ndim):
+                    for i in range(n):
+                        all_w[j, k, i] = tmp_w[k, i]
 
     if save_all:
-        return np.asarray(t), np.asarray(all_w)
+        return np.asarray(t), np.asarray(all_w).transpose(0,2,1)
     else:
-        return np.asarray(t[-1:]), np.asarray(tmp_w)
+        return np.asarray(t[-1:]), np.array(tmp_w.T, copy=False)
 
 
 # -------------------------------------------------------------------------------------

gala/potential/hamiltonian/chamiltonian.pyx

@@ -342,6 +342,7 @@ class Hamiltonian(CommonBase):
                     save_all=save_all,
                     err_if_fail=int(Integrator_kwargs.get('err_if_fail', 1)),
                     log_output=int(Integrator_kwargs.get('log_output', 0)),
+                    nbatch=Integrator_kwargs.get('nbatch', 100),
                 )
             else:
                 raise ValueError(f"Cython integration not supported for '{Integrator!r}'")

gala/potential/hamiltonian/src/chamiltonian.cpp

@@ -1,5 +1,6 @@
 #include <stddef.h>
 #include <math.h>
+#include "chamiltonian.h"
 #include "potential/src/cpotential.h"
 #include "frame/src/cframe.h"
 
@@ -34,6 +35,27 @@ void hamiltonian_gradient(CPotential *p, CFrameType *fr, double t, double *qp, d
     }
 }
 
+void hamiltonian_gradient_T(CPotential *p, CFrameType *fr, size_t n, double t, double *qp_T, double *dH_T) {
+    // qp_T: shape (n_dim, n)
+    // dH_T: shape (n_dim, n)
+
+    int ndim = p->n_dim;
+
+    // Initialize dH_T to zeros
+    for (int i = 0; i < 2 * ndim * n; i++) {
+        dH_T[i] = 0.0;
+    }
+
+    // Call gradient functions directly with transposed data
+    c_gradient(p, n, t, qp_T, dH_T + ndim * n);  // Write to momentum part
+    (fr->gradient)(t, (fr->parameters), qp_T, ndim, n, dH_T, NULL);  // Write to position part
+
+    // Negate the momentum derivatives
+    for (int i = 0; i < n * ndim; i++) {
+        dH_T[ndim * n + i] *= -1;  // pdot = -dH/dq
+    }
+}
+
 void hamiltonian_hessian(CPotential *p, CFrameType *fr, double t, double *qp, double *d2H) {
     int i;
 

gala/potential/hamiltonian/src/chamiltonian.h

@@ -3,4 +3,5 @@
 
 extern double hamiltonian_value(CPotential *p, CFrameType *fr, double t, double *q);
 extern void hamiltonian_gradient(CPotential *p, CFrameType *fr, double t, double *q, double *grad);
+extern void hamiltonian_gradient_T(CPotential *p, CFrameType *fr, size_t n, double t, double *q, double *grad);
 extern void hamiltonian_hessian(CPotential *p, CFrameType *fr, double t, double *q, double *hess);

gala/potential/potential/builtin/exp_fields.cc

@@ -191,7 +191,7 @@ double exp_value(double t, double *pars, double *q, int n_dim, void* state) {
   return field[5];
 }
 
-void exp_gradient_single(double t, double *__restrict__ pars, double6ptr q, int n_dim, double6ptr grad, void *__restrict__ state){
+void exp_gradient(double t, double *__restrict__ pars, double *__restrict__ q_in, int n_dim, size_t N, double *__restrict__ grad_in, void *__restrict__ state){
   gala_exp::State *exp_state = static_cast<gala_exp::State *>(state);
 
   if (!exp_state->is_static) {
@@ -202,11 +202,16 @@ void exp_gradient_single(double t, double *__restrict__ pars, double6ptr q, int
 
   // TODO: ask Martin/Mike for a way to compute only the force/acceleration - we're wasting
   // computation time here by computing all fields
-  auto field = exp_state->basis->getFields(q[0], q[1], q[2]);
+  double6ptr q = double6ptr{q_in, N};
+  double6ptr grad = double6ptr{grad_in, N};
+
+  for(size_t i = 0; i < N; i++) {
+    auto field = exp_state->basis->getFields(q.x[i], q.y[i], q.z[i]);
 
-  grad[0] += -field[6];
-  grad[1] += -field[7];
-  grad[2] += -field[8];
+    grad.x[i] += -field[6];
+    grad.y[i] += -field[7];
+    grad.z[i] += -field[8];
+  }
 }
 
 double exp_density(double t, double *pars, double *q, int n_dim, void* state) {
@@ -242,6 +247,4 @@ double exp_density(double t, double *pars, double *q, int n_dim, void* state) {
 //   }
 // }
 
-DEFINE_VECTORIZED_GRADIENT(exp)
-
 #endif  // USE_EXP