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

# <a href='http://bempp.com' style='line-height:50px'><img src='bempp-logo.jpg' style='float:left;height:50px;margin-right:10px'> This tutorial forms part of the documentation for Bempp. Find out more at bempp.com</a>
# 
# This example is based on work by Paulo Fernández and <a href='http://www.mecanica.usm.cl/quienes-somos/10-profesores/42-christopher-cooper'>Christopher Cooper Villagrán</a>.

# # Solvation energy of a protein with a continuum electrostatic model
# 
# ## Background
# 
# The energy required to dissolve a molecule (that is, moving it from vacuum into a solvent) is known as the solvation energy. This quantity is of interest in several biomolecular applications. In a continuum electrostatic description, the dissolved biomolecule can be modeled with two dielectric regions interfaced by a molecular surface: an outer *solvent* region, usually comprised of water and salt ions, and an inner *protein* region, where the solvent cannot access, which contains point charges at the locations of the atoms.  
# 
# <img src="sketch_protein.png", style="width:300px">
# 
# The protein region ($\Omega^-$) is governed by a Poisson equation with a source term due to the partial charges of the atoms, whereas the solvent region ($\Omega^\text{+}$) uses a modified Helmholtz equation to account for the presence of salt. The electrostatic potentials in the two regions, $\phi^-$ and $\phi^+$, are given by
# $$ \Delta \phi^- (\mathbf{x}) = \sum_k q_k \delta (\mathbf{x}_k) \quad\text{in }\Omega^-,$$
# $$ \Delta \phi^+ (\mathbf{x}) - \kappa \phi^+ (\mathbf{x}) = 0  \quad\text{in }\Omega^+,$$
# $$ \phi^- = \phi^+\quad\text{on }\Gamma$$
# $$\epsilon^- \frac{\partial \phi^-}{\partial \nu} = \epsilon^+ \frac{\partial \phi^+}{\partial \nu}\quad\text{on }\Gamma$$
# where $q_k$ and $\mathbf{x}_k$ are the charge and position of each molecule atom, $\kappa$ is the inverse of the <a href='https://en.wikipedia.org/wiki/Debye_length' target='new'>Debye length</a>.
# 
# 
# 
# Following the work by <a href='http://onlinelibrary.wiley.com/doi/10.1002/jcc.540110911/abstract' target='new'>Yoon and Lenhoff</a>, the system of PDEs with interface conditions can be formulated in terms of boundary operators, as follows:
# 
# $$
# \begin{bmatrix}
#     \tfrac12\mathsf{Id} + \mathsf{K}_L &  -\mathsf{V}_L \\
#     \tfrac12\mathsf{Id} - \mathsf{K}_H &  \frac{\epsilon_1}{\epsilon_2}\mathsf{V}_H
# \end{bmatrix}
# \left\{
# \begin{matrix}
#     \phi^- \\
#     \frac{\partial \phi^-}{\partial \mathbf{\nu}}
# \end{matrix}
# \right\}
# = \left\{
# \begin{matrix}
#     \sum_k q_k \\
#     0
# \end{matrix}
# \right\}
# $$
# 
# where $\mathsf{V}_L$ and $\mathsf{K}_L$ are the <a href='https://bempp.com/2017/07/11/available_operators/'>single and double layer boundary operators</a> for the <a href='https://bempp.com/tag/laplace'>Laplace</a> kernel (Poisson equation) and $\mathsf{V}_H$ and $\mathsf{K}_H$ are the corresponding operators for the <a href='https://bempp.com/tag/modified-helmholtz'>modified Helmholtz</a> kernel (Poisson-Boltzmann or Yukawa equation).
# 
# From there, we can compute the *reaction potential* ($\phi_\text{reac} = \phi^--\phi_\text{Coulomb}$) at the locations of the atoms using
# 
# $$\phi_\text{reac} = \mathcal{K}_L\left[ \frac{\partial \phi}{\partial n} \right] - \mathcal{V}_L \left[ \phi \right], $$
# where $\mathcal{V}_L$ and $\mathcal{K}_L$ are the single and double layer potential operators for the Laplace kernel.
# We then obtain the solvation energy using
# 
# $$ \Delta G_\text{solv} = \tfrac12 \sum_{k=1}^{N_q} q_k \phi_\text{reac}(\mathbf{x}_k) $$

