B
    âêõb<s  ã               @   s¾   d Z ddlmZmZmZ yddlZW n ek
r<   dZY nX 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ZeeƒG d	d
„ d
eƒƒZeeƒG dd„ deƒƒZdS )až  
This module contains classes for reading and writing Amber NetCDF-style files,
including both restarts and trajectories. The NetCDF engine used here is pulled
from the scipy distribution, and depends *only* on numpy. It is available
through parmed.utils.netcdf.

This module contains objects relevant to Amber NetCDF files. The use() function
is responsible for selecting the API based on a default choice or user-selection
(the latter is really only helpful for development to ensure that all packages
work correctly---there is no difference from a user perspective). ALL
NetCDF-file manipulation that the parmed/amber package does should be
contained in this module.
é    )ÚdivisionÚprint_functionÚabsolute_importN)Ú__version__)ÚFileFormatType)Úunit)Únetcdf_file)Úadd_metaclassc               @   s*  e Zd ZdZedd„ ƒZd,dd„Zed-d	d
„ƒZedd„ ƒZ	e	Z
edd„ ƒZej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j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jd#d"„ ƒZed$d%„ ƒZd&d'„ Zd(d)„ Zd*d+„ ZdS ).ÚNetCDFRestartz- Class to read or write NetCDF restart files c          	   C   sŽ   |   d¡s|   d¡s|   d¡r"dS yt| ddd}W n ttfk
rL   dS X z0y|j ¡ dkrddS W n tk
rz   dS X dS | ¡  X d	S )
az   Identifies the file type as an Amber NetCDF restart file

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

        Returns
        -------
        is_fmt : bool
            True if it is an Amber NetCDF restart file. False otherwise

        Notes
        -----
        Remote NetCDF files cannot be loaded
        zhttp://zhttps://zftp://FÚr)ÚmmapÚAMBERRESTARTTN)Ú
startswithÚ
NetCDFFileÚ	TypeErrorÚOSErrorÚConventionsÚdecodeÚAttributeErrorÚclose)ÚfilenameÚf© r   ú7lib/python3.7/site-packages/parmed/amber/netcdffiles.pyÚ	id_format#   s    
zNetCDFRestart.id_formatr   c             C   sT   d| _ | d¡r,tdk	r,tj||dd| _n$| d¡r@t d¡ t||dd| _dS )a  
        Opens a NetCDF File. The main constructor should never be called
        directly.  The alternative "open_old" and "open_new" constructors should
        be called instead for reading existing or writing new NetCDF files,
        respectively.
        FÚwNÚNETCDF3_64BIT)ÚformatzˆCould not find netCDF4 module. Falling back on scipy implementation, which can significantly slow down simulations if used as a reporter)r   )Úclosedr   ÚncÚDatasetÚ_ncfileÚwarningsÚwarnr   )ÚselfÚfnameÚmoder   r   r   Ú__init__G   s    

zNetCDFRestart.__init__Ú Nc	             C   sä  |dk	r||d dkr*d}	|dkrzt dƒ‚q€|d dkrrd}	|dkrJt dƒ‚x|D ]}
|
d	krPt d
ƒ‚qPW t|ƒ}q€t dƒ‚nd}	| |dƒ}|j}t|ƒ|_t|ƒ|_d|_d|_t|ƒ|_	d|_
d|_ttƒ|_| dd¡ | d|¡ d|_||_|j	|_	|r8| dd¡ | dd¡ | dd¡ d|_d|_d|_|	dkrT| d|¡ ||_| dd¡ | ddd¡}d|_| ddd¡}t td ƒ¡|dd…< |jr0| ddd!¡}t td"ƒ¡|d< t td#ƒ¡|d< t td$ƒ¡|d%< | ddd&¡}d'\|d< |d< |d%< | d(dd&¡}d)|_| d*dd+¡}d,|_| d-dd.¡}d)|_|jr†| d/dd.¡}d0|_t d1¡|_d1|_tdk	r†| d2¡ |	dkr®| d3dd¡}d4|_||d< n2|	dkrà| d5d6d7¡ | d8d6d7¡}||dd…< |S )9aÏ  
        Opens a new NetCDF file and sets the attributes

        Parameters
        ----------
        fname : str
            Name of the new file to open (overwritten)
        natom : int
            The number of atoms in the system
        box : bool
            Whether unit cell information is written or not
        vels : bool
            Whether velocity information is written or not
        title : str=''
            The title to write to the NetCDF restart file
        remd : str=None
            None -- No REMD information is written
            'T[emperature]' -- target temperature (or pH) will be written
            'M[ulti-D]' -- remd_dimtypes will be written
        remd_dimtypes : iterable of int=None
            Array of exchange types for each group. The length will be the REMD
            dimension (if `remd` above is "M[ulti-D]")

        Notes
        -----
        `remd` is case-insensitive, and done based on first-letter matching
        Nr   ZtTÚTEMPERATUREz+temp must be specified for T-REMD restarts.ZmMÚMULTIzLremd_dimtypes must be given for multi-D REMD, and must have the same length.)é   é   z=remd_dimtypes only supports dimension types 1 and 3 currentlyz+remd must be None, T[emperature] or M[ulti]r   r   z1.0Ú
