Skip to content
GitLab
Commit bbd77c31 authored by jdurrant's avatar jdurrant
Browse files

Updates to javascript version.

parent 6ee2d41c
Expand all Hide whitespace changes
Inline Side-by-side
javascript_grid_gen/__target__/gridder.fake_rdkit.js
View file @ bbd77c31
......@@ -2,12 +2,12 @@
__ipow__,__irshift__,__isub__,__ixor__,__jsUsePyNext__,__jsmod__,__k__,__kwargtrans__,__le__,__lshift__,__lt__,__matmul__,__mergefields__,__mergekwargtrans__,__mod__,__mul__,__ne__,__neg__,__nest__,__or__,__pow__,__pragma__,__proxy__,__pyUseJsNext__,__rshift__,__setitem__,__setproperty__,__setslice__,__sort__,__specialattrib__,__sub__,__super__,__t__,__terminal__,__truediv__,__withblock__,__xor__,abs,all,any,assert,bool,bytearray,bytes,callable,chr,copy,deepcopy,delattr,dict,dir,divmod,enumerate,
filter,float,getattr,hasattr,input,int,isinstance,issubclass,len,list,map,max,min,object,ord,pow,print,property,py_TypeError,py_iter,py_metatype,py_next,py_reversed,py_typeof,range,repr,round,set,setattr,sorted,str,sum,tuple,zip}from"./org.transcrypt.__runtime__.js";import*as __module_re__ from"./re.js";__nest__(re,"",__module_re__);var __name__="gridder.fake_rdkit";export var element_names_with_two_letters=["BR","CL","BI","AS","AG","LI","MG","RH","ZN","MN"];export var element_to_atomic_num=dict({"H":1,
"HE":2,"LI":3,"BE":4,"B":5,"C":6,"N":7,"O":8,"F":9,"NE":10,"NA":11,"MG":12,"AL":13,"SI":14,"P":15,"S":16,"CL":17,"AR":18,"K":19,"CA":20});export var Point=__class__("Point",[object],{__module__:__name__,get __init__(){return __get__(this,function(self,coor){self.x=coor[0];self.y=coor[1];self.z=coor[2]})},get __getitem__(){return __get__(this,function(self,key){if(key==0)return self.x;if(key==1)return self.y;if(key==2)return self.z;return 0})}});export var Atom=__class__("Atom",[object],{__module__:__name__,
get __init__(){return __get__(this,function(self,coor,py_name,element,resname){self.coor=Point(coor);self.py_name=py_name;self.element=element;self.resname=resname})},get GetAtomicNum(){return __get__(this,function(self){return element_to_atomic_num[self.element]})}});export var Mol=__class__("Mol",[object],{__module__:__name__,get __init__(){return __get__(this,function(self){self.atoms=[]})},get add_atom(){return __get__(this,function(self,coor,py_name,element,resname){self.atoms.append(Atom(coor,
py_name,element,resname))})},get GetAtomWithIdx(){return __get__(this,function(self,i){return self.atoms[i]})},get GetNumAtoms(){return __get__(this,function(self){return len(self.atoms)})},get GetConformer(){return __get__(this,function(self){return self})},get GetAtomPosition(){return __get__(this,function(self,i){return self.atoms[i].coor})}});export var MolIterator=__class__("MolIterator",[object],{__module__:__name__,get __init__(){return __get__(this,function(self,mol){self.mols=[mol]})},get __iter__(){return __get__(this,
function(self){return self})},[Symbol.iterator](){return this.__iter__()},get __next__(){return __get__(this,function(self){if(len(self.mols)==0){var __except0__=StopIteration;__except0__.__cause__=null;throw __except0__;}return self.mols.py_pop()})},next:__jsUsePyNext__});export var Chem=__class__("Chem",[object],{__module__:__name__,get MolFromPDBFile(){return function(filetxt,sanitize){var mol=Mol();var lines_src=filetxt.py_split("\n");var lines=function(){var __accu0__=[];for(var l of lines_src)if(l.startswith("ATOM")||
l.startswith("HETATM"))__accu0__.append(l);return __accu0__}();var coors=function(){var __accu0__=[];for(var l of lines)__accu0__.append([float(l.__getslice__(30,38,1)),float(l.__getslice__(38,46,1)),float(l.__getslice__(46,54,1))]);return __accu0__}();var names=function(){var __accu0__=[];for(var l of lines)__accu0__.append(l.__getslice__(11,16,1).strip());return __accu0__}();var elements_prep=function(){var __accu0__=[];for(var n of names)__accu0__.append(n.upper().strip());return __accu0__}();
for(var [i,e]of enumerate(elements_prep)){for(var num of"0123456789")var e=e.py_replace(num,"");elements_prep[i]=e}var elements_prep=function(){var __accu0__=[];for(var e of elements_prep)__accu0__.append(e.__getslice__(0,2,1));return __accu0__}();var elements=[];for(var e of elements_prep)elements.append(Chem.name_to_element(e));var resnames=function(){var __accu0__=[];for(var l of lines)__accu0__.append(l.__getslice__(16,21,1).strip());return __accu0__}();for(var i=0;i<len(coors);i++)mol.add_atom(coors[i],
names[i],elements[i],resnames[i]);return mol}},get name_to_element(){return function(py_name){if(__in__(py_name,element_names_with_two_letters))return py_name;else return py_name.__getslice__(0,1,1)}},get SDMolSupplier(){return function(filetxt,sanitize){if(typeof sanitize=="undefined"||sanitize!=null&&sanitize.hasOwnProperty("__kwargtrans__"))var sanitize=false;var mol=Mol();var txt=filetxt;var atoms=re.findall("^ *?[\\-0-9]+?\\.[\\-0-9]+? *?[\\-0-9]+?\\.[\\-0-9]+? *?[\\-0-9]+?\\.[\\-0-9]+? *?[a-zA-Z]+? *?[\\-0-9]+? *?[\\-0-9]+? *?[\\-0-9]+? *?[\\-0-9]+? *?[\\-0-9]+? *?[\\-0-9]+? *?$",
get __init__(){return __get__(this,function(self,coor,py_name,element,resname){self.coor=Point(coor);self.py_name=py_name;self.element=element;self.resname=resname})},get GetAtomicNum(){return __get__(this,function(self){var elem=__in__(self.element,element_to_atomic_num)?element_to_atomic_num[self.element]:100;return elem})}});export var Mol=__class__("Mol",[object],{__module__:__name__,get __init__(){return __get__(this,function(self){self.atoms=[]})},get add_atom(){return __get__(this,function(self,
coor,py_name,element,resname){self.atoms.append(Atom(coor,py_name,element,resname))})},get GetAtomWithIdx(){return __get__(this,function(self,i){return self.atoms[i]})},get GetNumAtoms(){return __get__(this,function(self){return len(self.atoms)})},get GetConformer(){return __get__(this,function(self){return self})},get GetAtomPosition(){return __get__(this,function(self,i){return self.atoms[i].coor})}});export var MolIterator=__class__("MolIterator",[object],{__module__:__name__,get __init__(){return __get__(this,
function(self,mol){self.mols=[mol]})},get __iter__(){return __get__(this,function(self){return self})},[Symbol.iterator](){return this.__iter__()},get __next__(){return __get__(this,function(self){if(len(self.mols)==0){var __except0__=StopIteration;__except0__.__cause__=null;throw __except0__;}return self.mols.py_pop()})},next:__jsUsePyNext__});export var Chem=__class__("Chem",[object],{__module__:__name__,get MolFromPDBFile(){return function(filetxt,sanitize){var mol=Mol();var lines_src=filetxt.py_split("\n");
var lines=function(){var __accu0__=[];for(var l of lines_src)if(l.startswith("ATOM")||l.startswith("HETATM"))__accu0__.append(l);return __accu0__}();var coors=function(){var __accu0__=[];for(var l of lines)__accu0__.append([float(l.__getslice__(30,38,1)),float(l.__getslice__(38,46,1)),float(l.__getslice__(46,54,1))]);return __accu0__}();var names=function(){var __accu0__=[];for(var l of lines)__accu0__.append(l.__getslice__(11,16,1).strip());return __accu0__}();var elements_prep=function(){var __accu0__=
[];for(var n of names)__accu0__.append(n.upper().strip());return __accu0__}();for(var [i,e]of enumerate(elements_prep)){for(var num of"0123456789")var e=e.py_replace(num,"");elements_prep[i]=e}var elements_prep=function(){var __accu0__=[];for(var e of elements_prep)__accu0__.append(e.__getslice__(0,2,1));return __accu0__}();var elements=[];for(var e of elements_prep)elements.append(Chem.name_to_element(e));var resnames=function(){var __accu0__=[];for(var l of lines)__accu0__.append(l.__getslice__(16,
21,1).strip());return __accu0__}();for(var i=0;i<len(coors);i++)mol.add_atom(coors[i],names[i],elements[i],resnames[i]);return mol}},get name_to_element(){return function(py_name){if(__in__(py_name,element_names_with_two_letters))return py_name;else return py_name.__getslice__(0,1,1)}},get SDMolSupplier(){return function(filetxt,sanitize){if(typeof sanitize=="undefined"||sanitize!=null&&sanitize.hasOwnProperty("__kwargtrans__"))var sanitize=false;var mol=Mol();var txt=filetxt;var atoms=re.findall("^ *?[\\-0-9]+?\\.[\\-0-9]+? *?[\\-0-9]+?\\.[\\-0-9]+? *?[\\-0-9]+?\\.[\\-0-9]+? *?[a-zA-Z]+? *?[\\-0-9]+? *?[\\-0-9]+? *?[\\-0-9]+? *?[\\-0-9]+? *?[\\-0-9]+? *?[\\-0-9]+? *?$",
txt,re.MULTILINE);var atoms=function(){var __accu0__=[];for(var a of atoms)__accu0__.append(a.strip().py_split().__getslice__(0,4,1));return __accu0__}();for(var [x,y,z,py_name]of atoms)mol.add_atom([float(x),float(y),float(z)],py_name,Chem.name_to_element(py_name),"");var py_iter=MolIterator(mol);return py_iter}}});
//# sourceMappingURL=gridder.fake_rdkit.map
\ No newline at end of file
javascript_grid_gen/__target__/gridder.grid_util.js
View file @ bbd77c31
'use strict';var math={};var random={};import{AssertionError,AttributeError,BaseException,DeprecationWarning,Exception,IndexError,IterableError,KeyError,NotImplementedError,RuntimeWarning,StopIteration,UserWarning,ValueError,Warning,__JsIterator__,__PyIterator__,__Terminal__,__add__,__and__,__call__,__class__,__envir__,__eq__,__floordiv__,__ge__,__get__,__getcm__,__getitem__,__getslice__,__getsm__,__gt__,__i__,__iadd__,__iand__,__idiv__,__ijsmod__,__ilshift__,__imatmul__,__imod__,__imul__,__in__,
__init__,__ior__,__ipow__,__irshift__,__isub__,__ixor__,__jsUsePyNext__,__jsmod__,__k__,__kwargtrans__,__le__,__lshift__,__lt__,__matmul__,__mergefields__,__mergekwargtrans__,__mod__,__mul__,__ne__,__neg__,__nest__,__or__,__pow__,__pragma__,__proxy__,__pyUseJsNext__,__rshift__,__setitem__,__setproperty__,__setslice__,__sort__,__specialattrib__,__sub__,__super__,__t__,__terminal__,__truediv__,__withblock__,__xor__,abs,all,any,assert,bool,bytearray,bytes,callable,chr,copy,deepcopy,delattr,dict,dir,
divmod,enumerate,filter,float,getattr,hasattr,input,int,isinstance,issubclass,len,list,map,max,min,object,ord,pow,print,property,py_TypeError,py_iter,py_metatype,py_next,py_reversed,py_typeof,range,repr,round,set,setattr,sorted,str,sum,tuple,zip}from"./org.transcrypt.__runtime__.js";import{KDTree}from"./gridder.kdtrees._kdtree.js";import*as __module_random__ from"./random.js";__nest__(random,"",__module_random__);import*as __module_math__ from"./math.js";__nest__(math,"",__module_math__);var __name__=
"gridder.grid_util";export var GPU_DIM=8;export var coord_to_key=function(c){return 1E3*round(c[0]+1E3*(c[1]+c[2]*1E3))};export var mol_gridify=function(grid,atom_coords,atom_layers,layer_offset,num_layers_to_consider,width,res,center,rot){var r2=4;var half_width=width/2;var data=function(){var __accu0__=[];for(var [i,c]of enumerate(atom_coords))__accu0__.append([c.x,c.y,c.z]);return __accu0__}();var tree=KDTree.initialize(data,3,0,null);var layers=dict({});for(var [i,c]of enumerate(data))layers[coord_to_key(c)]=
atom_layers[i];for(var rotation_idx=0;rotation_idx<len(grid);rotation_idx++)for(var layer_idx=0;layer_idx<len(grid[rotation_idx]);layer_idx++){if(layer_idx<=layer_offset)continue;if(layer_idx>layer_offset+num_layers_to_consider)continue;print(rotation_idx,layer_idx);for(var x=0;x<len(grid[rotation_idx][layer_idx]);x++)for(var y=0;y<len(grid[rotation_idx][layer_idx][x]);y++)for(var z=0;z<len(grid[rotation_idx][layer_idx][y]);z++){var tx=x-half_width;var ty=y-half_width;var tz=z-half_width;var tx=tx*
"gridder.grid_util";export var GPU_DIM=8;export var coord_to_key=function(c){return 1E3*round(c[0]+1E3*(c[1]+c[2]*1E3))};export var mol_gridify=function(grid,atom_coords,atom_layers,layer_offset,num_layers_to_consider,width,res,center,rot){var r2=3.0625;var half_width=width/2;var data=function(){var __accu0__=[];for(var [i,c]of enumerate(atom_coords))__accu0__.append([c.x,c.y,c.z]);return __accu0__}();var tree=KDTree.initialize(data,3,0,null);var layers=dict({});for(var [i,c]of enumerate(data))layers[coord_to_key(c)]=
atom_layers[i];for(var rotation_idx=0;rotation_idx<len(grid);rotation_idx++)for(var layer_idx=0;layer_idx<len(grid[rotation_idx]);layer_idx++){if(layer_idx<layer_offset)continue;if(layer_idx>=layer_offset+num_layers_to_consider)continue;print(rotation_idx,layer_idx);for(var x=0;x<len(grid[rotation_idx][layer_idx]);x++)for(var y=0;y<len(grid[rotation_idx][layer_idx][x]);y++)for(var z=0;z<len(grid[rotation_idx][layer_idx][y]);z++){var tx=x-half_width;var ty=y-half_width;var tz=z-half_width;var tx=tx*
res;var ty=ty*res;var tz=tz*res;var aw=rot[0];var ax=rot[1];var ay=rot[2];var az=rot[3];var bx=tx;var by=ty;var bz=tz;var cw=-(ax*bx)-ay*by-az*bz;var cx=aw*bx+ay*bz-az*by;var cy=aw*by+az*bx-ax*bz;var cz=aw*bz+ax*by-ay*bx;var dx=cw*-ax+cx*aw+cy*-az-cz*-ay;var dy=cw*-ay+cy*aw+cz*-ax-cx*-az;var dz=cw*-az+cz*aw+cx*-ay-cy*-ax;var tx=dx+center.x;var ty=dy+center.y;var tz=dz+center.z;var pt=[tx,ty,tz];var closest_atom_coords=tree.nearest_neighbor(pt,1,[]);if(closest_atom_coords[0][1]>r2)continue;var closest_atom_coords=
tree.proximal_neighbor(pt,r2,[]);for(var atom_inf of closest_atom_coords){var ft=layers[coord_to_key(atom_inf[0])];if(ft==-1)continue;var d2=atom_inf[1];var v=math.exp(-2*d2/r2);grid[rotation_idx][layer_offset+ft][x][y][z]+=v}}}return grid};export var make_tensor=function(shape){var t=[];for(var i1=0;i1<shape[0];i1++){var t1=[];for(var i2=0;i2<shape[1];i2++){var t2=[];for(var i3=0;i3<shape[2];i3++){var t3=[];for(var i4=0;i4<shape[3];i4++){var t4=[];for(var i5=0;i5<shape[4];i5++)t4.append(0);t3.append(t4)}t2.append(t3)}t1.append(t2)}t.append(t1)}return t};
export var rand_rot=function(){var q=function(){var __accu0__=[];for(var i=0;i<4;i++)__accu0__.append(random.random());return __accu0__}();var l=math.sqrt(sum(function(){var __accu0__=[];for(var v of q)__accu0__.append(Math.pow(v,2));return __accu0__}()));var q=function(){var __accu0__=[];for(var v of q)__accu0__.append(v/l);return __accu0__}();return q};export var get_raw_batch=function(r_coords,r_types,p_coords,p_types,conn,num_samples,width,res){var num_samples=num_samples===null?3:num_samples;
var width=width===null?24:width;var res=res===null?.5:res;var parent_channels=len(set(p_types));var rec_channels=len(set(r_types));var num_samples=1;var B=num_samples;var T=rec_channels+parent_channels;var N=width;var grid=make_tensor(tuple([B,T,N,N,N]));var rot=rand_rot();var grid=mol_gridify(grid,p_coords,p_types,0,parent_channels,width,res,conn,rot);var grid=mol_gridify(grid,r_coords,r_types,parent_channels,rec_channels,width,res,conn,rot);return grid};
tree.proximal_neighbor(pt,r2,[]);for(var atom_inf of closest_atom_coords){var ft=layers[coord_to_key(atom_inf[0])];if(ft==-1)continue;var d2=atom_inf[1];var v=math.exp(-2*d2/r2);grid[rotation_idx][layer_offset+ft][x][y][z]+=v}}}return grid};export var flatten_tensor=function(grid,shape){var flat=[];for(var i1=0;i1<shape[0];i1++)for(var i2=0;i2<shape[1];i2++)for(var i3=0;i3<shape[2];i3++)for(var i4=0;i4<shape[3];i4++)for(var i5=0;i5<shape[4];i5++)flat.append(grid[i1][i2][i3][i4][i5]);return flat};
export var make_tensor=function(shape){var t=[];for(var i1=0;i1<shape[0];i1++){var t1=[];for(var i2=0;i2<shape[1];i2++){var t2=[];for(var i3=0;i3<shape[2];i3++){var t3=[];for(var i4=0;i4<shape[3];i4++){var t4=[];for(var i5=0;i5<shape[4];i5++)t4.append(0);t3.append(t4)}t2.append(t3)}t1.append(t2)}t.append(t1)}return t};export var rand_rot=function(){var q=[1,0,0,0];return q};export var get_raw_batch=function(r_coords,r_types,p_coords,p_types,conn,num_samples,width,res){var num_samples=num_samples===
null?1:num_samples;var width=width===null?24:width;var res=res===null?.5:res;var parent_channels=4;var rec_channels=5;var B=num_samples;var T=rec_channels+parent_channels;var N=width;var shape=tuple([B,T,N,N,N]);var grid=make_tensor(shape);var rot=rand_rot();var grid=mol_gridify(grid,p_coords,p_types,0,parent_channels,width,res,conn,rot);var grid=mol_gridify(grid,r_coords,r_types,parent_channels,rec_channels,width,res,conn,rot);return flatten_tensor(grid,shape)};
//# sourceMappingURL=gridder.grid_util.map
\ No newline at end of file
javascript_grid_gen/__target__/gridder.util.js
View file @ bbd77c31
'use strict';var math={};import{AssertionError,AttributeError,BaseException,DeprecationWarning,Exception,IndexError,IterableError,KeyError,NotImplementedError,RuntimeWarning,StopIteration,UserWarning,ValueError,Warning,__JsIterator__,__PyIterator__,__Terminal__,__add__,__and__,__call__,__class__,__envir__,__eq__,__floordiv__,__ge__,__get__,__getcm__,__getitem__,__getslice__,__getsm__,__gt__,__i__,__iadd__,__iand__,__idiv__,__ijsmod__,__ilshift__,__imatmul__,__imod__,__imul__,__in__,__init__,__ior__,
__ipow__,__irshift__,__isub__,__ixor__,__jsUsePyNext__,__jsmod__,__k__,__kwargtrans__,__le__,__lshift__,__lt__,__matmul__,__mergefields__,__mergekwargtrans__,__mod__,__mul__,__ne__,__neg__,__nest__,__or__,__pow__,__pragma__,__proxy__,__pyUseJsNext__,__rshift__,__setitem__,__setproperty__,__setslice__,__sort__,__specialattrib__,__sub__,__super__,__t__,__terminal__,__truediv__,__withblock__,__xor__,abs,all,any,assert,bool,bytearray,bytes,callable,chr,copy,deepcopy,delattr,dict,dir,divmod,enumerate,
filter,float,getattr,hasattr,input,int,isinstance,issubclass,len,list,map,max,min,object,ord,pow,print,property,py_TypeError,py_iter,py_metatype,py_next,py_reversed,py_typeof,range,repr,round,set,setattr,sorted,str,sum,tuple,zip}from"./org.transcrypt.__runtime__.js";import*as __module_math__ from"./math.js";__nest__(math,"",__module_math__);import{Chem}from"./gridder.fake_rdkit.js";var __name__="gridder.util";export var get_coords=function(mol){var conf=mol.GetConformer();var coords=function(){var __accu0__=
[];for(var i=0;i<conf.GetNumAtoms();i++)__accu0__.append(conf.GetAtomPosition(i));return __accu0__}();return coords};export var get_types=function(mol){return function(){var __accu0__=[];for(var i=0;i<mol.GetNumAtoms();i++)__accu0__.append(mol.GetAtomWithIdx(i).GetAtomicNum());return __accu0__}()};export var generate_fragments=function(mol,max_heavy_atoms,only_single_bonds){var max_heavy_atoms=max_heavy_atoms===null?0:max_heavy_atoms;var only_single_bonds=only_single_bonds===null?true:only_single_bonds;
var splits=[];var splits=[tuple([mol,null])];return splits};export var load_receptor=function(rec_path){var rec=Chem.MolFromPDBFile(rec_path,__kwargtrans__({sanitize:false}));var rec=remove_water(rec);var rec=remove_hydrogens(rec);return rec};export var remove_hydrogens=function(m){try{for(var atom of m.GetAtoms())atom.SetFormalCharge(0);var m=Chem.RemoveHs(m)}catch(__except0__){m.atoms=function(){var __accu0__=[];for(var a of m.atoms)if(a.element!="H")__accu0__.append(a);return __accu0__}()}return m};
export var remove_water=function(m){m.atoms=function(){var __accu0__=[];for(var a of m.atoms)if(!__in__(a.resname,["WAT","HOH","TIP","TIP3","OH2"]))__accu0__.append(a);return __accu0__}();var merged=m;return merged};export var combine_all=function(frags){if(len(frags)==0)return null;var c=frags[0];for(var f of frags.__getslice__(1,null,1))var c=Chem.CombineMols(c,f);return c};export var load_ligand=function(sdf){var lig=py_next(Chem.SDMolSupplier(sdf,__kwargtrans__({sanitize:false})));var lig=remove_water(lig);
var lig=remove_hydrogens(lig);var frags=generate_fragments(lig,null,null);return tuple([lig,frags])};export var mol_to_points=function(mol,atom_types){var atom_types=atom_types===null?[6,7,8,9,15,16,17,35,53]:atom_types;var coords=get_coords(mol);var types=get_types(mol);var layers=function(){var __accu0__=[];for(var k of types)__accu0__.append(__in__(k,atom_types)?atom_types.index(k):-1);return __accu0__}();var coords=function(){var __accu0__=[];for(var [i,c]of enumerate(coords))if(layers[i]!=-1)__accu0__.append(c);
return __accu0__}();var layers=function(){var __accu0__=[];for(var l of layers)if(l!=-1)__accu0__.append(l);return __accu0__}();return tuple([coords,layers])};export var get_connection_point=function(frag){var dummy_idx=get_types(frag).index(0);var coords=get_coords(frag)[dummy_idx];return coords};export var frag_dist_to_receptor_raw=function(coords,frag){var conn=get_connection_point(frag);var dists=[];for(var i=0;i<len(coords);i++){var coord=coords[i];var tmp=[coord[0]-conn.x,coord[1]-conn.y,coord[2]-
conn.z];var tmp=[Math.pow(tmp[0],2),Math.pow(tmp[1],2),Math.pow(tmp[2],2)];var s=sum(tmp);dists.append(s)}var min_dist=math.sqrt(min(dists));return min_dist};export var mol_array=function(mol){var coords=get_coords(mol);var types=get_types(mol);var arr=[];for(var [i,coor]of enumerate(coords))arr.append([coor.x,coor.y,coor.z,types[i]]);return arr};
[];for(var i=0;i<conf.GetNumAtoms();i++)__accu0__.append(conf.GetAtomPosition(i));return __accu0__}();return coords};export var get_atomic_nums=function(mol){return function(){var __accu0__=[];for(var i=0;i<mol.GetNumAtoms();i++)__accu0__.append(mol.GetAtomWithIdx(i).GetAtomicNum());return __accu0__}()};export var generate_fragments=function(mol,max_heavy_atoms,only_single_bonds){var max_heavy_atoms=max_heavy_atoms===null?0:max_heavy_atoms;var only_single_bonds=only_single_bonds===null?true:only_single_bonds;
var splits=[];var splits=[tuple([mol,null])];return splits};export var load_receptor=function(rec_path){var rec=Chem.MolFromPDBFile(rec_path,__kwargtrans__({sanitize:false}));var rec=remove_water(rec);var rec=remove_hydrogens(rec);return rec};export var remove_hydrogens=function(m){m.atoms=function(){var __accu0__=[];for(var a of m.atoms)if(a.element!="H")__accu0__.append(a);return __accu0__}();return m};export var remove_water=function(m){m.atoms=function(){var __accu0__=[];for(var a of m.atoms)if(!__in__(a.resname,
["WAT","HOH","TIP","TIP3","OH2"]))__accu0__.append(a);return __accu0__}();var merged=m;return merged};export var combine_all=function(frags){if(len(frags)==0)return null;var c=frags[0];for(var f of frags.__getslice__(1,null,1))var c=Chem.CombineMols(c,f);return c};export var load_ligand=function(sdf){var lig=py_next(Chem.SDMolSupplier(sdf,__kwargtrans__({sanitize:false})));var lig=remove_water(lig);var lig=remove_hydrogens(lig);var frags=generate_fragments(lig,null,null);return tuple([lig,frags])};
export var mol_to_points=function(mol,atom_types,note_sulfur){if(typeof note_sulfur=="undefined"||note_sulfur!=null&&note_sulfur.hasOwnProperty("__kwargtrans__"))var note_sulfur=true;var atom_types=atom_types===null?[6,7,8,16]:atom_types;var coords=get_coords(mol);var atomic_nums=get_atomic_nums(mol);var layers=[];for(var t of atomic_nums)if(t==1)layers.append(-1);else if(t==6)layers.append(0);else if(t==7)layers.append(1);else if(t==8)layers.append(2);else if(!note_sulfur)layers.append(3);else if(note_sulfur&&
t==16)layers.append(3);else if(note_sulfur)layers.append(4);var coords=function(){var __accu0__=[];for(var [i,c]of enumerate(coords))if(layers[i]!=-1)__accu0__.append(c);return __accu0__}();var layers=function(){var __accu0__=[];for(var l of layers)if(l!=-1)__accu0__.append(l);return __accu0__}();return tuple([coords,layers])};export var get_connection_point=function(frag){var dummy_idx=get_atomic_nums(frag).index(0);var coords=get_coords(frag)[dummy_idx];return coords};export var frag_dist_to_receptor_raw=
function(coords,frag){var conn=get_connection_point(frag);var dists=[];for(var i=0;i<len(coords);i++){var coord=coords[i];var tmp=[coord[0]-conn.x,coord[1]-conn.y,coord[2]-conn.z];var tmp=[Math.pow(tmp[0],2),Math.pow(tmp[1],2),Math.pow(tmp[2],2)];var s=sum(tmp);dists.append(s)}var min_dist=math.sqrt(min(dists));return min_dist};export var mol_array=function(mol){var coords=get_coords(mol);var types=get_atomic_nums(mol);var arr=[];for(var [i,coor]of enumerate(coords))arr.append([coor.x,coor.y,coor.z,
types[i]]);return arr};
//# sourceMappingURL=gridder.util.map
\ No newline at end of file
javascript_grid_gen/__target__/make_grid.js
View file @ bbd77c31
'use strict';import{AssertionError,AttributeError,BaseException,DeprecationWarning,Exception,IndexError,IterableError,KeyError,NotImplementedError,RuntimeWarning,StopIteration,UserWarning,ValueError,Warning,__JsIterator__,__PyIterator__,__Terminal__,__add__,__and__,__call__,__class__,__envir__,__eq__,__floordiv__,__ge__,__get__,__getcm__,__getitem__,__getslice__,__getsm__,__gt__,__i__,__iadd__,__iand__,__idiv__,__ijsmod__,__ilshift__,__imatmul__,__imod__,__imul__,__in__,__init__,__ior__,__ipow__,
__irshift__,__isub__,__ixor__,__jsUsePyNext__,__jsmod__,__k__,__kwargtrans__,__le__,__lshift__,__lt__,__matmul__,__mergefields__,__mergekwargtrans__,__mod__,__mul__,__ne__,__neg__,__nest__,__or__,__pow__,__pragma__,__proxy__,__pyUseJsNext__,__rshift__,__setitem__,__setproperty__,__setslice__,__sort__,__specialattrib__,__sub__,__super__,__t__,__terminal__,__truediv__,__withblock__,__xor__,abs,all,any,assert,bool,bytearray,bytes,callable,chr,copy,deepcopy,delattr,dict,dir,divmod,enumerate,filter,float,
getattr,hasattr,input,int,isinstance,issubclass,len,list,map,max,min,object,ord,pow,print,property,py_TypeError,py_iter,py_metatype,py_next,py_reversed,py_typeof,range,repr,round,set,setattr,sorted,str,sum,tuple,zip}from"./org.transcrypt.__runtime__.js";import{ligand_sdf,receptor_pdb}from"./gridder.example_data.js";import{Point}from"./gridder.fake_rdkit.js";import{get_raw_batch}from"./gridder.grid_util.js";import{get_connection_point,load_ligand,load_receptor,mol_to_points}from"./gridder.util.js";
var __name__="__main__";export var make_grid=function(receptor,ligand,grid_center){var rec=load_receptor(receptor);var __left0__=load_ligand(ligand);var lig=__left0__[0];var frags=__left0__[1];var __left0__=mol_to_points(rec,null);var rec_coords=__left0__[0];var rec_layers=__left0__[1];var frags=frags.__getslice__(0,1,1);for(var [parent,frag]of frags){var __left0__=mol_to_points(parent,null);var parent_coords=__left0__[0];var parent_layers=__left0__[1];var conn=Point(grid_center);var grid=get_raw_batch(rec_coords,
rec_layers,parent_coords,parent_layers,conn,null,null,null);print(grid)}};if(__name__=="__main__")make_grid(receptor_pdb,ligand_sdf,[.512,3.311,12.006]);
var __name__="__main__";export var get_test_data=function(){return tuple([receptor_pdb,ligand_sdf])};export var make_grid=function(receptor,ligand,grid_center){var rec=load_receptor(receptor);var l=load_receptor(ligand);var __left0__=tuple([null,[tuple([l,null])]]);var lig=__left0__[0];var frags=__left0__[1];var __left0__=mol_to_points(rec,null,__kwargtrans__({note_sulfur:true}));var rec_coords=__left0__[0];var rec_layers=__left0__[1];var frags=[frags[0]];for(var [parent,frag]of frags){var __left0__=
mol_to_points(parent,null,__kwargtrans__({note_sulfur:false}));var parent_coords=__left0__[0];var parent_layers=__left0__[1];var conn=Point(grid_center);var grid=get_raw_batch(rec_coords,rec_layers,parent_coords,parent_layers,conn,1,24,.75);return grid}};if(__name__=="__main__");
//# sourceMappingURL=make_grid.map
\ No newline at end of file
javascript_grid_gen/__target__/make_grid.project
View file @ bbd77c31
{"options": {"source": "make_grid.py", "anno": false, "alimod": false, "build": false, "complex": false, "docat": false, "dassert": false, "dcheck": false, "dextex": false, "dlog": false, "dmap": false, "dnostrip": false, "dstat": false, "dtree": false, "esv": null, "ecom": true, "fcall": false, "gen": false, "iconv": false, "jscall": false, "jskeys": false, "jsmod": false, "kwargs": false, "keycheck": false, "license": false, "map": false, "nomin": false, "opov": false, "parent": null, "run": false, "symbols": null, "sform": false, "tconv": false, "unit": null, "verbose": false, "x": null, "xreex": false, "xglobs": false, "xpath": null, "xtiny": false, "star": false}, "modules": [{"source": "/Users/jdurrant/anaconda/envs/Python36/lib/python3.6/site-packages/transcrypt/modules/org/transcrypt/__runtime__.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/__target__/org.transcrypt.__runtime__.js"}, {"source": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/make_grid.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/__target__/make_grid.js"}, {"source": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/gridder/example_data.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/__target__/gridder.example_data.js"}, {"source": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/gridder/fake_rdkit.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/__target__/gridder.fake_rdkit.js"}, {"source": "/Users/jdurrant/anaconda/envs/Python36/lib/python3.6/site-packages/transcrypt/modules/re/__init__.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/__target__/re.js"}, {"source": "/Users/jdurrant/anaconda/envs/Python36/lib/python3.6/site-packages/transcrypt/modules/re/translate.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/__target__/re.translate.js"}, {"source": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/gridder/grid_util.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/__target__/gridder.grid_util.js"}, {"source": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/gridder/kdtrees/_kdtree.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/__target__/gridder.kdtrees._kdtree.js"}, {"source": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/gridder/kdtrees/_kdtree_type.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/__target__/gridder.kdtrees._kdtree_type.js"}, {"source": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/gridder/kdtrees/_utils.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/__target__/gridder.kdtrees._utils.js"}, {"source": "/Users/jdurrant/anaconda/envs/Python36/lib/python3.6/site-packages/transcrypt/modules/math/__init__.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/__target__/math.js"}, {"source": "/Users/jdurrant/anaconda/envs/Python36/lib/python3.6/site-packages/transcrypt/modules/random/__init__.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/__target__/random.js"}, {"source": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/gridder/util.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor.grid_gen/__target__/gridder.util.js"}]}
\ No newline at end of file
{"options": {"source": "make_grid.py", "anno": false, "alimod": false, "build": false, "complex": false, "docat": false, "dassert": false, "dcheck": false, "dextex": false, "dlog": false, "dmap": false, "dnostrip": false, "dstat": false, "dtree": false, "esv": null, "ecom": true, "fcall": false, "gen": false, "iconv": false, "jscall": false, "jskeys": false, "jsmod": false, "kwargs": false, "keycheck": false, "license": false, "map": false, "nomin": false, "opov": false, "parent": null, "run": false, "symbols": null, "sform": false, "tconv": false, "unit": null, "verbose": false, "x": null, "xreex": false, "xglobs": false, "xpath": null, "xtiny": false, "star": false}, "modules": [{"source": "/Users/jdurrant/anaconda/envs/Python36/lib/python3.6/site-packages/transcrypt/modules/org/transcrypt/__runtime__.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/__target__/org.transcrypt.__runtime__.js"}, {"source": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/make_grid.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/__target__/make_grid.js"}, {"source": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/gridder/example_data.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/__target__/gridder.example_data.js"}, {"source": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/gridder/fake_rdkit.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/__target__/gridder.fake_rdkit.js"}, {"source": "/Users/jdurrant/anaconda/envs/Python36/lib/python3.6/site-packages/transcrypt/modules/re/__init__.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/__target__/re.js"}, {"source": "/Users/jdurrant/anaconda/envs/Python36/lib/python3.6/site-packages/transcrypt/modules/re/translate.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/__target__/re.translate.js"}, {"source": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/gridder/grid_util.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/__target__/gridder.grid_util.js"}, {"source": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/gridder/kdtrees/_kdtree.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/__target__/gridder.kdtrees._kdtree.js"}, {"source": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/gridder/kdtrees/_kdtree_type.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/__target__/gridder.kdtrees._kdtree_type.js"}, {"source": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/gridder/kdtrees/_utils.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/__target__/gridder.kdtrees._utils.js"}, {"source": "/Users/jdurrant/anaconda/envs/Python36/lib/python3.6/site-packages/transcrypt/modules/math/__init__.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/__target__/math.js"}, {"source": "/Users/jdurrant/anaconda/envs/Python36/lib/python3.6/site-packages/transcrypt/modules/random/__init__.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/__target__/random.js"}, {"source": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/gridder/util.py", "target": "/Users/jdurrant/Documents/Work/durrant_git/moiety_predictor/javascript_grid_gen/__target__/gridder.util.js"}]}
\ No newline at end of file
javascript_grid_gen/compile_to_javascript.sh
View file @ bbd77c31
source activate Python36
rm __target__/*
transcrypt --ecom make_grid.py
cp index.html __target__/
......
javascript_grid_gen/gridder/grid_util.py
View file @ bbd77c31
......@@ -116,7 +116,6 @@ def mol_gridify(
print(rotation_idx, layer_idx) # To indicate progress.
for x in range(len(grid[rotation_idx][layer_idx])):
for y in range(len(grid[rotation_idx][layer_idx][x])):
for z in range(len(grid[rotation_idx][layer_idx][y])):
......@@ -165,9 +164,7 @@ def mol_gridify(
# Get the closest atom.
pt = [tx, ty, tz]
closest_atom_coords = tree.nearest_neighbor(
pt, 1, []
)
closest_atom_coords = tree.nearest_neighbor(pt, 1, [])
# If the closest one is farther than r2 away,
# continue.
......@@ -209,6 +206,7 @@ def mol_gridify(
grid[rotation_idx][layer_offset + ft][x][y][z] += v
return grid
def flatten_tensor(grid, shape):
flat = []
for i1 in range(shape[0]):
......@@ -219,6 +217,7 @@ def flatten_tensor(grid, shape):
flat.append(grid[i1][i2][i3][i4][i5])
return flat
def make_tensor(shape):
"""Creates a tensor to store the grid data in.
......@@ -267,9 +266,9 @@ def get_raw_batch(
p_coords,
p_types,
conn,
num_samples, # =3,
width, # =24,
res, # =0.5,
num_samples, # =3, VAL: 1
width, # =24, VAL: 24
res, # =0.5, VAL: 0.75
):
"""Sample a raw batch with provided atom coordinates.
......@@ -305,9 +304,9 @@ def get_raw_batch(
# TODO: For debugging
# num_samples = 1
B = num_samples
T = rec_channels + parent_channels
N = width
B = num_samples # 1
T = rec_channels + parent_channels # 9
N = width # 24
shape = (B, T, N, N, N)
grid = make_tensor(shape)
......
javascript_grid_gen/gridder/util.py
View file @ bbd77c31
......@@ -223,15 +223,30 @@ def mol_to_points(mol, atom_types, note_sulfur=True):
elif t == 8:
# Oxygen
layers.append(2)
# Below differs depending on protein or not.
elif not note_sulfur:
# Not noting sulfur (e.g., ligand), but some other atom.
layers.append(3)
elif note_sulfur and t == 16:
# Noting sulfur (e.g., protein) and sulfur found.
layers.append(3)
elif note_sulfur:
# Noting sulfur (e.g., protein) but some other atom.
layers.append(4)
else:
# So supposed to note sulfur (protein)
if t == 16:
# Noting sulfur (e.g., protein) and sulfur found.
layers.append(3)
else:
# Noting sulfur (e.g., protein) but some other atom.
layers.append(4)
# elif not note_sulfur:
# # Not noting sulfur (e.g., ligand), but some other atom.
# layers.append(3)
# elif note_sulfur and t == 16:
# # Noting sulfur (e.g., protein) and sulfur found.
# layers.append(3)
# elif note_sulfur:
# # Noting sulfur (e.g., protein) but some other atom.
# layers.append(4)
# layers = [(atom_types.index(k) if k in atom_types else -1) for k in types]
......
javascript_grid_gen/jdd_version_grid.json 0 → 100644
View file @ bbd77c31
This diff is collapsed.
javascript_grid_gen/make_grid.py
View file @ bbd77c31
......@@ -5,8 +5,12 @@ Create a protein/ligand grid. Can compile to JavaScript using Transcrypt.
from gridder.util import load_receptor, load_ligand, mol_to_points, get_connection_point
from gridder.grid_util import get_raw_batch
from gridder.fake_rdkit import Point
# __pragma__ ('skip')
import json
# __pragma__ ('noskip')
"""?
from gridder.example_data import receptor_pdb, ligand_sdf
?"""
......@@ -26,6 +30,11 @@ from gridder.example_data import receptor_pdb, ligand_sdf
# - frag_mass: [mass]
# - frag_dist: [dist]
"""?
def get_test_data():
return receptor_pdb, ligand_sdf
?"""
def make_grid(receptor: str, ligand: str, grid_center: list) -> None:
"""Makes a grid from the receptor or ligand.
......@@ -46,13 +55,20 @@ def make_grid(receptor: str, ligand: str, grid_center: list) -> None:
# load ligand and receptor
rec = load_receptor(receptor)
# If SDF
lig, frags = load_ligand(ligand) # List of tuples, (parent, frag)
# If PDB
# l = load_receptor(ligand)
# lig, frags = None, [(l, None)] # to get pdb to work
# compute shared receptor coords and layers
rec_coords, rec_layers = mol_to_points(rec, None, note_sulfur=True)
# Only keep the first fragment. A JDD addition.
frags = frags[:1]
# frags = frags[:1]
frags = [frags[0]]
for parent, frag in frags:
......@@ -80,29 +96,38 @@ def make_grid(receptor: str, ligand: str, grid_center: list) -> None:
parent_layers,
conn,
1,
width=24,
res=0.75,
24, # width=
0.75, # res=
)
# __pragma__ ('skip')
print(json.dumps(grid))
# __pragma__ ('noskip')
"""?
# print(str(grid))
return grid
?"""
if __name__ == "__main__":
# __pragma__ ('skip')
make_grid(
# "./1b6l/1b6l_protein.pdb",
"11gs/11gs_protein.pdb",
# "./1b6l/1b6l_ligand.sdf",
"11gs/11gs_ligand.minus-grid0-frag.sdf",
# [0.512000, 3.311000, 12.006000],
[14.62, 9.944, 24.471],
print(
make_grid(
# "./1b6l/1b6l_protein.pdb",
"11gs/11gs_protein.pdb",
# "./1b6l/1b6l_ligand.sdf",
"11gs/11gs_ligand.minus-grid0-frag.sdf",
# [0.512000, 3.311000, 12.006000],
[14.62, 9.944, 24.471],
)
)
# __pragma__ ('noskip')
"""?
make_grid(
receptor_pdb,
ligand_sdf,
[14.62, 9.944, 24.471],
)
# make_grid(
# receptor_pdb,
# ligand_sdf,
# [14.62, 9.944, 24.471],
# )
?"""
javascript_grid_gen/notebook_explore_prob2.ipynb 0 → 100644
View file @ bbd77c31
This diff is collapsed.
leadopt.py
View file @ bbd77c31
......@@ -21,7 +21,7 @@ def main():
parser.add_argument('-r', '--receptor', required=True, help='Receptor file (.pdb)')
parser.add_argument('-l', '--ligand', required=True, help='Ligand file (.sdf)')
parser.add_argument('-n', '--num_samples', type=int, default=16, help='Number of random rotation samples to use per prediction (default: 16)')
parser.add_argument('-k', '--num_suggestions', type=int, default=25, help='Number of suggestions per fragment')
......@@ -37,11 +37,11 @@ def main():
# run infer step
res = infer_all(
model=m,
fingerprints=fingerprints,
smiles=smiles,
rec_path=args.receptor,
lig_path=args.ligand,
model=m,
fingerprints=fingerprints,
smiles=smiles,
rec_path=args.receptor,
lig_path=args.ligand,
num_samples=args.num_samples,
k=args.num_suggestions
)
......@@ -53,4 +53,4 @@ if __name__=='__main__':
main()
# leadopt.py -r my_receptor.pdb -l my_ligand.sdf
\ No newline at end of file
# leadopt.py -r my_receptor.pdb -l my_ligand.sdf
leadopt/grid_util.py
View file @ bbd77c31
......@@ -18,7 +18,7 @@ GPU_DIM = 8
def gpu_gridify(grid, atom_num, atom_coords, atom_layers, layer_offset,
batch_idx, width, res, center, rot):
"""Adds atoms to the grid in a GPU kernel.
This kernel converts atom coordinate information to 3d voxel information.
Each GPU thread is responsible for one specific grid point. This function
receives a list of atomic coordinates and atom layers and simply iterates
......@@ -56,12 +56,12 @@ def gpu_gridify(grid, atom_num, atom_coords, atom_layers, layer_offset,
rot: (x,y,z,y) rotation quaternion
"""
x,y,z = numba.cuda.grid(3)
# center around origin
tx = x - (width/2)
ty = y - (width/2)
tz = z - (width/2)
# scale by resolution
tx = tx * res
ty = ty * res
......@@ -72,53 +72,53 @@ def gpu_gridify(grid, atom_num, atom_coords, atom_layers, layer_offset,
ax = rot[1]
ay = rot[2]
az = rot[3]
bw = 0
bx = tx
by = ty
bz = tz
# multiply by rotation vector
cw = (aw * bw) - (ax * bx) - (ay * by) - (az * bz)
cx = (aw * bx) + (ax * bw) + (ay * bz) - (az * by)
cy = (aw * by) + (ay * bw) + (az * bx) - (ax * bz)
cz = (aw * bz) + (az * bw) + (ax * by) - (ay * bx)
# multiply by conjugate
# dw = (cw * aw) - (cx * (-ax)) - (cy * (-ay)) - (cz * (-az))
dx = (cw * (-ax)) + (cx * aw) + (cy * (-az)) - (cz * (-ay))
dy = (cw * (-ay)) + (cy * aw) + (cz * (-ax)) - (cx * (-az))
dz = (cw * (-az)) + (cz * aw) + (cx * (-ay)) - (cy * (-ax))
# apply translation vector
tx = dx + center[0]
ty = dy + center[1]
tz = dz + center[2]
i = 0
while i < atom_num:
# fetch atom
fx, fy, fz = atom_coords[i]
ft = atom_layers[i]
i += 1
# invisible atoms
if ft == -1:
continue
# fixed radius (^2)
r2 = 4
# quick cube bounds check
if abs(fx-tx) > r2 or abs(fy-ty) > r2 or abs(fz-tz) > r2:
continue
# compute squared distance to atom
d2 = (fx-tx)**2 + (fy-ty)**2 + (fz-tz)**2
# compute effect
v = math.exp((-2 * d2) / r2)
# add effect
if d2 < r2:
grid[batch_idx, layer_offset+ft, x, y, z] += v
......@@ -127,7 +127,7 @@ def gpu_gridify(grid, atom_num, atom_coords, atom_layers, layer_offset,
def mol_gridify(grid, atom_coords, atom_layers, layer_offset, batch_idx,
width, res, center, rot):
"""Wrapper around gpu_gridify.
(See gpu_gridify() for details)
"""
dw = ((width - 1) // GPU_DIM) + 1
......@@ -148,28 +148,28 @@ def make_tensor(shape):
"""
# get cuda context
ctx = numba.cuda.cudadrv.driver.driver.get_active_context()
# setup tensor on gpu
t = torch.zeros(size=shape, dtype=torch.float32).cuda()
memory = numba.cuda.cudadrv.driver.MemoryPointer(ctx, ctypes.c_ulong(t.data_ptr()), t.numel() * 4)
cuda_arr = numba.cuda.cudadrv.devicearray.DeviceNDArray(
t.size(),
[i*4 for i in t.stride()],
np.dtype('float32'),
gpu_data=memory,
t.size(),
[i*4 for i in t.stride()],
np.dtype('float32'),
gpu_data=memory,
stream=torch.cuda.current_stream().cuda_stream
)
return (t, cuda_arr)
def rand_rot():
"""Returns a random uniform quaternion rotation."""
q = np.random.normal(size=4) # sample quaternion from normal distribution
q = q / np.sqrt(np.sum(q**2)) # normalize
return q
def get_batch(data, rec_channels, parent_channels, batch_size=16, batch_set=None,
width=48, res=0.5, ignore_receptor=False, ignore_parent=False):
......@@ -201,16 +201,16 @@ def get_batch(data, rec_channels, parent_channels, batch_size=16, batch_set=None
# create a tensor with shared memory on the gpu
torch_grid, cuda_grid = make_tensor((batch_size, dim, width, width, width))
if batch_set is None:
batch_set = np.random.choice(len(data), size=batch_size, replace=False)
examples = [data[idx] for idx in batch_set]
for i in range(len(examples)):
example = examples[i]
rot = rand_rot()
if ignore_receptor:
mol_gridify(
cuda_grid,
......@@ -237,7 +237,7 @@ def get_batch(data, rec_channels, parent_channels, batch_size=16, batch_set=None
)
else:
mol_gridify(
cuda_grid,
cuda_grid,
example['p_coords'],
example['p_types'],
layer_offset=0,
......@@ -258,7 +258,7 @@ def get_batch(data, rec_channels, parent_channels, batch_size=16, batch_set=None
center=example['conn'],
rot=rot
)