get rid of the dependency of networkx in the eg

examples/openmm/cv_Q/abf.py

@@ -1,8 +1,5 @@
 #!/usr/bin/env python
 
-from itertools import product
-
-import networkx as nx
 import numpy as np
 import openmm as mm
 import openmm.app as app
@@ -17,64 +14,33 @@
 AUGC = ["A", "U", "G", "C"]
 
 
-def init_graph(top):
-    graph = nx.Graph()
-    rna_indices = []
-    full_id2rna_heavy_id = {}
-    index = 0
-    for atom in top.atoms():
-        if atom.residue.name in AUGC and atom.element.name != "hydrogen":
-            graph.add_node(index, name=atom.name, resid=atom.residue.index)
-            full_id2rna_heavy_id[atom.index] = index
-            rna_indices.append(atom.index)
-            index += 1
-
-    for bond in top.bonds():
-        if (
-            bond.atom1.residue.name in AUGC
-            and bond.atom1.element.name != "hydrogen"
-            and bond.atom2.residue.name in AUGC
-            and bond.atom2.element.name != "hydrogen"
-        ):
-            graph.add_edge(
-                full_id2rna_heavy_id[bond.atom1.index], full_id2rna_heavy_id[bond.atom2.index]
-            )
-
-    return graph, rna_indices
-
-
-def gen_dihedral_indices(graph):
-    dihedral_indices = []
-    for n1, n2 in nx.edge_dfs(graph):
-        neibor1 = list(graph.adj[n1])
-        neibor1.remove(n2)
-        # we want to exclude neightbor n2
-        # because we are trying to construct a dihedral pre-n1-n2-post,
-        # and we don't want pre is the same as n2
-        neibor2 = list(graph.adj[n2])
-        neibor2.remove(n1)
-        for pre, post in product(neibor1, neibor2):
-            if pre != post:  # exlucde ring dihedrals
-                dihedral_indices.append([pre, n1, n2, post])
-    return dihedral_indices
-
-
-def calc_include_ij(graph, dihedral_indices):
+def create_exclusions_from_bonds(particles, bonds, bond_cutoff=3):
     """
-    calculate all the atom pairs that are connected by 1 or 2 or 3 bonds.
-    Those pairs can't be included in the native contacts
-    return a matrix for all the atoms in the graph:
-        True means we can include the pair
-        False means we should exclude the pair
+    create exclusion from bond.
     """
-    n_atoms = len(graph)
-    exclude_ij = np.full((n_atoms, n_atoms), False)
-    for dihedral in dihedral_indices:
-        indices = np.array(np.meshgrid(dihedral, dihedral)).T.reshape(-1, 2)
-        i, j = indices[:, 0], indices[:, 1]
-        exclude_ij[i, j] = True
-
-    return ~exclude_ij
+    n_particles = max(particles) + 1
+    exclusions = [set() for _ in range(n_particles)]
+    bonded12 = [set() for _ in range(n_particles)]
+    for bond in bonds:
+        p1, p2 = bond
+        exclusions[p1].add(p2)
+        exclusions[p2].add(p1)
+        bonded12[p1].add(p2)
+        bonded12[p2].add(p1)
+
+    for level in range(bond_cutoff - 1):
+        current_exclusions = [exclusion.copy() for exclusion in exclusions]
+        for i in range(n_particles):
+            for j in current_exclusions[i]:
+                exclusions[j].update(bonded12[i])
+
+    final_exclusions = []
+    for i in range(len(exclusions)):
+        for j in exclusions[i]:
+            if j < i:
+                final_exclusions.append((j, i))
+
+    return final_exclusions
 
 
 step_size = 2 * unit.femtosecond
@@ -85,15 +51,34 @@ def calc_include_ij(graph, dihedral_indices):
 pdb = app.PDBFile("../../inputs/GAGA.box_0mM.pdb")
 positions = pdb.getPositions(asNumpy=True).value_in_unit(unit.nanometer)
 
