B
    b_                 @   s   d Z ddlmZmZmZ ddl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 dd	lmZ dd
lmZ ddlmZ ddlmZ ddlmZ ddlZdZde ZG dd deZ ee
G dd de!Z"G dd de"Z#G dd de"Z$dS )a  
This is a pure-python module for reading and writing ASCII Amber structure
files, like trajectories and restarts.  Because it is pure-python and
formatted-ASCII, parsing and writing these files are expected to be slow. Binary
alternatives (like DCD and NetCDF, provided in netcdffiles.py) are strongly
encouraged, but these are provided for more complete compatibility and for
instances where the prequisites may not be installed.
    )divisionprint_functionabsolute_import)ceilN)FileFormatType)genopen)add_metaclass)range)string_types)	Structure)Atom)Vec3gGzt4@   c               @   s   e Zd ZdZdS )_FileEOFz For control flow N)__name__
__module____qualname____doc__ r   r   4lib/python3.7/site-packages/parmed/amber/asciicrd.pyr      s   r   c               @   sF   e Zd ZdZdZdZdddZdd Zedd	 Z	d
d Z
dd ZdS )_AmberAsciiCoordinateFilea%  
    Abstract base class for interacting with ASCII coordinate files.
    Opens a new ASCII coordinate file and either parses it (loading
    everything into memory) or sets it up for writing.

    Parameters
    ----------
    fname : str
        File name to open
    natom : int
        Number of atoms in the system
    hasbox : bool
        Does the system have PBCs?
    mode : str={'r', 'w'}
        Whether to open this file for 'r'eading or 'w'riting
    title : str, optional
        Title to write to a new trajectory (when mode='w')
    Nrc             C   s$  |dkrd| _ n(|dkr&d| _ d| _ntdt| j t|trVt||| _d| _	n$t
|drn|| _d| _	ntd	| || _|| _| jd krtd
| jd | j | _| jd | j| j  | _| j	rd| _| j dkr|   n>| j dkr |d kr| jd| j  n| j| d  d S )Nr   oldwnewFz%s mode must be 'r' or 'w'TreadzUnsupported type for %szThis object must be subclassed   z%s

)_status	_writebox
ValueErrortyper   
isinstancer
   r   _file_own_handlehasattr	TypeErrornatomhasboxCRDS_PER_LINENotImplementedError_full_lines_per_frame_nextrasclosed_parsewriteDEFAULT_TITLErstrip)selffnamer'   r(   modetitler   r   r   __init__6   s8    





z"_AmberAsciiCoordinateFile.__init__c             C   s   t ddS )z Handles actual file parsing zvirtual method not overwrittenN)r*   )r2   r   r   r   r.   [   s    z _AmberAsciiCoordinateFile._parsec             C   s   dd | j d D tj S )Nc             S   s   g | ]}t | qS r   )r   ).0Zxyzr   r   r   
<listcomp>a   s    z7_AmberAsciiCoordinateFile.positions.<locals>.<listcomp>r   )_coordinatesuZ	angstroms)r2   r   r   r   	positions_   s    z#_AmberAsciiCoordinateFile.positionsc             C   s    | j r| jp| j  d| _dS )z Close the open file handler TN)r$   r-   r#   close)r2   r   r   r   r<   c   s    z_AmberAsciiCoordinateFile.closec             C   s.   y| j p| j  W n tk
r(   Y nX dS )z+ Make sure the open file handler is closed N)r-   r#   r<   AttributeError)r2   r   r   r   __del__i   s    z!_AmberAsciiCoordinateFile.__del__)r   N)r   r   r   r   r0   r)   r6   r.   propertyr;   r<   r>   r   r   r   r   r      s   
%r   c                   s   e Zd ZdZedd ZdZdZd  fd
d	Ze	d!ddZ
dd Zedd Zejdd Zedd Zejdd Zedd Zedd Zedd Zejdd Zejdd Zejdd Z  ZS )"AmberAsciiRestartas  
    Parser for the Amber ASCII inpcrd/restart file format

    Parameters
    ----------
        fname : str
            File name to open
        mode : str={'r', 'w'}
            Whether to open this file for 'r'eading or 'w'riting
        natom : int, optional
            Number of atoms in the system (necessary when mode='w')
        hasbox : bool, optional
            Does the system have PBCs? Necessary when mode='w'
        title : str, optional
            Title to write to a new trajectory (when mode='w')
        time : float, optional
            The time to write to the restart file in ps. Default is 0.
    c       	   	      s  t tr4td  fddtdD }   nFtdrztdrztdrz }fddtdD }| |d	  }t	|d
