<jupyter_start><jupyter_text><jupyter_code>import sys
print(sys.version)
!wget https://repo.anaconda.com/miniconda/Miniconda3-py37_4.10.3-Linux-x86_64.sh
!bash Miniconda3-py*.sh -bfp /usr/local
!conda config --set always_yes yes
!conda config --add channels conda-forge
!conda create -n openmm python=3.7 cudatoolkit=10.0 git jupyterlab numpy pandas scipy matplotlib ipympl rdkit openbabel openmm mdtraj pymbar pdbfixer parmed openff-toolkit openmoltools openmmforcefields
sys.path.append('/usr/local/envs/openmm/lib/python3.7/site-packages')
import openmm.testInstallation
openmm.testInstallation.main()
!ls drive/MyDrive/Colab\ Notebooks
import mdtraj as mdt
import numpy as np
from simtk.openmm.app import *
from simtk.openmm import *
from simtk.unit import *
from sys import stdout
prmtop = AmberPrmtopFile('drive/MyDrive/Colab Notebooks/nacl.parm7')
inpcrd = AmberInpcrdFile('drive/MyDrive/Colab Notebooks/bstate.ncrst')
system = prmtop.createSystem(nonbondedMethod=PME, nonbondedCutoff=1*nanometer,
        constraints=HBonds)
integrator = LangevinIntegrator(300*kelvin, 1/picosecond, 0.002*picoseconds)
simulation = Simulation(prmtop.topology, system, integrator)
simulation.context.setPositions(inpcrd.positions)
if inpcrd.boxVectors is not None:
    simulation.context.setPeriodicBoxVectors(*inpcrd.boxVectors)
class Walker(object):
    def __init__(self, coordinates, boxvectors, weight):
        self.coordinates = coordinates
        self.boxvectors = boxvectors
        self.weight = weight
        self.pc = None
        self.bin = None
        
    def copy(self):
        new = Walker(self.coordinates, self.boxvectors, self.weight)
        new.pc = self.pc
        new.bin = self.bin
        return new

class Stepper(object):
    def __init__(self, simulation, nsteps):
        self.simulation = simulation
        self.nsteps = nsteps
        
    def run(self, walkers):
        new_walkers = []
        for w in walkers:
            self.simulation.context.setPositions(w.coordinates)
            self.simulation.context.setPeriodicBoxVectors(*w.boxvectors)
            self.simulation.step(self.nsteps)
            state = self.simulation.context.getState(getPositions=True, enforcePeriodicBox=True)
            new_walkers.append(Walker(state.getPositions(), state.getPeriodicBoxVectors(), w.weight))
            
        return new_walkers
stepper = Stepper(simulation, 1000)
state = simulation.context.getState(getPositions=True)
n_reps = 5
initial_pos = state.getPositions()
initial_vecs = state.getPeriodicBoxVectors()
weight = 1.0 / n_reps
walkers = [Walker(initial_pos, initial_vecs, weight) for i in range(n_reps)]
new_walkers = stepper.run(walkers)
class ProgressCoordinator(object):
    def __init__(self, a1, a2):
        self.a1 = a1
        self.a2 = a2
        
    def run(self, walkers):
        if not isinstance(walkers, list):
            walkers = [walkers]
        for i, w in enumerate(walkers):
            crds = w.coordinates
            x1 = np.array(crds[self.a1] / nanometer)
            x2 = np.array(crds[self.a2] / nanometer)
            pc = np.sqrt(((x1-x2)*(x1-x2)).sum())
            walkers[i].pc = pc
        return walkers
    
class Binner(object):
    def __init__(self, edges):
        self.edges = edges
    
    def run(self, walkers):
        if not isinstance(walkers, list):
            walkers = [walkers]
        pcs = [w.pc for w in walkers]
        if None in pcs:
            raise TypeError('Error: missing progress coordinates...')
        bin_ids = np.digitize(pcs, self.edges)
        for i, bin_id in enumerate(bin_ids):
            walkers[i].bin = bin_id
        return walkers
    