AmberToolsÚParmEdÚspatialr,   ÚatomÚcell_spatialÚlabelé   Úcell_angularÚremd_dimensionÚtimer+   Úd)r6   Ú
picosecondÚc)r/   Úxyz)r4   r2   Úalphazbeta Úgammaé   )r1   )ÚaÚbr9   Úcell_lengthsÚangstromÚcell_angles)r4   ÚdegreeÚcoordinates)r0   r/   Ú
velocitieszangstrom/picosecondg®Gázt4@FÚtemp0ÚkelvinÚremd_indicesÚi)r5   Úremd_dimtype)Ú
ValueErrorÚlenr!   ÚboolÚhasboxÚhasvelsr   ÚConventionVersionÚstrÚtitleÚapplicationÚprogramr   ÚprogramVersionÚcreateDimensionr/   r0   r1   r2   r4   r5   ÚcreateVariableÚunitsÚnpÚasarrayÚlistZfloat32Úscale_factorÚvelocity_scaler   Úset_auto_maskandscale)Úclsr%   ÚnatomÚboxÚvelsrR   ÚremdZtempZremd_dimtypesZ	remd_typeZdtr5   ÚinstÚncfileÚvr   r   r   Úopen_newX   s¦    















zNetCDFRestart.open_newc             C   sæ   | |dƒ}|j }|j ¡ |_|j ¡ |_|j ¡ |_|j ¡ |_t|dƒrX|j ¡ |_nd|_t|dƒrv|j ¡ |_nd|_x*|j	D ] }|dkr’q„t
|||j	| ƒ q„W d|jk|_d|jkoÆd|jk|_|jrâ|jd }|j|_|S )	z¹
        Opens the NetCDF file and sets the global attributes that the file sets

        Parameters
        ----------
        fname : str
            Name of the file to read
        r   rS   NrR   r6   rE   r@   rB   )r!   r   r   rP   rT   rU   ÚhasattrrS   rR   Ú
dimensionsÚsetattrÚ	variablesrO   rN   r\   r]   )r_   r%   rd   re   Údimrb   r   r   r   Úopen_oldÕ   s.    



 