-rna_graph, rna_indices = init_graph(pdb.topology)
-dihedral_indices = gen_dihedral_indices(rna_graph)
-include_ij = calc_include_ij(rna_graph, dihedral_indices)
+rna_indices = []
+for atom in pdb.topology.atoms():
+    if atom.residue.name in AUGC and atom.element.name != "hydrogen":
+        rna_indices.append(atom.index)
+
+rna_bonds = []
+for bond in pdb.topology.bonds():
+    if (
+        bond.atom1.residue.name in AUGC
+        and bond.atom1.element.name != "hydrogen"
+        and bond.atom2.residue.name in AUGC
+        and bond.atom2.element.name != "hydrogen"
+    ):
+        rna_bonds.append([bond.atom1.index, bond.atom2.index])
+
+exclusions = create_exclusions_from_bonds(rna_indices, rna_bonds)
 
 rna_pos = positions.astype("float")[np.asarray(rna_indices)]
 contact_matrix = sd.squareform(sd.pdist(rna_pos)) < contact_cutoff
 contacts = np.transpose(np.nonzero(contact_matrix))
 rna_id_contacts = np.array(
-    [[rna_indices[i], rna_indices[j]] for i, j in contacts if i != j and include_ij[i, j]]
+    [
+        [rna_indices[i], rna_indices[j]]
+        for i, j in contacts
+        if i != j
+        and (rna_indices[i], rna_indices[j]) not in exclusions
+        and (rna_indices[j], rna_indices[i]) not in exclusions
+    ]
 )
 # notice that we need to get rid of self-self contact!
 indices = np.unique(rna_id_contacts)

examples/openmm/cv_Q/log

@@ -1,101 +1,101 @@
 #"Step","Time (ps)","Speed (ns/day)","Elapsed Time (s)","Time Remaining"
