Add i-PI driven dynamics example with BEC support

Wire up the Calculator → i-PI BEC data path: extract "apt" from
LoremBEC model output, reshape to (3*natoms, 3), and expose as "BEC"
for i-PI's driven_dynamics module. Implement LOREMCalculator as a
proper Calculator subclass and clean up the i-PI driver (remove buggy
compute_structure override that operated on JSON strings).

Add md-ipi example with input.xml (eda-nve, bec mode="driver"),
run.sh, start.xyz, and README.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: sirmarcel

examples/md-ipi/README.md

@@ -0,0 +1,37 @@
+# Driven dynamics with LOREM via i-PI
+
+This example shows how to run an applied electric field simulation using a LOREM BEC model and i-PI's driven dynamics (Electric Dipole Approximation).
+
+Unlike the standard i-PI efield example which reads fixed Born Effective Charges from a file, LOREM computes BECs on-the-fly from its learned model (`<bec mode="driver"/>`).
+
+## Prerequisites
+
+- Trained `LoremBEC` checkpoint (with `model/model.yaml`, `model/baseline.yaml`, `model/model.msgpack`)
+- i-PI installed (`pip install ipi`)
+- LOREM i-PI driver installed: `lorem-install-ipi-driver`
+
+## Files
+
+- `input.xml` — i-PI configuration for driven dynamics (EDA-NVE). Adapt the E-field parameters (`amp`, `freq`, `peak`, `sigma`), cell size, and `start.xyz` to your system.
+- `start.xyz` — Starting structure in atomic units. Replace with your system's geometry.
+- `run.sh` — Launch script. Set `MODEL_PATH` to your trained checkpoint folder.
+
+## Running
+
+```bash
+# Edit run.sh to set MODEL_PATH, then:
+bash run.sh
+```
+
+i-PI outputs will be written to `i-pi.*` files. Key outputs:
+- `i-pi.properties.out` — energy, kinetic energy, E-field strength over time
+- `i-pi.positions_*` — nuclear trajectories
+- `i-pi.bec{x,y,z}_*` — BEC tensor components over time
+
+## Adapting to your system
+
+1. Replace `start.xyz` with your starting geometry (atomic units).
+2. Set the cell size in `input.xml` (large for isolated systems, physical for periodic).
+3. Set `pbc='True'` in `<ffsocket>` if your system is periodic.
+4. Adjust E-field parameters to match your target excitation.
+5. Point `MODEL_PATH` in `run.sh` to your trained LoremBEC checkpoint.

examples/md-ipi/input.xml

@@ -0,0 +1,55 @@
+<simulation verbosity='high'>
+
+  <total_steps>500</total_steps>
+
+  <ffsocket mode='unix' name='driver' pbc='False'>
+    <address>lorem</address>
+  </ffsocket>
+
+  <output prefix='i-pi'>
+    <trajectory filename='forces'>forces</trajectory>                <!--      total forces (a.u.): nuclear + q_e Z @ E(t) -->
+    <trajectory filename='Eforces'>Eforces</trajectory>              <!--     driven forces (a.u.): q_e Z @ E(t) -->
+    <trajectory filename='momenta'>momenta</trajectory>              <!--   nuclear momenta (a.u.) -->
+    <trajectory filename='velocities'>velocities</trajectory>        <!--   nuclear velocities (a.u.) -->
+    <trajectory filename='positions'>positions</trajectory>          <!-- nuclear positions (a.u.) -->
+    <trajectory filename='becx'>becx</trajectory>
+    <trajectory filename='becy'>becy</trajectory>
+    <trajectory filename='becz'>becz</trajectory>
+    <trajectory filename='dip' extra_type='dipole'> extras </trajectory>
+    <properties filename='properties.out' verbosity='low' > [ step, time, conserved, kinetic_md, potential, Efield, Eenvelope ] </properties>
+  </output>
+
+  <prng>
+    <seed>10545</seed>
+  </prng>
+
+  <system>
+    <forces>
+      <force forcefield='driver' />
+    </forces>
+
+    <initialize nbeads='1'>
+      <file mode='xyz' units='atomic_unit' > start.xyz </file>
+      <velocities mode='thermal' units='kelvin'> 0 </velocities>         <!-- no initial velocities: excited by E-field only -->
+      <cell mode="manual" units="atomic_unit">[ 50,0,0, 0,50,0, 0,0,50 ]</cell>
+    </initialize>
+
+    <ensemble>
+      <temperature units='kelvin'>0</temperature>
+    </ensemble>
+
+    <motion mode='driven_dynamics'>
+      <driven_dynamics mode='eda-nve'>                  <!-- Electric Dipole Approximation (EDA) on top of NVE -->
+        <timestep units='femtosecond'> 1 </timestep>
+        <efield>
+          <amp> [ 0 , 5e-2 , 0 ] </amp>                <!-- amplitude of the electric field pulse in a.u. -->
+          <freq units="thz"> 115 </freq>                <!-- frequency of the electric field pulse in THz -->
+          <peak units="picosecond"> 0.2 </peak>         <!-- peak of the electric field pulse -->
+          <sigma units="picosecond"> 0.05 </sigma>      <!-- width (std of gaussian envelope) -->
+        </efield>
+        <bec mode="driver"/>                            <!-- BECs computed by LOREM, not read from file -->
+      </driven_dynamics>
+    </motion>
+
+  </system>
+</simulation>

