isayevlab / aimnet2 Goto Github PK

Python 100.00%

aimnet2's Introduction

__ Update 6/10/24 __ We release new code, suaitable for large molecules and perioric calculations. Old code available in the old branch. Models were re-compiled and are not compatible with the new code.

AIMNet2 Calculator: Fast, Accurate Molecular Simulations

This package integrates the powerful AIMNet2 neural network potential into your simulation workflows. AIMNet2 provides fast and reliable energy, force, and property calculations for molecules containing a diverse range of elements.

Key Features:

Accurate and Versatile: AIMNet2 excels at modeling neutral, charged, organic, and elemental-organic systems.
Flexible Interfaces: Use AIMNet2 through convenient calculators for popular simulation packages like ASE and PySisyphus.
Flexible Long-Range Interactions: Optionally employ the Dumped-Shifted Force (DSF) or Ewald summation Coulomb models for accurate calculations in large or periodic systems.

Getting Started

1. Installation

While package is in alpha stage and repository is private, please install into your conda envoronment manually with

# install requirements
conda install -y pytorch pytorch-cuda=12.1 -c pytorch -c nvidia 
conda install -y -c pyg pytorch-cluster
conda install -y -c conda-forge openbabel ase
## pysis requirements
conda install -y -c conda-forge autograd dask distributed h5py fabric jinja2 joblib matplotlib numpy natsort psutil pyyaml rmsd scipy sympy scikit-learn
# now should not do any pip installs
pip install git+https://github.com/eljost/pysisyphus.git
# finally, this repo
git clone [email protected]:zubatyuk/aimnet2calc.git
cd aimnet2calc
python setup.py install

2. Available interfaces

