# Copyright (c) 2016-2024 The Regents of the University of Michigan # Part of GSD, released under the BSD 2-Clause License. """Read and write HOOMD schema GSD files. :py:mod:`gsd.hoomd` reads and writes GSD files with the ``hoomd`` schema. * `HOOMDTrajectory` - Read and write hoomd schema GSD files. * `Frame` - Store the state of a single frame. * `ConfigurationData` - Store configuration data in a frame. * `ParticleData` - Store particle data in a frame. * `BondData` - Store topology data in a frame. * `open` - Open a hoomd schema GSD file. * `read_log` - Read log from a hoomd schema GSD file into a dict of time-series arrays. See Also: See :ref:`hoomd-examples` for full examples. """ import copy import json import logging import warnings from collections import OrderedDict import numpy try: import gsd except ImportError: gsd = None fl_imported = True try: import gsd.fl except ImportError: fl_imported = False logger = logging.getLogger('gsd.hoomd') class ConfigurationData: """Store configuration data. Use the `Frame.configuration` attribute of a to access the configuration. Attributes: step (int): Time step of this frame (:chunk:`configuration/step`). dimensions (int): Number of dimensions (:chunk:`configuration/dimensions`). When not set explicitly, dimensions will default to different values based on the value of :math:`L_z` in `box`. When :math:`L_z = 0` dimensions will default to 2, otherwise 3. User set values always take precedence. """ _default_value = OrderedDict() _default_value['step'] = numpy.uint64(0) _default_value['dimensions'] = numpy.uint8(3) _default_value['box'] = numpy.array([1, 1, 1, 0, 0, 0], dtype=numpy.float32) def __init__(self): self.step = None self.dimensions = None self._box = None @property def box(self): """((6, 1) `numpy.ndarray` of ``numpy.float32``): Box dimensions. [lx, ly, lz, xy, xz, yz]. See :chunk:`configuration/box`. """ return self._box @box.setter def box(self, box): self._box = box try: Lz = box[2] except TypeError: return else: if self.dimensions is None: self.dimensions = 2 if Lz == 0 else 3 def validate(self): """Validate all attributes. Convert every array attribute to a `numpy.ndarray` of the proper type and check that all attributes have the correct dimensions. Ignore any attributes that are ``None``. Warning: Array attributes that are not contiguous numpy arrays will be replaced with contiguous numpy arrays of the appropriate type. """ logger.debug('Validating ConfigurationData') if self.box is not None: self.box = numpy.ascontiguousarray(self.box, dtype=numpy.float32) self.box = self.box.reshape([6]) class ParticleData: """Store particle data chunks. Use the `Frame.particles` attribute of a to access the particles. Instances resulting from file read operations will always store array quantities in `numpy.ndarray` objects of the defined types. User created frames may provide input data that can be converted to a `numpy.ndarray`. See Also: `hoomd.State` for a full description of how HOOMD interprets this data. Attributes: N (int): Number of particles in the frame (:chunk:`particles/N`). types (tuple[str]): Names of the particle types (:chunk:`particles/types`). position ((*N*, 3) `numpy.ndarray` of ``numpy.float32``): Particle position (:chunk:`particles/position`). orientation ((*N*, 4) `numpy.ndarray` of ``numpy.float32``): Particle orientation. (:chunk:`particles/orientation`). typeid ((*N*, ) `numpy.ndarray` of ``numpy.uint32``): Particle type id (:chunk:`particles/typeid`). mass ((*N*, ) `numpy.ndarray` of ``numpy.float32``): Particle mass (:chunk:`particles/mass`). charge ((*N*, ) `numpy.ndarray` of ``numpy.float32``): Particle charge (:chunk:`particles/charge`). diameter ((*N*, ) `numpy.ndarray` of ``numpy.float32``): Particle diameter (:chunk:`particles/diameter`). body ((*N*, ) `numpy.ndarray` of ``numpy.int32``): Particle body (:chunk:`particles/body`). moment_inertia ((*N*, 3) `numpy.ndarray` of ``numpy.float32``): Particle moment of inertia (:chunk:`particles/moment_inertia`). velocity ((*N*, 3) `numpy.ndarray` of ``numpy.float32``): Particle velocity (:chunk:`particles/velocity`). angmom ((*N*, 4) `numpy.ndarray` of ``numpy.float32``): Particle angular momentum (:chunk:`particles/angmom`). image ((*N*, 3) `numpy.ndarray` of ``numpy.int32``): Particle image (:chunk:`particles/image`). type_shapes (tuple[dict]): Shape specifications for visualizing particle types (:chunk:`particles/type_shapes`). """ _default_value = OrderedDict() _default_value['N'] = numpy.uint32(0) _default_value['types'] = ['A'] _default_value['typeid'] = numpy.uint32(0) _default_value['mass'] = numpy.float32(1.0) _default_value['charge'] = numpy.float32(0) _default_value['diameter'] = numpy.float32(1.0) _default_value['body'] = numpy.int32(-1) _default_value['moment_inertia'] = numpy.array([0, 0, 0], dtype=numpy.float32) _default_value['position'] = numpy.array([0, 0, 0], dtype=numpy.float32) _default_value['orientation'] = numpy.array([1, 0, 0, 0], dtype=numpy.float32) _default_value['velocity'] = numpy.array([0, 0, 0], dtype=numpy.float32) _default_value['angmom'] = numpy.array([0, 0, 0, 0], dtype=numpy.float32) _default_value['image'] = numpy.array([0, 0, 0], dtype=numpy.int32) _default_value['type_shapes'] = [{}] def __init__(self): self.N = 0 self.position = None self.orientation = None self.types = None self.typeid = None self.mass = None self.charge = None self.diameter = None self.body = None self.moment_inertia = None self.velocity = None self.angmom = None self.image = None self.type_shapes = None def validate(self): """Validate all attributes. Convert every array attribute to a `numpy.ndarray` of the proper type and check that all attributes have the correct dimensions. Ignore any attributes that are ``None``. Warning: Array attributes that are not contiguous numpy arrays will be replaced with contiguous numpy arrays of the appropriate type. """ logger.debug('Validating ParticleData') if self.position is not None: self.position = numpy.ascontiguousarray(self.position, dtype=numpy.float32) self.position = self.position.reshape([self.N, 3]) if self.orientation is not None: self.orientation = numpy.ascontiguousarray( self.orientation, dtype=numpy.float32 ) self.orientation = self.orientation.reshape([self.N, 4]) if self.typeid is not None: self.typeid = numpy.ascontiguousarray(self.typeid, dtype=numpy.uint32) self.typeid = self.typeid.reshape([self.N]) if self.mass is not None: self.mass = numpy.ascontiguousarray(self.mass, dtype=numpy.float32) self.mass = self.mass.reshape([self.N]) if self.charge is not None: self.charge = numpy.ascontiguousarray(self.charge, dtype=numpy.float32) self.charge = self.charge.reshape([self.N]) if self.diameter is not None: self.diameter = numpy.ascontiguousarray(self.diameter, dtype=numpy.float32) self.diameter = self.diameter.reshape([self.N]) if self.body is not None: self.body = numpy.ascontiguousarray(self.body, dtype=numpy.int32) self.body = self.body.reshape([self.N]) if self.moment_inertia is not None: self.moment_inertia = numpy.ascontiguousarray( self.moment_inertia, dtype=numpy.float32 ) self.moment_inertia = self.moment_inertia.reshape([self.N, 3]) if self.velocity is not None: self.velocity = numpy.ascontiguousarray(self.velocity, dtype=numpy.float32) self.velocity = self.velocity.reshape([self.N, 3]) if self.angmom is not None: self.angmom = numpy.ascontiguousarray(self.angmom, dtype=numpy.float32) self.angmom = self.angmom.reshape([self.N, 4]) if self.image is not None: self.image = numpy.ascontiguousarray(self.image, dtype=numpy.int32) self.image = self.image.reshape([self.N, 3]) if self.types is not None and (not len(set(self.types)) == len(self.types)): msg = 'Type names must be unique.' raise ValueError(msg) class BondData: """Store bond data chunks. Use the `Frame.bonds`, `Frame.angles`, `Frame.dihedrals`, `Frame.impropers`, and `Frame.pairs` attributes to access the bond topology. Instances resulting from file read operations will always store array quantities in `numpy.ndarray` objects of the defined types. User created frames may provide input data that can be converted to a `numpy.ndarray`. See Also: `hoomd.State` for a full description of how HOOMD interprets this data. Note: *M* varies depending on the type of bond. `BondData` represents all types of topology connections. ======== === Type *M* ======== === Bond 2 Angle 3 Dihedral 4 Improper 4 Pair 2 ======== === Attributes: N (int): Number of bonds/angles/dihedrals/impropers/pairs in the frame (:chunk:`bonds/N`, :chunk:`angles/N`, :chunk:`dihedrals/N`, :chunk:`impropers/N`, :chunk:`pairs/N`). types (list[str]): Names of the particle types (:chunk:`bonds/types`, :chunk:`angles/types`, :chunk:`dihedrals/types`, :chunk:`impropers/types`, :chunk:`pairs/types`). typeid ((*N*,) `numpy.ndarray` of ``numpy.uint32``): Bond type id (:chunk:`bonds/typeid`, :chunk:`angles/typeid`, :chunk:`dihedrals/typeid`, :chunk:`impropers/typeid`, :chunk:`pairs/types`). group ((*N*, *M*) `numpy.ndarray` of ``numpy.uint32``): Tags of the particles in the bond (:chunk:`bonds/group`, :chunk:`angles/group`, :chunk:`dihedrals/group`, :chunk:`impropers/group`, :chunk:`pairs/group`). """ def __init__(self, M): self.M = M self.N = 0 self.types = None self.typeid = None self.group = None self._default_value = OrderedDict() self._default_value['N'] = numpy.uint32(0) self._default_value['types'] = [] self._default_value['typeid'] = numpy.uint32(0) self._default_value['group'] = numpy.array([0] * M, dtype=numpy.int32) def validate(self): """Validate all attributes. Convert every array attribute to a `numpy.ndarray` of the proper type and check that all attributes have the correct dimensions. Ignore any attributes that are ``None``. Warning: Array attributes that are not contiguous numpy arrays will be replaced with contiguous numpy arrays of the appropriate type. """ logger.debug('Validating BondData') if self.typeid is not None: self.typeid = numpy.ascontiguousarray(self.typeid, dtype=numpy.uint32) self.typeid = self.typeid.reshape([self.N]) if self.group is not None: self.group = numpy.ascontiguousarray(self.group, dtype=numpy.int32) self.group = self.group.reshape([self.N, self.M]) if self.types is not None and (not len(set(self.types)) == len(self.types)): msg = 'Type names must be unique.' raise ValueError(msg) class ConstraintData: """Store constraint data. Use the `Frame.constraints` attribute to access the constraints. Instances resulting from file read operations will always store array quantities in `numpy.ndarray` objects of the defined types. User created frames may provide input data that can be converted to a `numpy.ndarray`. See Also: `hoomd.State` for a full description of how HOOMD interprets this data. Attributes: N (int): Number of constraints in the frame (:chunk:`constraints/N`). value ((*N*, ) `numpy.ndarray` of ``numpy.float32``): Constraint length (:chunk:`constraints/value`). group ((*N*, *2*) `numpy.ndarray` of ``numpy.uint32``): Tags of the particles in the constraint (:chunk:`constraints/group`). """ def __init__(self): self.M = 2 self.N = 0 self.value = None self.group = None self._default_value = OrderedDict() self._default_value['N'] = numpy.uint32(0) self._default_value['value'] = numpy.float32(0) self._default_value['group'] = numpy.array([0] * self.M, dtype=numpy.int32) def validate(self): """Validate all attributes. Convert every array attribute to a `numpy.ndarray` of the proper type and check that all attributes have the correct dimensions. Ignore any attributes that are ``None``. Warning: Array attributes that are not contiguous numpy arrays will be replaced with contiguous numpy arrays of the appropriate type. """ logger.debug('Validating ConstraintData') if self.value is not None: self.value = numpy.ascontiguousarray(self.value, dtype=numpy.float32) self.value = self.value.reshape([self.N]) if self.group is not None: self.group = numpy.ascontiguousarray(self.group, dtype=numpy.int32) self.group = self.group.reshape([self.N, self.M]) class Frame: """System state at one point in time. Attributes: configuration (`ConfigurationData`): Configuration data. particles (`ParticleData`): Particles. bonds (`BondData`): Bonds. angles (`BondData`): Angles. dihedrals (`BondData`): Dihedrals. impropers (`BondData`): Impropers. pairs (`BondData`): Special pair. constraints (`ConstraintData`): Distance constraints. state (dict): State data. log (dict): Logged data (values must be `numpy.ndarray` or `array_like`) """ def __init__(self): self.configuration = ConfigurationData() self.particles = ParticleData() self.bonds = BondData(2) self.angles = BondData(3) self.dihedrals = BondData(4) self.impropers = BondData(4) self.constraints = ConstraintData() self.pairs = BondData(2) self.state = {} self.log = {} self._valid_state = [ 'hpmc/integrate/d', 'hpmc/integrate/a', 'hpmc/sphere/radius', 'hpmc/sphere/orientable', 'hpmc/ellipsoid/a', 'hpmc/ellipsoid/b', 'hpmc/ellipsoid/c', 'hpmc/convex_polyhedron/N', 'hpmc/convex_polyhedron/vertices', 'hpmc/convex_spheropolyhedron/N', 'hpmc/convex_spheropolyhedron/vertices', 'hpmc/convex_spheropolyhedron/sweep_radius', 'hpmc/convex_polygon/N', 'hpmc/convex_polygon/vertices', 'hpmc/convex_spheropolygon/N', 'hpmc/convex_spheropolygon/vertices', 'hpmc/convex_spheropolygon/sweep_radius', 'hpmc/simple_polygon/N', 'hpmc/simple_polygon/vertices', ] def validate(self): """Validate all contained frame data.""" logger.debug('Validating Frame') self.configuration.validate() self.particles.validate() self.bonds.validate() self.angles.validate() self.dihedrals.validate() self.impropers.validate() self.constraints.validate() self.pairs.validate() # validate HPMC state if self.particles.types is not None: NT = len(self.particles.types) else: NT = 1 if 'hpmc/integrate/d' in self.state: self.state['hpmc/integrate/d'] = numpy.ascontiguousarray( self.state['hpmc/integrate/d'], dtype=numpy.float64 ) self.state['hpmc/integrate/d'] = self.state['hpmc/integrate/d'].reshape([1]) if 'hpmc/integrate/a' in self.state: self.state['hpmc/integrate/a'] = numpy.ascontiguousarray( self.state['hpmc/integrate/a'], dtype=numpy.float64 ) self.state['hpmc/integrate/a'] = self.state['hpmc/integrate/a'].reshape([1]) if 'hpmc/sphere/radius' in self.state: self.state['hpmc/sphere/radius'] = numpy.ascontiguousarray( self.state['hpmc/sphere/radius'], dtype=numpy.float32 ) self.state['hpmc/sphere/radius'] = self.state['hpmc/sphere/radius'].reshape( [NT] ) if 'hpmc/sphere/orientable' in self.state: self.state['hpmc/sphere/orientable'] = numpy.ascontiguousarray( self.state['hpmc/sphere/orientable'], dtype=numpy.uint8 ) self.state['hpmc/sphere/orientable'] = self.state[ 'hpmc/sphere/orientable' ].reshape([NT]) if 'hpmc/ellipsoid/a' in self.state: self.state['hpmc/ellipsoid/a'] = numpy.ascontiguousarray( self.state['hpmc/ellipsoid/a'], dtype=numpy.float32 ) self.state['hpmc/ellipsoid/a'] = self.state['hpmc/ellipsoid/a'].reshape( [NT] ) self.state['hpmc/ellipsoid/b'] = numpy.ascontiguousarray( self.state['hpmc/ellipsoid/b'], dtype=numpy.float32 ) self.state['hpmc/ellipsoid/b'] = self.state['hpmc/ellipsoid/b'].reshape( [NT] ) self.state['hpmc/ellipsoid/c'] = numpy.ascontiguousarray( self.state['hpmc/ellipsoid/c'], dtype=numpy.float32 ) self.state['hpmc/ellipsoid/c'] = self.state['hpmc/ellipsoid/c'].reshape( [NT] ) if 'hpmc/convex_polyhedron/N' in self.state: self.state['hpmc/convex_polyhedron/N'] = numpy.ascontiguousarray( self.state['hpmc/convex_polyhedron/N'], dtype=numpy.uint32 ) self.state['hpmc/convex_polyhedron/N'] = self.state[ 'hpmc/convex_polyhedron/N' ].reshape([NT]) sumN = numpy.sum(self.state['hpmc/convex_polyhedron/N']) self.state['hpmc/convex_polyhedron/vertices'] = numpy.ascontiguousarray( self.state['hpmc/convex_polyhedron/vertices'], dtype=numpy.float32 ) self.state['hpmc/convex_polyhedron/vertices'] = self.state[ 'hpmc/convex_polyhedron/vertices' ].reshape([sumN, 3]) if 'hpmc/convex_spheropolyhedron/N' in self.state: self.state['hpmc/convex_spheropolyhedron/N'] = numpy.ascontiguousarray( self.state['hpmc/convex_spheropolyhedron/N'], dtype=numpy.uint32 ) self.state['hpmc/convex_spheropolyhedron/N'] = self.state[ 'hpmc/convex_spheropolyhedron/N' ].reshape([NT]) sumN = numpy.sum(self.state['hpmc/convex_spheropolyhedron/N']) self.state['hpmc/convex_spheropolyhedron/sweep_radius'] = ( numpy.ascontiguousarray( self.state['hpmc/convex_spheropolyhedron/sweep_radius'], dtype=numpy.float32, ) ) self.state['hpmc/convex_spheropolyhedron/sweep_radius'] = self.state[ 'hpmc/convex_spheropolyhedron/sweep_radius' ].reshape([NT]) self.state['hpmc/convex_spheropolyhedron/vertices'] = ( numpy.ascontiguousarray( self.state['hpmc/convex_spheropolyhedron/vertices'], dtype=numpy.float32, ) ) self.state['hpmc/convex_spheropolyhedron/vertices'] = self.state[ 'hpmc/convex_spheropolyhedron/vertices' ].reshape([sumN, 3]) if 'hpmc/convex_polygon/N' in self.state: self.state['hpmc/convex_polygon/N'] = numpy.ascontiguousarray( self.state['hpmc/convex_polygon/N'], dtype=numpy.uint32 ) self.state['hpmc/convex_polygon/N'] = self.state[ 'hpmc/convex_polygon/N' ].reshape([NT]) sumN = numpy.sum(self.state['hpmc/convex_polygon/N']) self.state['hpmc/convex_polygon/vertices'] = numpy.ascontiguousarray( self.state['hpmc/convex_polygon/vertices'], dtype=numpy.float32 ) self.state['hpmc/convex_polygon/vertices'] = self.state[ 'hpmc/convex_polygon/vertices' ].reshape([sumN, 2]) if 'hpmc/convex_spheropolygon/N' in self.state: self.state['hpmc/convex_spheropolygon/N'] = numpy.ascontiguousarray( self.state['hpmc/convex_spheropolygon/N'], dtype=numpy.uint32 ) self.state['hpmc/convex_spheropolygon/N'] = self.state[ 'hpmc/convex_spheropolygon/N' ].reshape([NT]) sumN = numpy.sum(self.state['hpmc/convex_spheropolygon/N']) self.state['hpmc/convex_spheropolygon/sweep_radius'] = ( numpy.ascontiguousarray( self.state['hpmc/convex_spheropolygon/sweep_radius'], dtype=numpy.float32, ) ) self.state['hpmc/convex_spheropolygon/sweep_radius'] = self.state[ 'hpmc/convex_spheropolygon/sweep_radius' ].reshape([NT]) self.state['hpmc/convex_spheropolygon/vertices'] = numpy.ascontiguousarray( self.state['hpmc/convex_spheropolygon/vertices'], dtype=numpy.float32 ) self.state['hpmc/convex_spheropolygon/vertices'] = self.state[ 'hpmc/convex_spheropolygon/vertices' ].reshape([sumN, 2]) if 'hpmc/simple_polygon/N' in self.state: self.state['hpmc/simple_polygon/N'] = numpy.ascontiguousarray( self.state['hpmc/simple_polygon/N'], dtype=numpy.uint32 ) self.state['hpmc/simple_polygon/N'] = self.state[ 'hpmc/simple_polygon/N' ].reshape([NT]) sumN = numpy.sum(self.state['hpmc/simple_polygon/N']) self.state['hpmc/simple_polygon/vertices'] = numpy.ascontiguousarray( self.state['hpmc/simple_polygon/vertices'], dtype=numpy.float32 ) self.state['hpmc/simple_polygon/vertices'] = self.state[ 'hpmc/simple_polygon/vertices' ].reshape([sumN, 2]) for k in self.state: if k not in self._valid_state: raise RuntimeError('Not a valid state: ' + k) class _HOOMDTrajectoryIterable: """Iterable over a HOOMDTrajectory object.""" def __init__(self, trajectory, indices): self._trajectory = trajectory self._indices = indices self._indices_iterator = iter(indices) def __next__(self): return self._trajectory[next(self._indices_iterator)] next = __next__ # Python 2.7 compatibility def __iter__(self): return type(self)(self._trajectory, self._indices) def __len__(self): return len(self._indices) class _HOOMDTrajectoryView: """A view of a HOOMDTrajectory object. Enables the slicing and iteration over a subset of a trajectory instance. """ def __init__(self, trajectory, indices): self._trajectory = trajectory self._indices = indices def __iter__(self): return _HOOMDTrajectoryIterable(self._trajectory, self._indices) def __len__(self): return len(self._indices) def __getitem__(self, key): if isinstance(key, slice): return type(self)(self._trajectory, self._indices[key]) return self._trajectory[self._indices[key]] class HOOMDTrajectory: """Read and write hoomd gsd files. Args: file (`gsd.fl.GSDFile`): File to access. Open hoomd GSD files with `open`. """ def __init__(self, file): if file.mode == 'ab': msg = 'Append mode not yet supported' raise ValueError(msg) self._file = file self._initial_frame = None # Used to cache positive results when chunks exist in frame 0. self._chunk_exists_frame_0 = {} logger.info('opening HOOMDTrajectory: ' + str(self.file)) if self.file.schema != 'hoomd': raise RuntimeError('GSD file is not a hoomd schema file: ' + str(self.file)) valid = False version = self.file.schema_version if version < (2, 0) and version >= (1, 0): valid = True if not valid: raise RuntimeError( 'Incompatible hoomd schema version ' + str(version) + ' in: ' + str(self.file) ) logger.info('found ' + str(len(self)) + ' frames') @property def file(self): """:class:`gsd.fl.GSDFile`: The file handle.""" return self._file def __len__(self): """The number of frames in the trajectory.""" return self.file.nframes def append(self, frame): """Append a frame to a hoomd gsd file. Args: frame (:py:class:`Frame`): Frame to append. Write the given frame to the file at the current frame and increase the frame counter. Do not write any fields that are ``None``. For all non-``None`` fields, scan them and see if they match the initial frame or the default value. If the given data differs, write it out to the frame. If it is the same, do not write it out as it can be instantiated either from the value at the initial frame or the default value. """ logger.debug('Appending frame to hoomd trajectory: ' + str(self.file)) frame.validate() # want the initial frame specified as a reference to detect if chunks # need to be written if self._initial_frame is None and len(self) > 0: self._read_frame(0) for path in [ 'configuration', 'particles', 'bonds', 'angles', 'dihedrals', 'impropers', 'constraints', 'pairs', ]: container = getattr(frame, path) for name in container._default_value: if self._should_write(path, name, frame): logger.debug('writing data chunk: ' + path + '/' + name) data = getattr(container, name) if name == 'N': data = numpy.array([data], dtype=numpy.uint32) if name == 'step': data = numpy.array([data], dtype=numpy.uint64) if name == 'dimensions': data = numpy.array([data], dtype=numpy.uint8) if name in ('types', 'type_shapes'): if name == 'type_shapes': data = [json.dumps(shape_dict) for shape_dict in data] wid = max(len(w) for w in data) + 1 b = numpy.array(data, dtype=numpy.dtype((bytes, wid))) data = b.view(dtype=numpy.int8).reshape(len(b), wid) self.file.write_chunk(path + '/' + name, data) # write state data for state, data in frame.state.items(): self.file.write_chunk('state/' + state, data) # write log data for log, data in frame.log.items(): self.file.write_chunk('log/' + log, data) self.file.end_frame() def truncate(self): """Remove all frames from the file.""" self.file.truncate() self._initial_frame = None def close(self): """Close the file.""" self.file.close() del self._initial_frame def _should_write(self, path, name, frame): """Test if we should write a given data chunk. Args: path (str): Path part of the data chunk. name (str): Name part of the data chunk. frame (:py:class:`Frame`): Frame data is from. Returns: False if the data matches that in the initial frame. False if the data matches all default values. True otherwise. """ container = getattr(frame, path) data = getattr(container, name) if data is None: return False if self._initial_frame is not None: initial_container = getattr(self._initial_frame, path) initial_data = getattr(initial_container, name) if numpy.array_equal(initial_data, data): logger.debug( 'skipping data chunk, matches frame 0: ' + path + '/' + name ) return False matches_default_value = False if name == 'types': matches_default_value = data == container._default_value[name] else: matches_default_value = numpy.array_equiv( data, container._default_value[name] ) if matches_default_value and not self._chunk_exists_frame_0.get( path + '/' + name, False ): logger.debug('skipping data chunk, default value: ' + path + '/' + name) return False return True def extend(self, iterable): """Append each item of the iterable to the file. Args: iterable: An iterable object the provides :py:class:`Frame` instances. This could be another HOOMDTrajectory, a generator that modifies frames, or a list of frames. """ for item in iterable: self.append(item) def _read_frame(self, idx): """Read the frame at the given index from the file. Args: idx (int): Frame index to read. Returns: `Frame` with the frame data Replace any data chunks not present in the given frame with either data from frame 0, or initialize from default values if not in frame 0. Cache frame 0 data to avoid file read overhead. Return any default data as non-writeable numpy arrays. """ if idx >= len(self): raise IndexError logger.debug('reading frame ' + str(idx) + ' from: ' + str(self.file)) if self._initial_frame is None and idx != 0: self._read_frame(0) frame = Frame() # read configuration first if self.file.chunk_exists(frame=idx, name='configuration/step'): step_arr = self.file.read_chunk(frame=idx, name='configuration/step') frame.configuration.step = step_arr[0] if idx == 0: self._chunk_exists_frame_0['configuration/step'] = True elif self._initial_frame is not None: frame.configuration.step = self._initial_frame.configuration.step else: frame.configuration.step = frame.configuration._default_value['step'] if self.file.chunk_exists(frame=idx, name='configuration/dimensions'): dimensions_arr = self.file.read_chunk( frame=idx, name='configuration/dimensions' ) frame.configuration.dimensions = dimensions_arr[0] if idx == 0: self._chunk_exists_frame_0['configuration/dimensions'] = True elif self._initial_frame is not None: frame.configuration.dimensions = ( self._initial_frame.configuration.dimensions ) else: frame.configuration.dimensions = frame.configuration._default_value[ 'dimensions' ] if self.file.chunk_exists(frame=idx, name='configuration/box'): frame.configuration.box = self.file.read_chunk( frame=idx, name='configuration/box' ) if idx == 0: self._chunk_exists_frame_0['configuration/box'] = True elif self._initial_frame is not None: frame.configuration.box = copy.copy(self._initial_frame.configuration.box) else: frame.configuration.box = copy.copy( frame.configuration._default_value['box'] ) # then read all groups that have N, types, etc... for path in [ 'particles', 'bonds', 'angles', 'dihedrals', 'impropers', 'constraints', 'pairs', ]: container = getattr(frame, path) if self._initial_frame is not None: initial_frame_container = getattr(self._initial_frame, path) container.N = 0 if self.file.chunk_exists(frame=idx, name=path + '/N'): N_arr = self.file.read_chunk(frame=idx, name=path + '/N') container.N = N_arr[0] if idx == 0: self._chunk_exists_frame_0[path + '/N'] = True elif self._initial_frame is not None: container.N = initial_frame_container.N # type names if 'types' in container._default_value: if self.file.chunk_exists(frame=idx, name=path + '/types'): tmp = self.file.read_chunk(frame=idx, name=path + '/types') tmp = tmp.view(dtype=numpy.dtype((bytes, tmp.shape[1]))) tmp = tmp.reshape([tmp.shape[0]]) container.types = list(a.decode('UTF-8') for a in tmp) if idx == 0: self._chunk_exists_frame_0[path + '/types'] = True elif self._initial_frame is not None: container.types = copy.copy(initial_frame_container.types) else: container.types = copy.copy(container._default_value['types']) # type shapes if 'type_shapes' in container._default_value and path == 'particles': if self.file.chunk_exists(frame=idx, name=path + '/type_shapes'): tmp = self.file.read_chunk(frame=idx, name=path + '/type_shapes') tmp = tmp.view(dtype=numpy.dtype((bytes, tmp.shape[1]))) tmp = tmp.reshape([tmp.shape[0]]) container.type_shapes = list( json.loads(json_string.decode('UTF-8')) for json_string in tmp ) if idx == 0: self._chunk_exists_frame_0[path + '/type_shapes'] = True elif self._initial_frame is not None: container.type_shapes = copy.copy( initial_frame_container.type_shapes ) else: container.type_shapes = copy.copy( container._default_value['type_shapes'] ) for name in container._default_value: if name in ('N', 'types', 'type_shapes'): continue # per particle/bond quantities if self.file.chunk_exists(frame=idx, name=path + '/' + name): container.__dict__[name] = self.file.read_chunk( frame=idx, name=path + '/' + name ) if idx == 0: self._chunk_exists_frame_0[path + '/' + name] = True else: if ( self._initial_frame is not None and initial_frame_container.N == container.N ): # read default from initial frame container.__dict__[name] = initial_frame_container.__dict__[ name ] else: # initialize from default value tmp = numpy.array([container._default_value[name]]) s = list(tmp.shape) s[0] = container.N container.__dict__[name] = numpy.empty(shape=s, dtype=tmp.dtype) container.__dict__[name][:] = tmp container.__dict__[name].flags.writeable = False # read state data for state in frame._valid_state: if self.file.chunk_exists(frame=idx, name='state/' + state): frame.state[state] = self.file.read_chunk( frame=idx, name='state/' + state ) # read log data logged_data_names = self.file.find_matching_chunk_names('log/') for log in logged_data_names: if self.file.chunk_exists(frame=idx, name=log): frame.log[log[4:]] = self.file.read_chunk(frame=idx, name=log) if idx == 0: self._chunk_exists_frame_0[log] = True elif self._initial_frame is not None: frame.log[log[4:]] = self._initial_frame.log[log[4:]] frame.log[log[4:]].flags.writeable = False # store initial frame if self._initial_frame is None and idx == 0: self._initial_frame = copy.deepcopy(frame) return frame def __getitem__(self, key): """Index trajectory frames. The index can be a positive integer, negative integer, or slice and is interpreted the same as `list` indexing. Warning: As you loop over frames, each frame is read from the file when it is reached in the iteration. Multiple passes may lead to multiple disk reads if the file does not fit in cache. """ if isinstance(key, slice): return _HOOMDTrajectoryView(self, range(*key.indices(len(self)))) if isinstance(key, int): if key < 0: key += len(self) if key >= len(self) or key < 0: raise IndexError() return self._read_frame(key) raise TypeError def __iter__(self): """Iterate over frames in the trajectory.""" return _HOOMDTrajectoryIterable(self, range(len(self))) def __enter__(self): """Enter the context manager.""" return self def __exit__(self, exc_type, exc_value, traceback): """Close the file when the context manager exits.""" self.file.close() def flush(self): """Flush all buffered frames to the file.""" self._file.flush() def open(name, mode='r'): # noqa: A001 - allow shadowing builtin open """Open a hoomd schema GSD file. The return value of `open` can be used as a context manager. Args: name (str): File name to open. mode (str): File open mode. Returns: `HOOMDTrajectory` instance that accesses the file **name** with the given **mode**. Valid values for ``mode``: +------------------+---------------------------------------------+ | mode | description | +==================+=============================================+ | ``'r'`` | Open an existing file for reading. | +------------------+---------------------------------------------+ | ``'r+'`` | Open an existing file for reading and | | | writing. | +------------------+---------------------------------------------+ | ``'w'`` | Open a file for reading and writing. | | | Creates the file if needed, or overwrites | | | an existing file. | +------------------+---------------------------------------------+ | ``'x'`` | Create a gsd file exclusively and opens it | | | for reading and writing. | | | Raise :py:exc:`FileExistsError` | | | if it already exists. | +------------------+---------------------------------------------+ | ``'a'`` | Open a file for reading and writing. | | | Creates the file if it doesn't exist. | +------------------+---------------------------------------------+ """ if not fl_imported: msg = 'file layer module is not available' raise RuntimeError(msg) if gsd is None: msg = 'gsd module is not available' raise RuntimeError(msg) gsdfileobj = gsd.fl.open( name=str(name), mode=mode, application='gsd.hoomd ' + gsd.version.version, schema='hoomd', schema_version=[1, 4], ) return HOOMDTrajectory(gsdfileobj) def read_log(name, scalar_only=False): """Read log from a hoomd schema GSD file into a dict of time-series arrays. Args: name (str): File name to open. scalar_only (bool): Set to `True` to include only scalar log values. The log data includes :chunk:`configuration/step` and all matching :chunk:`log/user_defined`, :chunk:`log/bonds/user_defined`, and :chunk:`log/particles/user_defined` quantities in the file. Returns: `dict` Note: `read_log` issues a `RuntimeWarning` when there are no matching ``log/`` quantities in the file. Caution: `read_log` requires that a logged quantity has the same shape in all frames. Use `open` and `Frame.log` to read files where the shape changes from frame to frame. To create a *pandas* ``DataFrame`` with the logged data: .. ipython:: python import pandas df = pandas.DataFrame(gsd.hoomd.read_log('log-example.gsd', scalar_only=True)) df """ if not fl_imported: msg = 'file layer module is not available' raise RuntimeError(msg) if gsd is None: msg = 'gsd module is not available' raise RuntimeError(msg) min_supported_numpy = 2 if int(numpy.version.version.split('.')[0]) < min_supported_numpy: msg = 'read_log requires numpy >= 2.0' raise RuntimeError(msg) with gsd.fl.open( name=str(name), mode='r', application='gsd.hoomd ' + gsd.version.version, schema='hoomd', schema_version=[1, 4], ) as gsdfileobj: logged_data_names = gsdfileobj.find_matching_chunk_names('log/') # Always log timestep associated with each log entry logged_data_names.insert(0, 'configuration/step') if len(logged_data_names) == 1: warnings.warn( 'No logged data in file: ' + str(name), RuntimeWarning, stacklevel=2 ) logged_data_dict = dict() for log in logged_data_names: log_exists_frame_0 = gsdfileobj.chunk_exists(frame=0, name=log) is_configuration_step = log == 'configuration/step' if log_exists_frame_0 or is_configuration_step: if is_configuration_step and not log_exists_frame_0: # handle default configuration step on frame 0 tmp = numpy.array([0], dtype=numpy.uint64) else: tmp = gsdfileobj.read_chunk(frame=0, name=log) # if chunk contains string, put it in the numpy array if isinstance(tmp, str): tmp = numpy.array([tmp], dtype=numpy.dtypes.StringDType) if scalar_only and not tmp.shape[0] == 1: continue if tmp.shape[0] == 1: logged_data_dict[log] = numpy.full( fill_value=tmp[0], shape=(gsdfileobj.nframes,), dtype=tmp.dtype, ) else: logged_data_dict[log] = numpy.tile( tmp, (gsdfileobj.nframes, *tuple(1 for _ in tmp.shape)) ) for idx in range(1, gsdfileobj.nframes): for key in logged_data_dict.keys(): if not gsdfileobj.chunk_exists(frame=idx, name=key): continue data = gsdfileobj.read_chunk(frame=idx, name=key) if ( not isinstance( logged_data_dict[key].dtype, numpy.dtypes.StringDType ) and len(logged_data_dict[key][idx].shape) == 0 ): logged_data_dict[key][idx] = data[0] else: logged_data_dict[key][idx] = data return logged_data_dict