Second step fails

[moradza]

I am following steps torchANI tutorial steps. The code fails in the second call to simulation function. Any source of error. I couldn't install mamba in my workstation, and I used conda with pytorch=1.12, pytorch=1.11 throws symbol error.

import openmmtools

ala2 = openmmtools.testsystems.AlanineDipeptideVacuum(constraints=None)

while ala2.system.getNumForces() > 0:
  ala2.system.removeForce(0)

assert ala2.system.getNumConstraints() == 0
assert ala2.system.getNumForces() == 0

atomic_numbers = [atom.element.atomic_number for atom in ala2.topology.atoms()]

import torch as pt
from torchani.models import ANI2x
from NNPOps import OptimizedTorchANI

class NNP(pt.nn.Module):

  def __init__(self, atomic_numbers):

    super().__init__()

    # Store the atomic numbers
    self.atomic_numbers = pt.tensor(atomic_numbers).unsqueeze(0)

    # Create an ANI-2x model
    self.model = ANI2x(periodic_table_index=True)

    # Accelerate the model
    self.model = OptimizedTorchANI(self.model, self.atomic_numbers)

  def forward(self, positions):

    # Prepare the positions
    positions = positions.unsqueeze(0).float() * 10 # nm --> Å

    # Run ANI-2x
    result = self.model((self.atomic_numbers, positions))

    # Get the potential energy
    energy = result.energies[0] * 2625.5 # Hartree --> kJ/mol

    return energy

nnp = NNP(atomic_numbers)

pos = pt.tensor(ala2.positions.tolist())
energy_1 = float(nnp(pos))


from openmmtorch import TorchForce

pt.jit.script(nnp).save('model.pt')
force = TorchForce('model.pt')

ala2.system.addForce(force)
assert ala2.system.getNumForces() == 1


import sys
from openmm import LangevinMiddleIntegrator, BrownianIntegrator
from openmm.app import Simulation, StateDataReporter
from openmm.unit import kelvin, picosecond, femtosecond

temperature = 298.15 * kelvin
frictionCoeff = 1 / picosecond
timeStep = 1 * femtosecond
integrator = LangevinMiddleIntegrator(temperature, frictionCoeff, timeStep)

simulation = Simulation(ala2.topology, ala2.system, integrator)
simulation.context.setPositions(ala2.positions)

reporter = StateDataReporter(file=sys.stdout, reportInterval=1, step=True, time=True, potentialEnergy=True, temperature=True)
simulation.reporters.append(reporter)

from openmm.unit import kilojoule_per_mole
simulation.step(1000)

Warning: importing 'simtk.openmm' is deprecated.  Import 'openmm' instead.
#"Step","Time (ps)","Potential Energy (kJ/mole)","Temperature (K)"
1,0.001,-1301525.4216868656,4.132715143549229
Traceback (most recent call last):
  File "/home/amoradzadeh/Work/OpenMM-Torch/example-openmm-torch.py", line 83, in <module>
    simulation.step(1000)
  File "/home/amoradzadeh/anaconda3/envs/tmp-openmm-torch/lib/python3.9/site-packages/openmm/app/simulation.py", line 141, in step
    self._simulate(endStep=self.currentStep+steps)
  File "/home/amoradzadeh/anaconda3/envs/tmp-openmm-torch/lib/python3.9/site-packages/openmm/app/simulation.py", line 210, in _simulate
    self.integrator.step(stepsToGo)
  File "/home/amoradzadeh/anaconda3/envs/tmp-openmm-torch/lib/python3.9/site-packages/openmm/openmm.py", line 6145, in step
    return _openmm.LangevinMiddleIntegrator_step(self, steps)
openmm.OpenMMException: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "<string>", line 57, in <backward op>
            self_scalar_type = self.dtype
            def backward(grad_output):
                grad_self = AD_sum_backward(grad_output, self_size, dim, keepdim).to(self_scalar_type) / AD_safe_size(self_size, dim)
                            ~~~~~~~~~~~~~~~ <--- HERE
                return grad_self, None, None, None
  File "<string>", line 24, in AD_sum_backward
            if not keepdim and len(sizes) > 0:
                if len(dims) == 1:
                    return grad.unsqueeze(dims[0]).expand(sizes)
                           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
                else:
                    res = AD_unsqueeze_multiple(grad, dims, len(sizes))
RuntimeError: expand(CUDADoubleType{[1, 1]}, size=[1]): the number of sizes provided (1) must be greater or equal to the number of dimensions in the tensor (2)

[raimis]