kst	|d	k rdS yt
|d }t|d	  W n& tk
r   dS  tk
r   Y nX y|dkrdS d}x|d
d D ]}|d	7 }||kr P xZtdD ]N}|d }||d  dkrLdS t|||d   ||d  dkr*dS q*W |d	7 }||krP x^tdD ]R}|d d }||d  dkrdS t|||d   ||d  dkrdS qW qW W n ttfk
r   dS X dS )a0   Identifies the file type as an Amber restart/inpcrd file

        Parameters
        ----------
        filename : str
            Name of the file to check format for

        Returns
        -------
        is_fmt : bool
            True if it is an Amber restart/inpcrd file. False otherwise
        r   c                s   g | ]}   qS r   )readline)r7   i)fr   r   r8      s    z/AmberAsciiRestart.id_format.<locals>.<listcomp>   rA   seektellc                s   g | ]}   qS r   )rA   )r7   rB   )filenamer   r   r8      s    r      Fr   Nr         .   
0123456789$   T)r"   r
   r   r	   r<   r%   rF   rE   splitlenintfloatr    
IndexError)	rG   linesZcurZwordsr'   rB   linejZj12r   )rC   rG   r   	id_format   s^    





  

  zAmberAsciiRestart.id_format   zrestart created by ParmEdr   r   N        c                s>   d| _ d| _d| _d| _t|| _tt| ||||| dS )z
        For restart files, natom and hasbox are determined automatically for
        mode='r', and can be determined at write-time when the coordinates are
        set.
        FN)	_coords_written_cell_lengths_written_cell_angles_written_vels_writtenrR   timesuperr@   r6   )r2   r3   r4   r'   r(   r5   r^   )	__class__r   r   r6      s    
