# [BUG] openmmml 1.8: AIMNet2 on a particle subset (ML/MM) + PBC → `AssertionError: Invalid shape` (1.7 works)

[Kamecs]

Dear OpenMM-ML Developer

Thank you for your providing nice code.
Following error was observed. Please confirm.

• openmmml 1.8 only. Registering an AIMNet2 force on a subset of atoms
  (ML/MM mechanical embedding, e.g. ligand-only) and running with periodic
  boundary conditions crashes at the first ML/MM step:

  AssertionError: Invalid shape      # aimnet.calculators.calculator.maybe_pad_dim0
  

• 1.7 does not have this bug. The identical system runs fine on a 1.7-pinned
  container.

• Regression: 1.8 removed 1.7's per-force position slicing / force
  zero-padding and added force.setParticles(atoms) — see Root cause.

• Current workaround: pin openmmml==1.7 for partial-atom ML/MM use.

Root cause

In 1.8, openmmml/models/aimnet2potential.py, _computeAIMNet2:

def _computeAIMNet2(state, model, numbers, charge, multiplicity, periodic):
    positions = torch.tensor(
        state.getPositions(asNumpy=True).value_in_unit(unit.angstrom),   # ← ALL particles
        dtype=torch.float32, device=numbers.device)
    args = {'coord': positions.unsqueeze(0), 'numbers': numbers, ...}

• state.getPositions() returns positions of all system particles (e.g. 40426)
• numbers is built in addForces from the atoms passed to the force →
  shape (1, 32) for a ligand-only force
• In aimnet's PBC path, coord (40426, 3) → make_nbmat → nbmat N=40427, then
  pad_input/maybe_pad_dim0 tries to pad numbers (32,) to N=40427:
  diff = 40395 → AssertionError: Invalid shape

1.7 avoided this: its _computeAIMNet2 sliced positions by the force's particle
indices and zero-padded the returned forces back to the full system size (no
setParticles). 1.8 removed the slicing/padding and added
force.setParticles(atoms), so the full-system positions no longer line up with
the ligand-only numbers.

Full-system ML (all atoms in one force) is unaffected — len(atoms) ==
numParticles masks the bug, which is likely why it went unnoticed.

Minimal reproduction

# System: protein + ligand + water/ions, 40426 particles, periodic box
# (octahedron / triclinic; any PBC works)
from openmmml import MLPotential

ligand_atoms = [a for a in top.atoms() if a.residue.name == "MOL"]   # 32 atoms
potential = MLPotential("aimnet2")
mm_system = build_mm_system(...)                                     # standard MM forces
potential.createMixedSystem(top, mm_system, ligand_atoms,
                            removeConstraints=False)
# run a few steps with the periodic box:
# → after "Switching to DSF Coulomb for PBC":
#   AssertionError: Invalid shape   (aimnet maybe_pad_dim0)

Suggested fix

_computeAIMNet2 has no record of which particles belong to the force (OpenMM
State does not index per-force). Two options:

Option 1 (minimal): pass the particle indices in at addForces time and
slice positions:

def _computeAIMNet2(state, model, numbers, charge, multiplicity, periodic,
                    particle_indices=None):
    all_pos = state.getPositions(asNumpy=True).value_in_unit(unit.angstrom)
    if particle_indices is not None:
        all_pos = all_pos[particle_indices]          # force's particles only
    positions = torch.tensor(all_pos, dtype=torch.float32, device=numbers.device)
    ...   # model returns forces of shape (len(particle_indices), 3),
          # matching the force registered via setParticles(atoms)

# in addForces:
force.addForce(partial(_computeAIMNet2, model=model, numbers=numbers,
                       ..., particle_indices=[a.index for a in atoms]))

Option 2 (restore 1.7 behavior): keep the full-system computation and
zero-pad/extract forces exactly as 1.7 did.

Questions for maintainers

1. Was _computeAIMNet2 only ever exercised with full-system ML (all atoms in
   one force)? If so, the bug is masked by len(atoms) == numParticles.
