FEAT: Implement batch calculations for ase calculators

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<double> R_data(static_cast<size_t>(nSystems) *
+                               static_cast<size_t>(nAtoms) * 3);
+    std::vector<int> atomicNrs_data(static_cast<size_t>(nSystems) *
+                                    static_cast<size_t>(nAtoms));
+    std::vector<double> boxes_data(static_cast<size_t>(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<size_t> R_shape = {static_cast<size_t>(nSystems),
+                                   static_cast<size_t>(nAtoms), 3};
+    py::array_t<double> R_np(R_shape, R_data.data());
+
+    std::vector<size_t> atomicNrs_shape = {static_cast<size_t>(nSystems),
+                                           static_cast<size_t>(nAtoms)};
+    py::array_t<int> atomicNrs_np(atomicNrs_shape, atomicNrs_data.data());
+
+    std::vector<size_t> boxes_shape = {static_cast<size_t>(nSystems), 3, 3};
+    py::array_t<double> boxes_np(boxes_shape, boxes_data.data());
+
+    std::tuple<py::array_t<double>, py::array_t<double>> py_result =
+        (*batch_calculate)(R_np, atomicNrs_np, boxes_np, calculator)
+            .cast<std::tuple<py::array_t<double>, py::array_t<double>>>();
+
+    // copy the results to the output arrays
+    py::array_t<double> E = std::get<0>(py_result);
+    auto buffer_E = E.request();
+    double *ptr_E = static_cast<double *>(buffer_E.ptr);
+    std::copy(ptr_E, ptr_E + buffer_E.size, energies);
+
+    py::array_t<double> F = std::get<1>(py_result);
+    auto buffer_F = F.request();
+    double *ptr_F = static_cast<double *>(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;
+}

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<py::object>
+      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;
 };

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<Matter>(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<double *> boxes(nSystems);
+  for (long i = 0; i < nSystems; ++i) {
+    boxes[i] = matter->getCell().data();
+  }
+
+  std::vector<double> calculated_energy(nSystems, 0.0);
+
+  std::vector<AtomMatrix> calculated_forces(nSystems,
+                                            MatrixXd::Zero(nAtoms, 3));
+  std::vector<double *> 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<AtomMatrix>, _1, _2, threshold);
+    REQUIRE(matEq(calculated_forces[i], expected_forces));
+  };
+}
+
 } // namespace tests

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

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