diff --git a/client/potentials/ASE/ASE.cpp b/client/potentials/ASE/ASE.cpp index 8db7414bc..a0423ba2d 100644 --- a/client/potentials/ASE/ASE.cpp +++ b/client/potentials/ASE/ASE.cpp @@ -58,6 +58,9 @@ ASE::ASE(const Parameters &a_params) calculator = py_module.attr("ase_calc")(); _calculate = py_module.attr("_calculate"); + if (py::hasattr(py_module, "batch_calculate")) { + batch_calculate = py_module.attr("batch_calculate"); + } } catch (const std::exception &e) { fprintf(stderr, @@ -104,3 +107,71 @@ void ASE::force(long nAtoms, const double *R, const int *atomicNrs, double *F, counter++; return; } + +void ASE::forceBatch(long nSystems, long nAtoms, const double *const *R, + const int *const *atomicNrs, double *const *forces, + double *energies, double *variances, + const double *const *boxes) { + if (!batch_calculate) { + for (int i = 0; i < nSystems; ++i) { + force(nAtoms, R[i], atomicNrs[i], forces[i], &energies[i], &variances[i], + boxes[i]); + } + return; + } + + try { + variances = nullptr; + + std::vector R_data(static_cast(nSystems) * + static_cast(nAtoms) * 3); + std::vector atomicNrs_data(static_cast(nSystems) * + static_cast(nAtoms)); + std::vector boxes_data(static_cast(nSystems) * 9); + + for (long i = 0; i < nSystems; ++i) { + std::copy(R[i], R[i] + nAtoms * 3, R_data.begin() + i * nAtoms * 3); + std::copy(atomicNrs[i], atomicNrs[i] + nAtoms, + atomicNrs_data.begin() + i * nAtoms); + std::copy(boxes[i], boxes[i] + 9, boxes_data.begin() + i * 9); + } + + std::vector R_shape = {static_cast(nSystems), + static_cast(nAtoms), 3}; + py::array_t R_np(R_shape, R_data.data()); + + std::vector atomicNrs_shape = {static_cast(nSystems), + static_cast(nAtoms)}; + py::array_t atomicNrs_np(atomicNrs_shape, atomicNrs_data.data()); + + std::vector boxes_shape = {static_cast(nSystems), 3, 3}; + py::array_t boxes_np(boxes_shape, boxes_data.data()); + + std::tuple, py::array_t> py_result = + (*batch_calculate)(R_np, atomicNrs_np, boxes_np, calculator) + .cast, py::array_t>>(); + + // copy the results to the output arrays + py::array_t E = std::get<0>(py_result); + auto buffer_E = E.request(); + double *ptr_E = static_cast(buffer_E.ptr); + std::copy(ptr_E, ptr_E + buffer_E.size, energies); + + py::array_t F = std::get<1>(py_result); + auto buffer_F = F.request(); + double *ptr_F = static_cast(buffer_F.ptr); + for (long i = 0; i < nSystems; ++i) { + std::copy(ptr_F + i * nAtoms * 3, ptr_F + (i + 1) * nAtoms * 3, + forces[i]); + } + } catch (py::error_already_set &e) { + fprintf(stderr, "ASE calculator: Python error: %s\n", e.what()); + exit(1); + } catch (const std::exception &e) { + fprintf(stderr, "ASE calculator: C++ exception: %s\n", e.what()); + exit(1); + } + + counter += nSystems; + return; +} diff --git a/client/potentials/ASE/ASE.h b/client/potentials/ASE/ASE.h index 86686cfa9..e90e1741c 100644 --- a/client/potentials/ASE/ASE.h +++ b/client/potentials/ASE/ASE.h @@ -24,6 +24,9 @@ class ASE : public Potential { py::object calculator; // Member to store the ASE calculator object py::object _calculate; // Member to store the Python function to calculate // forces and energy + std::optional + batch_calculate; // Member to store the Python function to calculate + // forces and energies of multiple structures at once public: ASE(const Parameters &a_params); @@ -40,4 +43,12 @@ class ASE : public Potential { [[nodiscard]] bool needsPerImageInstance() const noexcept override { return true; } + + [[nodiscard]] bool supportsBatchEvaluation() const noexcept override { + return true; + } + void forceBatch(long nSystems, long nAtoms, const double *const *positions, + const int *const *atomicNrs, double *const *forces, + double *energies, double *variances, + const double *const *boxes) override; }; diff --git a/client/unit_tests/ASEPotTest.cpp b/client/unit_tests/ASEPotTest.cpp index 7b6df2f17..749477299 100644 --- a/client/unit_tests/ASEPotTest.cpp +++ b/client/unit_tests/ASEPotTest.cpp @@ -39,6 +39,7 @@ class ASEPotTest { pot = eonc::helpers::makePotential(params.potential_options.potential, params); + REQUIRE(pot->supportsBatchEvaluation()); matter = std::make_shared(pot, params); const std::string confile("pos.con"); @@ -78,4 +79,47 @@ TEST_CASE_METHOD(ASEPotTest, "ASE LJ energy and forces match reference", REQUIRE(matEq(calculated_forces, expected_forces)); } +TEST_CASE_METHOD(ASEPotTest, "ASE batch energy and forces match reference", + "[PotTest][ASE]") { + // Reference: two Al atoms at (0,0,0) and (1.5,0,0) in 10x10x10 box + // LJ(epsilon=1, sigma=1, rc=10, smooth=False) via ASE + const double expected_energy = -0.339200131812; + AtomMatrix expected_forces(2, 3); + expected_forces.row(0) << 1.158021344587014, 0.0, 0.0; + expected_forces.row(1) << -1.158021344587014, 0.0, 0.0; + + int nSystems = 2; + int nAtoms = matter->numberOfAtoms(); + const double *const positions[] = {matter->getPositions().data(), + matter->getPositions().data()}; + const int *const atomicNrs[] = {matter->getAtomicNrs().data(), + matter->getAtomicNrs().data()}; + + std::vector boxes(nSystems); + for (long i = 0; i < nSystems; ++i) { + boxes[i] = matter->getCell().data(); + } + + std::vector calculated_energy(nSystems, 0.0); + + std::vector calculated_forces(nSystems, + MatrixXd::Zero(nAtoms, 3)); + std::vector calculated_force_ptrs(nSystems); + for (long i = 0; i < nSystems; ++i) { + calculated_force_ptrs[i] = calculated_forces[i].data(); + } + + pot->forceBatch(nSystems, nAtoms, positions, atomicNrs, + calculated_force_ptrs.data(), calculated_energy.data(), + nullptr, boxes.data()); + + for (int i = 0; i < nSystems; ++i) { + REQUIRE_THAT(calculated_energy[i], WithinAbs(expected_energy, threshold)); + + auto matEq = + std::bind(eonc::helpers::eigenEquality, _1, _2, threshold); + REQUIRE(matEq(calculated_forces[i], expected_forces)); + }; +} + } // namespace tests diff --git a/client/unit_tests/data/systems/ase_pot/ase_lj.py b/client/unit_tests/data/systems/ase_pot/ase_lj.py index b3fc22a02..65a3b7295 100644 --- a/client/unit_tests/data/systems/ase_pot/ase_lj.py +++ b/client/unit_tests/data/systems/ase_pot/ase_lj.py @@ -1,4 +1,5 @@ """ASE LJ calculator for eOn test suite.""" + from ase import Atoms from ase.calculators.lj import LennardJones @@ -7,6 +8,20 @@ def ase_calc(): return LennardJones(epsilon=1.0, sigma=1.0, rc=10.0, ro=0.0, smooth=False) +def batch_calculate(Rs, atomicNrs, boxes, calc): + num_structures = np.shape(Rs)[0] + energies = np.empty(num_structures) + forces = np.empty((num_structures, atomicNrs, 3)) + for structure_idx in range(num_structures): + energies[structure_idx], forces[structure_idx, :] = _calculate( + Rs[structure_idx, :, :], + atomicNrs[structure_idx, :], + boxes[structure_idx, :, :], + calc, + ) + return energies, forces + + def _calculate(R, atomicNrs, box, calc): system = Atoms(symbols=atomicNrs, positions=R, pbc=True, cell=box) system.calc = calc diff --git a/docs/source/user_guide/ase_pot.md b/docs/source/user_guide/ase_pot.md index 1a391cba5..1a32c046f 100644 --- a/docs/source/user_guide/ase_pot.md +++ b/docs/source/user_guide/ase_pot.md @@ -116,6 +116,39 @@ def ase_calc(): return MACECalculator(model_paths="/absolute/path/to/model.pt", device="cuda") ``` +## Batch calculations + +For some potentials, like many MLIPs, batch evaluation is more efficient than individually +calculating the energy and force of each structure. eOn supports this as well, by adding +a function called `batch_calculate`. +```{code-block} python +from ase import Atoms +from ase.calculators.lj import LennardJones + +def ase_calc(): + # --- customize this section --- + calc = LennardJones(epsilon=0.0103, sigma=3.40, rc=10.0, ro=0.0, smooth=True) + return calc + +def batch_calculate(Rs, atomicNrs, boxs, calc): + # --- customize this section --- + return energies, forces + +#======================================================================= +# DO NOT EDIT below this line +def _calculate(R, atomicNrs, box, calc): + system = Atoms(symbols=atomicNrs, positions=R, pbc=True, cell=box) + system.calc = calc + forces = system.get_forces() + energy = system.get_potential_energy() + return energy, forces +#======================================================================= +``` +Now, if running an NEB, dimer or other method that calculates the energies +and forces of multiple structures at a time, `batch_calculate` is called instead of +calling `_calculate` for each structure. The energies should be returned as a `np.ndarray`, +and the forces as a np array with shape `n_structures*n_atoms*3`. + ## Alternatives If you installed eOn from conda-forge or prefer not to build from source, the