Could you have the list of your conda environment?

[moradza]

Here, it is.

# packages in environment at /home/amoradzadeh/anaconda3/envs/openmm-torch:
#
# Name                    Version                   Build  Channel
_libgcc_mutex             0.1                 conda_forge    conda-forge
_openmp_mutex             4.5                       2_gnu    conda-forge
astunparse                1.6.3              pyhd8ed1ab_0    conda-forge
blosc                     1.21.1               h83bc5f7_3    conda-forge
brotlipy                  0.7.0           py38h0a891b7_1004    conda-forge
bzip2                     1.0.8                h7f98852_4    conda-forge
c-ares                    1.18.1               h7f98852_0    conda-forge
ca-certificates           2022.9.24            ha878542_0    conda-forge
certifi                   2022.9.24          pyhd8ed1ab_0    conda-forge
cffi                      1.15.1           py38h4a40e3a_0    conda-forge
cftime                    1.6.2            py38h26c90d9_0    conda-forge
charset-normalizer        2.1.1              pyhd8ed1ab_0    conda-forge
cryptography              37.0.1           py38h9ce1e76_0
cudatoolkit               11.7.0              hd8887f6_10    conda-forge
cudnn                     8.4.1.50             hed8a83a_0    conda-forge
curl                      7.85.0               h2283fc2_0    conda-forge
fftw                      3.3.10          nompi_hf0379b8_105    conda-forge
hdf4                      4.2.15               h9772cbc_4    conda-forge
hdf5                      1.12.2          nompi_h4df4325_100    conda-forge
idna                      3.4                pyhd8ed1ab_0    conda-forge
importlib-metadata        4.11.4           py38h578d9bd_0    conda-forge
intel-openmp              2022.1.0          h9e868ea_3769
jpeg                      9e                   h166bdaf_2    conda-forge
keyutils                  1.6.1                h166bdaf_0    conda-forge
krb5                      1.19.3               h08a2579_0    conda-forge
lark-parser               0.12.0             pyhd8ed1ab_0    conda-forge
ld_impl_linux-64          2.36.1               hea4e1c9_2    conda-forge
libblas                   3.9.0            16_linux64_mkl    conda-forge
libcblas                  3.9.0            16_linux64_mkl    conda-forge
libcurl                   7.85.0               h2283fc2_0    conda-forge
libedit                   3.1.20191231         he28a2e2_2    conda-forge
libev                     4.33                 h516909a_1    conda-forge
libffi                    3.4.2                h7f98852_5    conda-forge
libgcc-ng                 12.1.0              h8d9b700_16    conda-forge
libgfortran-ng            12.1.0              h69a702a_16    conda-forge
libgfortran5              12.1.0              hdcd56e2_16    conda-forge
libgomp                   12.1.0              h8d9b700_16    conda-forge
liblapack                 3.9.0            16_linux64_mkl    conda-forge
libllvm11                 11.1.0               he0ac6c6_4    conda-forge
libnetcdf                 4.8.1           nompi_h21705cb_104    conda-forge
libnghttp2                1.47.0               hff17c54_1    conda-forge
libnsl                    2.0.0                h7f98852_0    conda-forge
libprotobuf               3.20.1               h6239696_4    conda-forge
libsqlite                 3.39.4               h753d276_0    conda-forge
libssh2                   1.10.0               hf14f497_3    conda-forge
libstdcxx-ng              12.1.0              ha89aaad_16    conda-forge
libuuid                   2.32.1            h7f98852_1000    conda-forge
libzip                    1.9.2                hc929e4a_1    conda-forge
libzlib                   1.2.13               h166bdaf_4    conda-forge
llvmlite                  0.39.1           py38h38d86a4_0    conda-forge
lz4-c                     1.9.3                h9c3ff4c_1    conda-forge
lzo                       2.10              h516909a_1000    conda-forge
magma                     2.5.4                h6103c52_2    conda-forge
mdtraj                    1.9.7            py38h4b5026a_2    conda-forge
mkl                       2022.1.0           hc2b9512_224
mpiplus                   v0.0.1          pyhd8ed1ab_1003    conda-forge
nccl                      2.14.3.1             h0800d71_0    conda-forge
ncurses                   6.3                  h27087fc_1    conda-forge
netcdf4                   1.6.1           nompi_py38h2250339_100    conda-forge
ninja                     1.11.0               h924138e_0    conda-forge
nnpops                    0.2             cuda112py38hcdcc72c_5    conda-forge
nose                      1.3.7                   py_1006    conda-forge
numba                     0.56.3           py38h9a4aae9_0    conda-forge
numexpr                   2.8.3           mkl_py38h9924f95_0    conda-forge
numpy                     1.23.3           py38h3a7f9d9_0    conda-forge
ocl-icd                   2.3.1                h7f98852_0    conda-forge
ocl-icd-system            1.0.0                         1    conda-forge
openmm                    7.7.0            py38h4b4c9af_1    conda-forge
openmm-torch              0.8             cuda112py38hdf64721_2    conda-forge
openmmtools               0.21.5             pyhd8ed1ab_0    conda-forge
openssl                   3.0.5                h166bdaf_2    conda-forge
packaging                 21.3               pyhd8ed1ab_0    conda-forge
pandas                    1.5.0            py38h8f669ce_0    conda-forge
pdbfixer                  1.8.1              pyh6c4a22f_0    conda-forge
pip                       22.3               pyhd8ed1ab_0    conda-forge
pycparser                 2.21               pyhd8ed1ab_0    conda-forge
pymbar                    3.1.0            py38h71d37f0_0    conda-forge
pyopenssl                 22.0.0             pyhd8ed1ab_1    conda-forge
pyparsing                 3.0.9              pyhd8ed1ab_0    conda-forge
pysocks                   1.7.1              pyha2e5f31_6    conda-forge
pytables                  3.7.0            py38hf632491_2    conda-forge
python                    3.8.13          ha86cf86_0_cpython    conda-forge
python-dateutil           2.8.2              pyhd8ed1ab_0    conda-forge
python_abi                3.8                      2_cp38    conda-forge
pytorch                   1.11.0          cuda112py38habe9d5a_202    conda-forge
pytz                      2022.4             pyhd8ed1ab_0    conda-forge
pyyaml                    6.0              py38h0a891b7_4    conda-forge
readline                  8.1.2                h0f457ee_0    conda-forge
requests                  2.28.1             pyhd8ed1ab_1    conda-forge
scipy                     1.9.1            py38hea3f02b_0    conda-forge
setuptools                65.5.0             pyhd8ed1ab_0    conda-forge
six                       1.16.0             pyh6c4a22f_0    conda-forge
sleef                     3.5.1                h9b69904_2    conda-forge
snappy                    1.1.9                hbd366e4_1    conda-forge
sqlite                    3.39.4               h4ff8645_0    conda-forge
tk                        8.6.12               h27826a3_0    conda-forge
torchani                  2.2                pyh9f0ad1d_0    conda-forge
typing_extensions         4.4.0              pyha770c72_0    conda-forge
urllib3                   1.26.11            pyhd8ed1ab_0    conda-forge
wheel                     0.37.1             pyhd8ed1ab_0    conda-forge
xz                        5.2.6                h166bdaf_0    conda-forge
yaml                      0.2.5                h7f98852_2    conda-forge
zipp                      3.9.0              pyhd8ed1ab_0    conda-forge
zlib                      1.2.13               h166bdaf_4    conda-forge
zstd                      1.5.2                h6239696_4    conda-forge

