.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/load_from_neurochem.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_examples_load_from_neurochem.py>`
        to download the full example code

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_load_from_neurochem.py:


Construct Model From NeuroChem Files
====================================

This tutorial illustrates how to manually load model from `NeuroChem files`_.

.. _NeuroChem files:
    https://github.com/isayev/ASE_ANI/tree/master/ani_models

.. GENERATED FROM PYTHON SOURCE LINES 14-15

To begin with, let's first import the modules we will use:

.. GENERATED FROM PYTHON SOURCE LINES 15-21

.. code-block:: default

    import os
    import torch
    import torchani
    import ase









.. GENERATED FROM PYTHON SOURCE LINES 22-23

Now let's read constants from constant file and construct AEV computer.

.. GENERATED FROM PYTHON SOURCE LINES 23-31

.. code-block:: default

    try:
        path = os.path.dirname(os.path.realpath(__file__))
    except NameError:
        path = os.getcwd()
    const_file = os.path.join(path, '../torchani/resources/ani-1x_8x/rHCNO-5.2R_16-3.5A_a4-8.params')  # noqa: E501
    consts = torchani.neurochem.Constants(const_file)
    aev_computer = torchani.AEVComputer(**consts)








.. GENERATED FROM PYTHON SOURCE LINES 32-33

Now let's read self energies and construct energy shifter.

.. GENERATED FROM PYTHON SOURCE LINES 33-36

.. code-block:: default

    sae_file = os.path.join(path, '../torchani/resources/ani-1x_8x/sae_linfit.dat')  # noqa: E501
    energy_shifter = torchani.neurochem.load_sae(sae_file)








.. GENERATED FROM PYTHON SOURCE LINES 37-38

Now let's read a whole ensemble of models.

.. GENERATED FROM PYTHON SOURCE LINES 38-41

.. code-block:: default

    model_prefix = os.path.join(path, '../torchani/resources/ani-1x_8x/train')  # noqa: E501
    ensemble = torchani.neurochem.load_model_ensemble(consts.species, model_prefix, 8)  # noqa: E501








.. GENERATED FROM PYTHON SOURCE LINES 42-43

Or alternatively a single model.

.. GENERATED FROM PYTHON SOURCE LINES 43-46

.. code-block:: default

    model_dir = os.path.join(path, '../torchani/resources/ani-1x_8x/train0/networks')  # noqa: E501
    model = torchani.neurochem.load_model(consts.species, model_dir)








.. GENERATED FROM PYTHON SOURCE LINES 47-51

You can create the pipeline of computing energies:
(Coordinates) -[AEVComputer]-> (AEV) -[Neural Network]->
(Raw energies) -[EnergyShifter]-> (Final energies)
From using either the ensemble or a single model:

.. GENERATED FROM PYTHON SOURCE LINES 51-56

.. code-block:: default

    nnp1 = torchani.nn.Sequential(aev_computer, ensemble, energy_shifter)
    nnp2 = torchani.nn.Sequential(aev_computer, model, energy_shifter)
    print(nnp1)
    print(nnp2)





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Sequential(
      (0): AEVComputer()
      (1): Ensemble(
        (0-7): 8 x ANIModel(
          (H): Sequential(
            (0): Linear(in_features=384, out_features=160, bias=True)
            (1): CELU(alpha=0.1)
            (2): Linear(in_features=160, out_features=128, bias=True)
            (3): CELU(alpha=0.1)
            (4): Linear(in_features=128, out_features=96, bias=True)
            (5): CELU(alpha=0.1)
            (6): Linear(in_features=96, out_features=1, bias=True)
          )
          (C): Sequential(
            (0): Linear(in_features=384, out_features=144, bias=True)
            (1): CELU(alpha=0.1)
            (2): Linear(in_features=144, out_features=112, bias=True)
            (3): CELU(alpha=0.1)
            (4): Linear(in_features=112, out_features=96, bias=True)
            (5): CELU(alpha=0.1)
            (6): Linear(in_features=96, out_features=1, bias=True)
          )
          (N): Sequential(
            (0): Linear(in_features=384, out_features=128, bias=True)
            (1): CELU(alpha=0.1)
            (2): Linear(in_features=128, out_features=112, bias=True)
            (3): CELU(alpha=0.1)
            (4): Linear(in_features=112, out_features=96, bias=True)
            (5): CELU(alpha=0.1)
            (6): Linear(in_features=96, out_features=1, bias=True)
          )
          (O): Sequential(
            (0): Linear(in_features=384, out_features=128, bias=True)
            (1): CELU(alpha=0.1)
            (2): Linear(in_features=128, out_features=112, bias=True)
            (3): CELU(alpha=0.1)
            (4): Linear(in_features=112, out_features=96, bias=True)
            (5): CELU(alpha=0.1)
            (6): Linear(in_features=96, out_features=1, bias=True)
          )
        )
      )
      (2): EnergyShifter()
    )
    Sequential(
      (0): AEVComputer()
      (1): ANIModel(
        (H): Sequential(
          (0): Linear(in_features=384, out_features=160, bias=True)
          (1): CELU(alpha=0.1)
          (2): Linear(in_features=160, out_features=128, bias=True)
          (3): CELU(alpha=0.1)
          (4): Linear(in_features=128, out_features=96, bias=True)
          (5): CELU(alpha=0.1)
          (6): Linear(in_features=96, out_features=1, bias=True)
        )
        (C): Sequential(
          (0): Linear(in_features=384, out_features=144, bias=True)
          (1): CELU(alpha=0.1)
          (2): Linear(in_features=144, out_features=112, bias=True)
          (3): CELU(alpha=0.1)
          (4): Linear(in_features=112, out_features=96, bias=True)
          (5): CELU(alpha=0.1)
          (6): Linear(in_features=96, out_features=1, bias=True)
        )
        (N): Sequential(
          (0): Linear(in_features=384, out_features=128, bias=True)
          (1): CELU(alpha=0.1)
          (2): Linear(in_features=128, out_features=112, bias=True)
          (3): CELU(alpha=0.1)
          (4): Linear(in_features=112, out_features=96, bias=True)
          (5): CELU(alpha=0.1)
          (6): Linear(in_features=96, out_features=1, bias=True)
        )
        (O): Sequential(
          (0): Linear(in_features=384, out_features=128, bias=True)
          (1): CELU(alpha=0.1)
          (2): Linear(in_features=128, out_features=112, bias=True)
          (3): CELU(alpha=0.1)
          (4): Linear(in_features=112, out_features=96, bias=True)
          (5): CELU(alpha=0.1)
          (6): Linear(in_features=96, out_features=1, bias=True)
        )
      )
      (2): EnergyShifter()
    )




.. GENERATED FROM PYTHON SOURCE LINES 57-58

You can also create an ASE calculator using the ensemble or single model:

.. GENERATED FROM PYTHON SOURCE LINES 58-63

.. code-block:: default

    calculator1 = torchani.ase.Calculator(consts.species, nnp1)
    calculator2 = torchani.ase.Calculator(consts.species, nnp2)
    print(calculator1)
    print(calculator1)





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    <torchani.ase.Calculator object at 0x7f28a943c5b0>
    <torchani.ase.Calculator object at 0x7f28a943c5b0>




.. GENERATED FROM PYTHON SOURCE LINES 64-65

Now let's define a methane molecule

.. GENERATED FROM PYTHON SOURCE LINES 65-74

.. code-block:: default

    coordinates = torch.tensor([[[0.03192167, 0.00638559, 0.01301679],
                                 [-0.83140486, 0.39370209, -0.26395324],
                                 [-0.66518241, -0.84461308, 0.20759389],
                                 [0.45554739, 0.54289633, 0.81170881],
                                 [0.66091919, -0.16799635, -0.91037834]]],
                               requires_grad=True)
    species = consts.species_to_tensor(['C', 'H', 'H', 'H', 'H']).unsqueeze(0)
    methane = ase.Atoms(['C', 'H', 'H', 'H', 'H'], positions=coordinates.squeeze().detach().numpy())








.. GENERATED FROM PYTHON SOURCE LINES 75-76

Now let's compute energies using the ensemble directly:

.. GENERATED FROM PYTHON SOURCE LINES 76-82

.. code-block:: default

    energy = nnp1((species, coordinates)).energies
    derivative = torch.autograd.grad(energy.sum(), coordinates)[0]
    force = -derivative
    print('Energy:', energy.item())
    print('Force:', force.squeeze())





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Energy: -40.45902210758762
    Force: tensor([[ 0.031, -0.132, -0.053],
            [-0.129,  0.164, -0.077],
            [ 0.086, -0.043,  0.041],
            [ 0.027,  0.006,  0.038],
            [-0.014,  0.005,  0.051]])




.. GENERATED FROM PYTHON SOURCE LINES 83-84

And using the ASE interface of the ensemble:

.. GENERATED FROM PYTHON SOURCE LINES 84-88

.. code-block:: default

    methane.set_calculator(calculator1)
    print('Energy:', methane.get_potential_energy() / ase.units.Hartree)
    print('Force:', methane.get_forces() / ase.units.Hartree)





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Energy: -40.45902210758762
    Force: [[ 0.03062528 -0.13160461 -0.05265211]
     [-0.12927507  0.16388635 -0.07736803]
     [ 0.08563145 -0.04288925  0.04082095]
     [ 0.02681219  0.00601404  0.03809873]
     [-0.01379385  0.00459347  0.05110047]]




.. GENERATED FROM PYTHON SOURCE LINES 89-90

We can do the same thing with the single model:

.. GENERATED FROM PYTHON SOURCE LINES 90-99

.. code-block:: default

    energy = nnp2((species, coordinates)).energies
    derivative = torch.autograd.grad(energy.sum(), coordinates)[0]
    force = -derivative
    print('Energy:', energy.item())
    print('Force:', force.squeeze())

    methane.set_calculator(calculator2)
    print('Energy:', methane.get_potential_energy() / ase.units.Hartree)
    print('Force:', methane.get_forces() / ase.units.Hartree)




.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Energy: -40.46280037151413
    Force: tensor([[ 0.056, -0.127, -0.054],
            [-0.140,  0.155, -0.075],
            [ 0.075, -0.037,  0.039],
            [ 0.024,  0.002,  0.033],
            [-0.016,  0.007,  0.057]])
    Energy: -40.46280037151413
    Force: [[ 0.05614968 -0.12697159 -0.05413406]
     [-0.1400753   0.15523373 -0.07533834]
     [ 0.0753209  -0.0374411   0.03949669]
     [ 0.02423098  0.00235087  0.03344228]
     [-0.01562626  0.00682809  0.05653343]]





.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 2.026 seconds)


.. _sphx_glr_download_examples_load_from_neurochem.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example




    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: load_from_neurochem.py <load_from_neurochem.py>`

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: load_from_neurochem.ipynb <load_from_neurochem.ipynb>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_