for_devon.py

import numpy as np
import MDAnalysis as mda
from MDAnalysis.analysis import align
from scipy.spatial.transform import Rotation 
import matplotlib.pyplot as plt

import argparse

'''
Create a synthetic 2D image from a pdb. Here a reference pdb is used to 
create the projection, however, if one does not want to use a reference,
then one should use the same pdb for both the reference and the system.

Different projections can be generated by providing a different quaternion vector.
I left them out of the argparse, because it is common to use sine or cosine function to
define them, so it would be more of a trouble to put them as a flag than to edit the code.
'''

class img_generator:

    def __init__(self, ref_pdb, system_pdb):

        self.ref_pdb = ref_pdb
        self.ref_universe = mda.Universe(self.ref_pdb)

        self.system_pdb = system_pdb
        self.system_universe = mda.Universe(self.system_pdb)
        self.align_system()
        
        
    ##\brief Perform rmsd alignment using a reference    
    def align_system(self):
        
        #Set origin as the cente rof mass
        system_ca = self.system_universe.select_atoms('name CA').positions - \
                    self.system_universe.atoms.center_of_mass()

        reference_ca = self.ref_universe.select_atoms('name CA').positions - \
                       self.ref_universe.atoms.center_of_mass()

        rot_mat, _ = align.rotation_matrix(system_ca, reference_ca)

        # Rotate the coordinates
        self.system_atoms = np.dot(rot_mat, system_ca.T)
        self.n_atoms = self.system_atoms.shape[1]

    ##\brief Rotate the coordinates using a quaternion vector q = w + ix +jy + zk
    #\param q: quaternion vector q, it should be provided as q = [x, y, z, w].
    def rotate_system(self, q):

        r_matrix = Rotation.from_quat(q).as_matrix()

        #Perform rotation
        self.system_atoms = np.matmul(r_matrix, self.system_atoms)


    ##\brief Generate an image by modeling atoms as gaussians and projecting them onto a grid
    #\param n_pixels 1D number of pixels. The image will be n_pixels x n_pixels.
    #\pixel_size Lenght of a pixel in real space (1 pixel = 1cm for example).
    #\sigma Standard deviation for the gaussians
    #\cutoff Neighbor cutoff for the gaussians. Two gaussians are considered as neighbors if their mean
    #        is at a distance less than sigma*cutoff
    def generate_image(self, n_pixels, pixel_size, sigma, cutoff):

        grid_min = -pixel_size * (n_pixels + 1)*0.5;
        grid_max = pixel_size * (n_pixels - 3)*0.5 + pixel_size;

        x_grid = np.arange(grid_min, grid_max, pixel_size);
        y_grid = np.arange(grid_min, grid_max, pixel_size);

        Ixy = np.zeros((n_pixels, n_pixels))

        for atom in range(self.n_atoms):
        
            #Values of the gaussians
            gauss_x = np.zeros((n_pixels))
            gauss_y = np.zeros((n_pixels))
            
            #Selected neighbors
            x_neigh = np.where(np.absolute(x_grid - self.system_atoms[0,:][atom]) <= cutoff*sigma)
            y_neigh = np.where(np.absolute(y_grid - self.system_atoms[1,:][atom]) <= cutoff*sigma)        
                
            gauss_x[x_neigh] = np.exp( -0.5 * ( ((x_grid[x_neigh] - self.system_atoms[0,:][atom]) /sigma)**2) )
            gauss_y[y_neigh] = np.exp( -0.5 * ( ((y_grid[y_neigh] - self.system_atoms[1,:][atom]) /sigma)**2) )
            
            Ixy = np.add(Ixy, gauss_x[:, None]*gauss_y)

        self.image = np.sqrt(2*np.pi) * sigma * Ixy
        

    def plot_image(self):

        fig, ax = plt.subplots()

        ax.imshow(self.image)
        plt.show()
   
   
def main():

    parser = argparse.ArgumentParser()

    parser.add_argument("--ref_pdb", help="/path/to/reference/pdb", required=True)
    parser.add_argument("--system_pdb", 
                        help="/path/to/system/pdb", 
                        required=True)
    parser.add_argument("--n_pixels", help="Number of pixels for the generated image", 
                        required=True, type=int)
    parser.add_argument("--pixel_size", help="Pixel size (relation between image and real space)", 
                        required=True, type=float)
    parser.add_argument("--sigma", help="Standard diviation for the gaussians", 
                        default=1, type=float)
    parser.add_argument("--cutoff", help="Neighbor cutoff for atoms", default=3, type=float)


    args = parser.parse_args()
    
    #Parse flags
    ref_pdb = args.ref_pdb
    system_pdb = args.system_pdb
    n_pixels = args.n_pixels
    pixel_size = args.pixel_size
    sigma = args.sigma
    cutoff = args.cutoff

    #Create quaternion vector
    Q = [0, 0, np.sin(np.pi/4), np.cos(np.pi/4)]

    #Create the generator
    gen = img_generator(ref_pdb, system_pdb)
    
    
    #Rotate the atoms
    gen.rotate_system(Q)

    #Generate the image
    gen.generate_image(n_pixels, pixel_size, sigma, cutoff)

    #Plot the image
    gen.plot_image()


if __name__ == '__main__':
    main()