zAmberAsciiRestart.__init__Fc             C   s`   | |}|rXt  }x$t|jD ]}|jt ddd qW |j|_|jd |_|j|_|S |S d S )NZXXXr   )ZresnameZresnum)r   r	   r'   Zadd_atomr   boxcoordinatesr9   )clsrG   Z	structurer2   obj_r   r   r   parse   s    zAmberAsciiRestart.parsec          	   C   sR  | j  }| j   |d  | _t|d   d | _yt|d   d | _	W n t
k
rt   d| _	Y nX x|d  s|  qxW t|tt| jd d krd | _| _nt|tt| jd d krd	| _d| _nzt|tdt| jd  d krd| _d	| _nHt|tdt| jd  d krLd	 | _| _ntd
t|| jf t| jdf| _| jrt| jdf| _| jrtd| _td| _d}|tt| jd  }d}xt||D ]}|| }t|dd }t|dd }t|dd }	|||	g| j|< |d7 }y4t|dd }
t|dd }t|dd }W n tk
rv   Y nX |
||g| j|< |d7 }qW | jd| jdf| _|}| jr|tt| jd  }d}xt||D ]}|| }t|dd t }t|dd t }t|dd t }	|||	g| j|< |d7 }y@t|dd t }
t|dd t }t|dd t }W n tk
r   Y nX |
||g| j|< |d7 }qW |}| jd| jdf| _| jrN|| }t|dd t|dd t|dd g| jdd< t|dd t|dd t|dd g| jdd< dS )a  
        This method parses the data out of the ASCII restart file and creates
        self._coordinates and self._velocities as np.ndarray(natom*3) arrays

        This method is called automatically for 'old' restart files and should
        not be called by external callers
        r   r   g        g       @rH   Fr   Tz8Badly formatted restart file. Has %d lines for %d atoms.rI      rN   0   <   H   N)r#   	readlinesr<   stripr5   rQ   rO   r'   rR   r^   rS   poprP   r   r(   hasvelsRuntimeErrornpzerosr9   _velocities_cell_lengths_cell_anglesr	   r    reshapeVELSCALE)r2   rT   	startlineZendlineidxrB   rU   Zx1Zy1Zz1Zx2Zy2Zz2r   r   r   r.      s    	

$$zAmberAsciiRestart._parsec             C   s"   | j dkrt| dstd| jS )Nr   r9   zCoordinates not yet set)r   r%   rp   r9   )r2   r   r   r   rb   E  s    zAmberAsciiRestart.coordinatesc          	   C   s2  | j dkrtdtj|dd }| jdkrVt|d| j krVtdt|| jf | jrdtdt|d | _|	d	| jdf| _
| jd
| j| jf  d}d}xdt| jD ]V}|d }| j||| ||d  ||d  f  |d7 }|d dkr| jd qW | jd dkr(| jd d| _d S )Nr   z(Cannot set coordinates on an old restartF)copyr   r   zGot %d coordinates for %d atomsz&Coordinates have already been written.rg   z
%5d%15.7e
z%12.7f%12.7f%12.7fr   rH   r   T)r   rp   rq   arrayravelr'   rP   r    rZ   rv   r9   r#   r/   r^   r	   )r2   stuffnumwritfmtrB   i3r   r   r   rb   K  s,    
*  c             C   s,   | j dkrt| dstd| js&d S | jS )Nr   rs   zVelocities not set yet)r   r%   rp   ro   rs   )r2   r   r   r   
velocitiesd  s
    zAmberAsciiRestart.velocitiesc          	   C   s6  | j dkrtdtj|dd }| js2td| js>| jrFtd| jrTtdt	|d| j
 kr|td	t	|| j
f |d
| j
df| _d}d}xpt| j
D ]b}|d }| j||| t ||d  t ||d  t f  |d7 }|d dkr| jd qW | j
d dkr&| jd d | _| _d S )Nr   z'Cannot set velocities on an old restartF)rz   z)Coordinates must be set before velocitiesz.Velocities must be written before the box infozCan only write velocities oncer   zGot %d velocities for %d atoms.rg   z%12.7f%12.7f%12.7fr   r   rH   r   T)r   rp   rq   r{   r|   rZ   r[   r\   r]   rP   r'   r    rv   rs   r	   r#   r/   ONEVELSCALEro   )r2   r}   r   r~   rB   r   r   r   r   r   l  s4    
  c             C   s,   | j dkrt| dstd| js&d S | jS )Nr   rt   zCell lengths not yet available)r   r%   rp   r(   rt   )r2   r   r   r   cell_lengths  s
    zAmberAsciiRestart.cell_lengthsc             C   s,   | j dkrt| dstd| js&d S | jS )Nr   ru   zCell angles not yet available)r   r%   rp   r(   ru   )r2   r   r   r   cell_angles  s
    zAmberAsciiRestart.cell_anglesc             C   s   | j dkr&t| drt| ds&td| js0dS td}| j| j }}|d |d |d	 g|dd
< |d |d |d	 g|d
d< |S )z' Combined cell lengths and cell angles r   rt   ru   zCell parameters not yet setNrX   r   r   rH   r   )r   r%   rp   r(   rq   rr   r   r   )r2   ra   ZlengthsZanglesr   r   r   ra     s    

zAmberAsciiRestart.boxc             C   s~   | j dkrtd| js td| jr.tdt|dkrBtdtj|dd| _| j	
d	|d
 |d |d f  d| _d S )Nr   z&Cannot set cell lengths on old restartz&Coordinates must be written before boxz Can only write cell lengths oncer   z#Expected 3 numbers for cell lengthsF)rz   z%12.7f%12.7f%12.7fr   r   rH   T)r   rp   rZ   r[   rP   r    rq   r{   rt   r#   r/   )r2   r}   r   r   r   r     s    
"c             C   s   | j dkrtd| js td| js.td| jr<tdt|dkrPtdtj|dd	| _	| j
d
|d |d |d f  d| _d S )Nr   z%Cannot set cell angles on old restartz&Coordinates must be written before boxz%Must write cell lengths before angleszCan only write cell angles oncer   z"Expected 3 numbers for cell anglesF)rz   z%12.7f%12.7f%12.7f
r   r   rH   T)r   rp   rZ   r[   r\   rP   r    rq   r{   ru   r#   r/   )r2   r}   r   r   r   r     s    
"c             C   s    |dd | _ |dd | _dS )z. Writes both the cell lengths and cell angles Nr   )r   r   )r2   r}   r   r   r   ra     s    )r   r   NNrY   )F)r   r   r   r   staticmethodrW   r)   r0   r6   classmethodrf   r.   r?   rb   setterr   r   r   ra   __classcell__r   r   )r`   r   r@   r   s&   @ _r@   c               @   sd   e Zd ZdZdZdZdZedd Zdd Z	e
d	d
 Ze
