TorchProteinLibrary 0.2 Nucleic Acid Extension

.gitignore

@@ -14,6 +14,14 @@
 *.sty
 *.gif
 *.txt
+*.crg
+*.siz
+*.cube
+*.vscode
+build/*
+dist/*
+*/talled
+*.mp4
 
 #Python 
 *.pyc

.gitmodules

@@ -0,0 +1,4 @@
+[submodule "cusplibrary"]
+	path = cusplibrary
+	url = https://github.com/lupoglaz/cusplibrary.git
+	branch = develop

Benchmark/volumemul.py

@@ -0,0 +1,64 @@
+from core import ModuleBenchmark
+import torch
+from torch.autograd import profiler
+from TorchProteinLibrary import VolumeMul
+
+import numpy as np
+import seaborn as sea
+from matplotlib import pylab as plt
+import pandas as pd
+from tqdm import tqdm
+import pickle as pkl
+
+class VolumeMulBenchmark(ModuleBenchmark):
+	def __init__(self, device='gpu', num_sequences=32, seq_length=350):
+		super().__init__(device)
+		self.vol1 = 
+
+	def prepare(self):
+		self.angles.requires_grad_()
+		self.coords = self.module(self.angles, self.length)
+		self.grad_coords = torch.ones(self.coords.size(), dtype=torch.float, device='cuda')
+
+	def run_forward(self):
+		self.angles.detach()
+		self.coords = self.module(self.angles, self.length)
+
+	def run_backward(self):
+		self.coords.backward(self.grad_coords)
+
+def test_length( interval=(60, 20, 300), num_measurements=1, output_filename='length.dat'):
+	time = []
+	direction = []
+	length = []
+	for n in range(num_measurements):
+		for i in tqdm(range(interval[0], interval[2], interval[1])):
+			mb = Angles2BackboneBenchmark(seq_length=i)
+			dur_fwd = mb.measure_forward()
+			time.append(dur_fwd/1000.0)
+			direction.append('Forward')
+			length.append(i)
+			dur_bwd = mb.measure_backward()
+			time.append(dur_bwd/1000.0)
+			direction.append('Backward')
+			length.append(i)
+
+	data = pd.DataFrame({  'Time': time, 
+							'Direction': direction, 
+							'Length': length
+						})
+	data.to_pickle(output_filename)
+
+if __name__=='__main__':
+
+	test_length(interval=(60, 20, 700), output_filename='Data/ReducedModelTime.pkl', num_measurements=10)
+
+	data = pd.read_pickle('Data/ReducedModelTime.pkl')
+	sea.set_style("whitegrid")
+	sea.set_context("paper", font_scale=1.5, rc={"lines.linewidth": 1.5})
+	g1 = sea.relplot(x="Length", y="Time", hue='Direction', kind="line", style="Direction", height=6, aspect=1.3, markers=True, data=data)
+	plt.ylabel("Time, ms")
+	plt.xlabel("Sequence length, amino-acids")
+	sea.despine()
+	# plt.show()
+	plt.savefig("Fig/ReducedModelTime.png")

Layers/FullAtomModel/Angles2Coords/angles2coords_interface.cpp

