<jupyter_start><jupyter_text>Implements the Torus grand tour. See:* https://en.wikipedia.org/wiki/Grand_Tour_(data_visualisation)* https://epubs.siam.org/doi/abs/10.1137/0906011<jupyter_code>def torus_coordinates(k, p, scale):
    """
    alpha: R --> T^N
    returns the coordinates of a trajectory in a N-dimensional torus
    given a 'time' of k*scale*sqrt(2)
    k : frame
    p : dimension of data
    scale : controls the granularity
    """    
    N = int(2*p-3)
    primes = np.array([2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 
                       41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 
                       89, 97, 101, 103, 107, 109, 113, 127, 131, 
                       137, 139, 149, 151, 157, 163, 167, 173, 179, 
                       181, 191, 193, 197, 199, 211, 223, 227, 229, 
                       233, 239, 241, 251, 257, 263, 269, 271, 277,
                       281, 283, 293, 307, 311, 313, 317, 331, 337, 
                       347, 349, 353, 359, 367, 373, 379, 383, 389, 397])

    assert N <= len(primes), 'not implemented this many dimensions'
    lambdas = np.sqrt(primes[:N])
    step = np.sqrt(2)*scale
    coordinates = lambdas*step*k % 2*np.pi
    return coordinates

def angles(coordinates, p):
    """
    See appendix of the data. This maps the torus coordinates to rotations. 
    For N coordinates there are 2p-3 rotations. 
    Each rotation corresponds to a pair of axes of our data. 
    parameters
    ----------
    coordinates : the torus coordinates at a given frame
    p : the dimenion of our data
    returns
    -------
    angles : the rotations in radians
    ix : the dimensions of our data
    """
    angles = []
    ix = []
    for i in range(2):
        for j in range(i+1, p):
            angles.append(coordinates[j]-coordinates[i]) 
            ix.append([i, j])
            
    angles = np.array(angles)
    angles = angles % 2*np.pi
    return angles, ix

def Rij(theta, i,j, p=3):
    """
    Elements of SO(p).  Rotates i, j through angle theta.  I'm assuming this is what the paper means by: 
    "We let Rij(0) denote the element of SO(p) which rotates the standard basis vector e_i through an angle 
    0 towards the standard basis vector e_j inside the i, j coordinate 2-plane of R_p, leaving 
    fixed the orthogonal complement of this 2-plane."
    parameters
    ----------
    theta : the angle
    i, j : the indices
    p : the dimension
    returns
    -------
    R : p-dimensional rotation matrix. 
    """
    rot = np.array([[np.cos(theta), -np.sin(theta)], 
                    [np.sin(theta), np.cos(theta)]])
    R = np.eye(p)
    R[np.ix_([i, j], [i, j])] = rot
    return R
    

def rotations(k, p=3, scale=1e-2):
    """
    f : T^N --> SO(p)
    combines all the rotations
    parameters
    ----------
    k : the frame index
    p : the dimension
    fac : a scaling factor to control granularity
    returns
    -------
    Rotation matrix
    """
    R = np.eye(p)
    coords = torus_coordinates(k, p, scale=scale)
    thetas, ix = angles(coords, p)
    
    for i in range(thetas.shape[0]):
        R = Rij(thetas[i], ix[i][0],ix[i][1], p=p) @ R
    
    assert np.allclose(R @ R.T, np.eye(p)), "Rotation matrix isn't orthogonal"
    
    return R
    
def grand_tour_stop(frame, X, scale=1e-2, out_dim=2):
    """
    returns a stop on the grand tour.  
    parameters
    ----------
    frame: the 'stop' on the grand tour. 
    X : p-dimensional data
    scale : how close each 'stop' is to each other. 
    out_dim : we usually want a 2D output, but can output 3
    returns
    -------
    Xnew : a 2/3 dimensional data vector. 
    """
    assert out_dim in [2, 3], "Output dimension must be 2 or 3"
    
    p = X.shape[1]
    R = rotations(frame, p, scale=scale)
    Xnew = X @ R
    Xnew = Xnew[:, :out_dim]
    return Xnew
    
    
def tour_free_energy(X, nframes=30, nbins=100, out_file='grand_tour', out_dim=2, cols=None, scale=1e-2, fps=25, extent=2, contf_kwargs={}, cont_kwargs={}):
    from pyemma.plots.plots2d import get_histogram, _to_free_energy
    """
    Creates a grand tour gif
    """
    if out_dim !=2: 
        raise NotImplementedError("Haven't done 3d projection yet")
        
    fig, ax = plt.subplots()
    
    # This will set the extent
    xx = np.linspace(-extent, extent, nbins)
    yy = np.linspace(-extent, extent, nbins)
    
    # add the artists in a list. 
    artists = []
    for i, frame in enumerate(range(nframes)):
        if i % 10 == 0:
            print(i)
        Xnew = grand_tour_stop(frame, X, scale=scale, out_dim=out_dim)
        x, y, z = get_histogram(Xnew[:, 0], Xnew[:, 1], avoid_zero_count=True, nbins=[xx, yy])
        f = _to_free_energy(z, minener_zero=True)    
#         im1 = ax.contourf(x, y, f, **contf_kwargs)
        im = ax.contour(x, y, f, **cont_kwargs)
        artists.append(im1.collections)

    
    ani = ArtistAnimation(fig, artists)
    writergif = PillowWriter(fps=fps) 