# ## Example
# 
# ### Molecular structure and mesh
# 
# As an example, we will calculate the solvation energy of the bovine pancreatic trypsin inhibitor. The molecular crystal structure is available the protein data bank ([PDB code 5PTI](https://www.rcsb.org/pdb/explore.do?structureId=5PTI)), and we obtained the atomic radii and charges with [pdb2pqr](http://www.poissonboltzmann.org/), resulting in a `.pqr` file.
# 
# There are several molecular surface definitions which can be used. In this case, we use the *solvent-excluded surface* (SES), which is the result of rolling a spherical probe of the size of a water molecule around the protein, and tracking the contact points. We mesh the SES with [`msms`](http://mgl.scripps.edu/people/sanner/html/msms_home.html), and reformat the result to a `.msh` file. 
# 
# First, let's import the required libraries and mesh and read the `.pqr` file for the charges.

# In[1]:


import bempp.api
bempp.api.set_ipython_notebook_viewer()

grid = bempp.api.import_grid('5pti_d1.msh')
grid.plot()

import numpy as np
q, x_q = np.array([]), np.empty((0,3))

ep_in = 4.
ep_ex = 80.
k = 0.125

# Read charges and coordinates from the .pqr file
molecule_file = open('5pti.pqr', 'r').read().split('\n')
for line in molecule_file:
    line = line.split()
    if len(line)==0 or line[0]!='ATOM': continue
    q = np.append( q, float(line[8]))
    x_q = np.vstack(( x_q, np.array(line[5:8]).astype(float) ))


# Next, we compute the potential due to the charges on the boundary, required for the right-hand side

# In[2]:


def charges_fun(x, n, domain_index, result):
    global q, x_q, ep_in
    result[:] = np.sum(q/np.linalg.norm( x - x_q, axis=1 ))/(4*np.pi*ep_in)

def zero(x, n, domain_index, result):
    result[:] = 0
    
dirichl_space = bempp.api.function_space(grid, "DP", 0)
neumann_space = bempp.api.function_space(grid, "DP", 0)

charged_grid_fun = bempp.api.GridFunction(dirichl_space, fun=charges_fun)
zero_grid_fun = bempp.api.GridFunction(neumann_space, fun=zero)


# Next, we generate the $2\times 2$ block matrix with the single and double layer operators of the Laplace and modified Helmholtz kernels.

# In[3]:


# Define Operators
from bempp.api.operators.boundary import sparse, laplace, modified_helmholtz
identity = sparse.identity(dirichl_space, dirichl_space, dirichl_space)
slp_in   = laplace.single_layer(neumann_space, dirichl_space, dirichl_space)
dlp_in   = laplace.double_layer(dirichl_space, dirichl_space, dirichl_space)
slp_out  = modified_helmholtz.single_layer(neumann_space, dirichl_space, dirichl_space, k)
dlp_out  = modified_helmholtz.double_layer(dirichl_space, dirichl_space, dirichl_space, k)

# Matrix Assembly
blocked = bempp.api.BlockedOperator(2, 2)
blocked[0, 0] = 0.5*identity + dlp_in
blocked[0, 1] = -slp_in
blocked[1, 0] = 0.5*identity - dlp_out
blocked[1, 1] = (ep_in/ep_ex)*slp_out


# We now use `gmres` to solve the system. 

# In[4]:


sol, info, it_count = bempp.api.linalg.gmres(blocked, [charged_grid_fun, zero_grid_fun],
                                            use_strong_form=True, return_iteration_count=True, tol=1e-3)

print("The linear system was solved in {0} iterations".format(it_count))


# The two following `GridFunction` calls store the calculated boundary potential data (separated by $\phi$ and $\frac{\partial \phi}{\partial n}$) for visualization purposes. 

# In[5]:


solution_dirichl, solution_neumann = sol
solution_dirichl.plot()


# To compute $\phi_\text{reac}$ at the atoms' locations, we use `operators.potential` to then multiply by the charge and add to compute $\Delta G_\text{solv}$. The `332.064` term is just a unit conversion constant.

# In[6]:


slp_q = bempp.api.operators.potential.laplace.single_layer(neumann_space, x_q.transpose())
dlp_q = bempp.api.operators.potential.laplace.double_layer(dirichl_space, x_q.transpose())
phi_q = slp_q * solution_neumann - dlp_q * solution_dirichl

# total dissolution energy applying constant to get units [kcal/mol]
total_energy = 2 * np.pi * 332.064 * np.sum(q * phi_q).real
print("Total solvation energy: {:7.2f} [kcal/Mol]".format(total_energy))