def take_snapshot(wrapped_context): # context = wrapped_context.context.context # extra indirection for OpenMM context_view = wrapped_context.view # positions = asarray(dlext.positions(context_view)) forces = asarray(dlext.forces(context_view)) ids = asarray(dlext.atom_ids(context_view)) # velocities = asarray(dlext.velocities(context_view)) if is_on_gpu(context_view): vel_mass = velocities else: inverse_masses = asarray(dlext.inverse_masses(context_view)) vel_mass = (velocities, inverse_masses.reshape((-1, 1))) # box_vectors = context.getSystem().getDefaultPeriodicBoxVectors() a = box_vectors[0].value_in_unit(unit.nanometer) b = box_vectors[1].value_in_unit(unit.nanometer) c = box_vectors[2].value_in_unit(unit.nanometer) H = ((a[0], b[0], c[0]), (a[1], b[1], c[1]), (a[2], b[2], c[2])) origin = (0.0, 0.0, 0.0) dt = context.getIntegrator().getStepSize() / unit.picosecond # OpenMM doesn't have images return Snapshot(positions, vel_mass, forces, ids, None, Box(H, origin), dt)
def update_snapshot(snapshot, sysview, location=default_location()): # positions = asarray(positions_types(sysview, location, AccessMode.Read)) vel_mass = asarray(velocities_masses(sysview, location, AccessMode.Read)) forces = asarray(net_forces(sysview, location, AccessMode.ReadWrite)) ids = asarray(rtags(sysview, location, AccessMode.Read)) imgs = asarray(images(sysview, location, AccessMode.Read)) # return Snapshot(positions, vel_mass, forces, ids, imgs, snapshot.box, snapshot.dt)
def python_update(positions, vel_mass, rtags, imgs, forces): positions = asarray(positions) vel_mass = asarray(vel_mass) ids = asarray(rtags) images = asarray(imgs) forces = asarray(forces) snap = Snapshot(positions=positions, vel_mass=vel_mass, forces=forces, ids=ids, images=images, box=self.pybox, dt=self.dt) self.state = update(snap, self.state) self.bias(snap, self.state) if self.callback: self.callback(snap, self.state, 0)
def take_snapshot(wrapped_context, location=default_location()): # context = wrapped_context.context sysview = wrapped_context.view # positions = asarray(positions_types(sysview, location, AccessMode.Read)) vel_mass = asarray(velocities_masses(sysview, location, AccessMode.Read)) forces = asarray(net_forces(sysview, location, AccessMode.ReadWrite)) ids = asarray(rtags(sysview, location, AccessMode.Read)) imgs = asarray(images(sysview, location, AccessMode.Read)) # box = sysview.particle_data().getGlobalBox() L = box.getL() xy = box.getTiltFactorXY() xz = box.getTiltFactorXZ() yz = box.getTiltFactorYZ() lo = box.getLo() H = ((L.x, xy * L.y, xz * L.z), (0.0, L.y, yz * L.z), (0.0, 0.0, L.z)) origin = (lo.x, lo.y, lo.z) dt = context.integrator.dt # return Snapshot(positions, vel_mass, forces, ids, imgs, Box(H, origin), dt)