def __init__(self, *args, attrmap=None, **kwargs): super().__init__(**kwargs) self.attrmap = attrmap self.trajectory = Trajectory() self.dumpattrs = {} self.dumpfiles = list(flatten([glob(f) for f in args[0].split()])) if len(self.dumpfiles) == 0: raise ValueError('No dump file specified.') self.read()
class DUMPReader(StructureIO): """Class for reading `LAMMPS dump` file format. Parameters ---------- *args : :class:`~python:list` :class:`~python:list` of one or more LAMMPS dump files. attrmap : class:`~python:dict` Python :class:`~python:dict` mapping custom dump attributes to :class:`~sknano.core.atoms.Atom` attributes. Examples -------- Using LAMMPS, one often saves molecular dynamics atom trajectories to one or more custom dump files containing per-atom attributes of of each atom at a given timestep. Often the per-atom attributes are calculated from custom LAMMPS compute commands that one may want to map to a specific :class:`~sknano.core.atoms.Atom` attribute. For example, consider a dump file containing per-atom dump attributes *id*, *mol*, *type*, *x*, *y*, *z*, *vx*, *vy*, *vz*, *c_atom_ke*, *c_atom_pe*, and *c_atom_CN*, where *c_atom_ke*, *c_atom_pe*, and *c_atom_CN* are the LAMMPS *compute-IDs* from custom compute commands in a LAMMPS input script, which calculate per-atom values for kinetic energy, potential energy, and coordination number, respectively. To map these compute-IDs to corresponding :class:`~sknano.core.atoms.Atom` attributes :attr:`~sknano.core.atom.EnergyAtom.ke`, :attr:`~sknano.core.atom.EnergyAtom.pe`, and :attr:`~sknano.core.atom.CNAtom.CN`), respectively, one would use the following :class:`~python:dict` mapping:: >>> from sknano.io import DUMPReader >>> compute_attrmap = {'c_atom_pe': 'pe', 'c_atom_ke': 'ke', ... 'c_atom_CN': 'CN'} >>> dumps = DUMPReader('dump.*', attrmap=compute_attrmap) >>> print(repr(dumps.dumpattrs2str())) 'id mol type x y z vx vy vz c_atom_ke c_atom_pe c_atom_CN' >>> print(repr(dumps.atomattrs2str())) 'id mol type x y z vx vy vz ke pe CN' """ def __init__(self, *args, attrmap=None, **kwargs): super().__init__(**kwargs) self.attrmap = attrmap self.trajectory = Trajectory() self.dumpattrs = {} self.dumpfiles = list(flatten([glob(f) for f in args[0].split()])) if len(self.dumpfiles) == 0: raise ValueError('No dump file specified.') self.read() def __getattr__(self, name): try: return getattr(self.trajectory, name) except AttributeError: return super().__getattr__(name) def __iter__(self): return iter(self.trajectory) def __getitem__(self, index): return self.trajectory[index] def read(self): """Read all snapshots from each dump file.""" for dumpfile in self.dumpfiles: with zopen(dumpfile) as f: snapshot = self.read_snapshot(f) while snapshot is not None: self.trajectory.append(snapshot) print(snapshot.timestep, end=' ') sys.stdout.flush() snapshot = self.read_snapshot(f) print() self.trajectory.sort(key=attrgetter('timestep')) self.trajectory.cull() print("read {:d} snapshots".format(self.Nsnaps)) self.trajectory.time_selection.all() self.trajectory.t0_snapshot = self.trajectory[0] if self.dumpattrs: print('Dumped Atom attributes: {}'.format(self.dumpattrs2str())) else: print('No dump column assignments') if 'x' not in self.dumpattrs or 'y' not in self.dumpattrs or \ 'z' not in self.dumpattrs: print('dump scaling status unknown') elif self.Nsnaps > 0: if self.scale_original: self.unscale() elif self.scale_original is None: print('dump scaling status unknown') else: print('dump is already unscaled') def read_snapshot(self, f): try: snapshot = Snapshot(self.trajectory) f.readline() snapshot.timestep = int(f.readline().strip().split()[0]) f.readline() snapshot.Natoms = int(f.readline().strip()) snapshot.atom_selection = np.zeros(snapshot.Natoms, dtype=bool) item = f.readline().strip() try: snapshot.boxstr = item.split('BOUNDS')[1].strip() except IndexError: snapshot.boxstr = '' snapshot.triclinic = False snapshot.bounding_box = Cuboid() if 'xy' in snapshot.boxstr: snapshot.triclinic = True for dim, tilt_factor in zip(('x', 'y', 'z'), ('xy', 'xz', 'yz')): bounds = f.readline().strip().split() setattr(snapshot, dim + 'lo', float(bounds[0])) setattr(snapshot, dim + 'hi', float(bounds[1])) setattr(snapshot.bounding_box, dim + 'min', getattr(snapshot, dim + 'lo')) setattr(snapshot.bounding_box, dim + 'max', getattr(snapshot, dim + 'hi')) try: setattr(snapshot, tilt_factor, float(bounds[2])) except IndexError: setattr(snapshot, tilt_factor, 0.0) if not self.dumpattrs: xflag = yflag = zflag = None attrs = f.readline().strip().split()[2:] for i, attr in enumerate(attrs): if attr in ('x', 'xu', 'xs', 'xsu'): self.dumpattrs['x'] = i if attr in ('x', 'xu'): xflag = False else: xflag = True elif attr in ('y', 'yu', 'ys', 'ysu'): self.dumpattrs['y'] = i if attr in ('y', 'yu'): yflag = False else: yflag = True elif attr in ('z', 'zu', 'zs', 'zsu'): self.dumpattrs['z'] = i if attr in ('z', 'zu'): zflag = False else: zflag = True else: self.dumpattrs[attr] = i self.scale_original = None if all([flag is False for flag in (xflag, yflag, zflag)]): self.scale_original = False if all([flag for flag in (xflag, yflag, zflag)]): self.scale_original = True self.atomattrs = \ sorted(self.dumpattrs, key=self.dumpattrs.__getitem__) self.attr_dtypes = [attr_dtypes[attr] if attr in attr_dtypes else float for attr in self.atomattrs] if self.attrmap is not None: self.remap_atomattr_names(self.attrmap) self.attrmap = None self.unknown_attrs = \ {attr: self.atomattrs.index(attr) for attr in set(self.atomattrs) - set(dir(Atom()))} [self.atomattrs.remove(attr) for attr in self.unknown_attrs] else: f.readline() snapshot.atomattrs = self.atomattrs snapshot.attr_dtypes = self.attr_dtypes atoms = \ np.zeros((snapshot.Natoms, len(self.atomattrs)), dtype=float) for n in range(snapshot.Natoms): line = [float(value) for col, value in enumerate(f.readline().strip().split()) if col not in self.unknown_attrs.values()] atoms[n] = line snapshot.atoms = atoms return snapshot except IndexError: return None def remap_atomattr_names(self, attrmap): """Rename attributes in the :attr:`DUMPReader.atomattrs` list. Parameters ---------- attrmap : :class:`~python:dict` :class:`~python:dict` mapping atom attribute name to new attribute name. """ for k, v in attrmap.items(): try: self.atomattrs[self.atomattrs.index(k)] = v except ValueError: pass def scale(self): """Scale cartesian coordinates to fractional coordinates.""" xi = self.dumpattrs['x'] yi = self.dumpattrs['y'] zi = self.dumpattrs['z'] for snapshot in self.trajectory: atoms = snapshot.get_atoms(asarray=True) if atoms is not None: xlo = snapshot.xlo xhi = snapshot.xhi ylo = snapshot.ylo yhi = snapshot.yhi zlo = snapshot.zlo zhi = snapshot.zhi lx = xhi - xlo ly = yhi - ylo lz = zhi - zlo xy = snapshot.xy xz = snapshot.xz yz = snapshot.yz if np.allclose([xy, xz, yz], np.zeros(3)): atoms[:, xi] = (atoms[:, xi] - snapshot.xlo) / lx atoms[:, yi] = (atoms[:, yi] - snapshot.ylo) / ly atoms[:, zi] = (atoms[:, zi] - snapshot.zlo) / lz else: xlo = xlo - min((0.0, xy, xz, xy + xz)) xhi = xhi - max((0.0, xy, xz, xy + xz)) lx = xhi - xlo ylo = ylo - min((0.0, yz)) yhi = yhi - max((0.0, yz)) ly = yhi - ylo atoms[:, xi] = (atoms[:, xi] - snapshot.xlo) / lx + \ (atoms[:, yi] - snapshot.ylo) * xy / (lx * ly) + \ (atoms[:, zi] - snapshot.zlo) * (yz * xy - ly * xz) / \ (lx * ly * lz) atoms[:, yi] = (atoms[:, yi] - snapshot.ylo) / ly + \ (atoms[:, zi] - snapshot.zlo) * yz / (ly * lz) atoms[:, zi] = (atoms[:, zi] - snapshot.zlo) / lz def unscale(self): """Unscale fractional coordinates to cartesian coordinates.""" xi = self.dumpattrs['x'] yi = self.dumpattrs['y'] zi = self.dumpattrs['z'] for snapshot in self.trajectory: atoms = snapshot.get_atoms(asarray=True) if atoms is not None: xlo = snapshot.xlo xhi = snapshot.xhi ylo = snapshot.ylo yhi = snapshot.yhi zlo = snapshot.zlo zhi = snapshot.zhi lx = xhi - xlo ly = yhi - ylo lz = zhi - zlo xy = snapshot.xy xz = snapshot.xz yz = snapshot.yz if np.allclose([xy, xz, yz], np.zeros(3)): atoms[:, xi] = snapshot.xlo + atoms[:, xi] * lx atoms[:, yi] = snapshot.ylo + atoms[:, yi] * ly atoms[:, zi] = snapshot.zlo + atoms[:, zi] * lz else: xlo = xlo - min((0.0, xy, xz, xy + xz)) xhi = xhi - max((0.0, xy, xz, xy + xz)) lx = xhi - xlo ylo = ylo - min((0.0, yz)) yhi = yhi - max((0.0, yz)) ly = yhi - ylo atoms[:, xi] = snapshot.xlo + atoms[:, xi] * lx + \ atoms[:, yi] * xy + atoms[:, zi] * xz atoms[:, yi] = snapshot.ylo + atoms[:, yi] * ly + \ atoms[:, zi] * yz atoms[:, zi] = snapshot.zlo + atoms[:, zi] * lz def atomattrs2str(self): """Return a space-separated string of parsed atom attributes.""" return ' '.join(self.atomattrs) def dumpattrs2str(self): """Return a space-separated string of the dumped atom attributes.""" return ' '.join(sorted(self.dumpattrs, key=self.dumpattrs.__getitem__))
def test1(): traj = Trajectory() assert_equal(traj.Nsnaps, 0) print(traj)