examples/md-ipi/run.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+# Driven dynamics with LOREM BEC model via i-PI.
+#
+# Prerequisites:
+#   - i-PI installed (`pip install ipi`)
+#   - LOREM driver installed (`lorem-install-ipi-driver`)
+#   - Trained LoremBEC checkpoint
+#
+# Adapt MODEL_PATH and start.xyz / input.xml to your system.
+
+MODEL_PATH="/path/to/checkpoint"
+
+# Install LOREM driver into i-PI (idempotent)
+lorem-install-ipi-driver
+
+# Start i-PI server
+i-pi input.xml > i-pi.out &
+echo "i-PI started"
+
+# Wait for socket to be ready
+sleep 5
+
+# Start LOREM driver
+i-pi-driver -a lorem -u -m lorem \
+    -o model_path=${MODEL_PATH},template=start.xyz \
+    > driver.out &
+echo "LOREM driver started"
+
+wait
+echo "Simulation complete"

examples/md-ipi/start.xyz

@@ -0,0 +1,5 @@
+3
+positions in atomic_units (ground state water, PES from pswater)
+O  3.77946e+00  3.41870e+00  3.77945e+00
+H  3.77945e+00  4.52675e+00  5.21072e+00
+H  3.77945e+00  4.52675e+00  2.34819e+00

src/lorem/calculator.py

@@ -126,6 +126,15 @@ def calculate(
             elif key == "stress":
                 raise KeyError
 
+        # BEC passthrough: when model outputs "apt" (e.g. LoremBEC),
+        # expose as "BEC" in (3*natoms, 3) layout for i-PI compatibility
+        if "apt" in results:
+            apt = np.array(
+                results["apt"][self.batch.sr.atom_mask].reshape(-1, 3, 3),
+                dtype=np.float32,
+            )
+            actual_results["BEC"] = apt.reshape(-1, 3)
+
         if self.add_offset:
             energy_offset = np.sum(
                 [self.species_weights[Z] for Z in atoms.get_atomic_numbers()]
@@ -136,7 +145,7 @@ def calculate(
         return actual_results
 
     def get_property(self, name, atoms=None, allow_calculation=True):
-        if name not in self.implemented_properties:
+        if name not in self.implemented_properties and name != "BEC":
             raise PropertyNotImplementedError(f"{name} property not implemented")
 
         self.update(atoms)

src/lorem/dynamics/ase.py

@@ -1,4 +1,7 @@
-class LOREMCalculator:
-    @classmethod
-    def from_checkpoint(cls, model_path):
-        pass
+from lorem.calculator import Calculator
+
+
+class LOREMCalculator(Calculator):
+    """LOREM ASE calculator with BEC support for i-PI integration."""
+
+    pass

src/lorem/dynamics/ipi.py

@@ -21,11 +21,3 @@ def __init__(self, template, model_path, device="cpu", *args, **kwargs):
     def check_parameters(self):
         super().check_parameters()
         self.ase_calculator = LOREMCalculator.from_checkpoint(self.model_path)
-
-    def compute_structure(self, cell, pos):
-        pot_ipi, force_ipi, vir_ipi, extras = super().compute_structure(cell, pos)
-
-        if "BEC" not in extras and "apt" in extras:
-            extras["BEC"] = extras["apt"]
-
-        return pot_ipi, force_ipi, vir_ipi, extras

tests/test_calculator.py

@@ -0,0 +1,60 @@
+import numpy as np
+
+from ase.build import bulk
+
+from lorem.calculator import Calculator
+from lorem.models.bec import LoremBEC
+from lorem.models.mlip import Lorem
+
+
+def test_calculator_energy_forces():
+    model = Lorem(cutoff=5.0, num_features=8, num_spherical_features=2, num_radial=4)
+    calc = Calculator.from_model(model)
+
+    atoms = bulk("Ar") * [2, 2, 2]
+    calc.calculate(atoms)
+
+    assert "energy" in calc.results
+    assert "forces" in calc.results
+    assert calc.results["forces"].shape == (len(atoms), 3)
+    assert "BEC" not in calc.results  # Lorem (not LoremBEC) has no BEC
+
+
+def test_calculator_bec():
+    model = LoremBEC(cutoff=5.0, num_features=8, num_spherical_features=2, num_radial=4)
+    calc = Calculator.from_model(model)
+
+    atoms = bulk("Ar") * [2, 2, 2]
+    calc.calculate(atoms)
+
+    assert "energy" in calc.results
+    assert "forces" in calc.results
+    assert "BEC" in calc.results
+
+    natoms = len(atoms)
+    assert calc.results["BEC"].shape == (3 * natoms, 3)
+    assert calc.results["BEC"].dtype == np.float32
+
+
+def test_calculator_bec_get_property():
+    model = LoremBEC(cutoff=5.0, num_features=8, num_spherical_features=2, num_radial=4)
+    calc = Calculator.from_model(model)
+
+    atoms = bulk("Ar") * [2, 2, 2]
+    calc.calculate(atoms)
+
+    bec = calc.get_property("BEC", atoms)
+    assert bec.shape == (3 * len(atoms), 3)
+
+
+def test_lorem_calculator_import():
+    from lorem import LOREMCalculator
+
+    model = LoremBEC(cutoff=5.0, num_features=8, num_spherical_features=2, num_radial=4)
+    calc = LOREMCalculator.from_model(model)
+
+    atoms = bulk("Ar") * [2, 2, 2]
+    calc.calculate(atoms)
+
+    assert "BEC" in calc.results
+    assert calc.results["BEC"].shape == (3 * len(atoms), 3)