@@ -9,7 +9,10 @@ void Angles2Coords_forward(     torch::Tensor sequences,
                                 torch::Tensor input_angles, 
                                 torch::Tensor output_coords,
                                 torch::Tensor res_names,
-                                torch::Tensor atom_names
+                                torch::Tensor res_nums,
+                                torch::Tensor atom_names,
+                                int polymer_type,
+                                torch::Tensor chain_names
                         ){
     bool add_terminal = false;
     CHECK_CPU_INPUT_TYPE(sequences, torch::kByte);
@@ -23,60 +26,122 @@ void Angles2Coords_forward(     torch::Tensor sequences,
     }
 
     int batch_size = input_angles.sizes()[0];
-
-    #pragma omp parallel for
-    for(int i=0; i<batch_size; i++){
-
-        torch::Tensor single_sequence = sequences[i];
-        torch::Tensor single_atom_names = atom_names[i];
-        torch::Tensor single_res_names = res_names[i];
-        torch::Tensor single_angles = input_angles[i];
-        torch::Tensor single_coords = output_coords[i];
-
-        std::string seq = StringUtil::tensor2String(single_sequence);
-
-        int length = single_angles.sizes()[1];
-        int num_atoms = ProtUtil::getNumAtoms(seq, add_terminal);
-
-        if( single_coords.sizes()[0]<3*num_atoms){
-            ERROR("incorrect coordinates tensor length");
-        }
-
-        if( length<seq.length() || single_angles.sizes()[0]<8 ){
-            ERROR("incorrect angles tensor length");
-        }
-
-        if( single_res_names.sizes()[0]<seq.length() ){
-            ERROR("incorrect res names tensor length");
-        }
-
-        if( single_atom_names.sizes()[0]<seq.length() ){
-            ERROR("incorrect atom names tensor length");
-        }
-        torch::Tensor dummy_grad = torch::zeros_like(single_angles);
-        AT_DISPATCH_FLOATING_TYPES(single_angles.type(), "cConformation", ([&] {
-            cConformation<scalar_t> conf( seq, single_angles.data<scalar_t>(), dummy_grad.data<scalar_t>(), length, single_coords.data<scalar_t>());
-            for(int j=0; j<conf.groups.size(); j++){
-                for(int k=0; k<conf.groups[j]->atomNames.size(); k++){
-                    int idx = conf.groups[j]->atomIndexes[k];
-                    torch::Tensor single_atom_name = single_atom_names[idx];
-                    torch::Tensor single_res_name = single_res_names[idx];
-                    StringUtil::string2Tensor(ProtUtil::convertRes1to3(conf.groups[j]->residueName), single_res_name);
-                    StringUtil::string2Tensor(conf.groups[j]->atomNames[k], single_atom_name);
+
+    if(polymer_type == 0){
+
+        #pragma omp parallel for
+
+        for(int i=0; i<batch_size; i++){
+
+            torch::Tensor single_sequence = sequences[i];
+            torch::Tensor single_atom_names = atom_names[i];
+            torch::Tensor single_res_names = res_names[i];
+            torch::Tensor single_res_nums = res_nums[i];
+            torch::Tensor single_angles = input_angles[i];
+            torch::Tensor single_coords = output_coords[i];
+
+            std::string seq = StringUtil::tensor2String(single_sequence);
+
+            int length = single_angles.sizes()[1];
+
+            int num_atoms = ProtUtil::getNumAtoms(seq, add_terminal);
+
+            if( single_coords.sizes()[0]<3*num_atoms){
+
+                ERROR("incorrect coordinates tensor length");
+            }
+
+            if( length<seq.length() || single_angles.sizes()[0]<8 ){
+                ERROR("incorrect angles tensor length");
+            }
+
+            if( single_res_names.sizes()[0]<seq.length() ){
+                ERROR("incorrect res names tensor length");
+            }
+
+            if( single_atom_names.sizes()[0]<seq.length() ){
+                ERROR("incorrect atom names tensor length");
+            }
+            torch::Tensor dummy_grad = torch::zeros_like(single_angles);
+    //        Conformation and convertRes1to3
+            AT_DISPATCH_FLOATING_TYPES(single_angles.type(), "cConformation", ([&] {
+                cConformation<scalar_t> conf( seq, single_angles.data<scalar_t>(), dummy_grad.data<scalar_t>(), length, single_coords.data<scalar_t>(), polymer_type, chain_names);
+                for(int j=0; j<conf.groups.size(); j++){
+                    for(int k=0; k<conf.groups[j]->atomNames.size(); k++){
+                        int idx = conf.groups[j]->atomIndexes[k];
+                        torch::Tensor single_atom_name = single_atom_names[idx];
+                        torch::Tensor single_res_name = single_res_names[idx];
+                        single_res_nums[idx] = (int)conf.groups[j]->residueIndex;
+                        StringUtil::string2Tensor(ProtUtil::convertRes1to3(conf.groups[j]->residueName), single_res_name);
+                        StringUtil::string2Tensor(conf.groups[j]->atomNames[k], single_atom_name);
+                    }
                 }
-            }                    
-        }));
-        // cConformation<double> conf( seq, single_angles.data<double>(), dummy_grad.data<double>(),
-        //                     length, single_coords.data<double>());
-        //Output atom names and residue names
-
+            }));
+            // cConformation<double> conf( seq, single_angles.data<double>(), dummy_grad.data<double>(),
+            //                     length, single_coords.data<double>());
+            //Output atom names and residue names
+              }
+          }
+        if(polymer_type == 1 or polymer_type == 2){
+
+            #pragma omp parallel for
+
+            for(int i=0; i<batch_size; i++){
+
+                torch::Tensor single_sequence = sequences[i];
+                torch::Tensor single_atom_names = atom_names[i];
+                torch::Tensor single_res_names = res_names[i];
+                torch::Tensor single_res_nums = res_nums[i];
+                torch::Tensor single_angles = input_angles[i];
+                torch::Tensor single_coords = output_coords[i];
+
+                std::string seq = StringUtil::tensor2String(single_sequence);
+
+                int length = single_angles.sizes()[1];
+
+
+                int num_atoms = ProtUtil::getNumAtoms(seq, add_terminal, polymer_type, chain_names);
+
+                if( single_coords.sizes()[0]<3*num_atoms){
+                    ERROR("incorrect coordinates tensor length");
+                }
+
+                if( length<seq.length() || single_angles.sizes()[0]<10){
+                    ERROR("incorrect angles tensor length");
+                }
+
+                if( single_res_names.sizes()[0]<seq.length() ){
+                    ERROR("incorrect res names tensor length");
+                }
+
+                if( single_atom_names.sizes()[0]<seq.length() ){
+                    ERROR("incorrect atom names tensor length");
+                }
+                torch::Tensor dummy_grad = torch::zeros_like(single_angles);
+        //        Conformation and convertRes1to3
+
+                AT_DISPATCH_FLOATING_TYPES(single_angles.type(), "cConformation", ([&] {
+                    cConformation<scalar_t> conf( seq, single_angles.data<scalar_t>(), dummy_grad.data<scalar_t>(), length, single_coords.data<scalar_t>(), polymer_type, chain_names);
+                    for(int j=0; j<conf.groups.size(); j++){
+                        for(int k=0; k<conf.groups[j]->atomNames.size(); k++){
+                            int idx = conf.groups[j]->atomIndexes[k];
+                            torch::Tensor single_atom_name = single_atom_names[idx];
+                            torch::Tensor single_res_name = single_res_names[idx];
+                            single_res_nums[idx] = (int)conf.groups[j]->residueIndex;
+                            StringUtil::string2Tensor(ProtUtil::convertRes1to3(conf.groups[j]->residueName, polymer_type), single_res_name);
+                            StringUtil::string2Tensor(conf.groups[j]->atomNames[k], single_atom_name);}
+                    }
+                }));
+        }
     }
 }
 
 void Angles2Coords_backward(    torch::Tensor grad_atoms,
                                 torch::Tensor grad_angles,
                                 torch::Tensor sequences,
-                                torch::Tensor input_angles
+                                torch::Tensor input_angles,
+                                int polymer_type,
+                                torch::Tensor chain_names
                         ){
     bool add_terminal = false;
     CHECK_CPU_INPUT_TYPE(sequences, torch::kByte);
@@ -98,13 +163,15 @@ void Angles2Coords_backward(    torch::Tensor grad_atoms,
         torch::Tensor single_grad_atoms = grad_atoms[i];
 
         std::string seq = StringUtil::tensor2String(single_sequence);
-
+
+//      Get Num Atoms
         uint length = single_angles.sizes()[1];
-        int num_atoms = ProtUtil::getNumAtoms(seq, add_terminal);
+        int num_atoms = ProtUtil::getNumAtoms(seq, add_terminal, polymer_type, chain_names);
 
+//      Conformation
         torch::Tensor dummy_coords = torch::zeros({3*num_atoms}, torch::TensorOptions().dtype(grad_atoms.dtype()));
         AT_DISPATCH_FLOATING_TYPES(single_angles.type(), "cConformation", ([&] {
-            cConformation<scalar_t> conf(seq, single_angles.data<scalar_t>(), single_grad_angles.data<scalar_t>(),length, dummy_coords.data<scalar_t>());
+            cConformation<scalar_t> conf(seq, single_angles.data<scalar_t>(), single_grad_angles.data<scalar_t>(),length, dummy_coords.data<scalar_t>(), polymer_type, chain_names);
             conf.backward(conf.root, single_grad_atoms.data<scalar_t>());
         }));
         // cConformation<double> conf( seq, single_angles.data<double>(), single_grad_angles.data<double>(),
@@ -125,7 +192,8 @@ void Angles2Coords_save(    const char* sequence,
     }
     std::string aa(sequence);
     uint length = aa.length();
-    int num_atoms = ProtUtil::getNumAtoms(aa, add_terminal);
+    int polymer_type = 1; //"hard coded" need to change to arg
+    int num_atoms = ProtUtil::getNumAtoms(aa, add_terminal, polymer_type);
     torch::Tensor dummy_grad = torch::zeros_like(input_angles);
     torch::Tensor dummy_coords = torch::zeros({3*num_atoms}, torch::TensorOptions().dtype(input_angles.dtype()));
     AT_DISPATCH_FLOATING_TYPES(input_angles.type(), "cConformation.save", ([&]{
@@ -136,9 +204,9 @@ void Angles2Coords_save(    const char* sequence,
 
 }
 
-int getSeqNumAtoms( const char *sequence){
+int getSeqNumAtoms( const char *sequence, int polymer_type, torch::Tensor chain_names){
     bool add_terminal = false;
     std::string seq(sequence);
-    int num_atoms = ProtUtil::getNumAtoms(seq, add_terminal);
+    int num_atoms = ProtUtil::getNumAtoms(seq, add_terminal, polymer_type, chain_names);
     return num_atoms;
 }

Layers/FullAtomModel/Angles2Coords/angles2coords_interface.h

@@ -3,17 +3,22 @@ void Angles2Coords_forward(     at::Tensor sequences,
                                 at::Tensor input_angles, 
                                 at::Tensor output_coords,
                                 at::Tensor res_names,
-                                at::Tensor atom_names
+                                torch::Tensor res_nums,
+                                at::Tensor atom_names,
+                                int polymer_type,
+                                torch::Tensor chain_names
                         );
 
 void Angles2Coords_backward(    at::Tensor grad_atoms,
                                 at::Tensor grad_angles,
                                 at::Tensor sequences,
-                                at::Tensor input_angles
+                                at::Tensor input_angles,
+                                int polymer_type,
+                                torch::Tensor chain_names
                         );
 void Angles2Coords_save(    const char* sequence,
                             at::Tensor input_angles, 
                             const char* output_filename,
                             const char mode
                         );
-int getSeqNumAtoms( const char *sequence);
+int getSeqNumAtoms( const char *sequence, int polymer_type = 0, torch::Tensor chain_names = {});