class Recycler(object):
    def __init__(self, initial_coordinates, initial_boxvectors, target_pc, retrograde=False):
        self.initial_coordinates = initial_coordinates
        self.initial_boxvectors = initial_boxvectors
        self.target_pc = target_pc
        self.retrograde = retrograde
        
    def run(self, walkers):
        flux = 0.0
        if not isinstance(walkers, list):
            walkers = [walkers]
        for i in range(len(walkers)):
            if walkers[i].pc is None:
                raise TypeError('Error - missing progress coordinate')
            if ((not self.retrograde and walkers[i].pc > self.target_pc) or
                (self.retrograde and walkers[i].pc < self.target_pc)):
                walkers[i].coordinates = self.initial_coordinates
                walkers[i].boxvectors = self.initial_boxvectors
                walkers[i].pc = None
                walkers[i].bin = None
                flux += walkers[i].weight 
        return walkers, flux

class Bin(object):
    def __init__(self, index):
        self.index = index
        self.walkers = []
        
    def add(self, walkers):
        if not isinstance(walkers, list):
            walkers = [walkers]
        self.walkers += walkers
        
    def split_merge(self, target_size):
        if len(self.walkers) == target_size:
            return
        probs = np.array([w.weight for w in self.walkers])
        old_weight = probs.sum()
        ids = np.random.choice(range(len(self.walkers)), target_size, p=probs/old_weight)
        new_walkers = []
        for i in ids:
            new_walkers.append(self.walkers[i].copy())
        new_weight = np.array([w.weight for w in new_walkers]).sum()
        fac = old_weight / new_weight
        for i in range(len(new_walkers)):
            new_walkers[i].weight *= fac
        self.walkers = list(new_walkers)
        
class SplitMerger(object):
    def __init__(self, target_size):
        self.target_size = target_size
    
    def run(self, walkers):
        if not isinstance(walkers, list):
            walkers = [walkers]
            
        bins = {}
        for w in walkers:
            if not w.bin in bins:
                bins[w.bin] = Bin(w.bin)
            bins[w.bin].add(w)
        
        for bin in bins:
            bins[bin].split_merge(self.target_size)
        
        new_walkers = []
        for bin in bins:
            new_walkers += bins[bin].walkers
            
        return new_walkers
pc = ProgressCoordinator(0, 1)
new_walkers = pc.run(new_walkers)
splitmerger = SplitMerger(n_reps)
binner = Binner([0, 0.26, 0.28, 0.3, 0.32, 0.34, 0.36, 0.38, 0.4, 0.45, 0.5, 0.55, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5])
new_walkers = binner.run(new_walkers)
recycler = Recycler(initial_pos, initial_vecs, 0.26, retrograde=True)
new_walkers, flux = recycler.run(new_walkers)
print(flux)
new_walkers = splitmerger.run(new_walkers)
print([(w.bin, w.weight) for w in new_walkers])
n_cycles=5
for i in range(n_cycles):
    new_walkers = stepper.run(new_walkers)
    new_walkers = pc.run(new_walkers)
    new_walkers = binner.run(new_walkers)
    new_walkers, flux = recycler.run(new_walkers)
    if flux > 0.0:
        new_walkers = pc.run(new_walkers)
        new_walkers = binner.run(new_walkers)
    print(i, flux, len(new_walkers), min([w.pc for w in new_walkers]))
    new_walkers = splitmerger.run(new_walkers)
for i in range(n_cycles):
    new_walkers = stepper.run(new_walkers)
    new_walkers = pc.run(new_walkers)
    new_walkers = binner.run(new_walkers)
    new_walkers, flux = recycler.run(new_walkers)
    if flux > 0.0:
        new_walkers = pc.run(new_walkers)
        new_walkers = binner.run(new_walkers)
    print(i, flux, len(new_walkers), min([w.pc for w in new_walkers]))
    new_walkers = splitmerger.run(new_walkers)
for i in range(n_cycles):
    new_walkers = stepper.run(new_walkers)
    new_walkers = pc.run(new_walkers)
    new_walkers = binner.run(new_walkers)
    new_walkers, flux = recycler.run(new_walkers)
    if flux > 0.0:
        new_walkers = pc.run(new_walkers)
        new_walkers = binner.run(new_walkers)
    print(i, flux, len(new_walkers), min([w.pc for w in new_walkers]))
    new_walkers = splitmerger.run(new_walkers)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10778518425355302, 'fraction_numerical': 0.014463589485599296, 'mean_word_length': 2.7003257328990227, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 9, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '520352', 'license': 'Apache-2.0', 'n_tokens_mistral': 2808, 'n_tokens_neox': 2690, 'n_words': 566, 'path': 'NaClWE/NaClWE.ipynb', 'repo_name': 'CharlieLaughton/colabtools', 'size': '7951'}