zNetCDFRestart.open_oldc             C   s&   | j jd d d … }| d| jdf¡S )NrD   r+   r,   )r!   rk   Úreshaper0   )r$   Zcoordsr   r   r   rD   ü   s    zNetCDFRestart.coordinatesc             C   s:   t j|dd | jdf¡}|| jjd d d …< |  ¡  d S )NF)Úcopyr,   rD   )rY   Úarrayrn   r0   r!   rk   Úflush)r$   Ústuffr   r   r   rD     s    c             C   s<   d| j jkr8| j jd d d … }| d| jdf¡| j S d S )NrE   éÿÿÿÿr,   )r!   rk   rn   r0   r]   )r$   rb   r   r   r   rE     s    zNetCDFRestart.velocitiesc             C   s4   t  || jdf¡| j | jjd d d …< |  ¡  d S )Nr,   rE   )rY   rn   r0   r]   r!   rk   rq   )r$   rr   r   r   r   rE     s    (c             C   s$   d| j jkr | j jd d d … S d S )Nr@   )r!   rk   )r$   r   r   r   r@     s    zNetCDFRestart.cell_lengthsc             C   s&   t  |¡| jjd d d …< |  ¡  d S )Nr@   )rY   rZ   r!   rk   rq   )r$   rr   r   r   r   r@     s    c             C   s$   d| j jkr | j jd d d … S d S )NrB   )r!   rk   )r$   r   r   r   rB     s    zNetCDFRestart.cell_anglesc             C   s&   t  |¡| jjd d d …< |  ¡  d S )NrB   )rY   rZ   r!   rk   rq   )r$   rr   r   r   r   rB   "  s    c             C   s4   | j d k	r0| jd k	r0| j | j }}t ||f¡S d S )N)r@   rB   rY   Zconcatenate)r$   ZlengZangr   r   r   ra   '  s    zNetCDFRestart.boxc             C   s    |d d… | _ |dd … | _d S )Nr,   )r@   rB   )r$   rr   r   r   r   ra   -  s    c             C   s   | j jd  ¡ S )Nr6   )r!   rk   ÚgetValue)r$   r   r   r   r6   2  s    zNetCDFRestart.timec             C   s    t |ƒ| jjd d< |  ¡  d S )Nr6   r   )Úfloatr!   rk   rq   )r$   rr   r   r   r   r6   6  s    c             C   s   | j jd  ¡ S )NrF   )r!   rk   rt   )r$   r   r   r   rF   ;  s    zNetCDFRestart.temp0c             C   s   | j jd d d … S )NrH   )r!   rk   )r$   r   r   r   rH   ?  s    zNetCDFRestart.remd_indicesc             C   s*   t j|dd| jjd d d …< |  ¡  d S )NrI   )ÚdtyperH   )rY   rZ   r!   rk   rq   )r$   rr   r   r   r   rH   C  s    c             C   s   | j jd d d … S )NrJ   )r!   rk   )r$   r   r   r   rJ   H  s    zNetCDFRestart.remd_dimtypec             C   s   | j sd| _ | j ¡  d S )NT)r   r!   r   )r$   r   r   r   r   L  s    zNetCDFRestart.closec             C   s   | j pt| dƒo| j ¡  d S )Nr!   )r   rh   r!   r   )r$   r   r   r   Ú__del__Q  s    zNetCDFRestart.__del__c             C   s2   y| j  ¡  W n tk
r,   | j  ¡  Y nX d S )N)r!   rq   r   Úsync)r$   r   r   r   rq   T  s    zNetCDFRestart.flush)r   )r(   NNN)Ú__name__Ú
__module__Ú__qualname__Ú__doc__Ústaticmethodr   r'   Úclassmethodrg   rm   ÚparseÚpropertyrD   ÚsetterrE   r@   rB   ra   r6   rF   rH   rJ   r   rw   rq   r   r   r   r   r
      s6   $
 {%r
   c               @   s  e Zd ZdZedd„ ƒZd4dd„Zed5dd„ƒZedd„ ƒZ	e
dd„ ƒZdd„ Ze
dd„ ƒZdd„ Ze
dd„ ƒZdd„ Ze
dd„ ƒZeZd6dd„ZeZe
dd „ ƒZd!d"„ Ze
d#d$„ ƒZd%d&„ Ze
d'd(„ ƒZd)d*„ Ze
d+d,„ ƒZejd-d,„ ƒZd.d/„ Zd0d1„ Zd2d3„ Zd	S )7Ú
NetCDFTrajap   Class to read or write NetCDF restart files

    Parameters
    ----------
    fname : str
        Name of the file to open
    mode : str
        Mode to open in:
            - 'w' means write-mode
            - 'r' means read-mode

    Notes
    -----
    You should use the open_new and open_old alternative constructors instead of
    the default constructor
    c          	   C   sŽ   |   d¡s|   d¡s|   d¡r"dS yt| ddd}W n ttfk
rL   dS X z0y|j ¡ dkrddS W n tk
rz   dS X dS | ¡  X d	S )
a€   Identifies the file type as an Amber NetCDF trajectory file

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

        Returns
        -------
        is_fmt : bool
            True if it is an Amber NetCDF trajectory file. False otherwise

        Notes
        -----
        Remote NetCDF files cannot be loaded
        zhttp://zhttps://zftp://Fr   )r   ÚAMBERTN)r   r   r   r   r   r   r   r   )r   r   r   r   r   r   n  s    
zNetCDFTraj.id_formatr   c             C   sT   d| _ | d¡r,tdk	r,tj||dd| _n$| d¡r@t d¡ t||dd| _dS )z Opens a NetCDF File Fr   Nr   )r   zˆCould not find netCDF4 module. Falling back on scipy implementation, which can significantly slow down simulations if used as a reporter)r   )r   r   r   r    r!   r"   r#   r   )r$   r%   r&   r   r   r   r'   ’  s    

zNetCDFTraj.__init__TFNr(   c
             C   s  | |dƒ}