[raimis]

I have created an environment:

conda env create mmh/openmm-8-beta-linux
conda activate openmm-8-beta-linux
conda install -c conda-forge openmmtools

The script above works with problem:

python debug_84.py 
/shared2/raimis/opt/miniconda/envs/openmm-8-beta-linux/lib/python3.10/site-packages/torchani/__init__.py:55: UserWarning: Dependency not satisfied, torchani.ase will not be available
  warnings.warn("Dependency not satisfied, torchani.ase will not be available")
/shared2/raimis/opt/miniconda/envs/openmm-8-beta-linux/lib/python3.10/site-packages/torchani/resources/
Warning: importing 'simtk.openmm' is deprecated.  Import 'openmm' instead.
#"Step","Time (ps)","Potential Energy (kJ/mole)","Temperature (K)"
1,0.001,-1301525.3878454717,4.122359749003946
2,0.002,-1301527.570361078,11.412210638851885
3,0.003,-1301529.5611348676,19.057213393688563
4,0.004,-1301531.0187404307,25.971037213133815
5,0.005,-1301531.831716344,30.366503161425385
...

@moradza try to create a new environment as indicated with the latest (22.9.0) conda.

[raimis]

@moradza this might be a conda issue (#88 (comment)). Could you try to install with mamba?

[moradza]

@raimis I updated conda and tried installation with new environment, but the same error persists.
I am trying mamba installation, but my conda fails to install mamba package. I need some time to figure out what is wrong with it.