Double Well potential using OpenMM

!pip install -q condacolab
import condacolab
condacolab.install()

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/condacolab.py:26
     25 try:
---> 26     import google.colab
     27 except ImportError:

ModuleNotFoundError: No module named 'google'

During handling of the above exception, another exception occurred:

RuntimeError                              Traceback (most recent call last)
Cell In[1], line 2
      1 get_ipython().system('pip install -q condacolab')
----> 2 import condacolab
      3 condacolab.install()

File /opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/condacolab.py:28
     26     import google.colab
     27 except ImportError:
---> 28     raise RuntimeError("This module must ONLY run as part of a Colab notebook!")
     31 __version__ = "0.1.6"
     32 __author__ = "Jaime Rodríguez-Guerra <jaimergp@users.noreply.github.com>"

RuntimeError: This module must ONLY run as part of a Colab notebook!

%%capture
!conda install -c conda-forge openmm mdtraj parmed
!pip install py3dmol 

import numpy as np
import matplotlib.pyplot as plt
import openmm as mm

def run_simulation(n_steps=10000):
    "Simulate a single particle in the double well"
    
    system = mm.System()
    system.addParticle(1)# added particle with a unit mass
    force = mm.CustomExternalForce('2*(x-1)^2*(x+1)^2 + y^2')# defines the potential
    force.addParticle(0, [])
    system.addForce(force)
    integrator = mm.LangevinIntegrator(500, 1, 0.02) # Langevin integrator with 500K temperature, gamma=1, step size = 0.02
    context = mm.Context(system, integrator)
    context.setPositions([[0, 0, 0]])
    context.setVelocitiesToTemperature(500)
    x = np.zeros((n_steps, 3))
    
    for i in range(n_steps):
        x[i] = context.getState(getPositions=True).getPositions(asNumpy=True)._value
        integrator.step(1)
        
    return x

trajectory = run_simulation(25000)

ylabels = ['x', 'y']

for i in range(2):
    plt.subplot(2, 1, i+1)
    plt.plot(trajectory[:, i])
    plt.ylabel(ylabels[i])
plt.xlabel('Simulation time')
plt.show()

plt.hist2d(trajectory[:, 0], trajectory[:, 1], bins=(25, 25), cmap=plt.cm.jet)
plt.show()