fix DCD to read/write timestep information

Author: stefdoerr

moleculekit/fileformats/dcd/dcd.pyx

@@ -201,7 +201,20 @@ cdef class DCDTrajectoryFile:
         free(self.filename)
         self.close()
 
-    def _initialize_write(self, int n_atoms, int with_unitcell):
+    def read_header(self):
+        """
+        Returns
+        -------
+        istart : int
+            The starting step
+        nsavc : int
+            The number of steps between each saved frame
+        delta : float
+            The timestep of the simulation
+        """
+        return self.fh.istart, self.fh.nsavc, self.fh.delta
+
+    def _initialize_write(self, int n_atoms, int with_unitcell, int istart, int nsavc, double delta):
         """We don't actually want to open the dcd file during write mode
         until we know how many atoms the user wants to save. so this
         is delayed until the first call to write()
@@ -210,7 +223,7 @@ cdef class DCDTrajectoryFile:
 
         self.n_atoms = n_atoms
         self.with_unitcell = with_unitcell
-        self.fh = open_dcd_write(self.filename, "dcd", self.n_atoms, with_unitcell)
+        self.fh = open_dcd_write(self.filename, "dcd", self.n_atoms, with_unitcell, istart, nsavc, delta)
         if self.fh is NULL:
             raise IOError('There was an error opening the file: %s' % self.filename)
         self.is_open = True
@@ -435,7 +448,7 @@ cdef class DCDTrajectoryFile:
         # If we got some other status, thats a "real" error.
         raise IOError("Error: %s", ERROR_MESSAGES(status))
 
-    def write(self, xyz, cell_lengths=None, cell_angles=None):
+    def write(self, xyz, cell_lengths=None, cell_angles=None, istart=0, nsavc=1, delta=1.0):
         """write(xyz, cell_lengths=None, cell_angles=None)
 
         Write one or more frames of data to the DCD file
@@ -453,6 +466,12 @@ cdef class DCDTrajectoryFile:
             Organized analogously to cell_lengths. Gives the alpha, beta and
             gamma angles respectively. By convention for
             this trajectory format, the angles should be in units of degrees.
