diff --git a/package/MDAnalysis/coordinates/CRD.py b/package/MDAnalysis/coordinates/CRD.py index 79f88634971..952b63b4a17 100644 --- a/package/MDAnalysis/coordinates/CRD.py +++ b/package/MDAnalysis/coordinates/CRD.py @@ -110,7 +110,7 @@ def Writer(self, filename, **kwargs): return CRDWriter(filename, **kwargs) -class CRDWriter(base.WriterBase): +class CRDWriter(base.SingleFrameWriterBase): """CRD writer that implements the CHARMM CRD coordinate format. It automatically writes the CHARMM EXT extended format if there @@ -148,13 +148,16 @@ class CRDWriter(base.WriterBase): "NUMATOMS": "{0:5d}\n", } - def __init__(self, filename, **kwargs): + def __init__(self, filename, n_atoms=None, append=False, **kwargs): """ Parameters ---------- filename : str or :class:`~MDAnalysis.lib.util.NamedStream` name of the output file or a stream - + n_atoms : int (optional) + number of atoms in the coordinate file; + append : bool (optional) + append to an existing CRD file; default is to overwrite extended : bool (optional) By default, noextended CRD format is used [``False``]. However, extended CRD format can be forced by @@ -165,13 +168,14 @@ def __init__(self, filename, **kwargs): .. versionadded:: 2.2.0 """ - self.filename = util.filename(filename, ext='crd') + filename = util.filename(filename, ext='crd') + super().__init__(filename, n_atoms, append) self.crd = None # account for explicit crd format, if requested self.extended = kwargs.pop("extended", False) - def write(self, selection, frame=None): + def _write_next_frame(self, selection, frame=None): """Write selection at current trajectory frame to file. Parameters diff --git a/package/MDAnalysis/coordinates/DCD.py b/package/MDAnalysis/coordinates/DCD.py index 3b340b78917..ca182589ac8 100644 --- a/package/MDAnalysis/coordinates/DCD.py +++ b/package/MDAnalysis/coordinates/DCD.py @@ -338,6 +338,7 @@ class DCDWriter(base.WriterBase): def __init__(self, filename, n_atoms, + append=False, convert_units=True, step=1, dt=1, @@ -351,6 +352,9 @@ def __init__(self, filename of trajectory n_atoms : int number of atoms to be written + append : bool (optional) + append to an existing trajectory if True, default is to + overwrite an existing file: ``False`` convert_units : bool (optional) convert from MDAnalysis units to format specific units step : int (optional) @@ -376,11 +380,10 @@ def __init__(self, General writer arguments """ - self.filename = filename self._convert_units = convert_units if n_atoms is None: raise ValueError("n_atoms argument is required") - self.n_atoms = n_atoms + super().__init__(filename, n_atoms, append) self._file = DCDFile(self.filename, 'w') self.step = step self.dt = dt diff --git a/package/MDAnalysis/coordinates/FHIAIMS.py b/package/MDAnalysis/coordinates/FHIAIMS.py index ce5bf8259e7..a8ce679d90d 100644 --- a/package/MDAnalysis/coordinates/FHIAIMS.py +++ b/package/MDAnalysis/coordinates/FHIAIMS.py @@ -205,7 +205,7 @@ def Writer(self, filename, n_atoms=None, **kwargs): return FHIAIMSWriter(filename, n_atoms=n_atoms, **kwargs) -class FHIAIMSWriter(base.WriterBase): +class FHIAIMSWriter(base.SingleFrameWriterBase): """FHI-AIMS Writer. Single frame writer for the `FHI-AIMS`_ format. Writes geometry (3D and @@ -228,7 +228,8 @@ class FHIAIMSWriter(base.WriterBase): 'box_triclinic': "lattice_vector {box[0]:12.8f} {box[1]:12.8f} {box[2]:12.8f}\nlattice_vector {box[3]:12.8f} {box[4]:12.8f} {box[5]:12.8f}\nlattice_vector {box[6]:12.8f} {box[7]:12.8f} {box[8]:12.8f}\n" } - def __init__(self, filename, convert_units=True, n_atoms=None, **kwargs): + def __init__(self, filename, convert_units=True, n_atoms=None, + append=False, **kwargs): """Set up the FHI-AIMS Writer Parameters @@ -239,8 +240,8 @@ def __init__(self, filename, convert_units=True, n_atoms=None, **kwargs): number of atoms """ - self.filename = util.filename(filename, ext='.in', keep=True) - self.n_atoms = n_atoms + filename = util.filename(filename, ext='.in', keep=True) + super().__init__(filename, n_atoms, append) def _write_next_frame(self, obj): """Write selection at current trajectory frame to file. diff --git a/package/MDAnalysis/coordinates/GRO.py b/package/MDAnalysis/coordinates/GRO.py index aff46e5b86c..e5cbc714617 100644 --- a/package/MDAnalysis/coordinates/GRO.py +++ b/package/MDAnalysis/coordinates/GRO.py @@ -261,7 +261,7 @@ def Writer(self, filename, n_atoms=None, **kwargs): return GROWriter(filename, n_atoms=n_atoms, **kwargs) -class GROWriter(base.WriterBase): +class GROWriter(base.SingleFrameWriterBase): """GRO Writer that conforms to the Trajectory API. Will attempt to write the following information from the topology: @@ -311,7 +311,8 @@ class GROWriter(base.WriterBase): } fmt['xyz_v'] = fmt['xyz'][:-1] + "{vel[0]:8.4f}{vel[1]:8.4f}{vel[2]:8.4f}\n" - def __init__(self, filename, convert_units=True, n_atoms=None, **kwargs): + def __init__(self, filename, convert_units=True, n_atoms=None, + append=False, **kwargs): """Set up a GROWriter with a precision of 3 decimal places. Parameters @@ -326,7 +327,9 @@ def __init__(self, filename, convert_units=True, n_atoms=None, **kwargs): By default, all the atoms were reindexed to have a atom id starting from 1. [``True``] However, this behaviour can be turned off by specifying `reindex` ``=False``. - + append : bool (optional) + If ``True``, open file in append mode. [``False``] + Note ---- To use the reindex keyword, user can follow the two examples given @@ -344,13 +347,13 @@ def __init__(self, filename, convert_units=True, n_atoms=None, **kwargs): w.write(u.atoms) """ - self.filename = util.filename(filename, ext='gro', keep=True) - self.n_atoms = n_atoms + filename = util.filename(filename, ext='gro', keep=True) + super().__init__(filename, n_atoms, append) self.reindex = kwargs.pop('reindex', True) self.convert_units = convert_units # convert length and time to base units - def write(self, obj): + def _write_next_frame(self, obj): """Write selection at current trajectory frame to file. Parameters diff --git a/package/MDAnalysis/coordinates/H5MD.py b/package/MDAnalysis/coordinates/H5MD.py index dfb4a8210fe..b7aac858171 100644 --- a/package/MDAnalysis/coordinates/H5MD.py +++ b/package/MDAnalysis/coordinates/H5MD.py @@ -1050,7 +1050,8 @@ class H5MDWriter(base.WriterBase): @due.dcite(Doi("10.1016/j.cpc.2014.01.018"), description="Specifications of the H5MD standard", path=__name__, version='1.1') - def __init__(self, filename, n_atoms, n_frames=None, driver=None, + def __init__(self, filename, n_atoms, + append=False, n_frames=None, driver=None, convert_units=True, chunks=None, compression=None, compression_opts=None, positions=True, velocities=True, forces=True, timeunit=None, lengthunit=None, @@ -1060,14 +1061,13 @@ def __init__(self, filename, n_atoms, n_frames=None, driver=None, if not HAS_H5PY: raise RuntimeError("H5MDWriter: Please install h5py") - self.filename = filename if n_atoms == 0: raise ValueError("H5MDWriter: no atoms in output trajectory") + super().__init__(filename, n_atoms, append) self._driver = driver if self._driver == 'mpio': raise ValueError("H5MDWriter: parallel writing with MPI I/O " "is not currently supported.") - self.n_atoms = n_atoms self.n_frames = n_frames self.chunks = (1, n_atoms, 3) if chunks is None else chunks if self.chunks is False and self.n_frames is None: diff --git a/package/MDAnalysis/coordinates/MOL2.py b/package/MDAnalysis/coordinates/MOL2.py index 5b2e93de978..d0a6a17f9d5 100644 --- a/package/MDAnalysis/coordinates/MOL2.py +++ b/package/MDAnalysis/coordinates/MOL2.py @@ -287,17 +287,21 @@ class MOL2Writer(base.WriterBase): multiframe = True units = {'time': None, 'length': 'Angstrom'} - def __init__(self, filename, n_atoms=None, convert_units=True): + def __init__(self, filename, n_atoms=None, append=False, + convert_units=True): """Create a new MOL2Writer Parameters ---------- filename: str name of output file + append: bool (optional) + If ``True``, append to an existing file. If ``False``, overwrite + the file. Default is ``False``. convert_units: bool (optional) units are converted to the MDAnalysis base format; [``True``] """ - self.filename = filename + super().__init__(filename, n_atoms, append) self.convert_units = convert_units # convert length and time to base units self.frames_written = 0 diff --git a/package/MDAnalysis/coordinates/NAMDBIN.py b/package/MDAnalysis/coordinates/NAMDBIN.py index 2400a34fd77..1c179eb71bb 100644 --- a/package/MDAnalysis/coordinates/NAMDBIN.py +++ b/package/MDAnalysis/coordinates/NAMDBIN.py @@ -91,7 +91,7 @@ def Writer(self, filename, **kwargs): return NAMDBINWriter(filename, **kwargs) -class NAMDBINWriter(base.WriterBase): +class NAMDBINWriter(base.SingleFrameWriterBase): """Writer for NAMD binary coordinate files. @@ -100,7 +100,7 @@ class NAMDBINWriter(base.WriterBase): format = ['COOR', 'NAMDBIN'] units = {'time': None, 'length': 'Angstrom'} - def __init__(self, filename, n_atoms=None, **kwargs): + def __init__(self, filename, n_atoms=None, append=False, **kwargs): """ Parameters ---------- @@ -109,7 +109,8 @@ def __init__(self, filename, n_atoms=None, **kwargs): n_atoms : int number of atoms for the output coordinate """ - self.filename = util.filename(filename) + filename = util.filename(filename) + super().__init__(filename, n_atoms, append) def _write_next_frame(self, obj): """Write information associated with ``obj`` at current frame into diff --git a/package/MDAnalysis/coordinates/PDB.py b/package/MDAnalysis/coordinates/PDB.py index 48b9fad7409..a18977ba386 100644 --- a/package/MDAnalysis/coordinates/PDB.py +++ b/package/MDAnalysis/coordinates/PDB.py @@ -461,7 +461,7 @@ def close(self): self._pdbfile.close() -class PDBWriter(base.WriterBase): +class PDBWriter(base.SingleFrameWriterBase): """PDB writer that implements a subset of the `PDB 3.3 standard`_ . PDB format as used by NAMD/CHARMM: 4-letter resnames and segID are allowed, @@ -612,7 +612,8 @@ class PDBWriter(base.WriterBase): def __init__(self, filename, bonds="conect", n_atoms=None, start=0, step=1, remarks="Created by PDBWriter", - convert_units=True, multiframe=None, reindex=True): + convert_units=True, multiframe=None, reindex=True, + append=False): """Create a new PDBWriter Parameters @@ -643,7 +644,9 @@ def __init__(self, filename, bonds="conect", n_atoms=None, start=0, step=1, reindex: bool (optional) If ``True`` (default), the atom serial is set to be consecutive numbers starting at 1. Else, use the atom id. - + append: bool (optional) + If ``True``, append to an existing file. If ``False``, overwrite + the file. [``False``] """ # n_atoms = None : dummy keyword argument # (not used, but Writer() always provides n_atoms as the second argument) @@ -651,7 +654,7 @@ def __init__(self, filename, bonds="conect", n_atoms=None, start=0, step=1, # TODO: - remarks should be a list of lines and written to REMARK # - additional title keyword could contain line for TITLE - self.filename = filename + super().__init__(filename, n_atoms, append) # convert length and time to base units self.convert_units = convert_units self._multiframe = self.multiframe if multiframe is None else multiframe @@ -895,32 +898,7 @@ def _update_frame(self, obj): # hack... self.obj = obj self.ts = obj.universe.trajectory.ts - - def write(self, obj): - """Write object *obj* at current trajectory frame to file. - - *obj* can be a selection (i.e. a - :class:`~MDAnalysis.core.groups.AtomGroup`) or a whole - :class:`~MDAnalysis.core.universe.Universe`. - - The last letter of the :attr:`~MDAnalysis.core.groups.Atom.segid` is - used as the PDB chainID (but see :meth:`~PDBWriter.ATOM` for - details). - - Parameters - ---------- - obj - The :class:`~MDAnalysis.core.groups.AtomGroup` or - :class:`~MDAnalysis.core.universe.Universe` to write. - """ - - self._update_frame(obj) - self._write_pdb_header() - # Issue 105: with write() ONLY write a single frame; use - # write_all_timesteps() to dump everything in one go, or do the - # traditional loop over frames - self._write_next_frame(self.ts, multiframe=self._multiframe) - # END and CONECT records are written when file is being close()d + def write_all_timesteps(self, obj): """Write all timesteps associated with *obj* to the PDB file. @@ -965,7 +943,7 @@ def write_all_timesteps(self, obj): for framenumber in range(start, len(traj), step): traj[framenumber] - self._write_next_frame(self.ts, multiframe=True) + self._write_next_frame(obj, multiframe=True) # CONECT record is written when the file is being close()d self.close() @@ -973,7 +951,7 @@ def write_all_timesteps(self, obj): # Set the trajectory to the starting position traj[start] - def _write_next_frame(self, ts=None, **kwargs): + def _write_next_frame(self, obj, **kwargs): '''write a new timestep to the PDB file :Keywords: @@ -994,13 +972,17 @@ def _write_next_frame(self, ts=None, **kwargs): .. versionchanged:: 1.0.0 Renamed from `write_next_timestep` to `_write_next_frame`. ''' - if ts is None: - try: - ts = self.ts - except AttributeError: - errmsg = ("PBDWriter: no coordinate data to write to " - "trajectory file") - raise NoDataError(errmsg) from None + + self._update_frame(obj) + self._write_pdb_header() + + try: + ts = self.ts + except AttributeError: + errmsg = ("PBDWriter: no coordinate data to write to " + "trajectory file") + raise NoDataError(errmsg) from None + self._check_pdb_coordinates() self._write_timestep(ts, **kwargs) @@ -1360,3 +1342,21 @@ class MultiPDBWriter(PDBWriter): format = ['PDB', 'ENT'] multiframe = True # For Writer registration singleframe = False + + def write(self, obj, **kwargs): + """Write current timestep, using the supplied `obj`. + + Parameters + ---------- + obj : :class:`~MDAnalysis.core.groups.AtomGroup` or :class:`~MDAnalysis.core.universe.Universe` + write coordinate information associate with `obj` + + Note + ---- + The size of the `obj` must be the same as the number of atoms provided + when setting up the trajectory. + """ + + # overwrite the default behavior of PDBWriter + self._n_frames_written += 1 + return self._write_next_frame(obj, multiframe=True, **kwargs) \ No newline at end of file diff --git a/package/MDAnalysis/coordinates/PDBQT.py b/package/MDAnalysis/coordinates/PDBQT.py index 0e07008d984..12d3d581458 100644 --- a/package/MDAnalysis/coordinates/PDBQT.py +++ b/package/MDAnalysis/coordinates/PDBQT.py @@ -190,7 +190,7 @@ def Writer(self, filename, **kwargs): return PDBQTWriter(filename, **kwargs) -class PDBQTWriter(base.WriterBase): +class PDBQTWriter(base.SingleFrameWriterBase): """PDBQT writer that implements a subset of the PDB_ 3.2 standard and the PDBQT_ spec. .. _PDB: http://www.wwpdb.org/documentation/file-format-content/format32/v3.2.html @@ -212,14 +212,15 @@ class PDBQTWriter(base.WriterBase): units = {'time': None, 'length': 'Angstrom'} pdb_coor_limits = {"min": -999.9995, "max": 9999.9995} - def __init__(self, filename, **kwargs): - self.filename = util.filename(filename, ext='pdbqt') + def __init__(self, filename, n_atoms=None, append=False, **kwargs): + filename = util.filename(filename, ext='pdbqt') + super().__init__(filename, n_atoms, append) self.pdb = util.anyopen(self.filename, 'w') def close(self): self.pdb.close() - def write(self, selection, frame=None): + def _write_next_frame(self, selection, frame=None): """Write selection at current trajectory frame to file. Parameters diff --git a/package/MDAnalysis/coordinates/PQR.py b/package/MDAnalysis/coordinates/PQR.py index 1fd971b4d37..64d81273cbc 100644 --- a/package/MDAnalysis/coordinates/PQR.py +++ b/package/MDAnalysis/coordinates/PQR.py @@ -171,7 +171,7 @@ def Writer(self, filename, **kwargs): return PQRWriter(filename, **kwargs) -class PQRWriter(base.WriterBase): +class PQRWriter(base.SingleFrameWriterBase): """Write a single coordinate frame in whitespace-separated PQR format. Charges ("Q") are taken from the @@ -197,7 +197,7 @@ class PQRWriter(base.WriterBase): " {pos[2]:-8.3f} {charge:-7.4f} {radius:6.4f}\n") fmt_remark = "REMARK {0} {1}\n" - def __init__(self, filename, convert_units=True, **kwargs): + def __init__(self, filename, n_atoms=None, append=False, convert_units=True, **kwargs): """Set up a PQRWriter with full whitespace separation. Parameters @@ -210,11 +210,12 @@ def __init__(self, filename, convert_units=True, **kwargs): remark lines (list of strings) or single string to be added to the top of the PQR file """ - self.filename = util.filename(filename, ext='pqr') + filename = util.filename(filename, ext='pqr') + super().__init__(filename, n_atoms, append) self.convert_units = convert_units # convert length and time to base units self.remarks = kwargs.pop('remarks', "PQR file written by MDAnalysis") - def write(self, selection, frame=None): + def _write_next_frame(self, selection, frame=None): """Write selection at current trajectory frame to file. Parameters diff --git a/package/MDAnalysis/coordinates/TRJ.py b/package/MDAnalysis/coordinates/TRJ.py index 1c342ded68f..c7fa690531d 100644 --- a/package/MDAnalysis/coordinates/TRJ.py +++ b/package/MDAnalysis/coordinates/TRJ.py @@ -885,15 +885,15 @@ class NCDFWriter(base.WriterBase): 'velocity': 'Angstrom/ps', 'force': 'kcal/(mol*Angstrom)'} - def __init__(self, filename, n_atoms, remarks=None, convert_units=True, + def __init__(self, filename, n_atoms, append=False, + remarks=None, convert_units=True, velocities=False, forces=False, scale_time=None, scale_cell_lengths=None, scale_cell_angles=None, scale_coordinates=None, scale_velocities=None, scale_forces=None, **kwargs): - self.filename = filename if n_atoms == 0: raise ValueError("NCDFWriter: no atoms in output trajectory") - self.n_atoms = n_atoms + super().__init__(filename, n_atoms, append) # convert length and time to base units on the fly? self.convert_units = convert_units diff --git a/package/MDAnalysis/coordinates/TRZ.py b/package/MDAnalysis/coordinates/TRZ.py index ccdb7feafb7..0b4db3dccc4 100644 --- a/package/MDAnalysis/coordinates/TRZ.py +++ b/package/MDAnalysis/coordinates/TRZ.py @@ -397,7 +397,8 @@ class TRZWriter(base.WriterBase): units = {'time': 'ps', 'length': 'nm', 'velocity': 'nm/ps'} - def __init__(self, filename, n_atoms, title='TRZ', convert_units=True): + def __init__(self, filename, n_atoms, append=False, + title='TRZ', convert_units=True): """Create a TRZWriter Parameters @@ -406,18 +407,21 @@ def __init__(self, filename, n_atoms, title='TRZ', convert_units=True): name of output file n_atoms : int number of atoms in trajectory + append : bool (optional) + If ``True``, append to an existing file. If ``False``, overwrite + the file. Default is ``False``. title : str (optional) title of the trajectory; the title must be 80 characters or shorter, a longer title raises a ValueError exception. convert_units : bool (optional) units are converted to the MDAnalysis base format; [``True``] """ - self.filename = filename if n_atoms is None: raise ValueError("TRZWriter requires the n_atoms keyword") if n_atoms == 0: raise ValueError("TRZWriter: no atoms in output trajectory") - self.n_atoms = n_atoms + + super().__init__(filename, n_atoms, append) if len(title) > 80: raise ValueError("TRZWriter: 'title' must be 80 characters of shorter") diff --git a/package/MDAnalysis/coordinates/XDR.py b/package/MDAnalysis/coordinates/XDR.py index e2398aee84d..028ad6ce6d5 100644 --- a/package/MDAnalysis/coordinates/XDR.py +++ b/package/MDAnalysis/coordinates/XDR.py @@ -306,7 +306,8 @@ def Writer(self, filename, n_atoms=None, **kwargs): class XDRBaseWriter(base.WriterBase): """Base class for libmdaxdr file formats xtc and trr""" - def __init__(self, filename, n_atoms, convert_units=True, **kwargs): + def __init__(self, filename, n_atoms, append=False, + convert_units=True, **kwargs): """ Parameters ---------- @@ -314,11 +315,15 @@ def __init__(self, filename, n_atoms, convert_units=True, **kwargs): filename of trajectory n_atoms : int number of atoms to be written + append: bool (optional) + If ``True``, append to an existing file. If ``False``, overwrite + the file. Default is ``False``. convert_units : bool (optional) convert from MDAnalysis units to format specific units **kwargs : dict General writer arguments """ + super().__init__(filename, n_atoms, append) self.filename = filename self._convert_units = convert_units self.n_atoms = n_atoms diff --git a/package/MDAnalysis/coordinates/XTC.py b/package/MDAnalysis/coordinates/XTC.py index 3a1b6b439cb..63785ef4cc2 100644 --- a/package/MDAnalysis/coordinates/XTC.py +++ b/package/MDAnalysis/coordinates/XTC.py @@ -51,8 +51,8 @@ class XTCWriter(XDRBaseWriter): units = {'time': 'ps', 'length': 'nm'} _file = XTCFile - def __init__(self, filename, n_atoms, convert_units=True, - precision=3, **kwargs): + def __init__(self, filename, n_atoms, append=False, + convert_units=True, precision=3, **kwargs): """ Parameters ---------- @@ -60,13 +60,15 @@ def __init__(self, filename, n_atoms, convert_units=True, filename of the trajectory n_atoms : int number of atoms to write + append : bool (optional) + append to an existing trajectory if True convert_units : bool (optional) convert into MDAnalysis units precision : float (optional) set precision of saved trjactory to this number of decimal places. """ - super(XTCWriter, self).__init__(filename, n_atoms, convert_units, - **kwargs) + super().__init__(filename, n_atoms, append=append, + convert_units=convert_units, **kwargs) self.precision = precision def _write_next_frame(self, ag): diff --git a/package/MDAnalysis/coordinates/XYZ.py b/package/MDAnalysis/coordinates/XYZ.py index 20a2f75a886..2cf1b5d4f18 100644 --- a/package/MDAnalysis/coordinates/XYZ.py +++ b/package/MDAnalysis/coordinates/XYZ.py @@ -140,8 +140,8 @@ class XYZWriter(base.WriterBase): # these are assumed! units = {'time': 'ps', 'length': 'Angstrom'} - def __init__(self, filename, n_atoms=None, convert_units=True, - remark=None, **kwargs): + def __init__(self, filename, n_atoms=None, append=False, + convert_units=True, remark=None, **kwargs): """Initialize the XYZ trajectory writer Parameters @@ -155,6 +155,9 @@ def __init__(self, filename, n_atoms=None, convert_units=True, and that this file is used to store several different models instead of a single trajectory. If a number is provided each written TimeStep has to contain the same number of atoms. + append: bool (optional) + If ``True``, append to an existing file. If ``False``, overwrite + the file. Default is ``False``. convert_units : bool (optional) convert quantities to default MDAnalysis units of Angstrom upon writing [``True``] @@ -171,9 +174,8 @@ def __init__(self, filename, n_atoms=None, convert_units=True, an empty universe, please use ``add_TopologyAttr`` to add in the required elements or names. """ - self.filename = filename + super().__init__(filename, n_atoms, append) self.remark = remark - self.n_atoms = n_atoms self.convert_units = convert_units # can also be gz, bz2 @@ -200,24 +202,17 @@ def close(self): self._xyz.close() self._xyz = None - def write(self, obj): - """Write object `obj` at current trajectory frame to file. - - Atom elements (or names) in the output are taken from the `obj` or - default to the value of the `atoms` keyword supplied to the - :class:`XYZWriter` constructor. - - Parameters - ---------- - obj : Universe or AtomGroup - The :class:`~MDAnalysis.core.groups.AtomGroup` or - :class:`~MDAnalysis.core.universe.Universe` to write. + def _write_next_frame(self, obj): + """ + Write coordinate information in *ts* to the trajectory - .. versionchanged:: 2.0.0 - Deprecated support for Timestep argument has now been removed. - Use AtomGroup or Universe as an input instead. + .. versionchanged:: 1.0.0 + Print out :code:`remark` if present, otherwise use generic one + (Issue #2692). + Renamed from `write_next_timestep` to `_write_next_frame`. """ + # prepare the Timestep and extract atom names if possible # (The way it is written it should be possible to write # trajectories with frames that differ in atom numbers @@ -239,20 +234,11 @@ def write(self, obj): elif hasattr(obj, 'trajectory'): # For Universe only --- get everything ts = obj.trajectory.ts + else: + ts = None # update atom names self.atomnames = self._get_atoms_elements_or_names(atoms) - - self._write_next_frame(ts) - - def _write_next_frame(self, ts=None): - """ - Write coordinate information in *ts* to the trajectory - - .. versionchanged:: 1.0.0 - Print out :code:`remark` if present, otherwise use generic one - (Issue #2692). - Renamed from `write_next_timestep` to `_write_next_frame`. - """ + if ts is None: if not hasattr(self, 'ts'): raise NoDataError('XYZWriter: no coordinate data to write to ' diff --git a/package/MDAnalysis/coordinates/base.py b/package/MDAnalysis/coordinates/base.py index 6c24602ecd9..fa2c5c80d47 100644 --- a/package/MDAnalysis/coordinates/base.py +++ b/package/MDAnalysis/coordinates/base.py @@ -1473,9 +1473,41 @@ class WriterBase(IOBase, metaclass=_Writermeta): .. versionchanged:: 2.0.0 Deprecated :func:`write_next_timestep` has now been removed, please use :func:`write` instead. - + .. versionchanged:: 3.0.0 + Add append functionality to Writers. """ + def __init__(self, + filename, + n_atoms, + append=False): + """Set up single frame writer. + Parameters + ---------- + filename : str + filename of trajectory file + n_atoms : int + number of atoms in trajectory + append: bool (optional) + If ``True``, append to an existing file. If ``False``, overwrite + the file. Default is ``False``. + """ + self.filename = filename + self.n_atoms = n_atoms # number of atoms to be written + if append: + self.check_append_match() + self.append = append + self._n_frames_written = 0 + + def check_append_match(self): + """Check if the file to be appended matches the Writer. + + Raises + ------ + ValueError + If the file to be appended does not match the Writer. + """ + raise NotImplementedError("Writer does not support appending") def convert_dimensions_to_unitcell(self, ts, inplace=True): """Read dimensions from timestep *ts* and return appropriate unitcell. @@ -1493,7 +1525,7 @@ def convert_dimensions_to_unitcell(self, ts, inplace=True): lengths = self.convert_pos_to_native(lengths) return np.concatenate([lengths, angles]) - def write(self, obj): + def write(self, obj, **kwargs): """Write current timestep, using the supplied `obj`. Parameters @@ -1511,7 +1543,8 @@ def write(self, obj): Deprecated support for Timestep argument to write has now been removed. Use AtomGroup or Universe as an input instead. """ - return self._write_next_frame(obj) + self._n_frames_written += 1 + return self._write_next_frame(obj, **kwargs) def __del__(self): self.close() @@ -1545,6 +1578,61 @@ def has_valid_coordinates(self, criteria, x): x = np.ravel(x) return np.all(criteria["min"] < x) and np.all(x <= criteria["max"]) +class SingleFrameWriterBase(WriterBase): + """Base class for single frame writers. + + See :ref:`Trajectory API` definition in for the required attributes and + methods. + + .. versionadded:: 3.0.0 + Separate single frame writer as new base class. + """ + def __init__(self, + filename, + n_atoms, + append=False): + """Set up single frame writer. + + Parameters + ---------- + filename : str + filename of trajectory file + n_atoms : int + number of atoms in trajectory + append: bool (optional) + If ``True``, append to an existing file. If ``False``, overwrite + the file. Default is ``False``. + """ + self.filename = filename + self.n_atoms = n_atoms # number of atoms to be written + if append: + raise ValueError("Single frame writers do not support appending") + self.append = append + self._n_frames_written = 0 + + def write(self, obj, **kwargs): + """Write current timestep, using the supplied `obj`. + + Parameters + ---------- + obj : :class:`~MDAnalysis.core.groups.AtomGroup` or :class:`~MDAnalysis.core.universe.Universe` + write coordinate information associate with `obj` + + Note + ---- + The size of the `obj` must be the same as the number of atoms provided + when setting up the trajectory. + + .. versionchanged:: 2.0.0 + Deprecated support for Timestep argument to write has now been + removed. Use AtomGroup or Universe as an input instead. + """ + if self._n_frames_written >= 1: + raise ValueError("SingleFrameBaseWriter can only write one frame") + + self._n_frames_written += 1 + return self._write_next_frame(obj, **kwargs) + class SingleFrameReaderBase(ProtoReader): """Base class for Readers that only have one frame. diff --git a/package/MDAnalysis/coordinates/chemfiles.py b/package/MDAnalysis/coordinates/chemfiles.py index a7e2bd828c5..053109b4666 100644 --- a/package/MDAnalysis/coordinates/chemfiles.py +++ b/package/MDAnalysis/coordinates/chemfiles.py @@ -275,6 +275,7 @@ def __init__( self, filename, n_atoms=0, + append=False, mode="w", chemfiles_format="", topology=None, @@ -308,8 +309,7 @@ def __init__( raise RuntimeError( "Please install Chemfiles > {}" "".format(MIN_CHEMFILES_VERSION) ) - self.filename = filename - self.n_atoms = n_atoms + super().__init__(filename, n_atoms, append) if mode != "a" and mode != "w": raise IOError("Expected 'a' or 'w' as mode in ChemfilesWriter") self._file = chemfiles.Trajectory(self.filename, mode, chemfiles_format) diff --git a/package/MDAnalysis/coordinates/core.py b/package/MDAnalysis/coordinates/core.py index 45eb7659382..3cf261aaea2 100644 --- a/package/MDAnalysis/coordinates/core.py +++ b/package/MDAnalysis/coordinates/core.py @@ -85,7 +85,7 @@ def reader(filename, format=None, **kwargs): raise TypeError(errmsg) from None -def writer(filename, n_atoms=None, **kwargs): +def writer(filename, n_atoms=None, append=False, **kwargs): """Initialize a trajectory writer instance for *filename*. Parameters @@ -96,6 +96,9 @@ def writer(filename, n_atoms=None, **kwargs): n_atoms : int (optional) The number of atoms in the output trajectory; can be ommitted for single-frame writers. + append: bool (optional) + If ``True``, append to an existing file. If ``False``, overwrite + he file. Default is ``False``. multiframe : bool (optional) ``True``: write a trajectory with multiple frames; ``False`` only write a single frame snapshot; ``None`` first try to get @@ -123,4 +126,4 @@ def writer(filename, n_atoms=None, **kwargs): """ Writer = get_writer_for(filename, format=kwargs.pop('format', None), multiframe=kwargs.pop('multiframe', None)) - return Writer(filename, n_atoms=n_atoms, **kwargs) + return Writer(filename, n_atoms=n_atoms, append=append, **kwargs) diff --git a/package/MDAnalysis/coordinates/null.py b/package/MDAnalysis/coordinates/null.py index f27fc2088e8..52bca678568 100644 --- a/package/MDAnalysis/coordinates/null.py +++ b/package/MDAnalysis/coordinates/null.py @@ -51,8 +51,8 @@ class NullWriter(base.WriterBase): multiframe = True units = {'time': 'ps', 'length': 'Angstrom'} - def __init__(self, filename, **kwargs): - pass + def __init__(self, filename, n_atoms=None, append=False, **kwargs): + super().__init__(filename, n_atoms, append) def _write_next_frame(self, obj): try: diff --git a/testsuite/MDAnalysisTests/coordinates/base.py b/testsuite/MDAnalysisTests/coordinates/base.py index 5fd4c77a444..a028eb77214 100644 --- a/testsuite/MDAnalysisTests/coordinates/base.py +++ b/testsuite/MDAnalysisTests/coordinates/base.py @@ -650,6 +650,50 @@ def test_write_not_changing_ts(self, ref, universe, tmpdir): W.write(universe) assert_timestep_almost_equal(copy_ts, universe.trajectory.ts) + @pytest.mark.xfail(raises=NotImplementedError) + def test_write_append_to_trajectory(self, ref, universe, tmpdir): + outfile = 'write-append-to-trajectory.' + ref.ext + with tmpdir.as_cwd(): + with ref.writer(outfile, universe.atoms.n_atoms) as W: + for ts in universe.trajectory: + W.write(universe) + + with ref.writer(outfile, universe.atoms.n_atoms, + append=True) as W: + for ts in universe.trajectory: + W.write(universe) + copy = ref.reader(outfile) + assert_equal(len(universe.trajectory) * 2, len(copy)) + + with pytest.raises(ValueError, match='n_atoms'): + ref.writer(outfile, universe.atoms.n_atoms + 1, append=True) + copy = ref.reader(outfile) + assert_equal(len(universe.trajectory) * 2, len(copy)) + +class SingleFrameBaseWriterTest(BaseWriterTest): + """Test the base class for single frame writers""" + def test_write_append_to_trajectory(self, ref, universe, tmpdir): + outfile = 'write-single-frame.' + ref.ext + with tmpdir.as_cwd(): + with ref.writer(outfile, universe.atoms.n_atoms) as W: + W.write(universe) + with pytest.raises(ValueError, match="Single frame"): + with ref.writer(outfile, universe.atoms.n_atoms, + append=True) as W: + W.write(universe) + copy = ref.reader(outfile) + assert_equal(len(copy), 1) + + def test_write_multiple_frame(self, ref, universe, tmpdir): + outfile = 'write-single-frame.' + ref.ext + with tmpdir.as_cwd(): + with pytest.raises(ValueError, match="SingleFrameBaseWriter can"): + with ref.writer(outfile, universe.atoms.n_atoms) as W: + W.write(universe) + W.write(universe) + copy = ref.reader(outfile) + assert_equal(len(copy), 1) + def assert_timestep_equal(A, B, msg=''): """ assert that two timesteps are exactly equal and commutative @@ -711,3 +755,4 @@ def assert_timestep_almost_equal(A, B, decimal=6, verbose=True): if len(A.aux) > 0 and len(B.aux) > 0: assert_equal(A.aux, B.aux, err_msg='Auxiliary values do not match: ' 'A.aux = {}, B.aux = {}'.format(A.aux, B.aux)) + \ No newline at end of file diff --git a/testsuite/MDAnalysisTests/coordinates/test_crd.py b/testsuite/MDAnalysisTests/coordinates/test_crd.py index 623524a8a98..3d2925187dd 100644 --- a/testsuite/MDAnalysisTests/coordinates/test_crd.py +++ b/testsuite/MDAnalysisTests/coordinates/test_crd.py @@ -96,7 +96,16 @@ def test_write_EXT_read(self, u, outfile): len(u2.atoms.segments)), 'Equal number of segments expected in'\ 'both CRD formats' assert_allclose(cog1, cog2, rtol=1e-6, atol=0), 'Same centroid expected for both CRD formats' - + + def test_append_to_trajectory(self, u, outfile): + with pytest.raises(ValueError, match="Single frame"): + u.atoms.write(outfile, append=True) + + def test_write_multiple_frame(self, u, outfile): + with pytest.raises(ValueError, match="SingleFrameBaseWriter can"): + with u.trajectory.Writer(outfile) as W: + W.write(u) + W.write(u) class TestCRDWriterMissingAttrs(object): # All required attributes with the default value diff --git a/testsuite/MDAnalysisTests/coordinates/test_fhiaims.py b/testsuite/MDAnalysisTests/coordinates/test_fhiaims.py index ebdf0411769..02ca090015f 100644 --- a/testsuite/MDAnalysisTests/coordinates/test_fhiaims.py +++ b/testsuite/MDAnalysisTests/coordinates/test_fhiaims.py @@ -28,7 +28,7 @@ import numpy as np from MDAnalysisTests import make_Universe from MDAnalysisTests.coordinates.base import ( - _SingleFrameReader, BaseWriterTest) + _SingleFrameReader, SingleFrameBaseWriterTest) from MDAnalysisTests.datafiles import FHIAIMS from numpy.testing import (assert_equal, assert_array_almost_equal, @@ -189,7 +189,7 @@ def test_mixed_units(self, good_input_natural_units, good_input_mixed_units): "FHIAIMSReader failed to read positions in lattice units properly") -class TestFHIAIMSWriter(BaseWriterTest): +class TestFHIAIMSWriter(SingleFrameBaseWriterTest): prec = 6 ext = ".in" diff --git a/testsuite/MDAnalysisTests/coordinates/test_gro.py b/testsuite/MDAnalysisTests/coordinates/test_gro.py index e25bf969fc5..7c51c91d1c4 100644 --- a/testsuite/MDAnalysisTests/coordinates/test_gro.py +++ b/testsuite/MDAnalysisTests/coordinates/test_gro.py @@ -28,7 +28,7 @@ from MDAnalysis.transformations import translate from MDAnalysisTests import make_Universe from MDAnalysisTests.coordinates.base import ( - BaseReference, BaseReaderTest, BaseWriterTest, + BaseReference, BaseReaderTest, SingleFrameBaseWriterTest, ) from MDAnalysisTests.coordinates.reference import RefAdK from MDAnalysisTests.datafiles import ( @@ -188,7 +188,7 @@ def test_full_slice(self, ref, reader): assert_equal(frames, np.arange(u.trajectory.n_frames)) -class TestGROWriter(BaseWriterTest): +class TestGROWriter(SingleFrameBaseWriterTest): @staticmethod @pytest.fixture(scope='class') def ref(): @@ -301,7 +301,7 @@ def transformed(ref): return transformed -class TestGROWriterNoConversion(BaseWriterTest): +class TestGROWriterNoConversion(SingleFrameBaseWriterTest): @staticmethod @pytest.fixture(scope='class') def ref(): @@ -334,7 +334,7 @@ def ref(): return GROReaderIncompleteVelocitiesReference() -class TestGROWriterIncompleteVelocities(BaseWriterTest): +class TestGROWriterIncompleteVelocities(SingleFrameBaseWriterTest): @staticmethod @pytest.fixture(scope='class') def ref(): @@ -355,7 +355,7 @@ def ref(): return GROBZReference() -class TestGROBZ2Writer(BaseWriterTest): +class TestGROBZ2Writer(SingleFrameBaseWriterTest): @staticmethod @pytest.fixture(scope='class') def ref(): @@ -369,7 +369,7 @@ def __init__(self): self.topology = GRO_large -class TestGROLargeWriter(BaseWriterTest): +class TestGROLargeWriter(SingleFrameBaseWriterTest): @staticmethod @pytest.fixture(scope='class') def ref(): diff --git a/testsuite/MDAnalysisTests/coordinates/test_namdbin.py b/testsuite/MDAnalysisTests/coordinates/test_namdbin.py index d24835c0b66..14e6a4c2083 100644 --- a/testsuite/MDAnalysisTests/coordinates/test_namdbin.py +++ b/testsuite/MDAnalysisTests/coordinates/test_namdbin.py @@ -32,7 +32,7 @@ from MDAnalysisTests.datafiles import NAMDBIN, PDB_small from MDAnalysisTests.coordinates.base import (_SingleFrameReader, BaseReference, - BaseWriterTest) + SingleFrameBaseWriterTest) class TestNAMDBINReader(_SingleFrameReader): @@ -70,7 +70,7 @@ def __init__(self): self.container_format = True -class NAMDBINWriter(BaseWriterTest): +class NAMDBINWriter(SingleFrameBaseWriterTest): __test__ = True @staticmethod @pytest.fixture() diff --git a/testsuite/MDAnalysisTests/coordinates/test_writer_registration.py b/testsuite/MDAnalysisTests/coordinates/test_writer_registration.py index 839516070ed..5365ea28206 100644 --- a/testsuite/MDAnalysisTests/coordinates/test_writer_registration.py +++ b/testsuite/MDAnalysisTests/coordinates/test_writer_registration.py @@ -31,10 +31,6 @@ class MagicWriter(WriterBase): # this writer does the 'magic' format format = 'MAGIC' - def __init__(self, filename, n_atoms=None): - self.filename = filename - self.n_atoms = n_atoms - class MultiMagicWriter(MagicWriter): # this writer does the 'magic' and 'magic2' formats # but *only* supports multiframe writing.