-1,0.002,0,0.00018405914306640625,--
-2,0.004,0.14,1.2383630275726318,2:01
-3,0.006,0.141,2.4549551010131836,1:59
-4,0.008,0.204,2.5401570796966553,1:21
-5,0.01,0.258,2.6760189533233643,1:03
-6,0.012,0.313,2.764267683029175,0:51
-7,0.014,0.364,2.850210666656494,0:44
-8,0.016,0.412,2.936511993408203,0:38
-9,0.018000000000000002,0.458,3.021488666534424,0:34
-10,0.020000000000000004,0.501,3.1063854694366455,0:31
-11,0.022000000000000006,0.537,3.2184159755706787,0:28
-12,0.024000000000000007,0.575,3.3063764572143555,0:26
-13,0.02600000000000001,0.611,3.391796112060547,0:24
-14,0.02800000000000001,0.637,3.528839349746704,0:23
-15,0.030000000000000013,0.669,3.618229866027832,0:21
-16,0.032000000000000015,0.7,3.70469069480896,0:20
-17,0.034000000000000016,0.73,3.7896921634674072,0:19
-18,0.03600000000000002,0.758,3.8752076625823975,0:18
-19,0.03800000000000002,0.785,3.959963798522949,0:17
-20,0.04000000000000002,0.812,4.044415712356567,0:17
-21,0.04200000000000002,0.837,4.129042148590088,0:16
-22,0.044000000000000025,0.861,4.213843584060669,0:15
-23,0.04600000000000003,0.884,4.299461603164673,0:15
-24,0.04800000000000003,0.896,4.43713116645813,0:14
-25,0.05000000000000003,0.916,4.525743722915649,0:14
-26,0.05200000000000003,0.937,4.611554384231567,0:13
-27,0.054000000000000034,0.957,4.696833610534668,0:13
-28,0.056000000000000036,0.976,4.782029390335083,0:12
-29,0.05800000000000004,0.994,4.866597890853882,0:12
-30,0.06000000000000004,1.01,4.952362775802612,0:11
-31,0.06200000000000004,1.03,5.038139581680298,0:11
-32,0.06400000000000004,1.05,5.1235737800598145,0:11
-33,0.06600000000000004,1.06,5.209028005599976,0:10
-34,0.06800000000000005,1.08,5.294060468673706,0:10
-35,0.07000000000000005,1.09,5.378709316253662,0:10
-36,0.07200000000000005,1.1,5.5171427726745605,0:10
-37,0.07400000000000005,1.11,5.607499361038208,0:09
-38,0.07600000000000005,1.12,5.6947526931762695,0:09
-39,0.07800000000000006,1.14,5.780440092086792,0:09
-40,0.08000000000000006,1.15,5.865738153457642,0:09
-41,0.08200000000000006,1.16,5.9505109786987305,0:08
-42,0.08400000000000006,1.17,6.036832332611084,0:08
-43,0.08600000000000006,1.19,6.122369050979614,0:08
-44,0.08800000000000006,1.2,6.208016395568848,0:08
-45,0.09000000000000007,1.2,6.345700979232788,0:07
-46,0.09200000000000007,1.21,6.436450719833374,0:07
-47,0.09400000000000007,1.22,6.523990631103516,0:07
-48,0.09600000000000007,1.23,6.610021352767944,0:07
-49,0.09800000000000007,1.24,6.696079730987549,0:07
-50,0.10000000000000007,1.25,6.78236722946167,0:06
-51,0.10200000000000008,1.26,6.86974573135376,0:06
-52,0.10400000000000008,1.27,6.956451654434204,0:06
-53,0.10600000000000008,1.28,7.042772054672241,0:06
-54,0.10800000000000008,1.28,7.128982305526733,0:06
-55,0.11000000000000008,1.29,7.215255260467529,0:06
-56,0.11200000000000009,1.29,7.354937314987183,0:05
-57,0.11400000000000009,1.3,7.445296287536621,0:05
-58,0.11600000000000009,1.31,7.532633543014526,0:05
-59,0.11800000000000009,1.32,7.619118928909302,0:05
-60,0.12000000000000009,1.32,7.706413745880127,0:05
-61,0.1220000000000001,1.33,7.792268514633179,0:05
-62,0.1240000000000001,1.34,7.878641843795776,0:04
-63,0.12600000000000008,1.35,7.964428663253784,0:04
-64,0.12800000000000009,1.35,8.050028324127197,0:04
-65,0.1300000000000001,1.36,8.135518312454224,0:04
-66,0.1320000000000001,1.37,8.221822261810303,0:04
-67,0.1340000000000001,1.37,8.30713152885437,0:04
-68,0.1360000000000001,1.37,8.44601321220398,0:04
-69,0.1380000000000001,1.38,8.535516738891602,0:03
-70,0.1400000000000001,1.38,8.623224973678589,0:03
-71,0.1420000000000001,1.39,8.710314750671387,0:03
-72,0.1440000000000001,1.39,8.796385526657104,0:03
-73,0.1460000000000001,1.4,8.882198572158813,0:03
-74,0.1480000000000001,1.41,8.968223094940186,0:03
-75,0.1500000000000001,1.41,9.054269313812256,0:03
-76,0.1520000000000001,1.42,9.139775276184082,0:02
-77,0.1540000000000001,1.42,9.279144048690796,0:02
-78,0.1560000000000001,1.42,9.368648052215576,0:02
-79,0.1580000000000001,1.43,9.45512843132019,0:02
-80,0.16000000000000011,1.43,9.54301142692566,0:02
-81,0.16200000000000012,1.44,9.629714012145996,0:02
-82,0.16400000000000012,1.44,9.717389106750488,0:02
-83,0.16600000000000012,1.45,9.803569078445435,0:02
-84,0.16800000000000012,1.45,9.88936734199524,0:01
-85,0.17000000000000012,1.46,9.975153923034668,0:01
-86,0.17200000000000013,1.46,10.061434745788574,0:01
-87,0.17400000000000013,1.46,10.202369928359985,0:01
-88,0.17600000000000013,1.46,10.292763471603394,0:01
-89,0.17800000000000013,1.47,10.379234552383423,0:01
-90,0.18000000000000013,1.47,10.464409589767456,0:01
-91,0.18200000000000013,1.47,10.549905061721802,0:01
-92,0.18400000000000014,1.48,10.635522603988647,0:00
-93,0.18600000000000014,1.48,10.721351146697998,0:00
-94,0.18800000000000014,1.49,10.807268857955933,0:00
-95,0.19000000000000014,1.49,10.89450216293335,0:00
-96,0.19200000000000014,1.5,10.980336904525757,0:00
-97,0.19400000000000014,1.5,11.065917491912842,0:00
-98,0.19600000000000015,1.5,11.152062177658081,0:00
-99,0.19800000000000015,1.5,11.291614294052124,0:00
-100,0.20000000000000015,1.5,11.381214141845703,0:00
+1,0.002,0,0.00022101402282714844,--
+2,0.004,0.21,0.8227248191833496,1:20
+3,0.006,0.24,1.4425551891326904,1:09
+4,0.008,0.339,1.5272512435913086,0:48
+5,0.01,0.429,1.610468864440918,0:38
+6,0.012,0.51,1.693587064743042,0:31
+7,0.014,0.563,1.840118408203125,0:28
+8,0.016,0.628,1.9254417419433594,0:25
+9,0.018000000000000002,0.689,2.0073280334472656,0:22
+10,0.020000000000000004,0.744,2.0890250205993652,0:20
+11,0.022000000000000006,0.794,2.175332546234131,0:19
+12,0.024000000000000007,0.842,2.2585818767547607,0:18
+13,0.02600000000000001,0.885,2.3418662548065186,0:16
+14,0.02800000000000001,0.927,2.423513174057007,0:16
+15,0.030000000000000013,0.966,2.505436897277832,0:15
+16,0.032000000000000015,1,2.5873866081237793,0:14
+17,0.034000000000000016,1.04,2.6675479412078857,0:13
+18,0.03600000000000002,1.05,2.8018243312835693,0:13
+19,0.03800000000000002,1.08,2.8864688873291016,0:12
+20,0.04000000000000002,1.11,2.969536781311035,0:12
+21,0.04200000000000002,1.13,3.052172899246216,0:12
+22,0.044000000000000025,1.16,3.1340787410736084,0:11
+23,0.04600000000000003,1.18,3.2186665534973145,0:11
+24,0.04800000000000003,1.2,3.300657272338867,0:10
+25,0.05000000000000003,1.23,3.384861946105957,0:10
+26,0.05200000000000003,1.24,3.474282741546631,0:10
+27,0.054000000000000034,1.26,3.560129165649414,0:09
+28,0.056000000000000036,1.26,3.702094316482544,0:09
+29,0.05800000000000004,1.28,3.7889840602874756,0:09
+30,0.06000000000000004,1.29,3.876796007156372,0:09
+31,0.06200000000000004,1.31,3.9615395069122314,0:09
+32,0.06400000000000004,1.32,4.043964147567749,0:08
+33,0.06600000000000004,1.34,4.125448942184448,0:08
+34,0.06800000000000005,1.36,4.205911636352539,0:08
+35,0.07000000000000005,1.37,4.286661624908447,0:08
+36,0.07200000000000005,1.39,4.366098642349243,0:07
+37,0.07400000000000005,1.4,4.445927143096924,0:07
+38,0.07600000000000005,1.39,4.59551739692688,0:07
+39,0.07800000000000006,1.4,4.680671215057373,0:07
+40,0.08000000000000006,1.42,4.761897802352905,0:07
+41,0.08200000000000006,1.43,4.843828916549683,0:07
+42,0.08400000000000006,1.44,4.926450490951538,0:06
+43,0.08600000000000006,1.45,5.006367206573486,0:06
+44,0.08800000000000006,1.46,5.086201429367065,0:06
+45,0.09000000000000007,1.47,5.16817569732666,0:06
+46,0.09200000000000007,1.48,5.248043775558472,0:06
+47,0.09400000000000007,1.49,5.330909729003906,0:06
+48,0.09600000000000007,1.49,5.458130359649658,0:06
+49,0.09800000000000007,1.5,5.542110443115234,0:05
+50,0.10000000000000007,1.5,5.629708766937256,0:05
+51,0.10200000000000008,1.51,5.716973543167114,0:05
+52,0.10400000000000008,1.52,5.802804946899414,0:05
+53,0.10600000000000008,1.53,5.888056039810181,0:05
+54,0.10800000000000008,1.53,5.9713804721832275,0:05
+55,0.11000000000000008,1.54,6.0550971031188965,0:05
+56,0.11200000000000009,1.55,6.137209892272949,0:04
+57,0.11400000000000009,1.56,6.219832181930542,0:04
+58,0.11600000000000009,1.55,6.354286432266235,0:04
+59,0.11800000000000009,1.56,6.440207004547119,0:04
+60,0.12000000000000009,1.56,6.524070739746094,0:04
+61,0.1220000000000001,1.57,6.608417272567749,0:04
+62,0.1240000000000001,1.58,6.691519737243652,0:04
+63,0.12600000000000008,1.58,6.7744140625,0:04
+64,0.12800000000000009,1.59,6.857131004333496,0:03
+65,0.1300000000000001,1.59,6.941458702087402,0:03
+66,0.1320000000000001,1.6,7.025902986526489,0:03
+67,0.1340000000000001,1.6,7.111144304275513,0:03
+68,0.1360000000000001,1.61,7.19291877746582,0:03
+69,0.1380000000000001,1.61,7.321002721786499,0:03
+70,0.1400000000000001,1.61,7.408953428268433,0:03
+71,0.1420000000000001,1.61,7.494205713272095,0:03
+72,0.1440000000000001,1.62,7.577526569366455,0:02
+73,0.1460000000000001,1.62,7.660561800003052,0:02
+74,0.1480000000000001,1.63,7.744027853012085,0:02
+75,0.1500000000000001,1.63,7.830069541931152,0:02
+76,0.1520000000000001,1.64,7.915714740753174,0:02
+77,0.1540000000000001,1.64,7.99590539932251,0:02
+78,0.1560000000000001,1.63,8.14884090423584,0:02
+79,0.1580000000000001,1.64,8.235782146453857,0:02
+80,0.16000000000000011,1.64,8.321000337600708,0:02
+81,0.16200000000000012,1.64,8.404356241226196,0:01
+82,0.16400000000000012,1.65,8.487345218658447,0:01
+83,0.16600000000000012,1.65,8.571403980255127,0:01
+84,0.16800000000000012,1.66,8.65239667892456,0:01
+85,0.17000000000000012,1.66,8.736083507537842,0:01
+86,0.17200000000000013,1.67,8.817293643951416,0:01
+87,0.17400000000000013,1.67,8.90245246887207,0:01
+88,0.17600000000000013,1.66,9.061772108078003,0:01
+89,0.17800000000000013,1.66,9.149581909179688,0:01
+90,0.18000000000000013,1.67,9.23227310180664,0:01
+91,0.18200000000000013,1.67,9.316927671432495,0:00
+92,0.18400000000000014,1.67,9.403042554855347,0:00
+93,0.18600000000000014,1.68,9.48436975479126,0:00
+94,0.18800000000000014,1.68,9.56793761253357,0:00
+95,0.19000000000000014,1.68,9.651633739471436,0:00
+96,0.19200000000000014,1.69,9.734861850738525,0:00
+97,0.19400000000000014,1.69,9.81718397140503,0:00
+98,0.19600000000000015,1.69,9.898550987243652,0:00
+99,0.19800000000000015,1.7,9.981334209442139,0:00
+100,0.20000000000000015,1.69,10.123533964157104,0:00