#     writermp4 = FFMpegWriter(fps=fps)
    ani.save(f"{out_file}.mp4", writer=writergif)<jupyter_output><empty_output><jupyter_text>Cubes<jupyter_code>def ncube(dim=3, n=5):
    """
    dim : dimension
    n : number of points along each dimension
    """
    x = np.linspace(-1, 1, n)
    
    xs = np.meshgrid(*[x]*dim)
    X = np.concatenate([x.flatten()[:, None] for x in xs], axis=1)
    return X
# X = ncube(dim=3, n=10)
# create_gif(X, out_file='cube3d', cols=X[:, 0])
# X = ncube(dim=4, n=10)
# create_gif(X, out_file='cube4d', cols=X[:, 0])<jupyter_output><empty_output><jupyter_text>Pentapeptide Adapted from http://www.emma-project.org/latest/tutorials/notebooks/02-dimension-reduction-and-discretization.htmlCase-3:-another-molecular-dynamics-data-set-(pentapeptide)First get data and featurize<jupyter_code># Run this once. Also enable scrolling!
# pdb = mdshare.fetch('pentapeptide-impl-solv.pdb', working_directory='data')
# files = mdshare.fetch('pentapeptide-*-500ns-impl-solv.xtc', working_directory='data')
files = glob('data/*xtc')
pdb = glob('data/*pdb')[0]
feat = pyemma.coordinates.featurizer(pdb)
feat.add_backbone_torsions(cossin=True, periodic=False)
feat.add_sidechain_torsions(which='all', cossin=True, periodic=False)
data_torsions = pyemma.coordinates.load(files, features=feat)

feat.active_features = []
feat.add_distances_ca(periodic=False)
data_dists_ca = pyemma.coordinates.load(files, features=feat)<jupyter_output>/Users/robertarbon/opt/miniconda3/envs/msmopt/lib/python3.7/site-packages/mdtraj/geometry/dihedral.py:374: FutureWarning: arrays to stack must be passed as a "sequence" type such as list or tuple. Support for non-sequence iterables such as generators is deprecated as of NumPy 1.16 and will raise an error in the future.
  indices = np.vstack(x for x in indices if x.size)[id_sort]
/Users/robertarbon/opt/miniconda3/envs/msmopt/lib/python3.7/site-packages/pyemma/coordinates/data/featurization/angles.py:211: FutureWarning: arrays to stack must be passed as a "sequence" type such as list or tuple. Support for non-sequence iterables such as generators is deprecated as of NumPy 1.16 and will raise an error in the future.
  indices = np.vstack(valid.values())
100%|██████████| 25/25 [00:00<00:00, 39.54it/s]                     
getting output of FeatureReader:   0%|          | 0/25 [00:00<?, ?it/s]/Users/robertarbon/opt/miniconda3/envs/msmopt/lib/python3.7/site-packages/mdtraj/utils/unitcell.py:[...]<jupyter_text>Now do dimensionality reduction.  We'll stick with three dimensions at first<jupyter_code>dim = 4
lag = 20
stride = 100

vamp = pyemma.coordinates.vamp(data_torsions, lag=lag, dim=dim)
pca = pyemma.coordinates.pca(data_torsions, dim=dim)
tica = pyemma.coordinates.tica(data_torsions, lag=lag, dim=dim)

Xvamp = np.concatenate(vamp.get_output())
Xpca = np.concatenate(pca.get_output())
Xtica = np.concatenate(tica.get_output())
tour_free_energy(Xvamp, nframes=1200, out_file='vamp_wlall', extent=5, scale=1e-3, fps=20, 
                 cont_kwargs={'cmap': 'magma', 'levels': 10, 'lw': 2})<jupyter_output><empty_output><jupyter_text>Alanine dipeptide<jupyter_code>def get_colors(X, lag=1, k=100):
    clus = pyemma.coordinates.cluster_kmeans(X, k=k, max_iter=100)
    dtrajs = clus.dtrajs
    mm = pyemma.msm.estimate_markov_model(dtrajs, lag=lag)
    colors = mm.pi[dtrajs]
    return colors
trajs = [np.load('../../ala1/data/features/positions/{:03d}.npy'.format(i)) for i in range(750)]

lag = 5
dim = 4
stride = 100

tica = pyemma.coordinates.tica(trajs, lag=lag, dim=dim)
vamp = pyemma.coordinates.vamp(trajs, lag=lag, dim=dim)
pca  = pyemma.coordinates.pca(trajs, dim=dim)

Xvamp = np.concatenate(vamp.get_output())[::stride]
Xpca = np.concatenate(pca.get_output())[::stride]
Xtica = np.concatenate(tica.get_output())[::stride]

Vcols = get_colors(Xvamp)
Pcols = get_colors(Xpca) 
Tcols = get_colors(Xtica)
create_gif(Xtica, num_stops=1000, out_file='tica_ala', out_dim=2, cols=Tcols, scale=1e-2, fps=10, extent=3)
create_gif(Xvamp, num_stops=500, out_file='vamp_ala', out_dim=2, cols=Tcols, scale=1e-2, fps=10, extent=3)
create_gif(Xpca, num_stops=1000, out_file='pca_ala', out_dim=2, cols=Tcols, scale=1e-2, fps=10, extent=0.5)
pyemma.plots.plot_free_energy(Xvamp[:, 0],Xvamp[:, 1])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14075993091537134, 'fraction_numerical': 0.04447322970639033, 'mean_word_length': 3.114120781527531, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 22, 'https://': 2, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '122442', 'license': 'MIT', 'n_tokens_mistral': 3789, 'n_tokens_neox': 3366, 'n_words': 847, 'path': 'grand_tour/grand_tour.ipynb', 'repo_name': 'RobertArbon/msm_hyperparameter_optimisation', 'size': '9264'}