|
j }|dk	rd|d dkr,d|
_qj|d dkrZd|
_|dkrNtdƒ‚t|ƒ|
_qjtd	ƒ‚nd|
_t|ƒ|
_t|ƒ|
_t|ƒ|
_t|ƒ|
_	d
|_
d|_d|_d|_t|_d|_d
|
_
d|
_d|
_d|
_t|
_|j|
_| dd¡ | dd¡ | d|¡ |
jdkr| d|
j¡ dd|  |
_|
_|
_|
jrl| dd¡ | dd¡ | dd¡ d\|
_|
_|
_| ddd¡}t tdƒ¡|dd…< |
jrî| ddd¡}t tdƒ¡|dd…< | ddd¡}t tdƒtd ƒtd!ƒg¡|dd…< | d"d#d$¡}d%|_|
jr$| d&d#d'¡}d(|_d|
_|
jrf| d)d#d'¡}d*|_d+ |
_|_d|
_ t!dk	rf| "d,¡ |
j	rˆ| d-d#d'¡}d.|_d|
_#|
jr¾| d/d0d1¡}d(|_| d2d0d3¡}d4|_d|
_$|
jdkræ| d5d0d$¡}d6|_d|
_%n.|
jdkr| d7d8d9¡ | d:d8d;¡ d|
_%d|
_&|
S )<aê  
        Opens a new NetCDF file and sets the attributes

        Parameters
        ----------
        fname : str
            Name of the new file to open (overwritten)
        natom : int
            Number of atoms in the restart
        box : bool
            Indicates if cell lengths and angles are written to the NetCDF file
        crds : bool=True
            Indicates if coordinates are written to the NetCDF file
        vels : bool=False
            Indicates if velocities are written to the NetCDF file
        frcs : bool=False
            Indicates if forces are written to the NetCDF file
        remd : str=None
            'T[emperature]' if replica temperature is written
            'M[ulti]' if Multi-D REMD information is written
            None if no REMD information is written
        remd_dimension : int=None
            If remd above is 'M[ulti]', this is how many REMD dimensions exist
        title : str=''
            The title of the NetCDF trajectory file
        r   Nr   ZTtr)   ZMmr*   z-remd_dimension must be given for multi-D REMDz&remd must be T[emperature] or M[ultiD]rƒ   z1.0r-   r.   zParmEd-created trajectoryÚframer/   r,   r0   r5   r1   r4   r2   r3   )r,   r,   r3   r9   )r/   r:   )r1   Úabc)r4   r2   r;   zbeta r<   r6   r   )r„   r8   rD   )r„   r0   r/   rA   rE   zangstrom/picosecondg®Gázt4@FÚforceszkilocalorie/mole/angstromr@   r7   )r„   r1   rB   )r„   r4   rC   rF   rG   rH   rI   )r„   r5   rJ   )r5   )'r!   rc   rK   Úintr5   rM   rN   rO   ÚhascrdsÚhasfrcsr   rP   rS   rT   r   rU   rR   rV   r„   r/   r0   r1   r4   r2   rW   rY   rZ   r[   rX   Ú_last_crd_framer]   r\   Ú_last_vel_framer   r^   Ú_last_frc_frameÚ_last_box_frameÚ_last_remd_frameÚ_last_time_frame)r_   r%   r`   ra   Zcrdsrb   Zfrcsrc   r5   rR   rd   re   rf   r   r   r   rg   ž  s°    





$

zNetCDFTraj.open_newc             C   sî  | |dƒ}|j }|j ¡ |_|j ¡ |_|j ¡ |_|j ¡ |_t|dƒrX|j ¡ |_nd|_t|dƒrv|j ¡ |_nd|_x |j	D ]}t
|||j	| ƒ q„W d|jk|_d|jk|_d|jk|_d|jkoÔd	|jk|_|jrøt |jd dd… ¡|_|jrPy|jd j}W n tk
r*   d
}Y nX t |jd dd… ¡| |_||_|jrrt |jd dd… ¡|_|jdkrêd|jkr¤t|jd dd… ƒ|_nF|jr¼|jjd |_n.|jrÔ|jjd |_n|jrê|jjd |_|S )zÓ
        Opens the NetCDF file and sets the global attributes that the file sets

        Parameters
        ----------
        fname : str
            File name of the trajectory to open. It must exist
        r   rS   NrR   rD   rE   r†   r@   rB   r+   r6   r   )r!   r   r   rP   rT   rU   rh   rS   rR   ri   rj   rk   rˆ   rO   r‰   rN   rY   rp   Ú_coordinatesr\   r   Ú_velocitiesr]   Ú_forcesr„   rL   Úshape)r_   r%   rd   re   rl   Zscaler   r   r   rm     sP    





