Conformer batches

One vmapped solve over a stack of geometries: energies and forces for a whole conformer batch in a single call.

import jax; jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
import numpy as np
from dftax import Molecule, scf_batched
from dftax.energy.xc import PBE

# A 1-D H2 bond scan: 8 geometries differing only in the bond length.
mol = Molecule.from_xyz("H 0 0 0; H 0 0 1.4", "sto-3g")  # template atoms + basis
lengths = jnp.linspace(1.0, 2.4, 8)                       # Bohr
coords_batch = jnp.stack([jnp.array([[0.0, 0.0, 0.0], [0.0, 0.0, L]]) for L in lengths])

res = scf_batched(mol, coords_batch, PBE(), forces=True)  # batched KSResult
print("converged:", np.asarray(res.converged))
print(" R (Bohr)   E (Ha)        F_z on atom 1 (Ha/Bohr)")
for L, e, F in zip(lengths, res.e_tot, res.forces):
    print(f" {float(L):6.3f}   {float(e):+.6f}    {float(F[1, 2]):+.6f}")

imin = int(jnp.argmin(res.e_tot))
print(f"\nminimum-energy geometry in the scan: R = {float(lengths[imin]):.3f} Bohr")

E0725 14:06:39.126517 1262956 numa_hwloc.cc:121] Call to hwloc_set_cpubind() failed: Invalid argument [22]

converged: [ True  True  True  True  True  True  True  True]
 R (Bohr)   E (Ha)        F_z on atom 1 (Ha/Bohr)
  1.000   -1.094345    +0.377199
  1.200   -1.141672    +0.124290
  1.400   -1.152112    -0.005631
  1.600   -1.143585    -0.072202
  1.800   -1.125398    -0.105801
  2.000   -1.102416    -0.121867
  2.200   -1.077320    -0.127776
  2.400   -1.051732    -0.127258

minimum-energy geometry in the scan: R = 1.400 Bohr