%matplotlib inline
import numpy as np
from matplotlib.pylab import *

try: import seaborn  #optional, makes plots look nicer
except ImportError: pass

import moldesign as mdt
from moldesign import units as u

"""
Explanation: <span style="float:right">
<a href="http://moldesign.bionano.autodesk.com/" target="_blank" title="About">About</a>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
<a href="https://forum.bionano.autodesk.com/c/Molecular-Design-Toolkit" target="_blank" title="Forum">Forum</a>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
<a href="https://github.com/autodesk/molecular-design-toolkit/issues" target="_blank" title="Issues">Issues</a>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
<a href="http://bionano.autodesk.com/MolecularDesignToolkit/explore.html" target="_blank" title="Tutorials">Tutorials</a>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
<a href="http://autodesk.github.io/molecular-design-toolkit/" target="_blank" title="Documentation">Documentation</a></span>
</span>

<br>
<center><h1>Example 2: Using MD sampling to calculate UV-Vis spectra</h1> </center>

This notebook uses basic quantum chemical calculations to calculate the absorption spectra of a small molecule.

Author: Aaron Virshup, Autodesk Research<br>
Created on: September 23, 2016
Tags: excited states, CASSCF, absorption, sampling
End of explanation
"""


qmmol = mdt.from_name('benzene')
qmmol.set_energy_model(mdt.models.CASSCF, active_electrons=6,
                       active_orbitals=6, state_average=6, basis='sto-3g')

properties = qmmol.calculate()

"""
Explanation: Contents


Single point
Sampling
Post-processing
Create spectrum

Single point
Let's start with calculating the vertical excitation energy and oscillator strengths at the ground state minimum (aka Franck-Condon) geometry.
Note that the active space and number of included states here is system-specific.
End of explanation
"""


for fstate in xrange(1, len(qmmol.properties.state_energies)):
    excitation_energy = properties.state_energies[fstate] - properties.state_energies[0]
    
    print '--- Transition from S0 to S%d ---' % fstate    
    print 'Excitation wavelength: %s' % excitation_energy.to('nm', 'spectroscopy')
    print 'Oscillator strength: %s' % qmmol.properties.oscillator_strengths[0,fstate]

"""
Explanation: This cell print a summary of the possible transitions. 
Note: you can convert excitation energies directly to nanometers using Pint by calling energy.to('nm', 'spectroscopy').
End of explanation
"""


mdmol = mdt.Molecule(qmmol)
mdmol.set_energy_model(mdt.models.GAFF)
mdmol.minimize()

mdmol.set_integrator(mdt.integrators.OpenMMLangevin, frame_interval=250*u.fs,
                     timestep=0.5*u.fs, constrain_hbonds=False, remove_rotation=True,
                     remove_translation=True, constrain_water=False)
mdtraj = mdmol.run(5.0 * u.ps)

"""
Explanation: Sampling
Of course, molecular spectra aren't just a set of discrete lines - they're broadened by several mechanisms. We'll treat vibrations here by sampling the molecule's motion on the ground state at 300 Kelvin.
To do this, we'll sample its geometries as it moves on the ground state by:
 1. Create a copy of the molecule
 2. Assign a forcefield (GAFF2/AM1-BCC)
 3. Run dynamics for 5 ps, taking a snapshot every 250 fs, for a total of 20 separate geometries.
End of explanation
"""


post_traj = mdt.Trajectory(qmmol)
for frame in mdtraj:
    qmmol.positions = frame.positions
    qmmol.calculate()
    post_traj.new_frame()

"""
Explanation: Post-processing
Next, we calculate the spectrum at each sampled geometry.
End of explanation
"""


wavelengths_to_state = []
oscillators_to_state = []
for i in xrange(1, len(qmmol.properties.state_energies)):
    wavelengths_to_state.append( (post_traj.state_energies[:,i] - post_traj.potential_energy).to('nm', 'spectroscopy'))
    oscillators_to_state.append([o[0,i] for o in post_traj.oscillator_strengths])

    
for istate, (w,o) in enumerate(zip(wavelengths_to_state, oscillators_to_state)):
    plot(w,o, label='S0 -> S%d'%(istate+1),
         marker='o', linestyle='none')
xlabel('wavelength / nm'); ylabel('oscillator strength'); legend()

"""
Explanation: This cell plots the results - wavelength vs. oscillator strength at each geometry for each transition:
End of explanation
"""


from itertools import chain
all_wavelengths = u.array(list(chain(*wavelengths_to_state)))
all_oscs = u.array(list(chain(*oscillators_to_state)))
hist(all_wavelengths, weights=all_oscs, bins=50)
xlabel('wavelength / nm')

"""
Explanation: Create spectrum
We're finally ready to calculate a spectrum - we'll create a histogram of all calculated transition wavelengths over all states, weighted by the oscillator strengths.
End of explanation
"""