+        istart : int, optional
+            The starting step
+        nsavc : int, optional
+            The number of steps between written frames
+        delta : float, optional
+            The time-step of the simulation.
         """
         if str(self.mode) != 'w':
             raise ValueError('write() is only available when the file is opened in mode="w"')
@@ -478,7 +497,7 @@ cdef class DCDTrajectoryFile:
                                   warn_on_cast=False)
 
         if self._needs_write_initialization:
-            self._initialize_write(xyz.shape[1], (cell_lengths is not None) and (cell_angles is not None))
+            self._initialize_write(xyz.shape[1], (cell_lengths is not None) and (cell_angles is not None), istart, nsavc, delta)
         else:
             if not self.n_atoms == xyz.shape[1]:
                 raise ValueError('Number of atoms doesnt match')

moleculekit/fileformats/dcd/dcdlib.pxd

@@ -23,7 +23,7 @@ cdef extern from "include/dcdplugin.h":
     int read_next_timestep(dcdhandle *v, int natoms, molfile_timestep_t *ts)
     void close_file_read(dcdhandle *v)
 
-    dcdhandle* open_dcd_write(const char *path, const char *filetype, const int natoms, const int with_unitcell)
+    dcdhandle* open_dcd_write(const char *path, const char *filetype, const int natoms, const int with_unitcell, const int istart, const int nsavc, const double delta)
     int write_timestep(dcdhandle *v, molfile_timestep_t *ts)
     void close_file_write(dcdhandle *v)
     int dcd_nsets(dcdhandle* v)

moleculekit/fileformats/dcd/include/dcdplugin.h

@@ -68,7 +68,7 @@ void close_file_read(dcdhandle *v);
 int read_next_timestep(dcdhandle *v, int natoms, molfile_timestep_t *ts);
 
 dcdhandle* open_dcd_write(const char *path, const char *filetype, const int natoms,
-                          const int with_unitcell);
+                          const int with_unitcell, const int istart, const int nsavc, const double delta);
 int write_timestep(dcdhandle *v, const molfile_timestep_t *ts);
 void close_file_write(dcdhandle *v);
 int dcd_nsets(dcdhandle* v);

moleculekit/fileformats/dcd/src/dcdplugin.c

@@ -1144,13 +1144,11 @@ void close_file_read(dcdhandle *v)
 }
 
 dcdhandle *open_dcd_write(const char *path, const char *filetype, const int natoms,
-                          const int with_unitcell)
+                          const int with_unitcell, const int istart, const int nsavc, const double delta)
 {
   dcdhandle *dcd;
   fio_fd fd;
   int rc;
-  int istart, nsavc;
-  double delta;
   int charmm;
 
   if (fio_open(path, FIO_WRITE, &fd) < 0)
@@ -1163,10 +1161,6 @@ dcdhandle *open_dcd_write(const char *path, const char *filetype, const int nato
   memset(dcd, 0, sizeof(dcdhandle));
   dcd->fd = fd;
 
-  istart = 0;  /* starting timestep of DCD file                  */
-  nsavc = 1;   /* number of timesteps between written DCD frames */
-  delta = 1.0; /* length of a timestep                           */
-
   charmm = DCD_IS_CHARMM; /* charmm-formatted DCD file                */
   if (with_unitcell)
     charmm |= DCD_HAS_EXTRA_BLOCK;

moleculekit/readers.py

@@ -2523,15 +2523,14 @@ def SDFread(filename, frame=None, topoloc=None, mol_idx=None):
             "MoleculeKit will only read the first molecule from the SDF file."
         )
 
+    v3000 = False
     with openFileOrStringIO(filename, "r") as f:
         curr_mol = 0
         lines = []
         for line in f:
             lines.append(line)
             if "V3000" in line:
-                raise RuntimeError(
-                    "Moleculekit does not support parsing V3000 SDF files yet."
-                )
+                v3000 = True
             if line.strip() == "$$$$":
                 if curr_mol == mol_idx:
                     break
@@ -2544,61 +2543,119 @@ def SDFread(filename, frame=None, topoloc=None, mol_idx=None):
                 f"SDF file contains only {curr_mol} molecules. Cannot read the requested mol_idx {mol_idx}"
             )
 
-        topo = Topology()
-        coords = []
-        mol_start = 0
-
-        molname = lines[0].strip().split()
-        if len(molname):
-            molname = molname[0]
-        if len(molname) == 3:
-            resname = molname[:3].upper()
+        if v3000:
+            topo, coords = parseV3000SDF(lines, chargemap, bondmap)
         else:
-            resname = "MOL"
-
-        n_atoms = int(lines[mol_start + 3][:3])
-        n_bonds = int(lines[mol_start + 3][3:6])
+            topo = Topology()
+            coords = []
+            mol_start = 0
 
-        coord_start = mol_start + 4
-        bond_start = coord_start + n_atoms
-        bond_end = bond_start + n_bonds
-        for n in range(coord_start, bond_start):
-            line = lines[n]
-            coords.append(
-                [
-                    float(line[:10].strip()),
-                    float(line[10:20].strip()),
-                    float(line[20:30].strip()),
-                ]
-            )
-            atom_symbol = line[31:34].strip()
-            topo.record.append("HETATM")
-            topo.element.append(atom_symbol)
-            topo.name.append(atom_symbol)
-            topo.serial.append(n - coord_start)
-            topo.formalcharge.append(chargemap[line[36:39].strip()])
-            topo.resname.append(resname)
+            molname = lines[0].strip().split()
+            if len(molname):
+                molname = molname[0]
+            if len(molname) == 3:
+                resname = molname[:3].upper()
+            else:
+                resname = "MOL"
 
-        for n in range(bond_start, bond_end):
-            line = lines[n]
-            idx1 = line[:3].strip()
-            idx2 = line[3:6].strip()
-            bond_type = line[6:9].strip()
-            topo.bonds.append([int(idx1) - 1, int(idx2) - 1])
-            topo.bondtype.append(bondmap[bond_type])
+            n_atoms = int(lines[mol_start + 3][:3])
+            n_bonds = int(lines[mol_start + 3][3:6])
 
-        for line in lines[bond_end:]:
-            if line.strip() == "$$$$":
-                break
-            if line.startswith("M  CHG"):  # These charges are more correct
-                pairs = line.strip().split()[3:]
-                for cc in range(0, len(pairs), 2):
-                    topo.formalcharge[int(pairs[cc]) - 1] = int(pairs[cc + 1])
+            coord_start = mol_start + 4
+            bond_start = coord_start + n_atoms
+            bond_end = bond_start + n_bonds
+            for n in range(coord_start, bond_start):
+                line = lines[n]
+                coords.append(
+                    [
+                        float(line[:10].strip()),
+                        float(line[10:20].strip()),
+                        float(line[20:30].strip()),
+                    ]
+                )
+                atom_symbol = line[31:34].strip()
+                topo.record.append("HETATM")
+                topo.element.append(atom_symbol)
+                topo.name.append(atom_symbol)
+                topo.serial.append(n - coord_start)
+                topo.formalcharge.append(chargemap[line[36:39].strip()])
+                topo.resname.append(resname)
+
+            for n in range(bond_start, bond_end):
+                line = lines[n]
+                idx1 = line[:3].strip()
+                idx2 = line[3:6].strip()
+                bond_type = line[6:9].strip()
+                topo.bonds.append([int(idx1) - 1, int(idx2) - 1])
+                topo.bondtype.append(bondmap[bond_type])
+
+            for line in lines[bond_end:]:
+                if line.strip() == "$$$$":
+                    break
+                if line.startswith("M  CHG"):  # These charges are more correct
+                    pairs = line.strip().split()[3:]
+                    for cc in range(0, len(pairs), 2):
+                        topo.formalcharge[int(pairs[cc]) - 1] = int(pairs[cc + 1])
 
     traj = Trajectory(coords=np.vstack(coords))
     return MolFactory.construct(topo, traj, filename, frame)
 
 
+def parseV3000SDF(lines, chargemap, bondmap):
+    topo = Topology()
+    coords = []
+
+    molname = lines[0].strip().split()
+    if len(molname):
+        molname = molname[0]
+    if len(molname) == 3:
+        resname = molname[:3].upper()
+    else:
+        resname = "MOL"
+
+    section = None
+    for line in lines[4:]:
+        line = line.strip()
+        if line == "$$$$":
+            break
+        if "v30 end" in line.lower() or "m  end" in line.lower():
+            section = None
+            continue
+        if "begin ctab" in line.lower():
+            section = "ctab"
+            continue
+        if "begin atom" in line.lower():
+            section = "atom"
+            continue
+        if "begin bond" in line.lower():
+            section = "bond"
+            continue
+
+        if section is None:
+            continue
+        if section == "ctab":
+            # M  V30 COUNTS 49 53 0 0 0
+            if "v30 counts" in line.lower():
+                n_atoms, n_bonds = map(int, line.split()[3:5])
+        if section == "atom":
+            # M  V30 1 C 42.4990 45.8550 38.9350 0
+            pieces = line.split()[2:]
+            topo.record.append("HETATM")
+            topo.element.append(pieces[1])
+            topo.name.append(pieces[1])
+            topo.serial.append(int(pieces[0]))
+            topo.formalcharge.append(chargemap[pieces[5]])
+            topo.resname.append(resname)
+            coords.append([float(pieces[2]), float(pieces[3]), float(pieces[4])])
+        if section == "bond":
+            # M  V30 bond_idx bond_type atom1 atom2
+            pieces = line.split()[3:]
+            topo.bonds.append([int(pieces[1]) - 1, int(pieces[2]) - 1])
+            topo.bondtype.append(bondmap[pieces[0]])
+
+    return topo, coords
+
+
 def sdf_generator(sdffile):
     """Generates Molecule objects from an SDF file"""
     from io import StringIO
@@ -2884,29 +2941,40 @@ def NETCDFread(
 
 
 def DCDread(filename, frame=None, topoloc=None, stride=None, atom_indices=None):
-    from moleculekit.dcd import load_dcd
+    from moleculekit.fileformats.utils import cast_indices
+    from moleculekit.dcd import DCDTrajectoryFile
+
+    atom_indices = cast_indices(atom_indices)
+    with DCDTrajectoryFile(str(filename)) as f:
+        if frame is not None:
+            f.seek(frame)
+            n_frames = 1
+        else:
+            n_frames = None
+        xyz, cell_lengths, cell_angles = f.read(
+            n_frames=n_frames, stride=stride, atom_indices=atom_indices
+        )
+        istart, nsavc, delta = f.read_header()
+        # Timestep conversion factor found in OpenMM
+        delta = np.round(delta * 0.04888821, decimals=8)
+        steps = np.arange(istart, (nsavc * xyz.shape[0]) + 1, nsavc)
+        if stride is not None:
+            steps = steps[::stride]
 
-    xyz, cell_lengths, cell_angles = load_dcd(
-        filename, stride=stride, atom_indices=atom_indices, frame=frame
-    )
     xyz = np.transpose(xyz, (1, 2, 0))
     if cell_lengths is not None:
         cell_lengths = cell_lengths.T
     if cell_angles is not None:
         cell_angles = cell_angles.T
 
-    if stride is None:
-        stride = 1
-    initial = 0
-    time = (stride * np.arange(xyz.shape[2])) + initial
     return MolFactory.construct(
         None,
         Trajectory(
             coords=xyz,
             box=cell_lengths,
             boxangles=cell_angles,
-            step=range(xyz.shape[2]),
-            time=time * 1000,  # ps to fs
+            step=steps,
+            time=steps * delta * 1000,  # ps to fs
         ),
         filename,
         frame,

moleculekit/writers.py

@@ -666,8 +666,25 @@ def DCDwrite(mol, filename):
     else:
         boxangles = np.zeros((n_frames, 3), dtype=np.float32)
 
+    try:
+        istart = int(mol.step[0])
+        nsavc = int(mol.step[1] - mol.step[0])
+        fstep = mol.fstep * 1000  # ns to ps
+        delta = fstep / nsavc / 0.04888821  # Conversion factor found in OpenMM
+    except Exception:
+        istart = 0
+        nsavc = 1
+        delta = 1.0
+
     with DCDTrajectoryFile(filename, "w") as fh:
-        fh.write(xyz, cell_lengths=box, cell_angles=boxangles)
+        fh.write(
+            xyz,
+            cell_lengths=box,
+            cell_angles=boxangles,
+            istart=istart,
+            nsavc=nsavc,
+            delta=delta,
+        )
 
 
 def BINPOSwrite(mol, filename):