diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml index f06a87cc1..210f57642 100644 --- a/.github/workflows/benchmarks.yml +++ b/.github/workflows/benchmarks.yml @@ -130,14 +130,14 @@ jobs: echo "$PWD/install/lib/" > "${SITE_PACKAGES}/pyEXP.pth" - name: Install package and dependencies - env: - OMP_NUM_THREADS: 1 run: | export GALA_EXP_PREFIX=$PWD/EXP uv pip install --system -ve .[test] - name: Run the benchmarks uses: CodSpeedHQ/action@v4 + env: + OMP_NUM_THREADS: 1 with: mode: instrumentation run: uv run pytest tests/benchmarks -n auto --codspeed --durations=0 -v diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index ac76a85aa..1e2792fca 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -39,7 +39,8 @@ jobs: !contains(github.event.pull_request.labels.*.name, 'skip tests') && !contains(needs.check_skip_flags.outputs.head-commit-message, '[skip tests]') }} env: - EXP_REF: ${{ github.event_name == 'schedule' && 'devel' || 'v7.8.5' }} + # Run tests against a tagged EXP version, unless this is the devel branch or a PR into devel, in which case test against EXP devel + EXP_REF: ${{ (github.ref == 'refs/heads/devel' || (github.event_name == 'pull_request' && github.event.pull_request.base.ref == 'devel')) && 'devel' || 'v7.9.1' }} MACOSX_DEPLOYMENT_TARGET: "15.0" strategy: fail-fast: true @@ -190,40 +191,28 @@ jobs: so_ext="so" fi PYTHON_LIBRARY=${PYTHON_LIBRARY_PATH}/libpython${PYTHON_VERSION}.${so_ext} + GALA_EXP_PREFIX=$HOME/EXP-install/ + echo "GALA_EXP_PREFIX=${GALA_EXP_PREFIX}" >> $GITHUB_ENV cmake -G Ninja -B build \ - -DCMAKE_INSTALL_RPATH=$PWD/install/lib \ - --install-prefix $PWD/install \ - -DENABLE_PYEXP=on \ + --install-prefix $GALA_EXP_PREFIX \ + -DENABLE_PYEXP_ONLY=on \ + -DCMAKE_INSTALL_RPATH=${GALA_EXP_PREFIX}/lib \ -DPYTHON_EXECUTABLE=${PYTHON_EXEC} \ -DPYTHON_LIBRARY=${PYTHON_LIBRARY} \ -DCMAKE_CXX_FLAGS="${CXXFLAGS} -flto=auto" \ -DCMAKE_EXE_LINKER_FLAGS="${LDFLAGS}" \ -DCMAKE_SHARED_LINKER_FLAGS="${LDFLAGS}" - # Build EXP - # We are avoiding doing a CMake install as a hack to speed up the build - cmake --build build -v -t libexpui.${so_ext} libexputil.${so_ext} libyaml-cpp.${so_ext} pyEXP - - # The EXP CMake install doesn't offer a way to install just the shared libraries, so - # we are manually copying them to the install dir here. - # Users should do a proper "cmake --install" instead! - mkdir -p install/lib - if [[ "$RUNNER_OS" == "macOS" ]]; then - cp -P $(find build -name "*.so*" -o -name "*.dylib*") install/lib/ - else - cp -P $(find build -name "*.so*") install/lib/ - fi + cmake --build build -v + cmake --install build # Add the pyEXP dir as a pyEXP.pth file in Python site-packages: SITE_PACKAGES=$(python -c "import site; print(site.getsitepackages()[0])") - echo "$PWD/install/lib/" > "${SITE_PACKAGES}/pyEXP.pth" + echo "$GALA_EXP_PREFIX/lib/python${PYTHON_VERSION}/site-packages/" > "${SITE_PACKAGES}/pyEXP.pth" - name: Install package and dependencies run: | - if [[ "${{ matrix.gala-exp }}" == "1" ]]; then - export GALA_EXP_PREFIX=$PWD/EXP - fi uv pip install --system -ve .[${{ matrix.python-extras }}] env: GALA_NOGSL: ${{ matrix.gala-nogsl }} @@ -247,6 +236,7 @@ jobs: env: GALA_FORCE_EXP_TEST: ${{ matrix.gala-exp }} GALA_FORCE_PYEXP_TEST: ${{ matrix.gala-exp }} + OMP_NUM_THREADS: 1 run: | if [[ "${{ matrix.gala-exp }}" == "1" ]]; then uv run pytest tests/potential/potential/test_exp.py -n auto -ra --cov --cov-report=xml --cov-report=term-missing diff --git a/CHANGES.rst b/CHANGES.rst index 3f3f397b5..ee9304ff9 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -30,6 +30,9 @@ New Features with the xy-plane, the stream and progenitor are centered at (0, 0), and the stream primarily extends in the x direction (leading tail at positive x and trailing tail at negative x). +- EXP: constructing potentials from pyEXP objects is now supported via + ``gala.potential.PyEXPPotential``. +- EXP: force evaluation with ``gala.potential.EXPPotential`` should now be much faster. Bug fixes --------- @@ -71,6 +74,11 @@ API changes ``'v2'``, etc.). The old classes are deprecated and will be removed in a future release. +Build changes +------------- +- EXP: the instructions to build Gala against EXP have changed. Only the EXP install + dir is now used. + Other ----- diff --git a/docs/tutorials/exp.rst b/docs/tutorials/exp.rst index b9ee9b9fa..2ee120cde 100644 --- a/docs/tutorials/exp.rst +++ b/docs/tutorials/exp.rst @@ -10,10 +10,10 @@ simulation snapshots. This requires: #. building EXP, #. building Gala with EXP support, -#. and setting up a `~gala.potential.potential.EXPPotential` object using the user's EXP config and - coefficient files. +#. and setting up a `~gala.potential.potential.EXPPotential` or `~gala.potential.potential.PyEXPPotential` + object using a user-provided basis and coefficients. -Note that EXP support currently requires building Gala (and EXP) from source. +Note that EXP support currently requires building Gala from source. Additionally, this workflow has only been tested on Linux and MacOS with the setups seen in the `GitHub actions test config file `_. @@ -23,10 +23,12 @@ Building EXP ------------ The `EXP documentation `_ -is the authoritative source on how to build EXP. Currently, the only Gala-specific -addition to the instructions is that Gala expects the ``build`` directory to be present -in the EXP root directory. The ``install`` directory will be looked for in the EXP root -directory too, or one can set ``GALA_EXP_LIB_PATH`` (see below). +is the best place to read about how to build EXP. Gala doesn't have any special +instructions for the EXP build, except that the user must actually "install" EXP, +rather than just build it. This is demonstrated below. + +Gala is compatible with EXP version >= 7.9.1. If you encounter build issues, double +check the EXP version. To install EXP's dependencies, here is one recipe that we have found to work on Ubuntu 24.04:: @@ -36,48 +38,47 @@ To install EXP's dependencies, here is one recipe that we have found to work on Here is another recipe using modules that has been found to work on Flatiron Institute's rusty cluster:: - module load modules/2.3 cmake gcc openmpi hdf5 libtirpc eigen fftw git python + module load modules/2.4 cmake gcc openmpi hdf5 libtirpc eigen fftw git python uv EXP also builds on Mac by installing the dependencies with Homebrew:: - brew install cmake eigen fftw hdf5 open-mpi git ninja + brew install cmake eigen@3 fftw hdf5 open-mpi git ninja After installing the dependencies, one can download and build EXP on Linux with:: git clone --recursive https://github.com/EXP-code/EXP.git cd EXP - cmake -G Ninja -B build -DCMAKE_INSTALL_RPATH=$PWD/install/lib --install-prefix $PWD/install + cmake -G Ninja -B build -DENABLE_MINIMAL=on -DCMAKE_INSTALL_RPATH="$PWD/install/lib" -DCMAKE_BUILD_TYPE=Release --install-prefix $PWD/install cmake --build build cmake --install build +``-DENABLE_MINIMAL=on`` is optional but will make the build go faster. One can replace +this with ``-DENABLE_PYEXP_ONLY=on`` if one wants a minimal build with PyEXP. + +In this case, we installed EXP to the ``EXP/install/`` directory, but this can be any +directory. This will become the ``GALA_EXP_PREFIX`` directory in the next step. + For a full example of how to build EXP on Mac, see `this build recipe `_. -Note that building pyEXP is not required. However, some tests will use pyEXP if it is -present. +Note that building pyEXP is only necessary if one wants to use ``PyEXPPotential``. +Additionally, some tests will use pyEXP if it is present. ------------------------------ Building Gala with EXP support ------------------------------ -Building Gala with the ``GALA_EXP_PREFIX`` environment variable set to the EXP root dir +Building Gala with the ``GALA_EXP_PREFIX`` environment variable set to the EXP install dir will trigger compilation of the Gala's EXP Cython extensions. For example:: git clone https://github.com/adrn/gala.git cd gala - export GALA_EXP_PREFIX=/path/to/EXP - -If you build and install EXP following the instructions above, the EXP libraries will be -located in ``EXP/install/lib`` and the Gala build process knows to look there by default. If -you installed EXP to a different location, you can set the ``GALA_EXP_LIB_PATH`` -environment variable to point to the lib directory of the EXP install:: - - # Only do this if the install location is not $GALA_EXP_PREFIX/install - # export GALA_EXP_LIB_PATH=/path/to/EXP-install/lib + export GALA_EXP_PREFIX=/path/to/EXP/install/ -That is, ``GALA_EXP_LIB_PATH`` can be set if the CMake ``--install-prefix`` was set to a -location other than ``GALA_EXP_PREFIX/install``. ``GALA_EXP_LIB_PATH`` should be the -directory that contains the ``.so`` or ``.dylib`` files. +If you build and install EXP following the instructions above, the EXP installation will be +located in ``EXP/install/``. If you installed EXP to a different location, you can set the +``GALA_EXP_PREFIX`` to that location. In either case, ``GALA_EXP_PREFIX`` must be the directory +that contains the subdirectories ``lib`` and ``include``. Now you can run the Gala build. For example, using uv:: @@ -90,8 +91,14 @@ Or using venv:: . .venv/bin/activate python -m pip install -ve . -In either case, the pip output should show a message like ``Gala: installing with EXP -support``. +In either case, the output should show a message like ``Gala: installing with EXP support``. + +Note that in previous versions of Gala, the ``GALA_EXP_PREFIX`` was supposed to point to the +EXP repo root, rather than the EXP installation directory. This is no longer the case. The +EXP repo and build directories are not needed to build Gala with EXP support. + +Likewise, ``GALA_EXP_LIB_PATH`` was used in past Gala versions but not anymore. + ---------------------------------- Running Gala with an EXP potential @@ -142,6 +149,55 @@ integrate and plot an orbit: orbit = gp.Hamiltonian(exp_pot).integrate_orbit(w0, dt=1 * u.Myr, t1=0, t2=6 * u.Gyr) fig = orbit.plot(units=u.kpc, linestyle="-", alpha=0.5, label="orbit in m12m") + +----------------------------------- +Running Gala with a pyEXP potential +----------------------------------- + +If you are using +`pyEXP `_ +and have ``pyEXP.basis.BiorthBasis`` and ``pyEXP.coefs.Coefs`` objects (or any object +that subclasses them), you can use those to construct a Gala +`~gala.potential.potential.PyEXPPotential` object. + +Using ``PyEXPPotential``, the previous example would look like: + +.. code-block:: python + + import os + + import astropy.units as u + import pyEXP + + import gala.potential as gp + from gala.units import SimulationUnitSystem + + exp_units = SimulationUnitSystem(mass=1e12 * u.Msun, length=10 * u.kpc, G=1) + + # Construct the pyEXP basis + oldcwd = os.getcwd() + os.chdir("data") + with open("m12m-basis.yml") as fp: + basis = pyEXP.basis.Basis.factory(fp.read()) + os.chdir(oldcwd) + + # Construct the pyEXP coefs + coefs = pyEXP.coefs.Coefs.factory("data/m12m-coef.hdf5") + + pyexp_pot = gp.PyEXPPotential( + units=exp_units, + basis=basis, + coefs=coefs, + ) + + +Note that ``PyEXPPotential`` is missing some parameters, like ``snapshot_index``, that +``EXPPotential`` supports. This is because the intended workflow is for the user to construct +and modify the pyEXP basis and coefs objects using standard pyEXP methods and then pass those +objects to Gala. Otherwise, there should be no behavior or performance difference in using an +``EXPPotential`` or ``PyEXPPotential``. + + ----- Units ----- @@ -158,19 +214,20 @@ arbitrary, but it can be used to set physical scales to the simulations. Time Evolution -------------- -An `~gala.potential.potential.EXPPotential` may be time-evolving or static. If the coefficient -file has only one snapshot, the potential will be static. Likewise, if ``tmin``/``tmax`` -are passed such that only one snapshot from the coefs falls within that range, the +An `~gala.potential.potential.EXPPotential` or `~gala.potential.potential.PyEXPPotential` +may be time-evolving or static. If the coefficients only have snapshot, the potential +will be static. Likewise, for ``EXPPotential``, if ``tmin``/``tmax`` are passed such that +only one snapshot from the coefs falls within that range, the potential will be static. For the examples below, we use hypothetical files ``config.yml`` and ``coefs.h5`` that contain coefficients for multiple snapshots. -One can always check if an ``EXPPotential`` is static with: +One can always check if an ``EXPPotential`` or ``PyEXPPotential`` is static with: .. code-block:: python exp_pot.static -One can also "freeze" make a multi-snapshot potential (i.e. make it static) by selecting +One can also "freeze" a multi-snapshot ``EXPPotential`` (i.e. make it static) by selecting a single snapshot with the ``snapshot_index`` parameter: .. code-block:: python @@ -182,14 +239,17 @@ a single snapshot with the ``snapshot_index`` parameter: snapshot_index=0, ) +The equivalent for the pyEXP interface is to pass ``PyEXPPotential`` a coefs object that +only contains one snapshot. + For time-evolving potentials, if one tries to evaluate the potential outside of the -time range stored in the coefficients file (even indirectly, such as during an +time range stored in the coefficients (even indirectly, such as during an orbit integration), a C++ exception will be triggered, which will be raised to the user as a Python exception. The Python exception will contain the error message from C++. For example: ``RuntimeError: FieldWrapper::interpolator: time t=11.73 is out of bounds: [0.0195404, 11.724]``. -If the coefficients file stores a very large time range but the user is only interested +In ``EXPPotential``, if the coefficients store a very large time range but the user is only interested in a smaller range, one can specify ``tmin`` and/or ``tmax`` to load a smaller subset of the coefficient data (for memory efficiency): @@ -206,7 +266,7 @@ the coefficient data (for memory efficiency): Note that, as mentioned above, subsequently using a time outside this range will result in a Python exception. Or more precisely: using a time outside the range of snapshots that this ``tmin``/``tmax`` caused to be loaded will cause such an error. One can check the loaded -range of snapshots with: +range of snapshots (both ``EXPPotential`` and ``PyEXPPotential``) with: .. code-block:: python @@ -244,19 +304,35 @@ repo root: Composite Potentials -------------------- -`~gala.potential.potential.EXPPotential` fully supports composite potentials, including +`~gala.potential.potential.EXPPotential` and `~gala.potential.potential.PyEXPPotential` +fully support composite potentials, including mixing static and time-evolving potentials. The potentials will be combined at the C level as a :class:`~gala.potential.potential.CCompositePotential` when possible. See :ref:`_compositepotential` for more info. +-------------------------- +Performance Considerations +-------------------------- + +Within a timestep, the EXP force evaluation is parallelized with OpenMP threads across +orbits. With enough orbits (perhaps 1000 or more), you can expect to see a performance +benefit from using multiple threads. The number of OpenMP threads can be controlled +with standard OpenMP mechanisms, such as setting the ``OMP_NUM_THREADS`` environment +variable. + +Note that :class:`~gala.integrate.DOPRI853Integrator` batches the orbits into small +sets for performance, so EXP only sees the batch size at any given time and may not be +able to parallelize this well. One can use the ``nbatch`` integrator kwarg to tune the +batch size. + ----------- Limitations ----------- -The `~gala.potential.potential.EXPPotential` currently has the following limitations: +`~gala.potential.potential.EXPPotential` and `~gala.potential.potential.PyEXPPotential` +currently has the following limitations: * Hessian evaluation is not supported. * Pickling, saving, and loading is not supported. -* Performance may currently not be as high as native Gala potentials .. TODO (adrn): any other notable limitations? @@ -264,4 +340,5 @@ The `~gala.potential.potential.EXPPotential` currently has the following limitat API --- -See :class:`~gala.potential.potential.EXPPotential` for the complete API documentation. +See :class:`~gala.potential.potential.EXPPotential` and :class:`~gala.potential.potential.PyEXPPotential` +for the complete API documentation. diff --git a/setup.py b/setup.py index eb7e6a4a4..82085a07f 100755 --- a/setup.py +++ b/setup.py @@ -97,10 +97,8 @@ gsl_version = os.environ.get("GALA_GSL_VERSION", None) gsl_prefix = os.environ.get("GALA_GSL_PREFIX", None) -# The root directory of EXP (i.e., the repository root) +# The EXP installation prefix. This directory should contain 'include' and 'lib' subdirs. exp_prefix = os.environ.get("GALA_EXP_PREFIX", None) -# The path to the built/installed EXP libraries -exp_lib_path = os.environ.get("GALA_EXP_LIB_PATH", None) try: import pybind11 @@ -521,31 +519,23 @@ def base_cfg(): if "cyexp" in ext.name: if exp_prefix is not None: - if exp_lib_path is None: - # NOTE: this assumes user installed EXP to $GALA_EXP_PREFIX/install - lib_path_tmp = os.path.join(exp_prefix, "install", "lib") - if os.path.exists(lib_path_tmp): - exp_lib_path = lib_path_tmp - else: - msg = ( - "GALA_EXP_LIB_PATH not set, and no EXP libraries found in " - f"{lib_path_tmp}. Please set GALA_EXP_LIB_PATH to the path " - "where EXP libraries are installed." - ) - raise RuntimeError(msg) - - print(f"Gala: installing with EXP libraries at {exp_lib_path}") + exp_lib_path = os.path.join(exp_prefix, "lib") + if not os.path.exists(exp_lib_path): + msg = ( + f"No EXP libraries found in {exp_lib_path}. " + "Please set GALA_EXP_PREFIX to the directory that contains the 'lib' and 'include' " + "subdirectories of your EXP installation." + ) + raise RuntimeError(msg) ext.include_dirs.append(pybind11.get_include()) if extra_incl_flags is not None: ext.extra_compile_args.extend(extra_incl_flags) - if "exp" not in ext.libraries: - # TODO: we're compiling against installed EXP libraries, - # but headers from the source, because EXP doesn't install - # its headers. It would also need to install its vendored - # headers. + ext.extra_compile_args.extend(["-fopenmp"]) + ext.extra_link_args.extend(["-fopenmp"]) + if "exp" not in ext.libraries: ext.libraries.extend( ( "exputil", @@ -555,17 +545,7 @@ def base_cfg(): ) ext.library_dirs.append(exp_lib_path) ext.runtime_library_dirs.append(exp_lib_path) - ext.include_dirs.extend( - ( - os.path.join(exp_prefix, "include"), - # TODO: requires build in $GALA_EXP_PREFIX/build - os.path.join(exp_prefix, "build"), - os.path.join(exp_prefix, "expui"), - os.path.join(exp_prefix, "extern", "HighFive", "include"), - os.path.join(exp_prefix, "extern", "yaml-cpp", "include"), - ) - ) - + ext.include_dirs.append(os.path.join(exp_prefix, "include")) else: # Skip cyexp extension if EXP is not found continue diff --git a/src/gala/potential/potential/builtin/core.py b/src/gala/potential/potential/builtin/core.py index 04b1c5d6c..bf42972a0 100644 --- a/src/gala/potential/potential/builtin/core.py +++ b/src/gala/potential/potential/builtin/core.py @@ -43,7 +43,7 @@ ) if EXP_ENABLED: - from gala.potential.potential.builtin.cyexp import EXPWrapper + from gala.potential.potential.builtin.cyexp import EXPWrapper, PyEXPWrapper from ..core import PotentialBase, _potential_docstring from ..cpotential import CPotentialBase @@ -71,6 +71,7 @@ "NullPotential", "PlummerPotential", "PowerLawCutoffPotential", + "PyEXPPotential", "SatohPotential", "SphericalSplinePotential", "StonePotential", @@ -1732,3 +1733,122 @@ def tmax_exp(self) -> u.Quantity: The actual, loaded maximum time for which the potential is defined. """ return self.c_instance.tmax * self.parameters["snapshot_time_unit"] + + +@format_doc(common_doc=_potential_docstring) +class PyEXPPotential(CPotentialBase, EXP_only=True): + r""" + Calls the EXP code for the potential, using the pyEXP objects that the + user provides. + + This potential will usually be constructed with + :class:`~gala.units.SimulationUnitSystem` units. See the tutorial for more + information. + + .. note:: + + This potential requires EXP and pyEXP to be installed, and Gala must have been + built and installed with EXP support enabled. + See https://gala.adrian.pw/en/latest/tutorials/exp.html for more information. + + Parameters + ---------- + basis : `pyEXP.basis.BiorthBasis` + A pyEXP BiorthBasis object + coefs : `pyEXP.coefs.Coefs` + A pyEXP Coefs object + {common_doc} + + Attributes + ---------- + static : bool + Whether the potential is in static, i.e. fixed-time, mode. + tmin_exp, tmax_exp : `~astropy.units.Quantity` + The actual, loaded minimum and maximum time for which the potential is defined. + """ + + basis = PotentialParameter( + "basis", physical_type=None, python_only=True, convert=None + ) + coefs = PotentialParameter( + "coefs", physical_type=None, python_only=True, convert=None + ) + snapshot_time_unit = PotentialParameter( + "snapshot_time_unit", + physical_type=None, + default=None, + python_only=True, + convert=None, + ) + + def __init__(self, *args, **kwargs): + if "units" not in kwargs: + raise ValueError( + "Must specify a `units` keyword argument to initialize a PyEXPPotential " + "(most likely a SimulationUnitSystem with G=1)." + ) + + PotentialBase.__init__(self, *args, **kwargs) + + if self.parameters["snapshot_time_unit"] is None: + self.parameters["snapshot_time_unit"] = self.units["time"] + + # This hackery handles the situation where the snapshot time unit is different + # from the EXP (G=1) unit system that the coefficients/basis are in: + factor = 1 / ( + u.Quantity(1.0, self.parameters["snapshot_time_unit"]) + .decompose(self.units) + .value + ) + + try: + basis_capsule = self.parameters["basis"].get_shared_ptr_capsule() + except AttributeError as e: + raise ValueError( + "The `basis` parameter must be a pyEXP BiorthBasis object from pyEXP >= 7.9.1" + ) from e + + try: + coefs_capsule = self.parameters["coefs"].get_shared_ptr_capsule() + except AttributeError as e: + raise ValueError( + "The `coefs` parameter must be a pyEXP Coefs object from pyEXP >= 7.9.1" + ) from e + + self._setup_wrapper( + basis_capsule=basis_capsule, + coefs_capsule=coefs_capsule, + snapshot_time_factor=factor, + ) + + if EXP_ENABLED: + Wrapper = PyEXPWrapper + + def hessian(self, *args, **kwargs): + """ + Not implemented yet. + """ + raise NotImplementedError( + "Computing Hessian matrices for EXP potentials is not supported." + ) + + @property + def static(self) -> bool: + """ + Whether the potential is in static, i.e. fixed-time, mode. + """ + return self.c_instance.static + + @property + def tmin_exp(self) -> u.Quantity: + """ + The actual, loaded minimum time for which the potential is defined. + """ + return self.c_instance.tmin * self.parameters["snapshot_time_unit"] + + @property + def tmax_exp(self) -> u.Quantity: + """ + The actual, loaded maximum time for which the potential is defined. + """ + return self.c_instance.tmax * self.parameters["snapshot_time_unit"] diff --git a/src/gala/potential/potential/builtin/cyexp.pyx b/src/gala/potential/potential/builtin/cyexp.pyx index 7449c0d61..26af7d01a 100644 --- a/src/gala/potential/potential/builtin/cyexp.pyx +++ b/src/gala/potential/potential/builtin/cyexp.pyx @@ -6,19 +6,35 @@ # cython: language_level=3 # cython: language=c++ # cython: c_string_type=unicode, c_string_encoding=utf8 +# cython: cpp_locals=True +# cython: initializedcheck=True import numpy as np cimport numpy as np np.import_array() from libcpp.string cimport string +from libcpp.memory cimport shared_ptr from libcpp cimport bool as cbool +from cpython.pycapsule cimport PyCapsule_GetPointer + from ..cpotential cimport CPotentialWrapper from ..cpotential cimport densityfunc, energyfunc, gradientfunc, hessianfunc from ...._cconfig cimport USE_EXP +cdef extern from "EXP/Coefficients.H" namespace "CoefClasses": + cdef cppclass Coefs: + pass + ctypedef shared_ptr[Coefs] CoefsPtr + +cdef extern from "EXP/BiorthBasis.H" namespace "BasisClasses": + cdef cppclass Basis: + pass + ctypedef shared_ptr[Basis] BasisPtr + + cdef extern from "potential/potential/builtin/exp_fields.h" namespace "gala_exp": cdef cppclass State: double tmin @@ -35,6 +51,12 @@ cdef extern from "potential/potential/builtin/exp_fields.h" namespace "gala_exp" double snapshot_time_factor ) except + nogil + State pyexp_init( + BasisPtr *basis_ptr, + CoefsPtr *coefs_ptr, + double snapshot_time_factor + ) except + nogil + cdef extern from "potential/potential/builtin/exp_fields.h": # Note: the 'except +' annotations here don't actually do anything, since these functions # are not (currently) called directly from Cython/Python. But they serve as a reminder that @@ -95,3 +117,47 @@ cdef class EXPWrapper(CPotentialWrapper): @property def tmax(self): return self.exp_state.tmax + + +cdef class PyEXPWrapper(CPotentialWrapper): + cdef State exp_state + + def __init__( + self, G, parameters, q0, R, + basis_capsule, coefs_capsule, snapshot_time_factor + ): + cdef BasisPtr *basis_ptr + cdef CoefsPtr *coefs_ptr + + self.init( + [G], + np.ascontiguousarray(q0), + np.ascontiguousarray(R) + ) + + if USE_EXP == 1: + basis_ptr = PyCapsule_GetPointer(basis_capsule, "BiorthBasis_shared_ptr") + coefs_ptr = PyCapsule_GetPointer(coefs_capsule, "Coefs_shared_ptr") + + self.exp_state = pyexp_init( + basis_ptr, + coefs_ptr, + snapshot_time_factor + ) + self.cpotential.state[0] = &self.exp_state + self.cpotential.value[0] = (exp_value) + self.cpotential.density[0] = (exp_density) + self.cpotential.gradient[0] = (exp_gradient) + + + @property + def static(self): + return self.exp_state.is_static + + @property + def tmin(self): + return self.exp_state.tmin + + @property + def tmax(self): + return self.exp_state.tmax diff --git a/src/gala/potential/potential/builtin/exp_fields.cc b/src/gala/potential/potential/builtin/exp_fields.cc index 96f4f3f6b..45ea0080e 100644 --- a/src/gala/potential/potential/builtin/exp_fields.cc +++ b/src/gala/potential/potential/builtin/exp_fields.cc @@ -10,33 +10,83 @@ namespace fs = std::filesystem; // EXP headers -#include -#include -#include +#include +#include +#include +#include #include "exp_fields.h" #include "src/vectorization.h" namespace gala_exp { +State pyexp_init( + BasisClasses::BasisPtr *basis_ptr, + CoefClasses::CoefsPtr *coefs_ptr, + double snapshot_time_factor +) { + if (!basis_ptr) { + throw std::runtime_error("pyexp_init: basis pointer is null"); + } + + if (!coefs_ptr) { + throw std::runtime_error("pyexp_init: coefs pointer is null"); + } + + if (!*basis_ptr) { + throw std::runtime_error("pyexp_init: basis is null"); + } + + if (!*coefs_ptr) { + throw std::runtime_error("pyexp_init: coefs is null"); + } + + auto biorth_basis( + std::dynamic_pointer_cast( + *basis_ptr + ) + ); + if (!biorth_basis) { + throw std::runtime_error("pyEXP Basis must be a BiorthBasis."); + } + + return { biorth_basis, *coefs_ptr, snapshot_time_factor, -1 }; +} + State exp_init( const std::string &config_fn, const std::string &coeffile, int stride, double tmin, double tmax, int snapshot_index, double snapshot_time_factor) { YAML::Node yaml = YAML::LoadFile(std::string(config_fn)); - BasisClasses::BasisPtr basis; + auto load_basis = [](auto yaml, auto config_fn) -> auto { - // change the cwd to the directory of the config file - // so that relative paths in the config file work - // TODO: this is not thread-safe, threads share a cwd - ScopedChdir cd(fs::path(config_fn).parent_path()); - basis = BasisClasses::Basis::factory(yaml); - } + BasisClasses::BasisPtr base_basis; + { + // change the cwd to the directory of the config file + // so that relative paths in the config file work + // TODO: this is not thread-safe, threads share a cwd + ScopedChdir cd(fs::path(config_fn).parent_path()); + + base_basis = BasisClasses::Basis::factory(yaml); + } - if (!basis) { + if (!base_basis) { + std::ostringstream error_msg; + error_msg << "Failed to load basis from config file: " << config_fn; + throw std::runtime_error(error_msg.str()); + } + return base_basis; + }; + + auto biorth_basis( + std::dynamic_pointer_cast( + load_basis(yaml, config_fn) + ) + ); + if (!biorth_basis) { std::ostringstream error_msg; - error_msg << "Failed to load basis from config file: " << config_fn; + error_msg << "Basis in config file " << config_fn << " must be a BiorthBasis."; throw std::runtime_error(error_msg.str()); } @@ -49,6 +99,17 @@ State exp_init( throw std::runtime_error(error_msg.str()); } + try { + // Turn the "pure virtual" error in a more informative message + // TODO: is there a better way to "validate" the Coefs object? + coefs->Times(); + } catch (const std::runtime_error& e) { + std::ostringstream error_msg; + error_msg << "Failed to load coefficients from file: " << coeffile + << ". Error: " << e.what(); + throw std::runtime_error(error_msg.str()); + } + if(coefs->Times().empty()) { std::ostringstream error_msg; error_msg << "No times in coeffile=" << coeffile @@ -58,6 +119,34 @@ State exp_init( throw std::runtime_error(error_msg.str()); } + return { biorth_basis, coefs, snapshot_time_factor, snapshot_index }; +} + +State::State( + BiorthBasisPtr basis_, + CoefClasses::CoefsPtr coefs_, + double snapshot_time_factor_, + int snapshot_index) + : basis(basis_), + coefs(coefs_), + snapshot_time_factor(snapshot_time_factor_) { + + try { + // Turn the "pure virtual" error in a more informative message + // TODO: is there a better way to "validate" the Coefs object? + coefs->Times(); + } catch (const std::runtime_error& e) { + std::ostringstream error_msg; + error_msg << "Failed to fetch Times from Coefs object. " + << "Is this a valid, non-empty Coefs instance? " + << "Error: " << e.what(); + throw std::runtime_error(error_msg.str()); + } + + if(coefs->Times().empty()) { + throw std::runtime_error("No times in coefficients."); + } + if (coefs->Times().size() == 1 && snapshot_index < 0) { // If there is only one loaded snapshot in the coefs, // we treat it as static @@ -65,6 +154,7 @@ State exp_init( } bool is_static = false; + double tmin, tmax; if (snapshot_index >= 0) { const auto& times = coefs->Times(); @@ -95,7 +185,9 @@ State exp_init( } } - return { basis, coefs, tmin, tmax, is_static, snapshot_time_factor }; + this->is_static = is_static; + this->tmin = tmin; + this->tmax = tmax; } // Linear interpolator on coefficients. Higher order interpolation @@ -200,18 +292,21 @@ void exp_gradient(double t, double *__restrict__ pars, double *__restrict__ q_in ); } - // TODO: ask Martin/Mike for a way to compute only the force/acceleration - we're wasting - // computation time here by computing all fields 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]); + Eigen::Map eigen_x(q.x, N); + Eigen::Map eigen_y(q.y, N); + Eigen::Map eigen_z(q.z, N); + + auto& allaccel = exp_state->basis->getAccel(eigen_x, eigen_y, eigen_z); - grad.x[i] += -field[6]; - grad.y[i] += -field[7]; - grad.z[i] += -field[8]; + for(size_t i = 0; i < N; i++) { + grad.x[i] -= allaccel(i, 0); + grad.y[i] -= allaccel(i, 1); + grad.z[i] -= allaccel(i, 2); } + } double exp_density(double t, double *pars, double *q, int n_dim, void* state) { diff --git a/src/gala/potential/potential/builtin/exp_fields.h b/src/gala/potential/potential/builtin/exp_fields.h index da3780323..de9e342fe 100644 --- a/src/gala/potential/potential/builtin/exp_fields.h +++ b/src/gala/potential/potential/builtin/exp_fields.h @@ -4,19 +4,27 @@ #include #include -#include -#include +#include +#include namespace gala_exp { +using BiorthBasisPtr = shared_ptr; + class State { public: - BasisClasses::BasisPtr basis; + BiorthBasisPtr basis; CoefClasses::CoefsPtr coefs; + double snapshot_time_factor; double tmin; double tmax; bool is_static; - double snapshot_time_factor; + + State( + BiorthBasisPtr basis_, + CoefClasses::CoefsPtr coefs_, + double snapshot_time_factor_, + int snapshot_index); }; State exp_init( @@ -29,6 +37,12 @@ State exp_init( double snapshot_time_factor ); +State pyexp_init( + BasisClasses::BasisPtr *basis_ptr, + CoefClasses::CoefsPtr *coefs_ptr, + double snapshot_time_factor +); + CoefClasses::CoefStrPtr interpolator(double t, CoefClasses::CoefsPtr coefs); } diff --git a/tests/potential/potential/EXP-field-basis.yml b/tests/potential/potential/EXP-field-basis.yml new file mode 100644 index 000000000..82450b2be --- /dev/null +++ b/tests/potential/potential/EXP-field-basis.yml @@ -0,0 +1,7 @@ + +--- +# dummy field basis +id: field +parameters : + modelname: EXP-Hernquist.model +... diff --git a/tests/potential/potential/test_exp.py b/tests/potential/potential/test_exp.py index 1f4322a4b..48c605e2d 100644 --- a/tests/potential/potential/test_exp.py +++ b/tests/potential/potential/test_exp.py @@ -13,7 +13,7 @@ import gala.dynamics as gd import gala.potential as gp -from gala.potential.potential.builtin import EXPPotential +from gala.potential.potential.builtin import EXPPotential, PyEXPPotential from gala.units import SimulationUnitSystem from gala.util import chdir @@ -33,12 +33,14 @@ raise ImportError("pyEXP is required to run pyEXP tests") from e -EXP_CONFIG_FILE = str(this_path / "EXP-Hernquist-basis.yml") -EXP_SINGLE_COEF_FILE = str(this_path / "EXP-Hernquist-single-coefs.hdf5") -EXP_MULTI_COEF_FILE = str(this_path / "EXP-Hernquist-multi-coefs.hdf5") -EXP_MULTI_COEF_SNAPSHOT_TIME_FILE = str( - this_path / ("EXP-Hernquist-multi-coefs-snap-time-Gyr.hdf5") +EXP_CONFIG_FILE = this_path / "EXP-Hernquist-basis.yml" +EXP_FIELD_CONFIG_FILE = this_path / "EXP-field-basis.yml" # dummy +EXP_SINGLE_COEF_FILE = this_path / "EXP-Hernquist-single-coefs.hdf5" +EXP_MULTI_COEF_FILE = this_path / "EXP-Hernquist-multi-coefs.hdf5" +EXP_MULTI_COEF_SNAPSHOT_TIME_FILE = ( + this_path / "EXP-Hernquist-multi-coefs-snap-time-Gyr.hdf5" ) +EXP_UNITS = SimulationUnitSystem(mass=1.25234e11 * u.Msun, length=3.845 * u.kpc, G=1) # global pytest marker to skip tests if EXP is not enabled pytestmark = pytest.mark.skipif( @@ -49,12 +51,12 @@ # See: generate_exp.py, which generates the basis and coefficients for these tests -class EXPTestBase(PotentialTestBase): +# base for EXP and PyEXP tests +class CommonEXPTestBase(PotentialTestBase): tol = 1e-1 # increase tolerance for gradient test - exp_units = SimulationUnitSystem( - mass=1.25234e11 * u.Msun, length=3.845 * u.kpc, G=1 - ) + exp_units = EXP_UNITS + _tmp = gd.PhaseSpacePosition( pos=[-8, 0.0, 0.0] * u.kpc, vel=[0.0, 180, 0.0] * u.km / u.s, @@ -67,20 +69,6 @@ class EXPTestBase(PotentialTestBase): num_dx = 1e-3 skip_hessian = True - def setup_method(self): - assert os.path.exists(self.EXP_CONFIG_FILE), "EXP config file does not exist" - assert os.path.exists(self.EXP_COEF_FILE), "EXP coef file does not exist" - - self.potential = EXPPotential( - config_file=self.EXP_CONFIG_FILE, - coef_file=self.EXP_COEF_FILE, - # TODO: this is making the multi-coef test actually static! - # Need to fix the orbit integration then remove this - # snapshot_index=0, - units=self.exp_units, - ) - return super().setup_method() - # TODO: deepcopy is not implemented for EXPPotential @pytest.mark.skip(reason="Not implemented for EXP") def test_unitsystem(self): @@ -122,11 +110,11 @@ def test_orbit_integration(self, *args, **kwargs): ) @pytest.mark.skipif( - not FORCE_PYEXP_TEST, + not HAVE_PYEXP, reason="requires pyEXP", ) def test_pyexp(self): - """Test EXPPotential against pyEXP""" + """Test against pyEXP""" gala_test_x = [1.0, 2.0, -3.0] * u.kpc exp_test_x = gala_test_x.to_value(self.exp_units["length"]) @@ -135,7 +123,7 @@ def test_pyexp(self): config_str = fp.read() with chdir(os.path.dirname(self.EXP_CONFIG_FILE)): exp_basis = pyEXP.basis.Basis.factory(config_str) - exp_coefs = pyEXP.coefs.Coefs.factory(self.EXP_COEF_FILE) + exp_coefs = pyEXP.coefs.Coefs.factory(str(self.EXP_COEF_FILE)) # Use a snapshot time so that we don't have to rebuild the interpolation # functionality @@ -158,6 +146,41 @@ def test_pyexp(self): assert u.allclose(exp_grad, gala_grad) +class EXPTestBase(CommonEXPTestBase): + def setup_method(self): + assert os.path.exists(self.EXP_CONFIG_FILE), "EXP config file does not exist" + assert os.path.exists(self.EXP_COEF_FILE), "EXP coef file does not exist" + + self.potential = EXPPotential( + config_file=self.EXP_CONFIG_FILE, + coef_file=self.EXP_COEF_FILE, + units=self.exp_units, + ) + return super().setup_method() + + +@pytest.mark.skipif( + not HAVE_PYEXP, + reason="requires pyEXP", +) +class PyEXPTestBase(CommonEXPTestBase): + def setup_method(self): + assert os.path.exists(self.EXP_CONFIG_FILE), "EXP config file does not exist" + assert os.path.exists(self.EXP_COEF_FILE), "EXP coef file does not exist" + + with open(self.EXP_CONFIG_FILE) as fp, chdir(self.EXP_CONFIG_FILE.parent): + basis = pyEXP.basis.Basis.factory(fp.read()) + + coefs = pyEXP.coefs.Coefs.factory(str(self.EXP_COEF_FILE)) + + self.potential = PyEXPPotential( + basis=basis, + coefs=coefs, + units=self.exp_units, + ) + return super().setup_method() + + class TestEXPSingle(EXPTestBase): EXP_CONFIG_FILE = EXP_CONFIG_FILE EXP_COEF_FILE = EXP_SINGLE_COEF_FILE @@ -168,6 +191,16 @@ class TestEXPMulti(EXPTestBase): EXP_COEF_FILE = EXP_MULTI_COEF_FILE +class TestPyEXPSingle(PyEXPTestBase): + EXP_CONFIG_FILE = EXP_CONFIG_FILE + EXP_COEF_FILE = EXP_SINGLE_COEF_FILE + + +class TestPyEXPMulti(PyEXPTestBase): + EXP_CONFIG_FILE = EXP_CONFIG_FILE + EXP_COEF_FILE = EXP_MULTI_COEF_FILE + + def test_exp_unit_tests(): pot_single = EXPPotential( config_file=EXP_CONFIG_FILE, @@ -238,6 +271,38 @@ def test_exp_unit_tests(): assert u.allclose(pot_multi.tmax_exp, 2.0 * u.Gyr) +@pytest.mark.skipif(not HAVE_PYEXP, reason="requires pyEXP") +def test_pyexp_unit_tests(): + """Test PyEXPPotential static/dynamic behavior""" + units = EXPTestBase.exp_units + + with open(EXP_CONFIG_FILE) as fp, chdir(EXP_CONFIG_FILE.parent): + basis = pyEXP.basis.Basis.factory(fp.read()) + + coefs_single = pyEXP.coefs.Coefs.factory(str(EXP_SINGLE_COEF_FILE)) + coefs_multi = pyEXP.coefs.Coefs.factory(str(EXP_MULTI_COEF_FILE)) + + pot_single = PyEXPPotential(basis=basis, coefs=coefs_single, units=units) + pot_multi = PyEXPPotential(basis=basis, coefs=coefs_multi, units=units) + + assert pot_single.static is True + assert pot_multi.static is False + + test_x = [8.0, 0, 0] * u.kpc + assert u.allclose( + pot_single.energy(test_x, t=0 * u.Gyr), + pot_single.energy(test_x, t=1.4 * u.Gyr), + ) + assert not u.allclose( + pot_multi.energy(test_x, t=0 * u.Gyr), + pot_multi.energy(test_x, t=1.4 * u.Gyr), + ) + + # check tmin/tmax + assert u.allclose(pot_multi.tmin_exp, 0.0 * u.Gyr) + assert u.allclose(pot_multi.tmax_exp, 2.0 * u.Gyr) + + def test_multi_different_snapshot_time_unit(): pot_multi = EXPPotential( config_file=EXP_CONFIG_FILE, @@ -254,6 +319,28 @@ def test_multi_different_snapshot_time_unit(): assert u.allclose(pot_multi.tmax_exp, 1.0 * u.Gyr) +@pytest.mark.skipif(not HAVE_PYEXP, reason="requires pyEXP") +def test_pyexp_multi_different_snapshot_time_unit(): + """Test PyEXPPotential with different snapshot time units""" + units = EXPTestBase.exp_units + + with open(EXP_CONFIG_FILE) as fp, chdir(EXP_CONFIG_FILE.parent): + basis = pyEXP.basis.Basis.factory(fp.read()) + + coefs = pyEXP.coefs.Coefs.factory(str(EXP_MULTI_COEF_SNAPSHOT_TIME_FILE)) + + pot_multi = PyEXPPotential( + basis=basis, coefs=coefs, units=units, snapshot_time_unit=u.Gyr + ) + x = [8.0, 0, 0] * u.kpc + val0 = pot_multi.energy(x, t=0.0 * u.Gyr) + val1 = pot_multi.energy(x, t=1.0 * u.Gyr) + assert np.isclose(val1 / val0, 3.0) # see: generate_exp.py + + assert u.allclose(pot_multi.tmin_exp, 0.0 * u.Gyr) + assert u.allclose(pot_multi.tmax_exp, 1.0 * u.Gyr) + + def test_cython_exceptions(): """Test various exceptions propagated from C++""" units = SimulationUnitSystem(mass=1e11 * u.Msun, length=2.5 * u.kpc, G=1) @@ -299,19 +386,82 @@ def test_cython_exceptions(): ) -def test_composite(): - """Test that EXPPotential can be used in a CompositePotential""" - units = SimulationUnitSystem(mass=1.25234e11 * u.Msun, length=3.845 * u.kpc, G=1) - pot_single = EXPPotential( - config_file=EXP_CONFIG_FILE, - coef_file=EXP_SINGLE_COEF_FILE, - units=units, +@pytest.mark.skipif(not HAVE_PYEXP, reason="requires pyEXP") +def test_pyexp_exceptions(): + """Test various exceptions for PyEXPPotential""" + units = SimulationUnitSystem(mass=1e11 * u.Msun, length=2.5 * u.kpc, G=1) + + with open(EXP_CONFIG_FILE) as fp, chdir(EXP_CONFIG_FILE.parent): + basis = pyEXP.basis.Basis.factory(fp.read()) + coefs = pyEXP.coefs.Coefs.factory(str(EXP_MULTI_COEF_FILE)) + + # Test with None + with pytest.raises(ValueError, match="BiorthBasis"): + PyEXPPotential(basis=None, coefs=None, units=units) + + # Test with a real Coefs object that is empty + empty_coefs = pyEXP.coefs.Coefs(type="empty", verbose=False) + with pytest.raises(RuntimeError, match="Coefs"): + PyEXPPotential(basis=basis, coefs=empty_coefs, units=units) + + # Test with a non-BiorthBasis + with open(EXP_FIELD_CONFIG_FILE) as fp, chdir(EXP_FIELD_CONFIG_FILE.parent): + field_basis = pyEXP.basis.FieldBasis(fp.read()) + with pytest.raises(ValueError, match="BiorthBasis"): + PyEXPPotential(basis=field_basis, coefs=coefs, units=units) + + # Test with valid objects but runtime errors + with open(EXP_CONFIG_FILE) as fp, chdir(EXP_CONFIG_FILE.parent): + basis = pyEXP.basis.Basis.factory(fp.read()) + + pot = PyEXPPotential(basis=basis, coefs=coefs, units=units) + with pytest.raises(RuntimeError, match="time"): + pot.energy([0, 0, 0], t=float(0xBAD)) + + +def _make_exp_pot(config_fn, coef_fn): + return EXPPotential( + config_file=config_fn, + coef_file=coef_fn, + units=EXP_UNITS, ) - pot_multi = EXPPotential( - config_file=EXP_CONFIG_FILE, - coef_file=EXP_MULTI_COEF_FILE, - units=units, + + +def _make_pyexp_pot(config_fn, coef_fn): + return PyEXPPotential( + basis=_load_pyexp_basis(config_fn), + coefs=pyEXP.coefs.Coefs.factory(str(coef_fn)), + units=EXP_UNITS, ) + + +def _load_pyexp_basis(config_file): + """Helper to load pyEXP basis for parametrized tests""" + if not HAVE_PYEXP: + return None + with open(config_file) as fp, chdir(config_file.parent): + return pyEXP.basis.Basis.factory(fp.read()) + + +potentials_parametrize = pytest.mark.parametrize( + "make_pot", + [ + pytest.param(_make_exp_pot, id="exp"), + pytest.param( + _make_pyexp_pot, + id="pyexp", + marks=pytest.mark.skipif(not HAVE_PYEXP, reason="requires pyEXP"), + ), + ], +) + + +@potentials_parametrize +def test_composite_parametrized(make_pot): + """Test that both EXPPotential and PyEXPPotential can be used in a CompositePotential""" + pot_single = make_pot(EXP_CONFIG_FILE, EXP_SINGLE_COEF_FILE) + + pot_multi = make_pot(EXP_CONFIG_FILE, EXP_MULTI_COEF_FILE) composite_pot = pot_single + pot_multi assert isinstance( composite_pot, gp.potential.ccompositepotential.CCompositePotential @@ -355,6 +505,30 @@ def test_composite(): assert np.all(np.isfinite(orbit.t.value)) +@potentials_parametrize +def test_replace_units(make_pot): + """Test that replace_units works for both EXPPotential and PyEXPPotential""" + pot = make_pot(EXP_CONFIG_FILE, EXP_SINGLE_COEF_FILE) + + new_units = SimulationUnitSystem( + mass=EXP_UNITS["mass"] * 2.0, + length=EXP_UNITS["length"], + G=1.0, + ) + pot_replaced = pot.replace_units(new_units) + + assert pot_replaced.units == new_units + assert pot_replaced is not pot + + x = [1.0, 2.0, 3.0] * u.kpc + e1 = pot.energy(x) + + x_new = x.to_value(new_units["length"]) * new_units["length"] + e2 = pot_replaced.energy(x_new) + + assert u.isclose(e1, e2 / 2.0) + + def test_paths(): """ Test relative and absolute file paths @@ -374,49 +548,128 @@ def test_paths(): ) -def test_replace_units(): - """Test that replace_units works for EXPPotential""" +def test_replicate(): + """Test that replicate works for EXPPotential""" units = SimulationUnitSystem(mass=1e11 * u.Msun, length=2.5 * u.kpc, G=1) pot = EXPPotential( config_file=EXP_CONFIG_FILE, - coef_file=EXP_SINGLE_COEF_FILE, + coef_file=EXP_MULTI_COEF_FILE, units=units, + snapshot_index=0, ) - new_units = SimulationUnitSystem(mass=2e11 * u.Msun, length=2.5 * u.kpc, G=1) - pot_replaced = pot.replace_units(new_units) + pot_replicated = pot.replicate(snapshot_index=1) - assert pot_replaced.units == new_units - assert pot_replaced is not pot # should be a new instance + assert pot_replicated.units == pot.units + assert pot_replicated.parameters["snapshot_index"] == 1 + assert pot.parameters["snapshot_index"] == 0 + assert pot_replicated is not pot # should be a new instance - # Check that the energy at a point is the same in both unit systems + # Check that the energy at a point is not the same in both instances x = [1.0, 2.0, 3.0] * u.kpc e1 = pot.energy(x) - e2 = pot_replaced.energy(x.to_value(new_units["length"]) * new_units["length"]) - assert u.isclose(e1, e2 / 2.0) + e2 = pot_replicated.energy(x) + assert not u.isclose(e1, e2) -def test_replicate(): - """Test that replicate works for EXPPotential""" +@pytest.mark.xfail(reason="replicate not supported by PyEXP") +@pytest.mark.skipif(not HAVE_PYEXP, reason="requires pyEXP") +def test_pyexp_replicate(): + """Test that replicate works for PyEXPPotential using coef_file""" units = SimulationUnitSystem(mass=1e11 * u.Msun, length=2.5 * u.kpc, G=1) - pot = EXPPotential( - config_file=EXP_CONFIG_FILE, - coef_file=EXP_MULTI_COEF_FILE, + with open(EXP_CONFIG_FILE) as fp, chdir(EXP_CONFIG_FILE.parent): + basis = pyEXP.basis.Basis.factory(fp.read()) + coefs = pyEXP.coefs.Coefs.factory(str(EXP_SINGLE_COEF_FILE)) + + pot = PyEXPPotential( + basis=basis, + coefs=coefs, units=units, - snapshot_index=0, ) - pot_replicated = pot.replicate(snapshot_index=1) + pot_replicated = pot.replicate(coef_file=str(EXP_MULTI_COEF_FILE)) assert pot_replicated.units == pot.units - assert pot_replicated.parameters["snapshot_index"] == 1 - assert pot.parameters["snapshot_index"] == 0 + assert Path(pot_replicated.parameters["coef_file"]) == Path(EXP_MULTI_COEF_FILE) + assert Path(pot.parameters["coef_file"]) == Path(EXP_SINGLE_COEF_FILE) assert pot_replicated is not pot # should be a new instance - # Check that the energy at a point is the not same in both instances x = [1.0, 2.0, 3.0] * u.kpc - e1 = pot.energy(x) - e2 = pot_replicated.energy(x) + e1 = pot.energy(x, t=0 * u.Gyr) + e2 = pot_replicated.energy(x, t=0 * u.Gyr) assert not u.isclose(e1, e2) + + +@pytest.mark.skipif(not HAVE_PYEXP, reason="requires pyEXP") +def test_exp_pyexp_consistency_single(): + """Test that EXPPotential and PyEXPPotential give the same results""" + + # Create EXPPotential + exp_pot = EXPPotential( + config_file=EXP_CONFIG_FILE, + coef_file=EXP_SINGLE_COEF_FILE, + units=EXP_UNITS, + ) + + # Create PyEXPPotential with same data + with open(EXP_CONFIG_FILE) as fp, chdir(EXP_CONFIG_FILE.parent): + basis = pyEXP.basis.Basis.factory(fp.read()) + coefs = pyEXP.coefs.Coefs.factory(str(EXP_SINGLE_COEF_FILE)) + pyexp_pot = PyEXPPotential(basis=basis, coefs=coefs, units=EXP_UNITS) + + x = [1.0, 2.0, 3.0] * u.kpc + + # Compare energy + exp_energy = exp_pot.energy(x) + pyexp_energy = pyexp_pot.energy(x) + assert u.allclose(exp_energy, pyexp_energy) + + # Compare density + exp_density = exp_pot.density(x) + pyexp_density = pyexp_pot.density(x) + assert u.allclose(exp_density, pyexp_density) + + # Compare gradient + exp_gradient = exp_pot.gradient(x) + pyexp_gradient = pyexp_pot.gradient(x) + assert u.allclose(exp_gradient, pyexp_gradient) + + +@pytest.mark.skipif(not HAVE_PYEXP, reason="requires pyEXP") +def test_exp_pyexp_consistency_multi(): + """Test time-dependent consistency between EXPPotential and PyEXPPotential.""" + exp_dynamic = EXPPotential( + config_file=EXP_CONFIG_FILE, + coef_file=EXP_MULTI_COEF_FILE, + units=EXP_UNITS, + ) + + with open(EXP_CONFIG_FILE) as fp, chdir(EXP_CONFIG_FILE.parent): + basis = pyEXP.basis.Basis.factory(fp.read()) + coefs = pyEXP.coefs.Coefs.factory(str(EXP_MULTI_COEF_FILE)) + + pyexp_dynamic = PyEXPPotential( + basis=basis, + coefs=coefs, + units=EXP_UNITS, + ) + + assert exp_dynamic.static is False + assert pyexp_dynamic.static is False + + x = [2.5, -1.5, 0.4] * u.kpc + times = [0.0, 1.4] * u.Gyr + + for t in times: + exp_energy = exp_dynamic.energy(x, t=t) + pyexp_energy = pyexp_dynamic.energy(x, t=t) + exp_density = exp_dynamic.density(x, t=t) + pyexp_density = pyexp_dynamic.density(x, t=t) + exp_gradient = exp_dynamic.gradient(x, t=t) + pyexp_gradient = pyexp_dynamic.gradient(x, t=t) + + assert u.allclose(exp_energy, pyexp_energy) + assert u.allclose(exp_density, pyexp_density) + assert u.allclose(exp_gradient, pyexp_gradient) diff --git a/tests/potential/potential/test_time_interpolated.py b/tests/potential/potential/test_time_interpolated.py index a8e64f764..c3c3da187 100644 --- a/tests/potential/potential/test_time_interpolated.py +++ b/tests/potential/potential/test_time_interpolated.py @@ -478,6 +478,9 @@ def test_all_builtin_potentials_time_interpolated(pot_cls_name): params_const["r_c"] = params_time["r_c"] = 1.0 params_const["r_h"] = params_time["r_h"] = 10.0 + elif pot_cls_name in ("EXPPotential", "PyEXPPotential"): + pytest.skip(f"{pot_cls_name} uses its own interpolation") + pot_const = pot_cls(**params_const, units=galactic) pot_time = gp.TimeInterpolatedPotential( potential_cls=pot_cls,