2. Is force.setParticles(atoms) actually called on the 1.8 AIMNet2
   addForces path? A runtime patch implementing Option 1 returned forces of
   shape (32, 3) for the 32-atom force, but OpenMM rejected it with
   `PythonForce: The forces must be returned in a NumPy array of shape
   (# particles, 3) — implying the force is registered against all
   particles, not the 32. (A follow-up variant that zero-padded forces to the
   full system size segfaulted right after the first successful force
   computation — see Appendix.)
3. Do the other backends (torchmdnet, ani2x) share the full-system-positions
   issue?
4. Can the 1.7 slicing/padding behavior (or an equivalent for
   setParticles-registered partial forces) be restored in 1.8 / main?

Environment

	affected (1.8)	unaffected (1.7)
openmmml	1.8	1.7
OpenMM	8.6.1.dev	8.5.2.dev
aimnet	0.2.0	0.2.0
Platform	CUDA (RTX 2080 Ti)	CUDA (RTX 2080 Ti)

System: mechanical-embedding ML/MM, ligand-only AIMNet2 force (32 atoms) in a
40426-particle octahedron (triclinic) box, createMixedSystem(...,
removeConstraints=False).

---

Appendix: investigation history (2026-09-18/19)

Diagnostic output (instrumented 1.8 run, diag_shape.log)

[diag-compute] ENTER periodic=True pos=(40426, 3) numbers=(1, 32) charge=(1,) mult=(1,) box=(3, 3)
[diag-pad_input] #1 IN  nbmat=(40427, 69) {'coord': (40426, 3), 'numbers': (32,), ...}
[diag-pad] maybe_pad_dim0 RAISED: a.shape=(32,) N=40427 diff(N - a0)=40395

Runtime patch attempted on 1.8 (does NOT work — reference only)

A runtime patch wrapping MLPotential.addForces to slice
state.getPositions() by the force's particle indices (Option 1 above)
resolves the aimnet assert, but on 1.8 the simulation still does not run:
early versions hit the PythonForce shape rejection quoted in Question 2, and
the final version (sliced positions + zero-padded full-system forces,
shape/layout verified) crashes with a hard segfault immediately after the
first force computation succeeds (suspected OpenMM PythonForce C++ layer /
GIL interaction; unresolved).

import openmmml.models.aimnet2potential as _a2

_orig_addForces = _a2.MLPotential.addForces

def _wrapped_addForces(self, system, atoms, molecule):
    particle_indices = [atom.index for atom in atoms]

    def _wrapped(state, model, numbers, charge, multiplicity, periodic):
        import torch
        all_pos = state.getPositions(asNumpy=True).value_in_unit(unit.angstrom)
        positions = torch.tensor(all_pos[particle_indices], dtype=torch.float32,
                                 device=numbers.device)
        args = {'coord': positions.unsqueeze(0), 'numbers': numbers,
                'charge': charge, 'mult': multiplicity}
        if periodic:
            cell = torch.tensor(
                state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(unit.angstrom),
                dtype=torch.float32, device=numbers.device)
            args['cell'] = cell
        result = model(args, forces=True)
        energyScale = (unit.ev / unit.item).conversion_factor_to(unit.kilojoules_per_mole)
        energy = float(energyScale * result["energy"].sum().detach())
        forces = (10.0 * energyScale * result["forces"]).detach().cpu().numpy()[0]
        return energy, forces

    _a2._computeAIMNet2 = _wrapped
    return _orig_addForces(self, system, atoms, molecule)

_a2.MLPotential.addForces = _wrapped_addForces

Resolution status

2026-09-19: confirmed 1.7 works end-to-end on GPU (passes
"Switching to DSF Coulomb for PBC", prod at 4 fs stable) → we pin
openmmml==1.7 for ML/MM use until an upstream fix lands.

Best regards.

[peastman]

Can you provide a test that reproduces the problem? The "minimal reproduction" is just a code fragment that doesn't run and doesn't demonstrate anything.

state.getPositions() returns positions of all system particles (e.g. 40426)

That is incorrect. If you call setParticles() on the force, the State passed to the Python code contains only the positions of the particles you specified.

The test case for a mixed system with AIMNet2 passes, which makes me skeptical that this bug actually exists.

See our AI policy. This report was obviously AI generated, but that wasn't disclosed.