dd Ze
dd Zdd Zdd ZdS )
AmberMdcrdz
    A class to parse Amber ASCII trajectory files. This is *much* slower than
    parsing NetCDF files (or the equivalent parsing done in a compiled language
    like C or C++). For large trajectories, this may be significant.
    )r'   r(   
   ztrajectory created by ParmEdc          	      s   t | d  fddtdD }   yxztdD ]n}|d7 }x`tdD ]T}|d }|| |d  d	krnd
S t|| ||d   || |d  dkrJd
S qJW q4W W n ttfk
r   d
S X dS )a   Identifies the file type as an Amber mdcrd file

        Parameters
        ----------
        filename : str
            Name of the file to check format for

        Returns
        -------
        is_fmt : bool
            True if it is an Amber mdcrd file. False otherwise
        r   c                s   g | ]}   qS r   )rA   )r7   rB   )rC   r   r   r8     s    z(AmberMdcrd.id_format.<locals>.<listcomp>rD   rJ   r   r      rK   F   rM   T)r   r	   r<   rR   rS   r    )rG   rT   rB   rV   Zj8r   )rC   r   rW     s     
 zAmberMdcrd.id_formatc                s  | j     | _td| _td| _tdd| j	 d}tdd| j
 d}yBx: rt| jd } sxt td}d}| j    st xLt| jD ]>} st  fdd|D |||d < |d7 }| j   qW | j
r fdd|D |||| j
 < | jrj| j    s.t t dd |d< t dd	 |d
< t d	d |d< t| j|f| _t| j|f| _qXW W n6 tk
r   tdt Y n tk
r   Y nX | jd| jdf| _| jd| _| j   dS )af  
        This method parses the data out of the mdcrd file and creates
        self._coordinates as a list of np.ndarray(1, natom, 3) and
        self.cell_lengths as a list of np.ndarray(3) for each frame in the
        trajectory. This method is called automatically for 'old' trajectory
        files and should not be called by external callers.
        r   r   r   c                s    g | ]}t  ||d   qS )r   )rR   )r7   rV   )rU   r   r   r8     s    z%AmberMdcrd._parse.<locals>.<listcomp>r   c                s    g | ]}t  ||d   qS )r   )rR   )r7   rV   )rU   r   r   r8     s    N   r   rh   rH   zEUnexpected EOF in parsing mdcrd. natom and/or hasbox are likely wrongrg   )rg   r   )r#   rA   rm   r5   rq   Zndarrayr9   r   r	   r)   r,   rr   r'   r   StopIterationr+   r(   rR   Zconcatenate	_warningswarnRuntimeWarningrv   r<   )r2   ZmainiterZ	extraiterframeZcellry   rB   r   )rU   r   r.     sR    


 

   
 
zAmberMdcrd._parsec             C   s   | j S )N)r9   )r2   r   r   r   rb   ,  s    zAmberMdcrd.coordinatesc             C   s   | j jd S )Nr   )r9   shape)r2   r   r   r   r   0  s    zAmberMdcrd.framec             C   s   | j dkstd| jS )Nr   z Cannot access box of a new mdcrd)r   rp   r   )r2   r   r   r   ra   4  s    
zAmberMdcrd.boxc             C   s   | j dkstdy| }W n tk
r2   Y nX | jrBtdt|d| j kr\tdd}d}xPt| j	D ]B}|d }x(tdD ]}| j
d	|||    qW | j
d
 qpW | jr
|d | d }x.|| jd k r| j
d	||   |d7 }qW | j
d
 | j| _| j
  dS )a\  
        Prints 'stuff' (which must be either an iterable of 3*natom or have an
        attribute 'flatten' that converts it into an iterable of 3*natom) to the
        open file handler. Can only be called on a 'new' mdcrd, and adds these
        coordinates to the current end of the file.

        Parameters
        ----------
        stuff : array or iterable
            This must be an iterable of length 3*natom or a numpy array that can
            be flattened to a 3*natom-length array

        Raises
        ------
        If the coordinate file is an old one being parsed or if you are
        currently expected to provide unit cell dimensions, a RuntimeError is
        raised. If the provided coordinate data does not have length 3*natom, or
        cannot be ``flatten()``ed to create a 3*natom array, a ValueError is
        raised.
        r   z#Cannot print frames to an old mdcrdz*Box information not written for last framer   z3add_coordinates requires an array of length natom*3r   rg   r   z%8.3fr   r   N)r   rp   Zflattenr=   r   rP   r'   r    r	   r+   r#   r/   r,   r(   flush)r2   r}   rB   rV   Zi10Zextrar   r   r   add_coordinates:  s2    
zAmberMdcrd.add_coordinatesc             C   sj   | j dkstd| js tdt|dkr4td| jd|d |d |d	 f  d
| _| j  dS )a/  
        Prints 'stuff' (which must be a 3-element list, array.array, tuple, or
        np.ndarray) as the box lengths for this frame

        Parameters
        ----------
        stuff : array or iterable
            This must be an iterable of length 3 with the box lengths

        Raises
        ------
        If the coordinate file is an old one being parsed or if you are
        currently expected to provide coordinates, a RuntimeError is raised.
        raised. If the provided box lengths are not length 3, a ValueError is
        raised.
        r   z Cannot print box to an old mdcrdz(Should not be writing box info right nowr   z%add_box requires an array of length 3z%8.3f%8.3f%8.3f
r   r   rH   FN)r   rp   r   rP   r    r#   r/   r   )r2   r}   r   r   r   add_boxn  s    
"zAmberMdcrd.add_boxN)r   r   r   r   Z
extra_argsr)   r0   r   rW   r.   r?   rb   r   ra   r   r   r   r   r   r   r     s   !44r   )%r   Z
__future__r   r   r   Zmathr   Znumpyrq   Zparmed.formats.registryr   Zparmed.utils.ior   Zparmed.utils.sixr   Zparmed.utils.six.movesr	   r
   Zparmed.structurer   Zparmed.topologyobjectsr   Zparmed.unitZunitr:   Zparmed.vec3r   warningsr   rw   r   	Exceptionr   objectr   r@   r   r   r   r   r   <module>   s,   S  \