B
    bQ                 @   sj  d Z ddl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
Z
eeZG dd deZd	d
 Zdd ZG dd deZG dd deZdd ZG dd deZG dd deZG dd deZG dd deZG dd deZG dd deZG dd  d eZG d!d" d"eZG d#d$ d$eZd%d& ZG d'd( d(eZG d)d* d*eZ d+d, Z!G d-d. d.eZ"dS )/zk
This module contains classes regarding the Amoeba potential and loading in a
TINKER-based parameter file.
    N   )range)TinkerErrorTinkerWarning)OrderedDictc               @   s8   e Zd ZdZdd Zdd Zdd Zdd	 Zd
d ZdS )BookmarkedFilez: Allows setting a bookmark and rewinding to that bookmark c             O   s   t ||| _d S )N)open_stream)selfargskwargs r   :lib/python3.7/site-packages/parmed/tinker/parameterfile.py__init__   s    zBookmarkedFile.__init__c             C   s   |   | _d S )N)tellbookmark)r
   r   r   r   mark   s    zBookmarkedFile.markc             C   s   |  | j d S )N)seekr   )r
   r   r   r   rewind   s    zBookmarkedFile.rewindc             C   s   t | j|S )N)getattrr	   )r
   attrr   r   r   __getattr__   s    zBookmarkedFile.__getattr__c             C   s    y| j   W n   Y nX d S )N)r	   close)r
   r   r   r   __del__   s    zBookmarkedFile.__del__N)	__name__
__module____qualname____doc__r   r   r   r   r   r   r   r   r   r      s   r   c             C   s&   yt |  dS  tk
r    dS X d S )NTF)int
ValueError)thingr   r   r   _IS_INT#   s
    r!   c             C   s&   yt |  dS  tk
r    dS X d S )NTF)floatr   )r    r   r   r   	_IS_FLOAT*   s
    r#   c               @   sB   e Zd ZdZe ZdZdd Zedd Z	edd Z
d	d
 ZdS )
_ParamTypezI All parameter types. This caches the list of parameters for easy access  c             O   s   t ddS )z! Instantiates the parameter type zvirtual methodN)NotImplementedError)r
   r   r   r   r   r   r   8   s    z_ParamType.__init__c             C   sR   || j krDd| j|f }| j | |kr0|d7 }n|d7 }t|t || j |< d S )NzDuplicate %s type found: %sz [same parameters]z [different parameters])TypeList_param_typewarningswarnr   )clsobjkeyZwarnmsgr   r   r   register<   s    

z_ParamType.registerc             C   s   t  | _dS )zX
        Resets the cached list -- allows us to read in multiple parameter sets
        N)dictr'   )r+   r   r   r   resetG   s    z_ParamType.resetc             C   s   t | }t|t | s.tdt | t |f xNt|D ]B}|ds8|dkrPq8tt| |drbq8t| |t||kr8dS q8W dS )z' Make sure all attributes are the same z!cannot compare type %s to type %s_r'   __call__FT)type
isinstance	TypeErrordir
startswithhasattrr   )r
   otherr+   Zpropr   r   r   __eq__N   s       z_ParamType.__eq__N)r   r   r   r   r/   r'   r(   r   classmethodr.   r0   r:   r   r   r   r   r$   3   s   r$   c               @   s*   e Zd ZdZe ZdZdd Zdd ZdS )	_BondTypez Bond parameter type Zbondc             C   sP   t |t | }}t|t| | _| _dt||t||f }| | | d S )Nz%d-%d)r   r"   kreqminmaxr.   )r
   idx1idx2r=   r>   r-   r   r   r   r   d   s    z_BondType.__init__c             C   s   d| j | jf S )Nz<_BondType: k=%s; req=%s>)r=   r>   )r
   r   r   r   __repr__j   s    z_BondType.__repr__N)	r   r   r   r   r/   r'   r(   r   rC   r   r   r   r   r<   ^   s
   r<   c             O   s   | dkrt ||S t||S )z1 Factory that returns the appropriate angle type ZfF)_FourierAngleType
