B
    d L                 @   sH  d dl mZmZmZ d dlZd dlmZ d dlmZ d dlm	Z	 d dl
mZmZ d dlmZ d dlZd dlmZ d dlZd dlZd dlZd dlZd d	lmZ d2ddZd3ddZddddddddddddddddddd Zd!d" Zd4d$d%Zd5d&d'ZG d(d) d)ZG d*d+ d+Z G d,d- d-Z!d.d/ Z"e#d0krDe"ej$d1d  dS )6    )absolute_importdivisionprint_functionN)easy_pickle)dlite)Sorry)PropertyFileSidechainAngles)monomer_library)flex)rangeFc             C   s`   t jtjdd}|d kr\t jd}|d kr\t jtjdd}|d kr\| s\td|S )NZ	chem_dataZrotarama_dataZext_ref_filesa  Can't find chem_data/rotarama_data/ directory:
  Inside chem_data, please run
    svn --quiet --non-interactive --trust-server-cert export https://github.com/rlabduke/reference_data.git/trunk/Top8000/Top8000_rotamer_pct_contour_grids rotarama_data
    svn --quiet --non-interactive --trust-server-cert --force export https://github.com/rlabduke/reference_data.git/trunk/Top8000/Top8000_ramachandran_pct_contour_grids rotarama_data
    mmtbx.rebuild_rotarama_cache
  to resolve this problem.)libtbxenvfind_in_repositoriesospathjoinr   )optionalresult r   y/mnt/filia/a/genomebrowser/www/genomebrowser/fleming/tools/molprobity/modules/cctbx_project/mmtbx/rotamer/rotamer_eval.pyfind_rotarama_data_dir   s    r   c             C   s"   | d krt  } ttj| dS )Nzrotarama.dlite)r   r   	target_dbr   r   r   )rotarama_data_dirr   r   r   open_rotarama_dlite!   s    r   argasnaspcysglngluhisileleulysmetZphetyrproserthrtrpval)r   r   r   r   r   r    r!   r"   r#   r$   r%   pher&   r'   r(   r)   tyrr*   c             C   s*   t |dd }|d k	r|| dS |j| dS )NZget_comp_comp_id)comp_id)getattrZget_comp_comp_id_direct)residuemon_lib_srvZget_funcr   r   r   mon_lib_query<   s     
r1   Tc             C   s   g }x$|   D ]}||j   qW t| j  |}g }|sdxJ| D ]}|| qPW n.|d k	rx$| D ]}||j	   qvW g }	xT|D ]L}||kr|
dd}
| dkrd}
n| dkrd}
|
|kr|	| qW |	S )N*'ZO1PZOP1ZO2PZOP2)atomsappendnamestripupperr1   resname	atom_dictZnon_hydrogen_atomsZatom_idreplace)r/   r0   ignore_hydrogensZ	atom_listatomZmlqZreference_listatZnonmissingZ	atom_tempr   r   r   eval_residue_completenessB   s,    
r@   c          	   C   sf  g }|d krt j }xH|  D ]:}x2| D ]$}x| D ]}| }	 x| D ]}
g }|
 }|
jdkrd|j|	 |j
f }nd|j|	 |
j|j
f }g }x | D ]}|jdkr|j}qW t||d}|st|dkr<|| || |r<|| n$|| || |r<|| t|dkr`|| q`W qFW q4W q"W |S )N z
%2s%5s %3sz%2s%5s%1s%3sz CA )r/   r0   r   )r
   serverZmodelsZchainsZresidue_groups
conformersZonly_residueZaltlocidresidr9   r4   r6   xyzr@   lenr5   )Zpdb_hierarchyr0   r<   Zreport_whole_resZreturn_ca_posZmissing_atom_listZmodelchainresidue_grouprC   Z	conformeritemr/   keyZca_xyzr=   r?   r   r   r   eval_sidechain_completenessZ   sF    








rL   c               @   sN   e Zd Zi ZdddZdd Zdd Zd	d
 ZdddZdddZ	dd Z
dS )RotamerEvalN8000c             C   sx  |d krt d}|| _|d kr*tjj }|d kr8tj}|| _| jdkrTd| _d}nt	d|| _
|| _t | _tj}i | _x | D ]\}}| | j|< qW t }	tjtj|	d}
t|	d}xt D ]\}}| j|d k	rq|| d }|| d	 }|j|||	d
}|jr|
r6tjtj|	|s>t	dtjtj|	|jjd}|| j|< t|||< qW d S )NTrN   g~jth?z	rota8000-z1data_version given to RotamerEval not recognized.Z	NO_UPDATE)r   z.dataz.pickle)source_pathtarget_pathpath_prefixzchem_data/rotarama_data/*.pickle files are missing or out of date.
  Please run
    mmtbx.rebuild_rotarama_cache
  to resolve this problem.
)	file_name) r	   sidechain_anglesmmtbxr
   rB   sysstdoutdata_versionoutlier_thresholdr   logr0   	RotamerIDrot_idrM   aaTablesitemsr   r   r   existsr   r   
aminoAcidsget	pair_infoneeds_updater   loadtargetweakrefref)selfrS   r0   rY   rW   Z
fileprefixZmain_aaTablesaandt_weakrefrotamer_data_dirZ	no_updater   aafile	data_filepickle_filera   ndtr   r   r   __init__   sN    

 

zRotamerEval.__init__c             C   s&   | j | }|dkrdS ||S )zEvaluates the specified rotamer from 0.0 (worst) to 1.0 (best).

       Values below 0.003 are generally considered outliers.
       If the 3-letter amino acid name is not recognized, returns None.N)r\   r`   lowerZvalueAt)rg   ZaaName	chiAnglesrn   r   r   r   evaluate   s     zRotamerEval.evaluatec             C   s   i }|  }x|D ]x}||jd kr||jdddd krT||jddd= n,||jdddd kr||jddd= |||j< qW |S )NHD   )r4   r`   r6   r;   )rg   r/   r:   r4   r=   r   r   r   get_atom_dict   s    
zRotamerEval.get_atom_dictc             C   s   |  |}| jj||dS )N)resr:   )rv   rS   measureChiAngles)rg   r/   r:   r   r   r   
chi_angles   s    
zRotamerEval.chi_anglesc       
      C   s   ||g d dkst|d k	r4| }|j  }|dkr@d S |dkrLd}| |}y | jj||d}| 	||}W n t
k
r   d S X |d krd S | jj||dd}| j||}	|	dkr|d kst|	S |	d	kr|| jkrd
S || jk rdS |	S d S )Nru   glymser%   )rw   r:   F)symmetry	EXCEPTIONrA   ZUNCLASSIFIEDOUTLIER)countAssertionErrorr4   r9   rp   r7   rv   rS   rx   rr   	Exceptionr[   	wrap_chisidentifyrX   )
rg   r/   rI   r4   r9   r:   chisvaluer   rotamer_namer   r   r   evaluate_residue   s<    


zRotamerEval.evaluate_residuec       	      C   s   ddl m}m} |d k	r.| }|j  }|dks>|dkrBd S |dkrNd}| |}y | jj	||d}| 
||}W n tk
r   d S X |d krdS ||krd	S ||krd
S dS d S )Nr   )OUTLIER_THRESHOLDALLOWED_THRESHOLDrz   alar{   r%   )rw   r:   r~   ZFavoredZAllowed)Zmmtbx.validation.rotalyzer   r   r4   r9   rp   r7   rv   rS   rx   rr   r   )	rg   r/   r   r   r4   r9   r:   r   r   r   r   r   evaluate_residue_2   s2    

zRotamerEval.evaluate_residue_2c       
      C   s   |   }tjj}||jdkr|   }| jjd|j|   |d}|d ksh|j	d k	sh|j
d krld }|d k	rd}x8|D ]0\}}tt||  }	|	|k r~|	}|}q~W |S )NZcommon_amino_acidT)Zfine_samplingr-   
atom_names
sites_cartg    eA)r4   extract_xyziotbxZpdbZcommon_residue_names_get_classr9   r0   rotamer_iteratorextract_nameproblem_messagerotamer_infor   Zmeansqrtdot)
rg   r/   Zsites_cart_resultZ	get_classr   r   Zdist_minrZrotamer_sites_cartdr   r   r   nearest_rotamer_sites_cart  s*    


z&RotamerEval.nearest_rotamer_sites_cart)NNNrN   )NN)N)__name__
__module____qualname__r\   ro   rr   rv   ry   r   r   r   r   r   r   r   rM      s      
-	 
#
 rM   c               @   sD   e Zd Zi Zdd Zdd ZdddZdd	 ZdddZdd Z	dS )rZ   c             C   s   t | jdkr|  }| tj|d}x|D ]~}|d}|d  }|d  d}|d}|d }|d }	t	||	|}
g }|| jkr| j| }|
|
 || j|< q0W d S )Nr   zrotamer_names.props=ru   " )rG   namesfind_source_dirprocessr   r   r   splitr7   NamedRotr5   )rg   
source_dirZrota_names_listlineZ
split_lineaa_nameZrangesZ
name_splitrh   Zrot_namerotrotListr   r   r   ro   ;  s     





zRotamerID.__init__c             C   s`   |  }|dkrdS || jkr(td|| ||}| j| }x|D ]}||rD|jS qDW dS )Nr   r}   zUnknown residue name: %srA   )rp   r   r   r   containsr   )rg   r   r   r   r   r   r   r   r   r   P  s     




 
zRotamerID.identifyFc             C   s.   t jtjdd}|d kr*| s*td|S )NrT   rotamerz1Can't seem to find mmtbx/rotamer/ directory.
    )r   r   r   r   r   r   r   )r   r   r   r   r   r   [  s
    zRotamerID.find_source_dirc          
   C   s   g }yt |}W n6 tk
rF } zt|d  t  W d d }~X Y nX x,|D ]$}|dsN|dkrhqNqN|| qNW |  |S )Nz file not found#
)openImportErrorprintrU   exit
startswithr5   close)rg   ZfileLocZrotaListfer   r   r   r   r   c  s    
 zRotamerID.processTc             C   s   |  }g }x\tdt|D ]J}|| d k	r\||| d  || dk rf||  d7  < q|d  qW |dkr~| ||}|S )Nr   ih  T)rp   r   rG   r5   wrap_sym)rg   r   r   r|   r   ir   r   r   r   p  s    zRotamerID.wrap_chisc             C   sp   |  }|dks(|dks(|dks(|dkrlt|d }|| d k	rl|| d ||< || dk rl||  d7  < |S )Nr   r    r+   r,   ru      r   )rp   rG   )rg   r   r   r   r   r   r   r     s     zRotamerID.wrap_symN)F)T)
r   r   r   r   ro   r   r   r   r   r   r   r   r   r   rZ   7  s   

rZ   c               @   s$   e Zd Zdd Zdd Zdd ZdS )r   c             C   s&   || _ || _dd |dD | _d S )Nc             S   s   g | ]}t |qS r   )int).0br   r   r   
<listcomp>  s    z%NamedRot.__init__.<locals>.<listcomp>z, )r   r   r   bounds)rg   rh   r   r   r   r   r   ro     s    zNamedRot.__init__c             C   s   t | jd t | j S )Nr   )strr   r   )rg   r   r   r   __str__  s    zNamedRot.__str__c             C   sv   xpt dt| jdD ]Z}||d  d ks6| j| d kr:dS ||d  | j| k sj||d  | j|d  krdS qW dS )Nr      Fru   T)r   rG   r   )rg   Zanglesr   r   r   r   r     s     zNamedRot.containsN)r   r   r   ro   r   r   r   r   r   r   r     s   r   c          H   C   s  t ddd krtd nddlm} |  ddlm} d| k}t }|jd }t jd |ksft	|
d	gdgks|t	|dgdkst	|jd
krddddddddddddddddddddg}|dd |jdd D |st	||d d!gd"st	|jd
kst	xd#d$d%d&gd'fd(d)d*d+gd'fd,d-d.gd/fd0d1d2gd3fd4d5gd6fd7d8d9d:d;gd<fd=d>gd?fd@dAdBgdCfd=dDgdEfd7dFdGdHdIgdJfd=dKgdLfd,dMdNgdOfd=dPgdQfd@dRdSgdOfdTdUdVdWgdXfd4dYgdZfdTd[d\d]gd^fd_d`gdafd dbgd	fdcdddegdffd=dggdhfd,didjgdkfdTdldmdngdofdpd&dqgdofd4drgdsfd7dtdudvdwgdxfd7dydzd{d|gd}fd,d~dgdfd(dddgdfdcddgdfd7ddddgdfdTdddgdfd,ddgdfd_dgd-fd_dgdfdcddgdOfd(dddgdfd(dddgdfdddddgdfd@ddgd<fd,ddgdfd0d*dgdfd7ddddgdfd(dddgd'fd@ddgdfdTdddgdfdcddgdfdddtddgdfd=dgdfd@ddgdfd dgdfdcddgdfdddgdfdpddgdfd,ddgdfd(dddgdfd7ddddgdfdTdddgdfd dgdfd=dgdfd@ddgdfdddgdfd@ddgdfd4dgdfd@ddgdOfdddddgdOfd@ddgdfddddd gdfgDD ]>\}}}	d||| }
|rrt|d|
|	|
|	 f  qrW x(tj D ]\}}| d k	st	qW ~~x(tj D ]\}}| d kst	qW td d S (  NT)r   z:Skipping exercise(): rotarama_data directory not availabler   )rebuild_rotarama_cache)approx_equalz	--verboser*   g      ?rN   gUV.?gǞ?g MΚi?g9TJ4?g/#?geLfa?gg'?gzA?g]V%?g?g߿yq?g:vB?gq?gJ?ga,?gތ?g]X32?gniѲ?c             S   s   g | ]}|d  qS )i  r   )r   yr   r   r   r     s    zexercise.<locals>.<listcomp>   ZSER<   gjNM?ZMETgT@gfffffeg     0f@gffffff#@ZGLNg     d@g     @f@gZZILEgfffff&N@gLd@g3333338@ZPHEgYNgyX@gW@ZVALgyfgM@ZLYSg33333egffffff@g     egfffffeg33333sT@ZTHRg,S@g$@ZLEUgQgdg333333"@g̬Q@gfffff9@gifglfg33333bgfffffHg333333*@g33333OgT@gl_@geg        g     P@gfffffF@g33333s]g333333>@ZGLUgfffffRgdgya@g33333sA@g     @Og     @@@gyRg     @Kgffffff2gfffff=@ZPROg      =gfffffV@gA@ZASPgfffff&Tg3g     K@gLN@g33333S@g33333sNgLKg     A@g333333eg33333eg333333R@gL>@ZASNgfffffJ@gfffffe@g333333E@gQg̼egfffffff@glf@g     W@gffffffNg)e@g̜b@gfffffFVg@gQg     d@g     R@gf@g|e@gA@g     K@g33333bg     `P@g      ?g33333S@gfffffFa@g333333O@g̬dgffffff?gNg333333SgBgfffff&B@g333333Kg      d@g>@g?gffffff<g33333X@g33333`@gNg33333Ng     `fgdg      *@ǧJg|fg      <@gG@ZARGgfffff6d@g33333e@gfffffe@g̬Zg     4@g33333Qgd@g33333sHgMg=@ǧS@gfffffFW@g     Ngfffffe@g[gffffffMg@g33333eg)f@g     ]@gffffffIgfffffegL1@g     e@g̬a@ǧe@g333333@g     Sg333333f@gR@g33333Kǧ\@g      `g?gfffffM@g33333R@gfffff&Ngfffffffg     S@g33333M@g33333SR@g     @Rg     c@g333333U@ZTYRgfffffOgfffffY@gV@g33333cg      bg333333?gffffffQgf@gfffffL@gSgfffff6dǧbgH@gfffffd@gIc@g9R@g	@gRgfffff_@gYB@g333333?gLRg2@gLNg̬W@g     Eg\egffffff@ZHISgfffffFQg333333Vg333333T@ge@gyP@g33333E@g      f@gT@gfffff[gC@g̼`@gycg333333;@gcg      Wg     B@g?g33333Sbg33333c@g9W@g     Wg?d   z%4.1f %4.1f %4.1fOK)r   r   Zmmtbx.command_liner   runZlibtbx.test_utilsr   rM   r\   r   ZwhereIsZ	bin2indexrW   ZlookupTablerr   r]   )argsr   r   verboser   ZtblZaeqlZ	aminoAcidrq   Z	molpValueZr_evalrh   ri   r   r   r   exercise  s    

 
















"r   __main__ru   )F)N)T)NTFF)%
__future__r   r   r   libtbx.load_envr   r   r   libtbx.utilsr   Zmmtbx.rotamer.sidechain_anglesr   r	   rT   r
   Zmmtbx.monomer_library.serverZcctbx.array_familyr   re   rU   r   Z	iotbx.pdbr   	six.movesr   r   r   r_   r1   r@   rL   rM   rZ   r   r   r   argvr   r   r   r   <module>   sZ   


   
' 3[w
