diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 6d25adf82..4d98baea6 100644 --- a/.github/workflows/tests.yml +++ b/.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' diff --git a/gala/dynamics/mockstream/mockstream.pyx b/gala/dynamics/mockstream/mockstream.pyx index 75de2ebbf..264266efb 100644 --- a/gala/dynamics/mockstream/mockstream.pyx +++ b/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 diff --git a/gala/dynamics/nbody/nbody.pyx b/gala/dynamics/nbody/nbody.pyx index 23e7548fb..afcd0e8f1 100644 --- a/gala/dynamics/nbody/nbody.pyx +++ b/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) diff --git a/gala/integrate/cyintegrators/dop853.pxd b/gala/integrate/cyintegrators/dop853.pxd index 95eb1b8ad..a46034c3e 100644 --- a/gala/integrate/cyintegrators/dop853.pxd +++ b/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=? ) diff --git a/gala/integrate/cyintegrators/dop853.pyx b/gala/integrate/cyintegrators/dop853.pyx index 5ae62d47d..dc264cac0 100644 --- a/gala/integrate/cyintegrators/dop853.pyx +++ b/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 = ((hamiltonian.potential.c_instance)).cpotential CFrameType cf = ((hamiltonian.frame.c_instance)).cframe - # 0 below is for nbody - we ignore that in this test particle integration - w = dop853_helper( - cp, &cf, 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, 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) diff --git a/gala/integrate/cyintegrators/dopri/dop853.cpp b/gala/integrate/cyintegrators/dopri/dop853.cpp index ddd82d2f3..bee61b7d0 100644 --- a/gala/integrate/cyintegrators/dopri/dop853.cpp +++ b/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) { diff --git a/gala/integrate/cyintegrators/dopri/dop853.h b/gala/integrate/cyintegrators/dopri/dop853.h index 25703e643..4da216ad4 100644 --- a/gala/integrate/cyintegrators/dopri/dop853.h +++ b/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, diff --git a/gala/integrate/cyintegrators/ruth4.pxd b/gala/integrate/cyintegrators/ruth4.pxd index 044dadf1b..04a724034 100644 --- a/gala/integrate/cyintegrators/ruth4.pxd +++ b/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 diff --git a/gala/integrate/cyintegrators/ruth4.pyx b/gala/integrate/cyintegrators/ruth4.pyx index 3bbad45ff..d450fa5cf 100644 --- a/gala/integrate/cyintegrators/ruth4.pyx +++ b/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 = ((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) # ------------------------------------------------------------------------------------- diff --git a/gala/potential/hamiltonian/chamiltonian.pyx b/gala/potential/hamiltonian/chamiltonian.pyx index b6ac06f53..72633eab4 100644 --- a/gala/potential/hamiltonian/chamiltonian.pyx +++ b/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}'") diff --git a/gala/potential/hamiltonian/src/chamiltonian.cpp b/gala/potential/hamiltonian/src/chamiltonian.cpp index 980dce91a..629a7a895 100644 --- a/gala/potential/hamiltonian/src/chamiltonian.cpp +++ b/gala/potential/hamiltonian/src/chamiltonian.cpp @@ -1,5 +1,6 @@ #include #include +#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; diff --git a/gala/potential/hamiltonian/src/chamiltonian.h b/gala/potential/hamiltonian/src/chamiltonian.h index fdcf47131..88e63e75c 100644 --- a/gala/potential/hamiltonian/src/chamiltonian.h +++ b/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); diff --git a/gala/potential/potential/builtin/exp_fields.cc b/gala/potential/potential/builtin/exp_fields.cc index aafbaf9a1..96f4f3f6b 100644 --- a/gala/potential/potential/builtin/exp_fields.cc +++ b/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(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 diff --git a/gala/potential/potential/builtin/multipole.cpp b/gala/potential/potential/builtin/multipole.cpp index f1154af46..533f742ef 100644 --- a/gala/potential/potential/builtin/multipole.cpp +++ b/gala/potential/potential/builtin/multipole.cpp @@ -334,8 +334,7 @@ double mp_potential(double t, double *pars, double *q, int n_dim) { return val[0]; } -// TODO: de-scalarize -void mp_gradient_single(double t, double *__restrict__ pars, double6ptr q, int n_dim, double6ptr grad, void *__restrict__ state) { +void mp_gradient(double t, double *__restrict__ pars, double *__restrict__ q, int n_dim, size_t N, double *__restrict__ grad, void *__restrict__ state) { /* pars: - G (Gravitational constant) - lmax @@ -358,6 +357,28 @@ void mp_gradient_single(double t, double *__restrict__ pars, double6ptr q, int n Tlm[i] = pars[7 + 2*i]; } + mp_gradient_helper(double6ptr{q, N}, N, + G, M, r_s, + &Slm[0], &Tlm[0], + lmax, inner, double6ptr{grad, N}); +} + +void _mp_gradient_single(double t, double *__restrict__ pars, double6ptr q, int n_dim, double6ptr grad, void *__restrict__ state) { + // This is a helper used by axisym_cylspline_gradient, but not by mp_gradient. + + double G = pars[0]; + int lmax = (int)pars[1]; + int num_coeff = (int)pars[2]; + int inner = (int)pars[3]; + double M = pars[4]; + double r_s = pars[5]; + + double Slm[num_coeff], Tlm[num_coeff]; + for(int i=0; i= gridR[0]) && (Rasinh <= gridR[nR-1]) && - (zasinh >= gridz[0]) && (zasinh <= gridz[nz-1])) { // Use CylSpline + double6ptr q = double6ptr{q_in, N}; + double6ptr grad = double6ptr{grad_in, N}; - /* initialize interpolation */ - // TODO: define this in wrapper, make all CPotential's have a void - // pointer array to store things like this, all these functions then - // need to accept one more parameter (or is there a way to do optional - // args in C?), ??, profit. - gsl_spline2d_init(spline, gridR, gridz, gridPhi, nR, nz); + for(size_t i = 0; i < N; i++) { + double R = sqrt(q.x[i]*q.x[i] + q.y[i]*q.y[i]); + double Rasinh = asinh(R / Rscale); + double zasinh = asinh(q.z[i] / Rscale); - dPhi_dR = gsl_spline2d_eval_deriv_x(spline, Rasinh, zasinh, xacc, yacc); - dPhi_dR = dPhi_dR / (Rscale * cosh(Rasinh)); + if ((Rasinh >= gridR[0]) && (Rasinh <= gridR[nR-1]) && + (zasinh >= gridz[0]) && (zasinh <= gridz[nz-1])) { // Use CylSpline - dPhi_dz = gsl_spline2d_eval_deriv_y(spline, Rasinh, zasinh, xacc, yacc); - dPhi_dz = dPhi_dz / (Rscale * cosh(zasinh)); + double dPhi_dR = gsl_spline2d_eval_deriv_x(spline, Rasinh, zasinh, xacc, yacc); + dPhi_dR = dPhi_dR / (Rscale * cosh(Rasinh)); - if (logScaling) { - Phi = gsl_spline2d_eval(spline, Rasinh, zasinh, xacc, yacc); - Phi = -exp(Phi); - dPhi_dR = dPhi_dR * Phi; - dPhi_dz = dPhi_dz * Phi; - } + double dPhi_dz = gsl_spline2d_eval_deriv_y(spline, Rasinh, zasinh, xacc, yacc); + dPhi_dz = dPhi_dz / (Rscale * cosh(zasinh)); - if (R > 0) { - grad[0] = grad[0] + dPhi_dR * q[0] / R; - grad[1] = grad[1] + dPhi_dR * q[1] / R; - grad[2] = grad[2] + dPhi_dz; - } else { - grad[2] = grad[2] + dPhi_dz; - } + if (logScaling) { + double Phi = gsl_spline2d_eval(spline, Rasinh, zasinh, xacc, yacc); + Phi = -exp(Phi); + dPhi_dR = dPhi_dR * Phi; + dPhi_dz = dPhi_dz * Phi; + } - } else { // Use external Multipole - mp_gradient_single(t, &pars[5 + nR + nz + nR * nz], q, n_dim, grad, state); + if (R > 0) { + grad.x[i] += dPhi_dR * q.x[i] / R; + grad.y[i] += dPhi_dR * q.y[i] / R; + grad.z[i] += dPhi_dz; + } else { + grad.z[i] += dPhi_dz; + } + + } else { // Use external Multipole + _mp_gradient_single(t, &pars[5 + nR + nz + nR * nz], double6ptr{q_in + i, N}, n_dim, double6ptr{grad_in + i, N}, state); + } } + gsl_spline2d_free(spline); gsl_interp_accel_free(xacc); gsl_interp_accel_free(yacc); @@ -808,8 +827,4 @@ double axisym_cylspline_density(double t, double *pars, double *q, int n_dim) { return dens; } -DEFINE_VECTORIZED_GRADIENT(mp) -DEFINE_VECTORIZED_GRADIENT(mpetd) -DEFINE_VECTORIZED_GRADIENT(axisym_cylspline) - #endif // USE_GSL