#!/usr/bin/env python
# coding: utf-8

# # Find jump sites from Trajectory
# 
# This notebook shows how to find the jump sites from the trajectory directly.
# 
# The idea is the diffusing species naturally cluster on (meta-)stable sites in the simulation. These sites can be used as the jump sites to calculate transitions. 
# 
# The procedure is as follows: 
# 
# 1. Load trajectory
# 2. Convert trajectory into density volume
# 3. Find peaks in volume
# 4. Convert peaks into pymatgen structure
# 
# First, we load the data from a VASP simulation.

# In[1]:


from gemdat import Trajectory
from gemdat.utils import VASPRUN

# Use your own data:
# VASPRUN = 'path/to/your/vasprun.xml'

trajectory = Trajectory.from_vasprun(VASPRUN)


# Create a new trajectory object that only consists of the diffusing species, Lithium.

# In[2]:


# Create a new trajectory object that only consists of Lithium
diff_trajectory = trajectory.filter('Li')
diff_trajectory


# ### Converting a Trajectory to a Volume
# 
# To convert a `Trajectory` to a `Volume` object we use `Trajectory.to_volume()` method provided by gemdat. A volume in this context is a discretized grid version of the trajectory. Each grid voxel acts as a bin, and each coordinate in the trajectory is assigned to a bin. The number of of points in each bin is counted to create the density volume.
# 
# The resolution is the minimum size of the voxels in Ångstrom. 
# 
# Note that a voxel does not necessarily represent a cubic or even orthorhombic volume, and depends on the lattice parameters of the trajectory.

# In[3]:


vol = diff_trajectory.to_volume(resolution=0.3)


# ### Converting a Volume to Sites
# 
# We can directly convert a volume to a structure with the  `Volume.to_structure()` method. This uses a [watershed](https://scikit-image.org/docs/stable/auto_examples/segmentation/plot_watershed.html) algorithm to find the sites.
# 
# The `background_level` sets the sensitivity of the peak finder. Set it to a higher value if you get too many spurious peaks. Sometimes the generated sites need to be cleaned up a bit manually.

# In[4]:


sites = vol.to_structure(specie='Li', background_level=0.1)
sites[0:10]


# Verify the generated sites by plotting them on the volume.

# In[5]:


vol.plot_3d(structure=sites)


# Helper functions are available on the `Volume` object, for example, finding peaks or converting a volume to a vasp file. See [gemdat.Volume](https://gemdat.readthedocs.io/en/latest/api/gemdat_volume/) for more info. 
# 
# This is useful to read the volume in a dedicated viewer like [VESTA](http://www.jp-minerals.org/vesta/en/).

# In[6]:


# Write it to a vasp file for inspection in external tools
vol.to_vasp_volume(sites, filename='/tmp/example.vasp')


# ## Transitions
# 
# Using these sites, we can calculate the transitions and jumps of Lithium.

# In[7]:


transitions = trajectory.transitions_between_sites(sites=sites, floating_specie='Li')


# Note that this gives a warning that two sites (26, 27) are too close, so we delete one of them.

# In[8]:


del sites[26]

transitions = trajectory.transitions_between_sites(sites=sites, floating_specie='Li')


# Analyze and visualize jumps.

# In[9]:


jumps = transitions.jumps(minimal_residence=0)
jumps.plot_jumps_vs_time(n_parts=5)


# In[10]:


jumps.plot_jumps_3d()


# In[ ]: