ó žu_[c@s‰ddlZddlZddlZddlmZddlZddlZddlmZddl Z ddl m Z dZ dZ dZdZdZdZdd d gZid d 6d d6dd6dd6dd6dd6dd6dd6dd6dd6dd6dd 6d!d"6d#d$6d%d&6d'd(6d)d*6d+d,6d-d.6d/d06Zied1d2d3d4gƒd6ed1d2d3d4d5d6gƒd6ed1d2d3d4d5d7d8d9d:d;d<g ƒd-6ed1d2d3d4d5d=gƒd6ed1d2d3d4d5d>d?gƒd6ed1d2d3d4d5d7d8d@dAg ƒd6ed1d2d3d4d5d7dBd@gƒd)6ed1d2d3d4d5gƒd'6ed1d2d3d4d5d7dCdDgƒd+6ed1d2d3d4d5dEd?dCgƒd 6ed1d2d3d4d5dEd?gƒd6ed1d2d3d4d5d7dFdGgƒd!6ed1d2d3d4d5d7d8dHdIg ƒd6ed1d2d3d4d5d7dJdDdKdLg ƒd6ed1d2d3d4d5d7dFdMgƒd6ed1d2d3d4d5d7d8gƒd6ed1d2d3d4d5d7d8dHdLg ƒd 6ed1d2d3d4d5d7dCdDdKdNd:g ƒd%6ed1d2d3d4d5d7dCdDdOdNdPdQdRdSgƒd#6ed1d2d3d4d5d7dCdDdKdNd:dTg ƒd/6ZdUd]dV„ƒYZdW„ZdX„ZdY„ZdZ„Zd[„Zd\„ZdS(^iÿÿÿÿN(tSet(tspatial(tPDBListiiigB`åÐ"Ûá¿g\Âõ(\ç¿gV-²ó?i tILEitGLNitGLYtGLUitCYStHISitSERitLYSitPROitASNi tVALitTHRitASPitTRPitPHEi tALAi tMETi tLEUitARGitTYRtNtCAtCtOtCBtSGtCGtCDtNEtCZtNH1tNH2tOGtOG1tCG2tCEtNZtSDtCD1tCD2tCG1tOD1tOD2tOE1tOE2tND1tCE1tNE2tND2tCE2tNE1tCE3tCZ2tCZ3tCH2tOHtPDB_atomcBseZd„Zd„ZRS(cCsC||_||_||_||_||_||_||_dS(N(tatomtrestchain_IDtxtytztindex(tselft atom_typeR=R>R?R@RARB((sn/cstor/stanford/rbaltman/users/wtorng/DEEP_LEARNING/3DCNN/FSCNN_models/User_test/code/cut_box_site_test_pdb.pyt__init__/s      cCs|j|jkS(N(t__dict__(RCtother((sn/cstor/stanford/rbaltman/users/wtorng/DEEP_LEARNING/3DCNN/FSCNN_models/User_test/code/cut_box_site_test_pdb.pyt__eq__7s(t__name__t __module__RERH(((sn/cstor/stanford/rbaltman/users/wtorng/DEEP_LEARNING/3DCNN/FSCNN_models/User_test/code/cut_box_site_test_pdb.pyR;.s cCs#|j|ddƒ\}}||S(Ntki(tquery(t my_kd_treetcort PDB_entriestdti((sn/cstor/stanford/rbaltman/users/wtorng/DEEP_LEARNING/3DCNN/FSCNN_models/User_test/code/cut_box_site_test_pdb.pytfind_actual_pos:scCs:i}x-|D]%}|j|j|jg||js #cCsJ|d}tj|dƒtj|dƒ}tj|dƒtj|dƒ}tj||ƒ}d|jƒkr¹tj|dƒtj|dƒ}|tj|j|ƒƒ9}ntj||ƒ}|tjt|dƒƒ:}|tjt|dƒƒ:}|tjt|dƒƒ:}tj|||gdƒj}||gS(NRRRRitfloat16( tnumpytarraytcrosstkeystsigntdottsqrttsumtT(RTt referencetaxis_xt pseudo_axis_ytaxis_zt directiontaxis_yt transform((sn/cstor/stanford/rbaltman/users/wtorng/DEEP_LEARNING/3DCNN/FSCNN_models/User_test/code/cut_box_site_test_pdb.pytcenter_and_transformDs $$$cCsOxH|D]@}|j|jkr|j|jkr|j|jkrtSqWtS(N(R?R@RAtTruetFalse(t PDB_entrytPDB_listtelement((sn/cstor/stanford/rbaltman/users/wtorng/DEEP_LEARNING/3DCNN/FSCNN_models/User_test/code/cut_box_site_test_pdb.pytPDB_is_in_listWs 6cCs:i}g}g}g}g}d}d}g}xî|D]æ} | jƒ} | ddkr | dd!jdƒ} | dd!} | dd !} | d}| d jdƒ}t|ƒ}||f} t| d d !ƒt| d d!ƒt| dd!ƒg}|j|ƒ|j| ƒ|j| dƒ| |jƒkrg}td| d| d| d|dd|d d|dd|ƒ}|j|ƒ||| R?R@RAiRBtENDMDL(tsplittstriptinttfloattappendR[R;(t entry_listtID_dicttall_post all_linest all_atom_typeROt atom_indextmodel_IDtMODELStline1tlineR<R=R>tchaintres_notnew_postlRUtMODEL((sn/cstor/stanford/rbaltman/users/wtorng/DEEP_LEARNING/3DCNN/FSCNN_models/User_test/code/cut_box_site_test_pdb.pytgrab_PDBasD       6  ?  84 cECsP dGHd}d}d}tjj|ƒs9tj|ƒn|dkrN|}n|d|d|}td|dd ƒ} t} | d } | d } | d } | d }| d }| d }t|ƒ}t|ƒ}i}x~|D]v}|jƒ}|d }t|dƒ}t|d ƒ}t|d ƒ}||krMg||t j%|<|ƒ}?t j%|=|ƒ}@t j%|>|ƒ}A|?t&krFt&d}?n|@t&kr_t&d}@n|At&krxt&d}An|;j'd dkr¸|:d |?|@|Afd|:d |?|@|AfR=tresiName_to_labelR[RVRRhR]tbiast intersect1dtwherettolisttlent atom_densitytzerostnum_of_channelsty_dimtz_dimtrangetfloort num_3d_pixelR<tshapettheanotconfigtfloatXtndimagetfilterstgaussian_filtertNonetwritetdumptstrtclose(Et ptf_file_nametpdb_dirt target_RESt target_ATOMt site_namet pixel_sizetstdt output_dirtptf_IDtextracted_ptf_filetbox_sizet box_x_mint box_x_maxt box_y_mint box_y_maxt box_z_mint box_z_maxtptf_filet ptf_linest pdb_groupR€R‹tPDB_IDR?R@RAtdat_numt extracted_boxt all_ptf_linest pdb_file_nametpdb_fileR„R…RxRyRzR{RORMt actual_atomR>t res_atomsR=tlabelRTRaRgttransformed_postx_indexty_indextz_indext final_indextindRQtbox_oritnew_pos_in_boxt atom_countt thresholdtsamplegatatomstx_newty_newtz_newtbin_xtbin_ytbin_ztX_smoothtjt sample_time_t((sn/cstor/stanford/rbaltman/users/wtorng/DEEP_LEARNING/3DCNN/FSCNN_models/User_test/code/cut_box_site_test_pdb.pytcut_box_site_test_pdb“sð                               HPPP Y  '          ---1  '#'#  ((tsysRXt scipy.ndimageRœtsetsRtjsonR”RR­tBio.PDBRR¦R™R§R¨R«R¤RŸRžtlabel_atom_type_dictR;RRRVRhRnR†Rê(((sn/cstor/stanford/rbaltman/users/wtorng/DEEP_LEARNING/3DCNN/FSCNN_models/User_test/code/cut_box_site_test_pdb.pytsX      ’."(%%%"%(+%"(.74    2