AIMNet2ASECalculator Fails on Multiple Geometries Due to Shape Mismatch

[IliasChair]

Hello,

I have encountered an issue where the AIMNet2ASE calculator cannot handle multiple geometries within a single calculator instance. When attempting to compute the potential energy for different molecular structures sequentially, the following error occurs:

Couldn't find configuration file. Expected it at '[...]/.pysisyphusrc'.
Found model file: [...]/AIMNet2/aimnet2calc/assets/aimnet2/aimnet2_wb97m_0.jpt
Traceback (most recent call last):
  File "[...]/tmp.py", line 11, in <module>
    CH4_e = calc.get_potential_energy(CH4)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/ase/ase/calculators/abc.py", line 24, in get_potential_energy
    return self.get_property(name, atoms)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/ase/ase/calculators/calculator.py", line 538, in get_property
    self.calculate(atoms, [name], system_changes)
  File "[...]/AIMNet2/aimnet2calc/aimnet2ase.py", line 69, in calculate
    results = self.base_calc({
              ^^^^^^^^^^^^^^^^
  File "[...]/AIMNet2/aimnet2calc/calculator.py", line 59, in __call__
    return self.eval(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/AIMNet2/aimnet2calc/calculator.py", line 79, in eval
    data = self.prepare_input(data)
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/AIMNet2/aimnet2calc/calculator.py", line 99, in prepare_input
    data = self.pad_input(data)
           ^^^^^^^^^^^^^^^^^^^^
  File "[...]/AIMNet2/aimnet2calc/calculator.py", line 171, in pad_input
    data['numbers'] = maybe_pad_dim0(data['numbers'], N)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/AIMNet2/aimnet2calc/calculator.py", line 241, in maybe_pad_dim0
    assert _shape_diff == 0 or _shape_diff == 1, 'Invalid shape'
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: Invalid shape

Steps to Reproduce

Below I have provided a minimal example to reproduce this error:

from ase.atoms import Atoms
from aimnet2calc import AIMNet2ASE

calc = AIMNet2ASE()

h2o = Atoms("H2O", positions=[(0, 0, 0), (0, 0, 1), (0, 0, -1)])
CH4 = Atoms("CH4", positions=[(0, 0, 0), (0, 0, 1), (0, 0, -1), (0, 0, -2), (0, 0, -3)])  # crashes here


h2o_e = calc.get_potential_energy(h2o)
CH4_e = calc.get_potential_energy(CH4)

print(h2o_e, CH4_e)

[IliasChair]

Related to the issue described above, I have discovered another concerning problem:

Issue 2

If the number of atoms stays the same, but the system composition changes, the calculator does not detect the change. Instead, it reuses the results from the first calculation, leading to incorrect outputs.

Example 1

from ase.atoms import Atoms
from aimnet2calc import AIMNet2ASE

calc = AIMNet2ASE()

H20 = Atoms("H2O", positions=[(0, 0, 0), (0, 0, 1), (0, 0, -1)])
CO2 = Atoms("CO2", positions=[(0, 0, 0), (0, 0, 1), (0, 0, -1)])

H2O_e = calc.get_potential_energy(H20)
CO2_e = calc.get_potential_energy(CO2)

print(H2O_e, CO2_e)

output: [-2074.40625758] [-2074.40625758]

E.g. the calculator did not notice the change from H20 to CO2. This type of silent failure would be especially hard for users to debug or even notice.

Example 2

However, if the element types remain the same and only the atomic positions change, the calculator correctly updates the energy:

from ase.atoms import Atoms
from aimnet2calc import AIMNet2ASE

calc = AIMNet2ASE()

H2O = Atoms("H2O", positions=[(0, 0, 0), (0, 0, 1), (0, 0, -1)])
H2O_2 = Atoms("H2O", positions=[(0, 0, 0), (0, 1, 1), (0, 0, 1)])

h2o_e = calc.get_potential_energy(H2O)
H2O_2_e = calc.get_potential_energy(H2O_2)

print(h2o_e, H2O_2_e)

out: [-2074.40625758] [-2080.84284107]

This means the calculator properly tracks coordinate updates but does not recognize system composition changes.