_AngleType)typecoder   r   r   r   r   get_angle_typeo   s    
rG   c               @   s(   e Zd Ze ZdZdddZdd ZdS )rE   ZangleNc       	      C   s   t |t |t |  }}}dt|||t||f }t|t| | _| _|d k	r`t|| _nd | _|d k	rzt|| _nd | _| | | d S )Nz%d-%d-%d)	r   r?   r@   r"   r=   theteqtheteq2theteq3r.   )	r
   rA   rB   idx3r=   rH   rI   rJ   r-   r   r   r   r   z   s    z_AngleType.__init__c             C   sH   d| j | jf }| jd k	r(|d| j 7 }| jd k	r@|d| j 7 }|d S )Nz<_AngleType: k=%s; theteq=%sz; theteq2=%sz; theteq3=%s>)r=   rH   rI   rJ   )r
   retvalr   r   r   rC      s    

z_AngleType.__repr__)NN)r   r   r   r/   r'   r(   r   rC   r   r   r   r   rE   u   s   
rE   c               @   s   e Zd Zdd Zdd ZdS )rD   c             C   sd   t |t |t |  }}}dt|||t||f }t|| _t|| _t|| _| | | d S )Nz%d-%d-%d)r   r?   r@   r"   r=   rH   periodicityr.   )r
   rA   rB   rK   r=   rH   rN   r-   r   r   r   r      s    


z_FourierAngleType.__init__c             C   s   d| j | j| jf S )Nz4<_FourierAngleType: k=%s; theteq=%s; periodicity=%s>)r=   rH   rN   )r
   r   r   r   rC      s    z_FourierAngleType.__repr__N)r   r   r   r   rC   r   r   r   r   rD      s   rD   c               @   s&   e Zd Ze ZdZdd Zdd ZdS )_StretchBendTypezstretch-bendc             C   s\   t |t |t |  }}}dt|||t||f }t|t| | _| _| | | d S )Nz%d-%d-%d)r   r?   r@   r"   k1k2r.   )r
   rA   rB   rK   rP   rQ   r-   r   r   r   r      s    z_StretchBendType.__init__c             C   s   d| j | jf S )Nz <_StretchBendType: k1=%s; k2=%s>)rP   rQ   )r
   r   r   r   rC      s    z_StretchBendType.__repr__N)r   r   r   r/   r'   r(   r   rC   r   r   r   r   rO      s   rO   c               @   s&   e Zd Ze ZdZdd Zdd ZdS )_UreyBradleyTypezurey-bradleyc             C   s\   t |t |t |  }}}dt|||t||f }t|t| | _| _| | | d S )Nz%d-%d-%d)r   r?   r@   r"   r=   r>   r.   )r
   rA   rB   rK   r=   r>   r-   r   r   r   r      s    z_UreyBradleyType.__init__c             C   s   d| j | jf S )Nz <_UreyBradleyType: k=%s; req=%s>)r=   r>   )r
   r   r   r   rC      s    z_UreyBradleyType.__repr__N)r   r   r   r/   r'   r(   r   rC   r   r   r   r   rR      s   rR   c               @   s&   e Zd Ze ZdZdd Zdd ZdS )_OPBendTypezout-of-plane bendingc             C   sZ   t |t |t |t |f\}}}}t|| _d||t||t||f }| | | d S )Nz%d-%d-%d-%d)r   r"   r=   r?   r@   r.   )r
   rA   rB   rK   idx4r=   r-   r   r   r   r      s    $
z_OPBendType.__init__c             C   s
   d| j  S )Nz<_OPBendType: k=%s>)r=   )r
   r   r   r   rC      s    z_OPBendType.__repr__N)r   r   r   r/   r'   r(   r   rC   r   r   r   r   rS      s   rS   c               @   s&   e Zd Ze ZdZdd Zdd ZdS )_DihedralTypeZdihedralc             G   s  t |t |t |t |f\}}}}||k sD||krV||k sD||krVd||||f }n(||ksn||kr~||kr~d||||f }g g g   | _| _| _xhtt|d D ]T}| jt||d   | jt||d d   | jt||d d   qW | | | d S )Nz%d-%d-%d-%d      r   )	r   r=   phaserN   r   lenappendr"   r.   )r
   rA   rB   rK   rT   r   r-   ir   r   r   r      s    $  z_DihedralType.__init__c             C   s   d| j | j| jf S )Nz'<_DihedralType: k=%r; phase=%r; per=%r>)r=   rX   rN   )r
   r   r   r   rC      s    
z_DihedralType.__repr__N)r   r   r   r/   r'   r(   r   rC   r   r   r   r   rU      s   rU   c               @   s&   e Zd Ze ZdZdd Zdd ZdS )_PiTorsionTypez
pi-torsionc             C   sD   t |t | }}dt||t||f }t|| _| | | d S )Nz%d-%d)r   r?   r@   r"   r=   r.   )r
   rA   rB   r=   r-   r   r   r   r      s    
z_PiTorsionType.__init__c             C   s
   d| j  S )Nz<_PiTorsionType: k=%s>)r=   )r
   r   r   r   rC      s    z_PiTorsionType.__repr__N)r   r   r   r/   r'   r(   r   rC   r   r   r   r   r\      s   r\   c               @   s.   e Zd Ze ZdZdd Zdd Zdd ZdS )	_TorsionTorsionTypeztorsion-torsionc             C   sH   dd |D }dt | }t|t| | _| _t | _| | | d S )Nc             s   s   | ]}t |V  qd S )N)r   ).0r[   r   r   r   	<genexpr>   s    z/_TorsionTorsionType.__init__.<locals>.<genexpr>z%d-%d-%d-%d-%d)tupler   nxnyr   potential_gridr.   )r
   indexesra   rb   r-   r   r   r   r      s
    z_TorsionTorsionType.__init__c             C   s   t || jt |t |f< d S )N)r"   rc   )r
   xyZ	potentialr   r   r   	add_point   s    z_TorsionTorsionType.add_pointc             C   s   d| j | jf S )Nz)<_TorsionTorsion: %d x %d potential grid>)ra   rb   )r
   r   r   r   rC     s    z_TorsionTorsionType.__repr__N)	r   r   r   r/   r'   r(   r   rg   rC   r   r   r   r   r]      s
   r]   c               @   s.   e Zd Ze ZdZdd Zdd Zdd ZdS )	_MultipoleType	multipolec             C   s4   dd |D }d |}t|g| _| | | d S )Nc             s   s   | ]}t |V  qd S )N)str)r^   r[   r   r   r   r_     s    z*_MultipoleType.__init__.<locals>.<genexpr>-)joinr"   potential_termsr.   )r
   rd   Zp1r-   r   r   r   r     s    
z_MultipoleType.__init__c             C   s
   d| j  S )Nz<_MultipoleType: terms=%r>)rm   )r
   r   r   r   rC     s    z_MultipoleType.__repr__c             C   s"   x|D ]}| j t| qW d S )N)rm   rZ   r"   )r
   ZtermsZtermr   r   r   	add_terms  s    
z_MultipoleType.add_termsN)	r   r   r   r/   r'   r(   r   rC   rn   r   r   r   r   rh     s
   rh   c             C   s6   t | } y
tj|  S  tk
r0   t| |||S X dS )zv
    Factory for getting an _AtomType, but making sure that only one instance of
    a particular type is created
    N)r   	_AtomTyper'   KeyError)indexatomic_numbermassvalencer   r   r   get_atom_type  s
    
ru   c               @   s@   e Zd ZdZe Zdd ZedddZdd Z	e
d	d
 ZdS )ro   z An atom type c             C   s6   t || _t || _t|| _t || _| tj|< d S )N)r   rq   rr   r"   rs   rt   ro   r'   )r
   rq   rr   rs   rt   r   r   r   r   )  s
    



z_AtomType.__init__Nc             C   s@   t jt|  }t||_t||_|d k	r6t||_nd |_d S )N)ro   r'   r   r"   sizeepsilon	reduction)rq   rv   rw   rx   instr   r   r   set_vdw_params0  s    

z_AtomType.set_vdw_paramsc             C   sB   d| j | j| j| jf }t| dr:|d| j| j| jf 7 }|d S )Nz,<_AtomType: idx=%d; elem=%d; mass=%s; val=%drv   z; size=%s; eps=%s; red=%srL   )rq   rr   rs   rt   r8   rv   rw   rx   )r
   rM   r   r   r   rC   :  s    
z_AtomType.__repr__c             C   s   t  | _d S )N)r/   r'   )r+   r   r   r   r0   B  s    z_AtomType.reset)N)r   r   r   r   r/   r'   r   staticmethodrz   rC   r;   r0   r   r   r   r   ro   %  s   	ro   c               @   s   e Zd ZdZe Zdd Zdd Zdd Zdd	 Z	d
d Z
dd Zdd Zdd Zdd ZeeedZeeedZee	edZee
edZeeedZeeedZeeedZdd Zdd Zedd ZdS )_Atomz An atom in a parameter set c             C   s*   || _ || _t||||| _| tj|< d S )N)namedescriptionru   r3   r|   AtomList)r
   rq   	typeindexr}   descriprr   rs   valr   r   r   r   I  s    z_Atom.__init__c             C   s   | j jS )N)r3   rq   )r
   r   r   r   
_typeindexP      z_Atom._typeindexc             C   s   | j jS )N)r3   rr   )r
   r   r   r   _atomic_numberQ  r   z_Atom._atomic_numberc             C   s   | j jS )N)r3   rr   )r
   r   r   r   _elementR  r   z_Atom._elementc             C   s   | j jS )N)r3   rt   )r
   r   r   r   _valenceS  r   z_Atom._valencec             C   s   | j jS )N)r3   rv   )r
   r   r   r   _sizeT  r   z_Atom._sizec             C   s   | j jS )N)r3   rw   )r
   r   r   r   _epsilonU  r   z_Atom._epsilonc             C   s   | j jS )N)r3   rx   )r
   r   r   r   
_reductionV  r   z_Atom._reductionc             C   s   t dd S )NzCannot set this attribute)r&   )r
   r   r   r   _blockedW  r   z_Atom._blocked)fgetfsetc             C   s(   t || _t || _dd |D | _d S )Nc             S   s   g | ]}t |qS r   )r   )r^   r[   r   r   r   
<listcomp>d  s    z,_Atom.set_polarizability.<locals>.<listcomp>)r"   polarizabilitytholeconnected_types)r
   r   r   r   r   r   r   set_polarizabilitya  s    

z_Atom.set_polarizabilityc             C   s>   d| j | j| jf }t| dr6|d| j| j| jf 7 }|d S )Nz<_Atom "%s": name=%s; type=%dr   z-; dipole pol=%s; thole=%s; connected atoms=%rrL   )r~   r}   r   r8   r   r   r   )r
   Zretstrr   r   r   rC   f  s    
z_Atom.__repr__c             C   s   t  | _d S )N)r/   r'   )r+   r   r   r   r0   n  s    z_Atom.resetN)r   r   r   r   r/   r   r   r   r   r   r   r   r   r   r   propertyr   rr   elementrt   rv   rw   rx   r   rC   r;   r0   r   r   r   r   r|   F  s*   r|   c               C   s\   t   t  t  t  t  t  t  t  t	  t
  t  dS )z
    Resets all of the TypeList instances (without destroying the data inside
    them) so we can load multiple parameter sets
    N)r<   r0   rE   rD   rO   rR   rS   rU   r\   r]   rh   ro   r   r   r   r   r0   t  s    r0   c               @   s:   e Zd ZdZedejZedZd	ddZ	dd Z
dS )
AmoebaParameterSetzV
    Contains all of the parameters found in an Amoeba parameter file from TINKER
    zJatom *(\d+) *(\d+) *([A-Za-z\-\+\*0-9]+) *"(.+)" *(\d+) *(\d+\.\d+) *(\d+)zangle([ 345fF])Nc             C   sf   t  | _tj| _tj| _tj| _t	j| _
tj| _tj| _tj| _tj| _tj| _|d k	rb| | d S )N)r/   atomsro   r'   Z
atom_typesr<   ZbondsrE   ZanglesrO   Zstretch_bendsrR   Zurey_bradleysrS   ZopbendsrU   Z	dihedralsr]   Ztorsion_torsionsrh   Z
multipolesload_parameter_file)r
   fnamer   r   r   r     s    zAmoebaParameterSet.__init__c          
   C   sx	  t  | _t|d}d}| dd}x|r|sd|krBd}P |d  }|d|d  }|sz| dd}q(| }t|d	krq(t	|d
 rt
|d
 | j|d  < n>t|d
 rt|d
 | j|d  < n|d
 | j|d  < | dd}q(W |stdx0| dd  dkrH| dd}qW xt| dd  dkr| j|}| \}}}	}
}}}tt
|||	|
|||| jt
|< | dd}qLW x0| dd  dkr| dd}qW xF| dd  dkr8tj| d
d   | dd}qW x0| dd  dkrj| dd}q<W xD| dd  dkrt| d
d   | dd}qnW | j|}x(|s| dd}| j|}qW x|rvy&t| d f| d
d   W n> tk
rT   tdt| d | d
d   Y nX | dd}| j|}qW |  x6| dd  dkr|r| dd}qW xR| dd  dkr
|r
|  t| d
d   | dd}qW |   | dd}x6| dd  dkrZ|rZ| dd}q&W xR| dd  dkr|r|  t!| d
d   | dd}q^W |   | dd}x6| dd  dkr|r| dd}qW xR| dd  dkrR|rR|  t"| d
d   | dd}qW |   | dd}x6| dd  dkr|r| dd}qnW xR| dd  dkr|r|  t#| d
d   | dd}qW |   | dd}x2| dd dkrB|rB| dd}qW xN| dd dkr|r|  t$| d
d   | dd}qFW |   | dd}x2| dd dkr|r| dd}qW x| dd dkr~|r~| }t%|d
d |d |d }| dd}x.| r`|j&|   | dd}q4W | dd}|  qW |   | dd}x6| dd  dkr|r| dd}qW x| dd  dkrt|rt| }t'|d
d |d }|(|   |(|   |(|   |(|   | dd}|  qW |   | dd}x2| dd d kr|r| dd}qW x| dd d k	rd|	rd| }t
|d
 }y(| j| )|d	 |d! |dd  W n2 t*k
	rN   | j| )|d	 |d! g  Y nX | dd}qW |+  t,  dS )"zm
        Parses a parameter file and loads all of the parameters found into data
        structures.
        rF	 zLiterature ReferencesT#Nr   rW   r   z&Could not find force field attributes.   zatom    zvdw zbond z%s, %s   zstrbnd 	   z	ureybrad zopbend    ztorsion zpitors ztortors    
   z
multipole z	polarize rV   )-r/   Z
attributesr   readlinereplacestriprq   splitrY   r!   r   lowerr#   r"   r   lstripatomrematchgroupsr|   r   ro   rz   r<   anglererG   r5   LOGGERdebugreprr   rO   r   rR   rS   rU   r\   r]   rg   rh   rn   r   
IndexErrorr   r0   )r
   r   fZdone_with_attributeslineZwordsZrematchZnumZtypenumr}   r   Zanumrs   r   Ztortorri   rq   r   r   r   r     s   
 &&"" "" "" ""   "" 
"z&AmoebaParameterSet.load_parameter_file)N)r   r   r   r   recompileIr   r   r   r   r   r   r   r   r     s   

r   )#r   ZloggingZutils.six.movesr   
exceptionsr   r   collectionsr   r   r)   Z	getLoggerr   r   objectr   r!   r#   r$   r<   rG   rE   rD   rO   rR   rS   rU   r\   r]   rh   ru   ro   r|   r0   r   r   r   r   r   <module>   s6   
	+!.