zNetCDFTraj.open_oldc             C   s   | j S )N)r   )r$   r   r   r   rD   R  s    zNetCDFTraj.coordinatesc             C   s^   t  |¡r| t j¡}tj|dd}t || jdf¡| jj	d | j
< |  j
d7  _
|  ¡  dS )a  
        Adds a new coordinate frame to the end of a NetCDF trajectory. This
        should only be called on objects created with the "open_new"
        constructor.

        Parameters
        ----------
        stuff : iterable of floats or distance Quantity
            This array of floats is converted into a numpy array of shape
            (natom, 3). It can be passed either in the 2-D format of
            [ [x1, y1, z1], [x2, y2, z2], ... ] or in the 1-D format of
            [x1, y1, z1, x2, y2, z2, ... ].
        r   )rv   r,   rD   r+   N)ÚuÚis_quantityÚvalue_in_unitÚ	angstromsrY   rZ   rn   r0   r!   rk   rŠ   rq   )r$   rr   r   r   r   Úadd_coordinatesV  s    
 zNetCDFTraj.add_coordinatesc             C   s   | j S )N)r‘   )r$   r   r   r   rE   l  s    zNetCDFTraj.velocitiesc             C   sf   t  |¡r| t jt j ¡}t |¡}t || jdf¡| j	 | j
jd | j< |  jd7  _|  ¡  dS )a  
        Adds a new velocities frame to the end of a NetCDF trajectory. This
        should only be called on objects created with the "open_new"
        constructor.

        Parameters
        ----------
        stuff : iterable of floats or distance/time Quantity
            This array of floats is converted into a numpy array of shape
            (natom, 3). It can be passed either in the 2-D format of
            [ [x1, y1, z1], [x2, y2, z2], ... ] or in the 1-D format of
            [x1, y1, z1, x2, y2, z2, ... ].
        r,   rE   r+   N)r”   r•   r–   rA   r8   rY   rZ   rn   r0   r]   r!   rk   r‹   rq   )r$   rr   r   r   r   Úadd_velocitiesp  s    

&zNetCDFTraj.add_velocitiesc             C   s   | j S )N)r’   )r$   r   r   r   r†   †  s    zNetCDFTraj.forcesc             C   sV   t  |¡r| t jt j ¡ t || jdf¡| jj	d | j
< |  j
d7  _
|  ¡  dS )a  
        Adds a new coordinate frame to the end of a NetCDF trajectory. This
        should only be called on objects created with the "open_new"
        constructor.

        Parameters
        ----------
        stuff : iterable of floats or energy/distance Quantity
            This array of floats is converted into a numpy array of shape
            (natom, 3). It can be passed either in the 2-D format of
            [ [x1, y1, z1], [x2, y2, z2], ... ] or in the 1-D format of
            [x1, y1, z1, x2, y2, z2, ... ].
        r,   r†   r+   N)r”   r•   r–   Zkilocalories_per_moler—   rY   rn   r0   r!   rk   rŒ   rq   )r$   rr   r   r   r   Ú
add_forcesŠ  s
    
 zNetCDFTraj.add_forcesc             C   sH   y.t  | jjd d d … | jjd d d … f¡S  tk
rB   d S X d S )Nr@   rB   )rY   Zhstackr!   rk   ÚKeyError)r$   r   r   r   Úcell_lengths_anglesŸ  s
    zNetCDFTraj.cell_lengths_anglesc                sò   dd„ ‰ t |ƒdkr$|dkr$tdƒ‚t |ƒdkr@|dk	r@tdƒ‚t |ƒdkrlt |ƒdksdt |ƒdkrltdƒ‚|dkrŽ‡ fd	d
„|dd… D ƒ}‡ fdd
„|dd… D ƒ}t |¡| jjd | j< t |¡| jjd | j< |  jd7  _|  ¡  dS )aë  
        Adds a new cell length and angle frame to the end of a NetCDF
        trajectory.  This should only be called on objects created with the
        "open_new" constructor.

        Parameters
        ----------
        lengths : array of 3 (or 6) floats (or Quantities)
            This should be a 1-D array of 3 or 6 elements.  If 6 elements,
            `angles` should be None (below) and the first 3 elements are the box
            lengths (angstroms) and the last 3 are the box angles (degrees).
        angles : 3-item iterable = None
            These are the box angles (if lengths contains only 3 elements) in
            degrees. Must be a 1-D array of 3 elements or None if lengths
            includes angles as well.
        c             S   s   t  | ¡r|  |¡S | S )N)r”   r•   r–   )ÚxZdesired_unitsr   r   r   Ústrip_unitsº  s    
 