ASE [https://wiki.fysik.dtu.dk/ase]

from aimnet2calc import AIMNet2ASE
calc = AIMNet2ASE('aimnet2')

To specify total molecular charge and spin multiplicity, use optional charge and mult keyword arguments, or set_charge and set_mult methods:

calc = AIMNet2ASE('aimnet2', charge=1)
atoms1.calc = calc
# calculations on atoms1 will be done with charge 1
....
atoms2.calc = calc
calc.set_charge(-2)
# calculations on atoms1 will be done with charge -2

PySisyphus [https://pysisyphus.readthedocs.io]

from aimnet2calc import AIMNet2PySis
calc = AIMNet2PySis('aimnet2')

This produces standard PySisyphus calculator.

Instead of Pysis command line utility, use aimnet2pysis. This registeres AIMNet2 calculator with PySisyphus. Example calc section for PySisyphus YAML files:

calc:
   type: aimnet              # use AIMNet2 calculator
   model: aimnet2_b973c      # use aimnet2_b973c_0.jpt model

3. Base calculator

from aimnet2calc import AIMNet2Calculator

Initialization

calc = AIMNet2Calculator('aimnet2')

will load default AIMNet2 model aimnet2_wb97m_0.jpt as defined at aimnet2calc/models.py . If file does not exist on the machine, it will be downloaded from aimnet-model-zoo repository.

calc = AIMNet2Calculator('/path/to_a/model.jpt')

will load model from the file.

Input structure

The calculator accepts a dictionary containig lists, numpy arrays, torch tensors, or anything that could be accepted by torch.as_tensor.

The input could be for a single molecule (dict keys and shapes):

coord: (B, N, 3)  # atomic coordinates in Angstrom
numbers (B, N)    # atomic numbers
charge (B,)       # molecular charge
mult (B,)         # spin multiplicity, optional

or for a concatenation of molecules:

coord: (N, 3)  # atomic coordinates in Angstrom
numbers (N,)    # atomic numbers
charge (B,)    # molecular charge
mult (B,)      # spin multiplicity, optional
mol_idx (N,)   # molecule index for each atom, should contain integers in increasing order, with (B-1) is the maximum number.

where B is the number of molecules, N is number of atoms.

Calling calculator

results = calc(data, forces=False, stress=False, hessian=False)

results would be a dictionary of PyTorch tensors containing energy, charges, and possibly forces, stress and hessian if requested.

4. Long range Coulomb model

By default, Coulomb energy is calculated in O(N^2) manner, e.g. pair interaction between every pair of atoms in system. For very large or periodic systems, O(N) Dumped-Shifted Force Coulomb model could be employed doi: 10.1063/1.2206581. With AIMNet2Calculator interface, switch between standard and DSF Coulomb implementations im AIMNet2 models:

# switch to O(N)
calc.set_lrcoulomb_method('dsf', cutoff=15.0, dsf_alpha=0.2)
# switch to O(N^2), not suitable for PBC
calc.set_lrcoulomb_method('simple')

aimnet2's People

Contributors

Stargazers

Watchers

Forkers

leelasd mirual liucmu pythonfz hoelzerc simon-boothroyd-genesis pablo-arantes jthorton okanders gaoshan2006 grimme-lab senppoa gomesgroup chrinide kangmg cationic jyothishjoy morrowchem

aimnet2's Issues

Energy Unit

Thank you for your open-source code. I'm a student and am not familiar with the DFT and ML force field. Could you tell me the unit of the potential energy of AIMNet2? Thank you very much.

Error loading the model on CPU

Hi,

When I try to run the model on Mac CPU, I am getting the following errors. Here is the command I used

python aimnet2_ase_opt.py ../models/aimnet2_wb97m-d3_0.jpt test.xyz out.xyz --charge 0 --traj trj.xyz

The error is as below

Loading AIMNet2 model from file ../models/aimnet2_wb97m-d3_0.jpt
Traceback (most recent call last):
  File "aimnet2_ase_opt.py", line 74, in <module>
    model = torch.jit.load(args.model, map_location=device)
  File "/opt/anaconda3/envs/delfta/lib/python3.7/site-packages/torch/jit/_serialization.py", line 161, in load
    cpp_module = torch._C.import_ir_module(cu, str(f), map_location, _extra_files)
RuntimeError: [enforce fail at inline_container.cc:222] . file not found: wb97m_gas_0/version

I am using PyTorch 1.8.0

Conda package

Hi, thanks for making these fantastic models open-source. Is there any interest in making AIMNet2 conda installable similar to torchani to make it easier to install and use? I am more than happy to help or do this.

License

Hi, I could not find any reference to the license of AIMNet2 models. Could you tell me what license you use?
Thanks in advance!

D3 correction not applied

I'm using the default aimnet2_b973c_0.jpt model and it seems like the D3 corrections are not being applied. What makes me think this is that I can get consistent energies/forces with b973c DFT data, but only if I subtract the D3 correction manually from the DFT data first.

Is there a setting I need to apply to activate these corrections that I'm missing?

RuntimeError - sparse_nb is empty

Hey,
thanks for the ongoing development of the AIMNet potentials!

While running a minimum example with the mol_single.xyz structure, I encountered a RuntimeError:

from aimnet2calc import AIMNet2ASE
import ase.io

atoms = ase.io.read(r'AIMNet2\test\mol_single.xyz')
calc = AIMNet2ASE('aimnet2')
atoms.calc = calc
atoms.get_potential_energy()

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
File c:\Users\Andreas\Documents\AIMNet2\test\my_test.py:9
      [7](file:///C:/Users/Andreas/Documents/AIMNet2/test/my_test.py:7) calc = AIMNet2ASE('aimnet2')
      [8](file:///C:/Users/Andreas/Documents/AIMNet2/test/my_test.py:8) atoms.calc = calc
----> [9](file:///C:/Users/Andreas/Documents/AIMNet2/test/my_test.py:9) atoms.get_potential_energy()

File c:\Users\Andreas\Anaconda3\envs\aimnet\lib\site-packages\ase\atoms.py:755, in Atoms.get_potential_energy(self, force_consistent, apply_constraint)
    [752](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/atoms.py:752)     energy = self._calc.get_potential_energy(
    [753](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/atoms.py:753)         self, force_consistent=force_consistent)
    [754](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/atoms.py:754) else:
--> [755](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/atoms.py:755)     energy = self._calc.get_potential_energy(self)
    [756](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/atoms.py:756) if apply_constraint:
    [757](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/atoms.py:757)     for constraint in self.constraints:

File c:\Users\Andreas\Anaconda3\envs\aimnet\lib\site-packages\ase\calculators\abc.py:24, in GetPropertiesMixin.get_potential_energy(self, atoms, force_consistent)
     [22](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/calculators/abc.py:22) else:
     [23](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/calculators/abc.py:23)     name = 'energy'
---> [24](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/calculators/abc.py:24) return self.get_property(name, atoms)

File c:\Users\Andreas\Anaconda3\envs\aimnet\lib\site-packages\ase\calculators\calculator.py:538, in BaseCalculator.get_property(self, name, atoms, allow_calculation)
    [535](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/calculators/calculator.py:535)     if self.use_cache:
    [536](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/calculators/calculator.py:536)         self.atoms = atoms.copy()
--> [538](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/calculators/calculator.py:538)     self.calculate(atoms, [name], system_changes)
    [540](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/calculators/calculator.py:540) if name not in self.results:
    [541](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/calculators/calculator.py:541)     # For some reason the calculator was not able to do what we want,
    [542](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/calculators/calculator.py:542)     # and that is OK.
    [543](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/calculators/calculator.py:543)     raise PropertyNotImplementedError(
    [544](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/calculators/calculator.py:544)         '{} not present in this ' 'calculation'.format(name)
    [545](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/ase/calculators/calculator.py:545)     )

File c:\Users\Andreas\Anaconda3\envs\aimnet\lib\site-packages\aimnet2calc\aimnet2ase.py:65, in AIMNet2ASE.calculate(self, atoms, properties, system_changes)
     [62](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/aimnet2ase.py:62) else:
     [63](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/aimnet2ase.py:63)     cell = None
---> [65](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/aimnet2ase.py:65) results = self.base_calc({
     [66](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/aimnet2ase.py:66)     'coord': torch.tensor(self.atoms.positions, dtype=torch.float32, device=self.base_calc.device),
     [67](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/aimnet2ase.py:67)     'numbers': self._t_numbers,
     [68](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/aimnet2ase.py:68)     'cell': cell,
     [69](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/aimnet2ase.py:69)     'mol_idx': self._t_mol_idx,
     [70](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/aimnet2ase.py:70)     'charge': self._t_charge,
     [71](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/aimnet2ase.py:71)     'mult': self._t_mult,
     [72](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/aimnet2ase.py:72) }, forces='forces' in properties, stress='stress' in properties)
     [73](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/aimnet2ase.py:73) for k, v in results.items():
     [74](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/aimnet2ase.py:74)     results[k] = v.detach().cpu().numpy()

File c:\Users\Andreas\Anaconda3\envs\aimnet\lib\site-packages\aimnet2calc\calculator.py:59, in AIMNet2Calculator.__call__(self, *args, **kwargs)
     [58](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/calculator.py:58) def __call__(self, *args, **kwargs):
---> [59](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/calculator.py:59)     return self.eval(*args, **kwargs)

File c:\Users\Andreas\Anaconda3\envs\aimnet\lib\site-packages\aimnet2calc\calculator.py:79, in AIMNet2Calculator.eval(self, data, forces, stress, hessian)
     [78](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/calculator.py:78) def eval(self, data: Dict[str, Any], forces=False, stress=False, hessian=False) -> Dict[str, Tensor]:
---> [79](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/calculator.py:79)     data = self.prepare_input(data)
     [80](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/calculator.py:80)     if hessian and data['mol_idx'][-1] > 0:
     [81](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/calculator.py:81)         raise NotImplementedError('Hessian calculation is not supported for multiple molecules')

File c:\Users\Andreas\Anaconda3\envs\aimnet\lib\site-packages\aimnet2calc\calculator.py:98, in AIMNet2Calculator.prepare_input(self, data)
     [96](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/calculator.py:96)         print('Switching to DSF Coulomb for PBC')
     [97](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/calculator.py:97)         self.set_lrcoulomb_method('dsf')
---> [98](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/calculator.py:98) data = self.make_nbmat(data)
     [99](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/calculator.py:99) data = self.pad_input(data)
    [100](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/calculator.py:100) return data

File c:\Users\Andreas\Anaconda3\envs\aimnet\lib\site-packages\aimnet2calc\calculator.py:161, in AIMNet2Calculator.make_nbmat(self, data)
    [159](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/calculator.py:159) else:
    [160](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/calculator.py:160)     if 'nbmat' not in data:
--> [161](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/calculator.py:161)         data['nbmat'] = nblist_torch_cluster(data['coord'], self.cutoff, data['mol_idx'], max_nb=128)
    [162](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/calculator.py:162)         if self.lr:
    [163](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/calculator.py:163)             if 'nbmat_lr' not in data:

File c:\Users\Andreas\Anaconda3\envs\aimnet\lib\site-packages\aimnet2calc\nblist.py:37, in nblist_torch_cluster(coord, cutoff, mol_idx, max_nb)
     [35](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/nblist.py:35) sparse_nb = radius_graph(coord, batch=mol_idx, r=cutoff, max_num_neighbors=max_nb).to(torch.int32)
     [36](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/nblist.py:36) print(sparse_nb)
---> [37](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/nblist.py:37) max_num_neighbors = torch.unique(sparse_nb[0], return_counts=True)[1].max().item()
     [38](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/nblist.py:38) if max_num_neighbors < max_nb:
     [39](file:///C:/Users/Andreas/Anaconda3/envs/aimnet/lib/site-packages/aimnet2calc/nblist.py:39)     break

RuntimeError: max(): Expected reduction dim to be specified for input.numel() == 0. Specify the reduction dim with the 'dim' argument.

The tensor, which is returned by the radius_graph() function in nblist.py line 35 is empty, resulting in the error.

Running the same code entirely on CPU works fine.
Any ideas what is causing this behavior? Or could this be a CUDA/GPU related problem of my setup?

Best,
Andreas

ps: I'm using the following packages:

pytorch                   2.3.1           py3.10_cuda12.1_cudnn8_0    pytorch
pytorch-cuda              12.1                 hde6ce7c_5    pytorch
torch-cluster             1.6.3+pt23cpu            pypi_0    pypi

Fine tuning AIMNet2

Hi, I'm interested in fine-tuning AIMNet2 on my particular system.

I want to make sure additional DFT reference calculations are consistent with those used to train the original model, but I can't find any training data labelled with your reference method to which I can compare.

Is there any chance you could include these labels in a small test set in this repo, if they're not publicly available already?

ASE calculator using cached energy even after atomic positions change

We have a workflow that uses an ASE calculator to explore a potential energy surface. When we swapped out the xTB ASE calculator that we were using to try the AIMNet2ASE calculator, we noticed that the calculator seemed to compute the energy correctly in the first invocation but then continue returning the same energy repeatedly even after updating the atomic positions. We dug around and it looks like AIMNet2ASE.set_atoms() is changing the atom positions but the cache doesn't get cleared so the next time the energy is requested it uses the value from the cache instead of recalculating at the new positions. We added a self.reset() call at the end of the AIMNet2ASE.do_reset() function. This at least superficially seems to resolve the issue such that the energy seems to be recalculated when the atomic positions change but use the cached value when called on the same atomic positions. Our workflow also seems to be operating as expected again. I'm not sure if this is a correct or general solution though.

Periodic Systems

First of all, this is a great model and I'm enjoying playing around with it.
I set up various different systems and it was able to run a geometry optimization for all of them.
Furthermore, I ran an MD simulation and e.g., in a Box of sulfuric acid and water I could see a proton transfer

I used the ASECalculator to run the MD simulations. The unwrapped simulation looks good, but when wrapping the particles back into the Box, some particles are too close together.
Here is an example for a system of water.

Is there a certain key that has to be set for periodic systems / are periodic systems possible?

conda installation failed

Hi,

Thanks for the fantastic project and repo!

I tried to follow the conda instructions and it found some incompatibilities doing this command conda install -y -c pyg pytorch-cluster

Could not solve for environment specs
The following packages are incompatible
├─ pin-1 is installable and it requires
│ └─ python 3.12.* , which can be installed;
└─ pytorch-cluster is not installable because there are no viable options
├─ pytorch-cluster 1.5.9 would require
│ └─ python >=3.6,<3.7.0a0 , which conflicts with any installable versions previously reported;
├─ pytorch-cluster [1.5.9|1.6.0|1.6.1] would require
│ └─ python >=3.7,<3.8.0a0 , which conflicts with any installable versions previously reported;
├─ pytorch-cluster [1.5.9|1.6.0|1.6.1|1.6.2|1.6.3] would require
│ └─ python >=3.8,<3.9.0a0 , which conflicts with any installable versions previously reported;
├─ pytorch-cluster [1.5.9|1.6.0|1.6.1|1.6.2|1.6.3] would require
│ └─ python >=3.9,<3.10.0a0 , which conflicts with any installable versions previously reported;
├─ pytorch-cluster [1.6.0|1.6.1|1.6.2|1.6.3] would require
│ └─ python >=3.10,<3.11.0a0 , which conflicts with any installable versions previously reported;
├─ pytorch-cluster [1.6.2|1.6.3] would require
│ └─ python >=3.11,<3.12.0a0 , which conflicts with any installable versions previously reported;
└─ pytorch-cluster 1.6.3 would require
└─ pytorch 2.2.* , which conflicts with any installable versions previously reported.

Do you already have a solution to fix this?

Thanks again!
Ciao
Ferri

MD calculations explode

Hi,

I tried using the AIMnet ASE calculator to run some MD calculations, but my molecule instantly explodes and the temperature is goes to thousands of kelvin. Does AIMnet not yet work for MD calculations or am I doing something wrong?

from aimnet2ase import AIMNet2Calculator
import torch
from ase.io import read
from ase import units
from ase.md.langevin import Langevin
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary, ZeroRotation

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

aimnet = torch.jit.load("aimnet2_b973c_ens.jpt", map_location=device)
calc = AIMNet2Calculator(aimnet)

atoms = read("Benzene.xyz")
print(atoms)
calc.do_reset()
calc.set_charge(0)
atoms.set_calculator(calc)

MaxwellBoltzmannDistribution(atoms, temperature_K=300)
Stationary(atoms)  # zero linear momentum
ZeroRotation(atoms)  # zero angular momentum

dyn = Langevin(atoms, timestep=5.0 * units.fs, temperature_K=300, friction=0.002,
    trajectory='md.traj', logfile='md.log')
dyn.run(2000)  # take 1000 steps

Questions about extending AIMnet2 to dimers or large molecules

Hi,

Thank you for your great AI molecular force field model. May I ask several questions about Aimnet2 (or more general, different AIFF approaches for organic molecules).

Because the training data set of AIFF usually contain single molecular conformations, when applying Aimnet2 to predict the binding energy between dimers, is there an out-of-distribution problem? Is there anything else we should pay attention to?
Is there anything we should pay attention to when applying AIFF (e.g., AIMnet2) to predict the energy of large molecules, such as peptides or proteins (suppose we don't care about computational cost)?

Bests,
Bin

AIMNet2 Performance near TS

I have tested a simple PES calculation for the Diels-Alder reaction using the following models:

AIMNet2
torchani models (ani-1x, ani-2x, ani-ccx, ani-xnr)
DFT calculations as the target theory for each ML potential

The results indicate that AIMNet2 shows good predictive performance even near the transition state (TS).

My question is, what might be the reason AIMNet2 predicts well near the TS compared to ANI models? Additionally, does AIMNet2 generally predict well near the TS?

In a past presentation about AIMNet2 video, a TS-targeted model called RxnAIMNet is mentioned. Could you please clarify if this is the same model as the current AIMNet2?

Sincerely,
Kangmg

RuntimeError when using new AIMNet2ASE calculator

I installed the new aimnet2 calculator mostly according to instructions (I had to do the first two lines of installation commands in one line during environment creation, otherwise I ran into problems solving the environment). When I now try the following example script:

from aimnet2calc import AIMNet2ASE
from ase import Atoms

calc = AIMNet2ASE('aimnet2')
atoms = Atoms('CO', positions=[(0, 0, 0), (0, 0, 1.1)])
atoms.calc = calc
atoms.get_potential_energy()

I'm running into a RuntimeError:

{
	"name": "RuntimeError",
	"message": "The following operation failed in the TorchScript interpreter.\nTraceback of TorchScript, serialized code (most recent call last):\n  File \"code/__torch__/aimnet/models/aimnet2.py\", line 70, in forward\n    data9 = (atomic_sum).forward(data8, )\n    data10 = (lrcoulomb).forward(data9, )\n    return (dftd3).forward(data10, )\n            ~~~~~~~~~~~~~~ <--- HERE\n  def prepare_data(self: __torch__.aimnet.models.aimnet2.AIMNet2,\n    data: Dict[str, Tensor]) -> Dict[str, Tensor]:\n  File \"code/__torch__/aimnet/modules.py\", line 290, in forward\n  def forward(self: __torch__.aimnet.modules.DFTD3,\n    data: Dict[str, Tensor]) -> Dict[str, Tensor]:\n    c6ij, d_ij, = (self)._calc_c6ij(data, )\n                   ~~~~~~~~~~~~~~~~ <--- HERE\n    r4r2 = self.r4r2\n    _70 = annotate(List[Optional[Tensor]], [data[\"numbers\"]])\n  File \"code/__torch__/aimnet/modules.py\", line 336, in _calc_c6ij\n    cnmax = self.cnmax\n    _84 = annotate(List[Optional[Tensor]], [numbers])\n    cn0 = torch.clamp(cn, None, torch.index(cnmax, _84))\n                                ~~~~~~~~~~~ <--- HERE\n    _85 = torch.unsqueeze(torch.unsqueeze(cn0, -1), -1)\n    cn_i = torch.unsqueeze(_85, -1)\n\nTraceback of TorchScript, original code (most recent call last):\n  File \"/home/roman/repo/AIMNET2NB/aimnet2nbmat/aimnet/models/aimnet2.py\", line 136, in forward\n    \n        for m in self.outputs.children():\n            data = m(data)\n                   ~ <--- HERE\n        return data\n  File \"/home/roman/repo/AIMNET2NB/aimnet2nbmat/aimnet/modules.py\", line 524, in forward\n    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:\n        c6ij, d_ij = self._calc_c6ij(data)\n                     ~~~~~~~~~~~~~~~ <--- HERE\n        rrij = 3 * ops.ij_product(self.r4r2[data['numbers']], data['nbmat_lr'])\n        r0ij = self.a1 * rrij.sqrt() + self.a2\n  File \"/home/roman/repo/AIMNET2NB/aimnet2nbmat/aimnet/modules.py\", line 508, in _calc_c6ij\n        rcov_ij = rcov_i.unsqueeze(-1) + rcov_j\n        cn = (1.0 / (1.0 + torch.exp(self.k1*((rcov_ij)/d_ij - 1.0)))).sum(dim=-1)\n        cn = torch.clamp(cn, max=self.cnmax[numbers])\n                                 ~~~~~~~~~~~~~~~~~~~ <--- HERE\n        cn_i = cn.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)\n        cn_j = ops.select_j(cn, data['nbmat_lr']).unsqueeze(-1).unsqueeze(-1)\nRuntimeError: tensors used as indices must be long, byte or bool tensors\n",
	"stack": "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m\n\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)\nCell \u001b[1;32mIn [6], line 4\u001b[0m\n\u001b[0;32m      2\u001b[0m atoms \u001b[39m=\u001b[39m Atoms(\u001b[39m'\u001b[39m\u001b[39mCO\u001b[39m\u001b[39m'\u001b[39m, positions\u001b[39m=\u001b[39m[(\u001b[39m0\u001b[39m, \u001b[39m0\u001b[39m, \u001b[39m0\u001b[39m), (\u001b[39m0\u001b[39m, \u001b[39m0\u001b[39m, \u001b[39m1.1\u001b[39m)])\n\u001b[0;32m      3\u001b[0m atoms\u001b[39m.\u001b[39mcalc \u001b[39m=\u001b[39m calc\n\u001b[1;32m----> 4\u001b[0m atoms\u001b[39m.\u001b[39;49mget_potential_energy()\n\nFile \u001b[1;32mc:\\ProgramData\\mambaforge\\envs\\aimnet2\\lib\\site-packages\\ase\\atoms.py:755\u001b[0m, in \u001b[0;36mAtoms.get_potential_energy\u001b[1;34m(self, force_consistent, apply_constraint)\u001b[0m\n\u001b[0;32m    752\u001b[0m     energy \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_calc\u001b[39m.\u001b[39mget_potential_energy(\n\u001b[0;32m    753\u001b[0m         \u001b[39mself\u001b[39m, force_consistent\u001b[39m=\u001b[39mforce_consistent)\n\u001b[0;32m    754\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m--> 755\u001b[0m     energy \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_calc\u001b[39m.\u001b[39;49mget_potential_energy(\u001b[39mself\u001b[39;49m)\n\u001b[0;32m    756\u001b[0m \u001b[39mif\u001b[39;00m apply_constraint:\n\u001b[0;32m    757\u001b[0m     \u001b[39mfor\u001b[39;00m constraint \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mconstraints:\n\nFile \u001b[1;32mc:\\ProgramData\\mambaforge\\envs\\aimnet2\\lib\\site-packages\\ase\\calculators\\abc.py:24\u001b[0m, in \u001b[0;36mGetPropertiesMixin.get_potential_energy\u001b[1;34m(self, atoms, force_consistent)\u001b[0m\n\u001b[0;32m     22\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m     23\u001b[0m     name \u001b[39m=\u001b[39m \u001b[39m'\u001b[39m\u001b[39menergy\u001b[39m\u001b[39m'\u001b[39m\n\u001b[1;32m---> 24\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mget_property(name, atoms)\n\nFile \u001b[1;32mc:\\ProgramData\\mambaforge\\envs\\aimnet2\\lib\\site-packages\\ase\\calculators\\calculator.py:538\u001b[0m, in \u001b[0;36mBaseCalculator.get_property\u001b[1;34m(self, name, atoms, allow_calculation)\u001b[0m\n\u001b[0;32m    535\u001b[0m     \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39muse_cache:\n\u001b[0;32m    536\u001b[0m         \u001b[39mself\u001b[39m\u001b[39m.\u001b[39matoms \u001b[39m=\u001b[39m atoms\u001b[39m.\u001b[39mcopy()\n\u001b[1;32m--> 538\u001b[0m     \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mcalculate(atoms, [name], system_changes)\n\u001b[0;32m    540\u001b[0m \u001b[39mif\u001b[39;00m name \u001b[39mnot\u001b[39;00m \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mresults:\n\u001b[0;32m    541\u001b[0m     \u001b[39m# For some reason the calculator was not able to do what we want,\u001b[39;00m\n\u001b[0;32m    542\u001b[0m     \u001b[39m# and that is OK.\u001b[39;00m\n\u001b[0;32m    543\u001b[0m     \u001b[39mraise\u001b[39;00m PropertyNotImplementedError(\n\u001b[0;32m    544\u001b[0m         \u001b[39m'\u001b[39m\u001b[39m{}\u001b[39;00m\u001b[39m not present in this \u001b[39m\u001b[39m'\u001b[39m \u001b[39m'\u001b[39m\u001b[39mcalculation\u001b[39m\u001b[39m'\u001b[39m\u001b[39m.\u001b[39mformat(name)\n\u001b[0;32m    545\u001b[0m     )\n\nFile \u001b[1;32mc:\\ProgramData\\mambaforge\\envs\\aimnet2\\lib\\site-packages\\aimnet2calc-0.0.1-py3.10.egg\\aimnet2calc\\aimnet2ase.py:65\u001b[0m, in \u001b[0;36mAIMNet2ASE.calculate\u001b[1;34m(self, atoms, properties, system_changes)\u001b[0m\n\u001b[0;32m     62\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m     63\u001b[0m     cell \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n\u001b[1;32m---> 65\u001b[0m results \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mbase_calc({\n\u001b[0;32m     66\u001b[0m     \u001b[39m'\u001b[39;49m\u001b[39mcoord\u001b[39;49m\u001b[39m'\u001b[39;49m: torch\u001b[39m.\u001b[39;49mtensor(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49matoms\u001b[39m.\u001b[39;49mpositions, dtype\u001b[39m=\u001b[39;49mtorch\u001b[39m.\u001b[39;49mfloat32, device\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mbase_calc\u001b[39m.\u001b[39;49mdevice),\n\u001b[0;32m     67\u001b[0m     \u001b[39m'\u001b[39;49m\u001b[39mnumbers\u001b[39;49m\u001b[39m'\u001b[39;49m: \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_t_numbers,\n\u001b[0;32m     68\u001b[0m     \u001b[39m'\u001b[39;49m\u001b[39mcell\u001b[39;49m\u001b[39m'\u001b[39;49m: cell,\n\u001b[0;32m     69\u001b[0m     \u001b[39m'\u001b[39;49m\u001b[39mmol_idx\u001b[39;49m\u001b[39m'\u001b[39;49m: \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_t_mol_idx,\n\u001b[0;32m     70\u001b[0m     \u001b[39m'\u001b[39;49m\u001b[39mcharge\u001b[39;49m\u001b[39m'\u001b[39;49m: \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_t_charge,\n\u001b[0;32m     71\u001b[0m     \u001b[39m'\u001b[39;49m\u001b[39mmult\u001b[39;49m\u001b[39m'\u001b[39;49m: \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_t_mult,\n\u001b[0;32m     72\u001b[0m }, forces\u001b[39m=\u001b[39;49m\u001b[39m'\u001b[39;49m\u001b[39mforces\u001b[39;49m\u001b[39m'\u001b[39;49m \u001b[39min\u001b[39;49;00m properties, stress\u001b[39m=\u001b[39;49m\u001b[39m'\u001b[39;49m\u001b[39mstress\u001b[39;49m\u001b[39m'\u001b[39;49m \u001b[39min\u001b[39;49;00m properties)\n\u001b[0;32m     73\u001b[0m \u001b[39mfor\u001b[39;00m k, v \u001b[39min\u001b[39;00m results\u001b[39m.\u001b[39mitems():\n\u001b[0;32m     74\u001b[0m     results[k] \u001b[39m=\u001b[39m v\u001b[39m.\u001b[39mdetach()\u001b[39m.\u001b[39mcpu()\u001b[39m.\u001b[39mnumpy()\n\nFile \u001b[1;32mc:\\ProgramData\\mambaforge\\envs\\aimnet2\\lib\\site-packages\\aimnet2calc-0.0.1-py3.10.egg\\aimnet2calc\\calculator.py:59\u001b[0m, in \u001b[0;36mAIMNet2Calculator.__call__\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m     58\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m__call__\u001b[39m(\u001b[39mself\u001b[39m, \u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs):\n\u001b[1;32m---> 59\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39meval(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\nFile \u001b[1;32mc:\\ProgramData\\mambaforge\\envs\\aimnet2\\lib\\site-packages\\aimnet2calc-0.0.1-py3.10.egg\\aimnet2calc\\calculator.py:84\u001b[0m, in \u001b[0;36mAIMNet2Calculator.eval\u001b[1;34m(self, data, forces, stress, hessian)\u001b[0m\n\u001b[0;32m     82\u001b[0m data \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mset_grad_tensors(data, forces\u001b[39m=\u001b[39mforces, stress\u001b[39m=\u001b[39mstress, hessian\u001b[39m=\u001b[39mhessian)\n\u001b[0;32m     83\u001b[0m \u001b[39mwith\u001b[39;00m torch\u001b[39m.\u001b[39mjit\u001b[39m.\u001b[39moptimized_execution(\u001b[39mFalse\u001b[39;00m):\n\u001b[1;32m---> 84\u001b[0m     data \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mmodel(data)\n\u001b[0;32m     85\u001b[0m data \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mget_derivatives(data, forces\u001b[39m=\u001b[39mforces, stress\u001b[39m=\u001b[39mstress, hessian\u001b[39m=\u001b[39mhessian)\n\u001b[0;32m     86\u001b[0m data \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mprocess_output(data)\n\nFile \u001b[1;32mc:\\ProgramData\\mambaforge\\envs\\aimnet2\\lib\\site-packages\\torch\\nn\\modules\\module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39m\u001b[39minput\u001b[39m, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[0;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n\n\u001b[1;31mRuntimeError\u001b[0m: The following operation failed in the TorchScript interpreter.\nTraceback of TorchScript, serialized code (most recent call last):\n  File \"code/__torch__/aimnet/models/aimnet2.py\", line 70, in forward\n    data9 = (atomic_sum).forward(data8, )\n    data10 = (lrcoulomb).forward(data9, )\n    return (dftd3).forward(data10, )\n            ~~~~~~~~~~~~~~ <--- HERE\n  def prepare_data(self: __torch__.aimnet.models.aimnet2.AIMNet2,\n    data: Dict[str, Tensor]) -> Dict[str, Tensor]:\n  File \"code/__torch__/aimnet/modules.py\", line 290, in forward\n  def forward(self: __torch__.aimnet.modules.DFTD3,\n    data: Dict[str, Tensor]) -> Dict[str, Tensor]:\n    c6ij, d_ij, = (self)._calc_c6ij(data, )\n                   ~~~~~~~~~~~~~~~~ <--- HERE\n    r4r2 = self.r4r2\n    _70 = annotate(List[Optional[Tensor]], [data[\"numbers\"]])\n  File \"code/__torch__/aimnet/modules.py\", line 336, in _calc_c6ij\n    cnmax = self.cnmax\n    _84 = annotate(List[Optional[Tensor]], [numbers])\n    cn0 = torch.clamp(cn, None, torch.index(cnmax, _84))\n                                ~~~~~~~~~~~ <--- HERE\n    _85 = torch.unsqueeze(torch.unsqueeze(cn0, -1), -1)\n    cn_i = torch.unsqueeze(_85, -1)\n\nTraceback of TorchScript, original code (most recent call last):\n  File \"/home/roman/repo/AIMNET2NB/aimnet2nbmat/aimnet/models/aimnet2.py\", line 136, in forward\n    \n        for m in self.outputs.children():\n            data = m(data)\n                   ~ <--- HERE\n        return data\n  File \"/home/roman/repo/AIMNET2NB/aimnet2nbmat/aimnet/modules.py\", line 524, in forward\n    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:\n        c6ij, d_ij = self._calc_c6ij(data)\n                     ~~~~~~~~~~~~~~~ <--- HERE\n        rrij = 3 * ops.ij_product(self.r4r2[data['numbers']], data['nbmat_lr'])\n        r0ij = self.a1 * rrij.sqrt() + self.a2\n  File \"/home/roman/repo/AIMNET2NB/aimnet2nbmat/aimnet/modules.py\", line 508, in _calc_c6ij\n        rcov_ij = rcov_i.unsqueeze(-1) + rcov_j\n        cn = (1.0 / (1.0 + torch.exp(self.k1*((rcov_ij)/d_ij - 1.0)))).sum(dim=-1)\n        cn = torch.clamp(cn, max=self.cnmax[numbers])\n                                 ~~~~~~~~~~~~~~~~~~~ <--- HERE\n        cn_i = cn.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)\n        cn_j = ops.select_j(cn, data['nbmat_lr']).unsqueeze(-1).unsqueeze(-1)\nRuntimeError: tensors used as indices must be long, byte or bool tensors\n"
}

Do you guys have any idea what causes this problem?

netrork > network

There's a typo in the README title which looks a bit funny.

GPU Memory Problem

Hi, thanks for the great contribution!

While trying to run the minimal example mol_single_opt.yml, I ran into the following memory issue while aimnet tries to build the nb list apparently:

Traceback (most recent call last):
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/bin/aimnet2pysis", line 33, in <module>
    sys.exit(load_entry_point('aimnet2calc', 'console_scripts', 'aimnet2pysis')())
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/Projects/aimnet2calc/aimnet2calc/aimnet2pysis.py", line 66, in run_pysis
    run.run()
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/pysisyphus/run.py", line 2022, in run
    run_result = run_from_dict(run_dict, **run_kwargs)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/pysisyphus/run.py", line 1966, in run_from_dict
    run_result = main(run_dict, restart, cwd, scheduler)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/pysisyphus/run.py", line 1489, in main
    opt_result = run_opt(
                 ^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/pysisyphus/drivers/opt.py", line 143, in run_opt
    opt.run()
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/pysisyphus/optimizers/Optimizer.py", line 725, in run
    step = self.optimize()
           ^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/pysisyphus/optimizers/RFOptimizer.py", line 74, in optimize
    energy, gradient, H, big_eigvals, big_eigvecs, resetted = self.housekeeping()
                                                              ^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/pysisyphus/optimizers/HessianOptimizer.py", line 457, in housekeeping
    gradient = self.geometry.gradient
               ^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/pysisyphus/Geometry.py", line 937, in gradient
    return -self.forces
            ^^^^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/pysisyphus/Geometry.py", line 903, in forces
    forces = self.cart_forces
             ^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/pysisyphus/Geometry.py", line 883, in cart_forces
    results = self.calculator.get_forces(self.atoms, self._coords)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/Projects/aimnet2calc/aimnet2calc/aimnet2pysis.py", line 49, in get_forces
    res = self.model(_in, forces=True)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/Projects/aimnet2calc/aimnet2calc/calculator.py", line 59, in __call__
    return self.eval(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/Projects/aimnet2calc/aimnet2calc/calculator.py", line 79, in eval
    data = self.prepare_input(data)
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/Projects/aimnet2calc/aimnet2calc/calculator.py", line 98, in prepare_input
    data = self.make_nbmat(data)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/Projects/aimnet2calc/aimnet2calc/calculator.py", line 161, in make_nbmat
    data['nbmat'] = nblist_torch_cluster(data['coord'], self.cutoff, data['mol_idx'], max_nb=128)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/Projects/aimnet2calc/aimnet2calc/nblist.py", line 35, in nblist_torch_cluster
    sparse_nb = radius_graph(coord, batch=mol_idx, r=cutoff, max_num_neighbors=max_nb).to(torch.int32)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/torch_cluster/radius.py", line 135, in radius_graph
    edge_index = radius(x, x, r, batch, batch,
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/torch_cluster/radius.py", line 82, in radius
    return torch.ops.torch_cluster.radius(x, y, ptr_x, ptr_y, r,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/torch/_ops.py", line 854, in __call__
    return self_._op(*args, **(kwargs or {}))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 22.00 GiB. GPU

I wouldn't have expected to run out of memory that quickly on an A100 40 GB device 😁 The earlier versions (I guess without these nb lists?) ran fine, also on smaller devices.

Do you have any advice on what to do to fix this? Run the nb list creation on CPU only?

Many thanks in advance!

Best,
Max

Dockerfile and requirements.txt for reproducibility

Hi,

I am opening this issue to add the requirements.txt and Dockerfile for containerizing this awesome repo :)

Docker build error

I cloned the repo, tried to build the docker, and got this error.

% docker build --pull --rm -f "Dockerfile_cpu" -t aimnet-box "."
[+] Building 2.7s (11/11) FINISHED                                                                                                         docker:desktop-linux
 => [internal] load .dockerignore                                                                                                                          0.0s
 => => transferring context: 2B                                                                                                                            0.0s
 => [internal] load build definition from Dockerfile_cpu                                                                                                   0.0s
 => => transferring dockerfile: 317B                                                                                                                       0.0s
 => [internal] load metadata for docker.io/condaforge/mambaforge:latest                                                                                    2.5s
 => CANCELED [1/7] FROM docker.io/condaforge/mambaforge:latest@sha256:615e43f72b49a6d46d3f93f7fadf7b7e22f6de934b8e1c46b01becb2816e8fc7                     0.0s
 => => resolve docker.io/condaforge/mambaforge:latest@sha256:615e43f72b49a6d46d3f93f7fadf7b7e22f6de934b8e1c46b01becb2816e8fc7                              0.0s
 => => sha256:615e43f72b49a6d46d3f93f7fadf7b7e22f6de934b8e1c46b01becb2816e8fc7 2.36kB / 2.36kB                                                             0.0s
 => => sha256:1461e0a1fa14431128dc95d921655fd6cd0b9147b4ec757c6d99e02776e82b47 676B / 676B                                                                 0.0s
 => [internal] load build context                                                                                                                          0.0s
 => => transferring context: 2B                                                                                                                            0.0s
 => CACHED [2/7] RUN mamba install -q -y -c conda-forge openbabel                                                                                          0.0s
 => ERROR [3/7] COPY requirements.txt .                                                                                                                    0.0s
 => CACHED [4/7] RUN pip install requirements.txt                                                                                                          0.0s
 => CACHED [5/7] WORKDIR /app                                                                                                                              0.0s
 => ERROR [6/7] COPY ../calculators/ calculators/                                                                                                          0.0s
 => ERROR [7/7] COPY ../models/ models/                                                                                                                    0.0s
------
 > [3/7] COPY requirements.txt .:
------
------
 > [6/7] COPY ../calculators/ calculators/:
------
------
 > [7/7] COPY ../models/ models/:
------
Dockerfile_cpu:7
--------------------
   5 |     WORKDIR /app
   6 |     COPY ../calculators/ calculators/
   7 | >>> COPY ../models/ models/
   8 |     ENTRYPOINT [ "python3" ,"/app/calculators/aimnet2_ase_opt.py"]
   9 |     
--------------------
ERROR: failed to solve: failed to compute cache key: failed to calculate checksum of ref bf78f52f-4126-42f4-a771-f8eba05fa16e::x0f2um24r0hbeap1ulgdu9sq0: "/models": not found

Running on MacOS Ventura 13.6 (22G120) with Docker 4.23.0 (120376)

Apple Silicon/ARM build

Hey folks, just got a new laptop and I'm struggling to get the .jpt files to run. Could this be because they're compiled for AMD64/x86, and would it be possible to get a compiled version for ARM/Apple Silicon in the future?

corin

Hessian computation fails

Hi again,

are there any special things to consider when requesting a Hessian?

I tried this with the mol_single_opt.yml example, requesting a Hessian computation after optimization:

calc:
  type: aimnet
  model: aimnet2_b973c

geom:
  type: dlc
  fn: mol_single.xyz

opt:
  type: rfo
  max_cycles: 50
  do_hess: True

After the optimization is done, I get the following traceback:

... started Hessian calculation
Traceback (most recent call last):
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/bin/aimnet2pysis", line 33, in <module>
    sys.exit(load_entry_point('aimnet2calc', 'console_scripts', 'aimnet2pysis')())
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/Projects/aimnet2calc/aimnet2calc/aimnet2pysis.py", line 66, in run_pysis
    run.run()
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/pysisyphus/run.py", line 2022, in run
    run_result = run_from_dict(run_dict, **run_kwargs)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/pysisyphus/run.py", line 1966, in run_from_dict
    run_result = main(run_dict, restart, cwd, scheduler)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/pysisyphus/run.py", line 1489, in main
    opt_result = run_opt(
                 ^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/pysisyphus/drivers/opt.py", line 213, in run_opt
    do_final_hessian(
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/pysisyphus/helpers.py", line 465, in do_final_hessian
    hessian = geom.cart_hessian
              ^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/pysisyphus/Geometry.py", line 959, in cart_hessian
    results = self.calculator.get_hessian(self.atoms, self._coords)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/Projects/aimnet2calc/aimnet2calc/aimnet2pysis.py", line 56, in get_hessian
    res = self.model(_in, forces=True, hessian=True)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/Projects/aimnet2calc/aimnet2calc/calculator.py", line 59, in __call__
    return self.eval(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/Projects/aimnet2calc/aimnet2calc/calculator.py", line 85, in eval
    data = self.get_derivatives(data, forces=forces, stress=stress, hessian=hessian)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/Projects/aimnet2calc/aimnet2calc/calculator.py", line 225, in get_derivatives
    data['hessian'] = self.calculate_hessian(data['forces'], self._saved_for_grad['coord'])
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/Projects/aimnet2calc/aimnet2calc/calculator.py", line 233, in calculate_hessian
    torch.autograd.grad(_f, coord, retain_graph=True)[0]
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/torch/autograd/__init__.py", line 412, in grad
    result = _engine_run_backward(
             ^^^^^^^^^^^^^^^^^^^^^
  File "/fs/home/cvsik/miniforge3/envs/aimnet2/lib/python3.12/site-packages/torch/autograd/graph.py", line 744, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
NotImplementedError: the derivative for '_cdist_backward' is not implemented.

Any idea how to get around the missing deriv. impl. in pytorch?
Thanks!

Struggling With Autograd

Hey all, congrats on a fantastic paper. I'm trying out some of the scripts included here, but I'm struggling to get autograd to work (for gradient/Hessian)...

Here is my minimal reproducible example, using water as a test case:

import torch
import numpy as np

model_name = "models/aimnet2_wb97m-d3_ens.jpt"
#model_name = "models/aimnet2_b973c_ens.jpt"

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = torch.jit.load(model_name, map_location=device)

input_atomic_numbers = [8, 1, 1]
input_coord = [
    [ 0.0000,  0.0000,  0.1178],
    [ 0.0000,  0.7555, -0.4712],
    [ 0.0000, -0.7555, -0.4712],
]
input_charge = 0

numbers = torch.as_tensor([input_atomic_numbers], device=device)
coord = torch.as_tensor(input_coord, dtype=torch.float, device=device).view(1, numbers.shape[1], 3)
charge = torch.as_tensor([input_charge], dtype=torch.float, device=device)
_in = dict(coord=coord, numbers=numbers, charge=charge)

with torch.jit.optimized_execution(False):
    _in['coord'].requires_grad_(True)
    _out = model(_in)
    e = _out['energy']
    f = torch.autograd.grad(e, _in['coord'])[0]

print(e)
print(f)

Energy works fine, but I get the following error for torch.autograd.grad():

Traceback (most recent call last):
  File "/Users/cwagen/.../example.py", line 27, in <module>
    f = torch.autograd.grad(e, _in['coord'])[0]
  File "/opt/miniconda3/envs/aimnet2_pytorch200/lib/python3.10/site-packages/torch/autograd/__init__.py", line 303, in grad
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

This is confusing, because when I print _in before the energy call, it seems to require grad...

{'coord': tensor([[[ 0.0000,  0.0000,  0.1178],
         [ 0.0000,  0.7555, -0.4712],
         [ 0.0000, -0.7555, -0.4712]]], requires_grad=True), 'numbers': tensor([[8, 1, 1]]), 'charge': tensor([0.])}

Here's the result of conda list - I also tried the latest version of Pytorch w/ Python 3.11 and I got the same error:

# packages in environment at /opt/miniconda3/envs/aimnet2_pytorch200:
#
# Name                    Version                   Build  Channel
bzip2                     1.0.8                h10d778d_5    conda-forge
ca-certificates           2023.11.17           h8857fd0_0    conda-forge
filelock                  3.13.1             pyhd8ed1ab_0    conda-forge
gmp                       6.3.0                h93d8f39_0    conda-forge
gmpy2                     2.1.2           py310hb691cb2_1    conda-forge
icu                       73.2                 hf5e326d_0    conda-forge
jinja2                    3.1.3              pyhd8ed1ab_0    conda-forge
libabseil                 20230802.1      cxx17_h048a20a_0    conda-forge
libblas                   3.9.0           20_osx64_openblas    conda-forge
libcblas                  3.9.0           20_osx64_openblas    conda-forge
libcxx                    16.0.6               hd57cbcb_0    conda-forge
libffi                    3.4.2                h0d85af4_5    conda-forge
libgfortran               5.0.0           13_2_0_h97931a8_1    conda-forge
libgfortran5              13.2.0               h2873a65_1    conda-forge
libhwloc                  2.9.3           default_h24e0189_1009    conda-forge
libiconv                  1.17                 hd75f5a5_2    conda-forge
liblapack                 3.9.0           20_osx64_openblas    conda-forge
libopenblas               0.3.25          openmp_hfef2a42_0    conda-forge
libprotobuf               4.24.4               h0ee05dc_0    conda-forge
libsqlite                 3.44.2               h92b6c6a_0    conda-forge
libuv                     1.46.0               h0c2f820_0    conda-forge
libxml2                   2.11.6               hc0ae0f7_0    conda-forge
libzlib                   1.2.13               h8a1eda9_5    conda-forge
llvm-openmp               17.0.6               hb6ac08f_0    conda-forge
markupsafe                2.1.3           py310h6729b98_1    conda-forge
mkl                       2022.2.1         h44ed08c_16952    conda-forge
mpc                       1.3.1                h81bd1dd_0    conda-forge
mpfr                      4.2.1                h0c69b56_0    conda-forge
mpmath                    1.3.0              pyhd8ed1ab_0    conda-forge
ncurses                   6.4                  h93d8f39_2    conda-forge
networkx                  3.2.1              pyhd8ed1ab_0    conda-forge
numpy                     1.26.3          py310h4bfa8fc_0    conda-forge
openssl                   3.2.0                hd75f5a5_1    conda-forge
pip                       23.3.2             pyhd8ed1ab_0    conda-forge
python                    3.10.13         h00d2728_1_cpython    conda-forge
python_abi                3.10                    4_cp310    conda-forge
pytorch                   2.0.0           cpu_mkl_py310hed029b9_104    conda-forge
readline                  8.2                  h9e318b2_1    conda-forge
scipy                     1.11.4          py310h3f1db6d_0    conda-forge
setuptools                69.0.3             pyhd8ed1ab_0    conda-forge
sleef                     3.5.1                h6db0672_2    conda-forge
sympy                     1.12            pypyh9d50eac_103    conda-forge
tbb                       2021.11.0            he51d815_0    conda-forge
tk                        8.6.13               h1abcd95_1    conda-forge
typing_extensions         4.9.0              pyha770c72_0    conda-forge
tzdata                    2023d                h0c530f3_0    conda-forge
wheel                     0.42.0             pyhd8ed1ab_0    conda-forge
xz                        5.2.6                h775f41a_0    conda-forge

I'm new to PyTorch, so let me know what I'm doing wrong!

Malformatted traj.xyz

The traj.xyz constructed in the example docker run -it --rm -v $(pwd):/app/ aimnet-box models/aimnet2_wb97m-d3_ens.jpt input.sdf output.sdf --charge 0 --traj traj.xyz is malformatted, i.e. not even human-readable.

Would expect a xyz-file format.