def main(argv=None): args = process_command_line(argv) tick = time.perf_counter() with open(args.centers, 'rb') as f: centers = concatenate_trjs( pickle.load(f), args.atoms, enspara.util.parallel.auto_nprocs()) logger.info('Loaded %s centers with %s atoms using selection "%s" ' 'in %.1f seconds.', len(centers), centers.n_atoms, args.atoms, time.perf_counter() - tick) assig, dist = reassign( args.topologies, args.trajectories, [args.atoms]*len(args.topologies), centers=centers, frac_mem=args.mem_fraction) mem_highwater = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss logger.info( "Finished reassignments in %.1f seconds. Process memory high-water " "mark was %.2f GB (VRAM size is %.2f GB).", time.perf_counter() - tick, (mem_highwater / 1024**2), psutil.virtual_memory().total / 1024**3) ra.save(args.distances, dist) ra.save(args.assignments, assig) logger.info("Wrote distances at %s.", args.distances) logger.info("Wrote assignments at %s.", args.assignments) return 0
def run(self): # load and concat trjs if self.mem_efficient: trj_lengths, xyzs = load_as_concatenated( filenames=self.trj_filenames, processes=self.n_procs, top=self.base_struct_md, atom_indices=self.atom_indices_vals) trjs_sub = md.Trajectory( xyzs, self.base_struct_md.atom_slice( self.atom_indices_vals).topology) else: trj_lengths, xyzs = load_as_concatenated( filenames=self.trj_filenames, processes=self.n_procs, top=self.base_struct_md) trjs = md.Trajectory(xyzs, self.base_struct_md.topology) trjs_sub = trjs.atom_slice(self.atom_indices_vals) # determine if rebuilding all msm stuff if self.build_full: base_struct_centers = self.base_struct_md.atom_slice( self.atom_indices_vals) base_struct_centers.save_pdb("./centers.pdb") self.base_struct_md.save_pdb("./prot_masses.pdb") init_centers = None else: init_centers = md.load("./data/centers.xtc", top="./centers.pdb") # fit data with base clustering object self.base_clust_obj.fit(trjs_sub, init_centers=init_centers) center_indices, distances, assignments, centers = \ self.base_clust_obj.result_.partition(trj_lengths) # save data ra.save("./data/assignments.h5", assignments) ra.save("./data/distances.h5", distances) trjs_sub = trjs_sub[self.base_clust_obj.center_indices_] trjs_sub.superpose(trjs_sub[0]) trjs_sub.save_xtc("./data/centers.xtc") if not self.mem_efficient: full_centers = trjs[self.base_clust_obj.center_indices_] full_centers.superpose(self.base_struct_md) full_centers.save_xtc("./data/full_centers.xtc") # save states n_states = len(self.base_clust_obj.center_indices_) unique_states = np.arange(n_states) if init_centers is not None: unique_states = unique_states[-(n_states - len(init_centers)):] np.save("./data/unique_states.npy", unique_states)
def write_assignments_and_distances_with_reassign(result, args): if args.subsample == 1: logger.debug("Subsampling was 1, not reassigning.") ra.save(args.distances, result.distances) ra.save(args.assignments, result.assignments) elif not args.no_reassign: logger.debug("Reassigning data from subsampling of %s", args.subsample) assig, dist = reassign(args.topologies, args.trajectories, args.atoms, centers=result.centers) ra.save(args.distances, dist) ra.save(args.assignments, assig) else: logger.debug("Got --no-reassign, not doing reassigment")
def main(argv=None): args = process_command_line(argv) try: features = ra.load(args.features, keys=...) except exception.DataInvalid: features = ra.load(args.features) logger.info("Loaded data from %s with shape %s", args.features, features.shape) if args.cluster_algorithm == 'khybrid': clustering = KHybrid(metric=args.cluster_distance, cluster_radius=args.cluster_radius, kmedoids_updates=args.kmedoids_updates) elif args.cluster_algorithm == 'kcenters': clustering = KCenters(cluster_radius=args.cluster_radius, metric=args.cluster_distance) logger.info("Clustering with %s", clustering) clustering.fit(features._data) result = clustering.result_.partition(features.lengths) del features ra.save(args.distances, result.distances) logger.info("Wrote distances with shape %s to %s", result.distances.shape, args.distances) ra.save(args.assignments, result.assignments) logger.info("Wrote assignments with shape %s to %s", result.assignments.shape, args.cluster_centers) ra.save(args.cluster_centers, result.centers) logger.info("Wrote cluster_centers with shape %s to %s", result.centers.shape, args.cluster_centers) pickle.dump(result.center_indices, open(args.center_indices, 'wb')) logger.info("Wrote %s center_indices with shape to %s", len(result.center_indices), args.center_indices) return 0
def main(argv=None): '''Run the driver script for this module. This code only runs if we're being run as a script. Otherwise, it's silent and just exposes methods.''' args = process_command_line(argv) top = md.load(args.topology).top atom_ids = top.select(args.selection) logging.info("Running with %s total workers.", MPI_SIZE) logging.info( "Loading trajectories [%s::%s]; selection == '%s' w/ " "subsampling %s", MPI_RANK, MPI_SIZE, args.selection, args.subsample) with timed("load_as_concatenated took %.2f sec", logging.info): global_lengths, my_xyz = mpi.io.load_as_striped( filenames=args.trajectories, top=top, atom_indices=atom_ids, stride=args.subsample, processes=args.processes) with timed("Turned over array in %.2f min", logging.info): xyz = my_xyz.copy() del my_xyz my_xyz = xyz logging.info("Loaded %s frames in %s trjs (%.2fG).", len(my_xyz), len(args.trajectories) // MPI_SIZE, my_xyz.data.nbytes / 1024**3) trjs = md.Trajectory(my_xyz, topology=top.subset(atom_ids)) logging.info( "Beginning kcenters clustering with memory footprint of %.2fG " "RAM (coords are %.2fG; total VRAM is %.2fG)", resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024**2, trjs.xyz.nbytes / 1024**3, psutil.virtual_memory().total / 1024**3) if len(args.cluster_radii) > 1: raise NotImplementedError( "Multiple cluster radii are not yet supported") tick = time.perf_counter() local_dists, local_assigs, local_ctr_inds = kcenters_mpi( trjs, md.rmsd, dist_cutoff=args.cluster_radii[0]) tock = time.perf_counter() logging.info( "Finished kcenters clustering using %.2fG RAM (coords are " "%.2fG) in %.2f min.", resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024**2, trjs.xyz.nbytes / 1024**3, (tock - tick) / 60) for i in range(args.kmedoids_iters): with timed("KMedoids iteration {i} took %.2f sec".format(i=i), logging.info): local_ctr_inds, local_dists, local_assigs = _kmedoids_pam_update( X=trjs, metric=md.rmsd, medoid_inds=local_ctr_inds, assignments=local_assigs, distances=local_dists, random_state=args.random_state) with timed("Reassembled dist and assign arrays in %.2f sec", logging.info): all_dists = mpi.ops.assemble_striped_ragged_array( local_dists, global_lengths) all_assigs = mpi.ops.assemble_striped_ragged_array( local_assigs, global_lengths) ctr_inds = mpi.ops.convert_local_indices(local_ctr_inds, global_lengths) ctr_inds = partition_indices(ctr_inds, global_lengths) if MPI_RANK == 0: logging.info("Dumping center indices to %s", args.center_indices) with open(args.center_indices, 'wb') as f: pickle.dump([(trj, frame * args.subsample) for trj, frame in ctr_inds], f) if args.distances: ra.save(args.distances, ra.RaggedArray(all_dists, lengths=global_lengths)) if args.assignments: ra.save(args.assignments, ra.RaggedArray(all_assigs, lengths=global_lengths)) centers = load_frames(args.trajectories, ctr_inds, stride=args.subsample, top=md.load(args.topology).top) with open(args.center_structures, 'wb') as f: pickle.dump(centers, f) logging.info("Wrote %s centers to %s", len(centers), args.center_structures) return 0
def main(argv=None): '''Run the driver script for this module. This code only runs if we're being run as a script. Otherwise, it's silent and just exposes methods.''' args = process_command_line(argv) targets = { os.path.basename(topf): "%s xtcs" % len(trjfs) for topf, trjfs in zip(args.topologies, args.trajectories) } logger.info("Beginning RMSD Clustering app. Operating on targets:\n%s", json.dumps(targets, indent=4)) tick = time.perf_counter() lengths, xyz, select_top = load(args.topologies, args.trajectories, selections=args.atoms, stride=args.subsample, processes=args.processes) logger.info( "Loading finished in %.1f s. Clustering using %s atoms matching '%s'.", round(time.perf_counter() - tick, 2), xyz.shape[1], args.atoms) clustering = args.Clusterer(metric=md.rmsd, n_clusters=args.n_clusters, cluster_radius=args.rmsd_cutoff) # md.rmsd requires an md.Trajectory object, so wrap `xyz` in # the topology. clustering.fit(md.Trajectory(xyz=xyz, topology=select_top)) logger.info("Clustered %s frames into %s clusters in %s seconds.", sum(lengths), len(clustering.centers_), clustering.runtime_) result = clustering.result_.partition(lengths) outdir = os.path.dirname(args.centers) logger.info("Saving cluster centers at %s", outdir) try: os.makedirs(outdir) except FileExistsError: pass centers = load_asymm_frames(result.center_indices, args.trajectories, args.topologies, args.subsample) with open(args.centers, 'wb') as f: pickle.dump(centers, f) if args.subsample == 1: logger.debug("Subsampling was 1, not reassigning.") ra.save(args.distances, result.distances) ra.save(args.assignments, result.assignments) if not args.no_reassign: logger.debug("Reassigning data from subsampling of %s", args.subsample) # overwrite temporary output with actual results assig, dist = reassign(args.topologies, args.trajectories, args.atoms, centers=result.centers) ra.save(args.distances, dist) ra.save(args.assignments, assig) else: logger.debug("Got --no-reassign, not doing reassigment") logger.info("Success! Data can be found in %s.", os.path.dirname(args.distances)) return 0
def entry_point(): if True: # filenames filenames = np.sort([ os.path.abspath(pathname) for pathname in glob.glob("./trajectories/*.xtc") ]) print("obtained filenames!") # load atom indices pdb = md.load("prot_masses.pdb") iis = pdb.topology.select("backbone and resid 72 to 87") # iis = np.loadtxt("./atom-indices-bb.dat", dtype=int) # topology filename prot_masses = "./prot_masses.pdb" prot_masses = md.load(prot_masses) # load trjs print("about to load!!") centers = prot_masses.atom_slice(iis) trj_lengths, xyzs = load_as_concatenated(filenames=filenames, atom_indices=iis, processes=48, top=prot_masses) trjs_sub = md.Trajectory(xyz=xyzs, topology=centers.top) del xyzs if True: # get subset n_clusters = 10000 #n_clusters = None dist_cutoff = 0.01 clusterer = cluster.KCenters(metric=md.rmsd, cluster_radius=dist_cutoff, n_clusters=n_clusters) #clusterer = cluster.KHybrid(metric=md.rmsd, cluster_radius=dist_cutoff, n_clusters=n_clusters, kmedoids_updates=2) clusterer.fit(trjs_sub) center_indices, distances, assignments, centers = \ clusterer.result_.partition(trj_lengths) ra.save("./data/assignments.h5", assignments) ra.save("./data/distances.h5", distances) trjs_sub[clusterer.center_indices_].save_xtc("./data/centers.xtc") np.save("./data/center_indices.npy", clusterer.center_indices_) print("Done clustering!") if True: lag_time = 10 # 20ps * 200 = 4 ns #lag_time = 1 # 20ps * 200 = 4 ns assignments = ra.load("./data/assignments.h5") unique_states = np.unique(np.concatenate(assignments)) b = partial(builders.normalize, prior_counts=1 / unique_states.shape[0]) msm_obj = MSM(lag_time=lag_time, method=b) msm_obj.fit(assignments) np.save("./data/tcounts.npy", msm_obj.tcounts_) np.save("./data/tprobs.npy", msm_obj.tprobs_) np.save("./data/populations.npy", msm_obj.eq_probs_) print("Done MSM!")