z7NetCDFTraj.add_cell_lengths_angles.<locals>.strip_unitsr,   Nz%Both lengths and angles are required.é   z!Angles can be provided only once.z-6 numbers expected -- 3 lengths and 3 angles.c                s   g | ]}ˆ |t jƒ‘qS r   )r”   Zdegrees)Ú.0r   )rž   r   r   ú
<listcomp>Ä  s    z6NetCDFTraj.add_cell_lengths_angles.<locals>.<listcomp>c                s   g | ]}ˆ |t jƒ‘qS r   )r”   r—   )r    r   )rž   r   r   r¡   Å  s    r@   rB   r+   )rL   rK   rY   rZ   r!   rk   r   rq   )r$   ZlengthsZanglesr   )rž   r   Úadd_cell_lengths_angles©  s    $z"NetCDFTraj.add_cell_lengths_anglesc             C   s   | j jd d d … S )Nr6   )r!   rk   )r$   r   r   r   r6   Ï  s    zNetCDFTraj.timec             C   sF   t  |¡r| t j¡}t|ƒ| jjd | j< |  jd7  _|  ¡  dS )zÈ Adds the time to the current frame of the NetCDF file

        Parameters
        ----------
        stuff : float or time-dimension Quantity
            The time to add to the current frame
        r6   r+   N)	r”   r•   r–   Zpicosecondsru   r!   rk   r   rq   )r$   rr   r   r   r   Úadd_timeÓ  s
    
 zNetCDFTraj.add_timec             C   s   | j jd d d … S )NrH   )r!   rk   )r$   r   r   r   rH   à  s    zNetCDFTraj.remd_indicesc             C   s6   t j|dd| jjd | j< |  jd7  _|  ¡  dS )z¸ Add REMD indices to the current frame of the NetCDF file

        Parameters
        ----------
        stuff : iterable of int
            The indices in each REMD dimension
        rI   )rv   rH   r+   N)rY   rZ   r!   rk   rŽ   rq   )r$   rr   r   r   r   Úadd_remd_indicesä  s    	zNetCDFTraj.add_remd_indicesc             C   s   | j jd d d … S )NrF   )r!   rk   )r$   r   r   r   rF   ñ  s    zNetCDFTraj.temp0c             C   sF   t  |¡r| t j¡}t|ƒ| jjd | j< |  jd7  _|  ¡  dS )zÛ The temperature to add to the current frame of the NetCDF file

        Parameters
        ----------
        stuff : float or temperature Quantity
            The temperature to add to the current NetCDF file
        rF   r+   N)	r”   r•   r–   rG   ru   r!   rk   rŽ   rq   )r$   rr   r   r   r   Ú	add_temp0õ  s
    
 zNetCDFTraj.add_temp0c             C   s   | j jd d d … S )NrJ   )r!   rk   )r$   r   r   r   rJ     s    zNetCDFTraj.remd_dimtypec             C   s*   t j|dd| jjd d d …< |  ¡  d S )NrI   )rv   rJ   )rY   rZ   r!   rk   rq   )r$   rr   r   r   r   rJ     s    c             C   s   | j s| j ¡  d| _ dS )z Closes the NetCDF file TN)r   r!   r   )r$   r   r   r   r     s    
zNetCDFTraj.closec             C   s   | j pt| dƒo| j ¡  d S )Nr!   )r   rh   r!   r   )r$   r   r   r   rw     s    zNetCDFTraj.__del__c             C   s   t d kr| j ¡  d S )N)r   r!   rq   )r$   r   r   r   rq     s    zNetCDFTraj.flush)r   )TFFNNr(   )N) ry   rz   r{   r|   r}   r   r'   r~   rg   rm   r€   rD   r˜   rE   r™   r†   rš   rœ   ra   r¢   Zadd_boxr6   r£   rH   r¤   rF   r¥   rJ   r   r   rw   rq   r   r   r   r   r‚   [  s8   $
 |6
$r‚   )r|   Z
__future__r   r   r   ZnetCDF4r   ÚImportErrorZnumpyrY   Zparmedr   Zparmed.formats.registryr   r   r”   Zparmed.utils.netcdfr   r   Zparmed.utils.sixr	   r"   Úobjectr
   r‚   r   r   r   r   Ú<module>   s$   
  =