From faf5a92816b46bc021897deae2e54c3870805251 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Fri, 15 Nov 2024 14:28:21 -0700 Subject: [PATCH 01/43] Updated functions and formatting --- fake_data.py | 261 ++++++++++++++++++++++++++------------------------- 1 file changed, 131 insertions(+), 130 deletions(-) diff --git a/fake_data.py b/fake_data.py index 97fc841..a9223b8 100644 --- a/fake_data.py +++ b/fake_data.py @@ -1,161 +1,162 @@ import numpy as np -from itertools import product +from itertools import product, combinations from typing import NamedTuple import random import argparse -parser = argparse.ArgumentParser() -parser.add_argument('num_resid', type=int, help='Number of residues') -args = parser.parse_args() - class HSQCPeak(NamedTuple): - H1: float - N15: float + H1: float + N15: float class NOEPeak(NamedTuple): - H1: float - N15: float - H2: float + H1: float + N15: float + H2: float class Protein(NamedTuple): - x: float - y: float - z: float + x: float + y: float + z: float # Sampled points from unit square/cube def sample_unit(n, num_sides, min=0, max=1): - return np.random.uniform(min, max, size=(n, num_sides)) + return np.random.uniform(min, max, size=(n, num_sides)) # Noise to sampled points -def add_noise(point, scale=0.1): - """ - Random selection from normal (gaussian) distribution of 'scale' width from 0 (center). - 'size' makes sure that it's the same shape as the point we are adding noises to. - ex. point = [0,1,2], noise = [1,1,1], noisy_point = [1,2,3] - """ - noise = np.random.normal(0, scale, size=point.shape) - noisy_point = point + noise - return np.clip(noisy_point, 0, 1) +def add_noise(point, scale=0.1, min=0, max=1): + """ + Random selection from normal (gaussian) distribution of 'scale' width from 0 (center). + 'size' makes sure that it's the same shape as the point we are adding noises to. + ex. point = [0,1,2], noise = [1,1,1], noisy_point = [1,2,3] + """ + noise = np.random.normal(0, scale, size=point.shape) + noisy_point = point + noise + + # fold over points outside of boundaries + for point in noisy_point: + for value in point: + if max < value: + value = (max-(value-max)) + elif value < min: + value = (min-value) + + return noisy_point def calc_dist(p1, p2): - return np.sqrt((p2[0] - p1[0])**2 + (p2[1] - p1[1])**2 + (p2[2] - p1[2])**2) + return np.linalg.norm((p2-p1)) def distance_noe(protein, shifts, cutoff): - """ - Grabs close coordinates and 'associated' HSQC shift peaks (by same index) to create NOES. - Adds gaussian noise to all points. + """ + Grabs close coordinates and 'associated' HSQC shift peaks (by same index) to create NOES. + Adds gaussian noise to all points. - **Can end up with no NOEs depending on the cutoff** - """ - noe = [] + **Can end up with no NOEs depending on the cutoff** + """ + noes = [] - for i, atom1 in enumerate(protein): - for j, atom2 in enumerate(protein): - if i != j: - dist = calc_dist(atom1, atom2) - if dist < cutoff: - # print(dist, shifts[i], shifts[j]) - val = len(noe) # Append both shifts to correct noe index - noe.append(list(shifts[i][:])) # H1, N1 - noe[val].append(shifts[j][0]) # H2 + for i, atom1 in enumerate(protein): + for j, atom2 in enumerate(protein): + if i != j: + dist = calc_dist(atom1, atom2) + if dist < cutoff: + # print(dist, shifts[i], shifts[j]) + noe = list(shifts[i][:]) # H1, N1 + noe.append(shifts[j][0]) # H2 + noes.append(noe) - noisy_noe = add_noise(np.array(noe), scale=0.01) + noisy_noe = add_noise(np.array(noes), scale=0.01) - return noisy_noe + return noisy_noe def generate_data(num_resid): - """ - Generates all fake data and orders it in lists of namedtuples. - """ - - # 3D structure [x,y,z] - protein = sample_unit(num_resid,3) - - # "Actual" shifts [H1,N1] - actual_shifts = sample_unit(num_resid,2) - - # Predicted shifts [H1,N1] - predicted_shifts = add_noise(actual_shifts) - - # NOES [H1,N1,H2] - noe = distance_noe(protein, actual_shifts, 0.5) # likely need to change dist cutoff and only accounting for actual_shifts right now - - # Probably not the best way to do this - # Lists of namedtuples (one object per residue) - coords = [Protein(x=resid[0], y=resid[1], z=resid[2]) for resid in protein] - actual_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in actual_shifts] - predicted_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in predicted_shifts] - noe = [NOEPeak(H1=shift[0], N15=shift[1], H2=shift[2]) for shift in noe] - - # Fake_data = namedtuple('FakeData', ['protein', 'actual_shifts', 'predicted_shifts', 'noe']) - # fake_data = Fake_data(protein=protein, actual_shifts=actual_shifts, predicted_shifts=predicted_shifts, noe=noe) - - # print(coords) - # print(actual_shifts) - # print(noe) - - return coords, actual_shifts, predicted_shifts, noe + """ + Generates all fake data and orders it in lists of namedtuples. + """ + # 3D structure [x,y,z] + protein = sample_unit(num_resid, num_sides=3) + # "Actual" shifts [H1,N1] + actual_shifts = sample_unit(num_resid, num_sides=2) + # Predicted shifts [H1,N1] + predicted_shifts = add_noise(actual_shifts) + # NOES [H1,N1,H2] + noes = distance_noe(protein, actual_shifts, cutoff=0.5) # likely need to change dist cutoff and only accounting for actual_shifts right now + + # Lists of namedtuples (one object per residue) + coords = [Protein(x=resid[0], y=resid[1], z=resid[2]) for resid in protein] + actual_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in actual_shifts] + predicted_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in predicted_shifts] + noes = [NOEPeak(H1=shift[0], N15=shift[1], H2=shift[2]) for shift in noes] + + # print(coords) + # print(actual_shifts) + # print(noe) + + return coords, actual_shifts, predicted_shifts, noes # match H1 shifts def matchH(hshift, hsqc, tolerance_h): - """ - Loops through the HSQC list and matches the second proton shift of the NOE to a HSQC shift based on H cutoff. - Returns the index of said shift. - """ - h_idx = [] - for i, shift in enumerate(hsqc): - if abs(hshift - shift.H1) <= tolerance_h: - h_idx.append(i) - return h_idx + """ + Loops through the HSQC list and matches the second proton shift of the NOE to a HSQC shift based on H cutoff. + Returns the index of said shift. + """ + h_index = [] + for i, shift in enumerate(hsqc): + if abs(hshift - shift.H1) <= tolerance_h: + h_index.append(i) + + return h_index # match H1 and N15 shifts def matchNH(hshift, nshift, hsqc, tolerance_h, tolerance_n): - """ - Loops through the HSQC list and matches the first proton and nitrogen shift of the NOE to a HSQC shift based on H and N cutoffs. - Returns the index of said shift. - """ - nh_idx = [] - for i, shift in enumerate(hsqc): - if abs(hshift - shift.H1) <= tolerance_h and abs(nshift - shift.N15) <= tolerance_n: - nh_idx.append(i) - return nh_idx - -def noe_combinations(noe, actual_shifts, tolerance_h=0.2, tolerance_n=0.2): - """ - Loops through the NOES to match the H1, N1 and H2 shifts to two HSQC shifts. - Lists the HSQC shifts separately (H1/N1 as contact1 and H2 as contact2) for each NOE. - Returns a lists of tuples of all possible shift index combinations for each NOE. - ex. [[(1,2),(1,3)],[(3,0)]] - 1 and 2 are a pair of potential shift indices for one NOE, 1 and 3 are another for the same NOE. - - Does not consider: - - Same shift assignments (nh_index and h_index could be the same for one NOE). - Filter these out initially, but itertools product can still combine one shift as a possible pair. - - - Shift pair repetitions (two shifts are included in both the nh_index and h_index). - Gives a double count ex. contact1/2 = [[3,0]], product = [[(0,3),(3,0)]] - """ - contact1 = [] - contact2 = [] - - for i, peak in enumerate(noe): - # matching H1 and N15 in NOE to shifts in HSQC - nh_idx = matchNH(peak.H1, peak.N15, actual_shifts, tolerance_h, tolerance_n) - # matching H2 in NOE to H1 shifts in HSQC - h_idx = matchH(peak.H2, actual_shifts, tolerance_h) - - if nh_idx != h_idx: # Filter out same shift assignments for one NOE - contact1.append(nh_idx) - contact2.append(h_idx) - - contacts = [[] for i in range(len(contact1))] # List to match number of NOES - for i in range(len(contact1)): - contacts[i] += list(product(contact1[i], contact2[i])) - - return contacts - -num_resid = args.num_resid -# These should be grouped to export into an environment -coords, actual_shifts, predicted_shifts, noe = generate_data(num_resid) -combinations = noe_combinations(noe, actual_shifts) -print(combinations) + """ + Loops through the HSQC list and matches the first proton and nitrogen shift of the NOE to a HSQC shift based on H and N cutoffs. + Returns the index of said shift. + """ + nh_index = [] + for i, shift in enumerate(hsqc): + if abs(hshift - shift.H1) <= tolerance_h and abs(nshift - shift.N15) <= tolerance_n: + nh_index.append(i) + + return nh_index + +def noe_combinations(noes, actual_shifts, tolerance_h=0.2, tolerance_n=0.2): + """ + Loops through the NOES to match the H1, N1 and H2 shifts to two HSQC shifts. + Lists the HSQC shifts separately (H1/N1 as contact1 and H2 as contact2) for each NOE. + Returns a lists of tuples of all possible shift index combinations for each NOE. + ex. [[(1,2),(1,3)],[(3,0)]] + 1 and 2 are a pair of potential shift indices for one NOE, 1 and 3 are another for the same NOE. + """ + contact1 = [] + contact2 = [] + for i, peak in enumerate(noes): + # matching H1 and N15 in NOE to shifts in HSQC + nh_index = matchNH(peak.H1, peak.N15, actual_shifts, tolerance_h, tolerance_n) + # matching H2 in NOE to H1 shifts in HSQC + h_index = matchH(peak.H2, actual_shifts, tolerance_h) + + contact1.append(nh_index) + contact2.append(h_index) + + # print(contact1) + # print(contact2) + + contacts = [] + for i, contact in enumerate(contact1): + prod = product(contact, contact2[i]) + # filter out same shift pairs and inverted duplicates (ex. (0,0) and [(0,1),(1,0)]) + pairs = set(tuple(sorted(combo)) for combo in prod if combo[0] != combo[1]) + contacts.append(list(pairs)) + + return contacts + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('num_resid', type=int, help='Number of residues') + args = parser.parse_args() + + num_resid = args.num_resid + # These should be grouped to export into an environment + coords, actual_shifts, predicted_shifts, noes = generate_data(num_resid) + combinations = noe_combinations(noes, actual_shifts) + print(combinations) From 10dff002d6776d40d3c2cfc9870cd88990704f03 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Thu, 21 Nov 2024 19:15:04 -0700 Subject: [PATCH 02/43] Split energy and data functions --- energy.py | 122 +++++++++++++++++++++++++++++++++++++++++++++++++++ fake_data.py | 68 ++++------------------------ 2 files changed, 130 insertions(+), 60 deletions(-) create mode 100644 energy.py diff --git a/energy.py b/energy.py new file mode 100644 index 0000000..6400674 --- /dev/null +++ b/energy.py @@ -0,0 +1,122 @@ +import numpy as np +from itertools import product +import math + +# match H1 shifts +def matchH(hshift, hsqc, tolerance_h): + """ + Loops through the HSQC list and matches the second proton shift of the NOE to a HSQC shift based on H cutoff. + Returns the index of said shift. + """ + h_index = [] + for i, shift in enumerate(hsqc): + if abs(hshift - shift.H1) <= tolerance_h: + h_index.append(i) + + return h_index + +# match H1 and N15 shifts +def matchNH(hshift, nshift, hsqc, tolerance_h, tolerance_n): + """ + Loops through the HSQC list and matches the first proton and nitrogen shift of the NOE to a HSQC shift based on H and N cutoffs. + Returns the index of said shift. + """ + nh_index = [] + for i, shift in enumerate(hsqc): + if abs(hshift - shift.H1) <= tolerance_h and abs(nshift - shift.N15) <= tolerance_n: + nh_index.append(i) + + return nh_index + +def noe_combinations(noes, actual_shifts, tolerance_h=0.2, tolerance_n=0.2): + """ + Loops through the NOES to match the H1, N1 and H2 shifts to two HSQC shifts. + Lists the HSQC shifts separately (H1/N1 as contact1 and H2 as contact2) for each NOE. + Returns a lists of tuples of all possible shift index combinations for each NOE. + ex. [[(1,2),(1,3)],[(3,0)]] + 1 and 2 are a pair of potential shift indices for one NOE, 1 and 3 are another for the same NOE. + """ + contact1 = [] + contact2 = [] + for i, peak in enumerate(noes): + # matching H1 and N15 in NOE to shifts in HSQC + nh_index = matchNH(peak.H1, peak.N15, actual_shifts, tolerance_h, tolerance_n) + # matching H2 in NOE to H1 shifts in HSQC + h_index = matchH(peak.H2, actual_shifts, tolerance_h) + + contact1.append(nh_index) + contact2.append(h_index) + + # print(contact1) + # print(contact2) + + contacts = [] + for i, contact in enumerate(contact1): + prod = product(contact, contact2[i]) + # filter out same shift pairs and inverted duplicates (ex. (0,0) and [(0,1),(1,0)]) + pairs = set(tuple(sorted(combo)) for combo in prod if combo[0] != combo[1]) + contacts.append(list(pairs)) + + return contacts + +class Energy(): + + def __init__(self, coords, actual_shifts, noes): + self.coords = coords + self.actual_shifts = actual_shifts + self.noes = noes + + def setup_noe_restraints(self): + combos = noe_combinations(self.noes, self.actual_shifts) + return combos + + def flat_bottom(self, x, tolerance): + """ + Defines the flat bottom equation with a flat region from zero to the distance tolerance. + """ + assert x > 0 + + # ends of flat bottom region + xmax = tolerance + xmin = 0 + + # if x <= tolerance return energy of zero + # if x > tolerance return energy relating to quadratic + y = np.piecewise(x, + [(xmin <= x) & (x <= xmax), x > xmax], + [lambda x: 0, + lambda x: x**2 - xmax*x]) + return y + + def energy_loop(self, restraints, assignments): + """ + Loops over the NOE restraints and then each possible shift option. + If these shift options have been assigned, the NOE is activated, if not the NOE restraint should be passed. + Activation means the coordinate indices are grabbed from the assignment dictionary and used in the flat_bottom function. + Energy is then summed across all activated NOEs. + """ + + total_energy = 0 + # assignments = {0: 0, 1 : 1, 2 : 2, 3 : 3} # this would be added to parameters as external dict + for restraint in restraints: + restraint_energy = math.inf + activated = False + for i, j in restraint: + if i in assignments.values() and j in assignments.values(): + activated = True + + k, l = assignments.get(i), assignments.get(j) + # distance to energy evaluation + dist = np.linalg.norm((np.array(self.coords[k]) - np.array(self.coords[l]))) + energy_value = self.flat_bottom(dist, tolerance=0.5) + # keep the smallest energy (distance) + restraint_energy = energy_value if energy_value < restraint_energy else restraint_energy + # add energy to the summation of the restraints + # problem with this adding the restraint energy each time in inner loop + total_energy += restraint_energy + + else: + # NOE not activated + break + + return total_energy diff --git a/fake_data.py b/fake_data.py index a9223b8..bc2ad93 100644 --- a/fake_data.py +++ b/fake_data.py @@ -1,9 +1,11 @@ import numpy as np -from itertools import product, combinations +from itertools import product from typing import NamedTuple import random import argparse +from energy import Energy + class HSQCPeak(NamedTuple): H1: float N15: float @@ -93,63 +95,6 @@ def generate_data(num_resid): return coords, actual_shifts, predicted_shifts, noes -# match H1 shifts -def matchH(hshift, hsqc, tolerance_h): - """ - Loops through the HSQC list and matches the second proton shift of the NOE to a HSQC shift based on H cutoff. - Returns the index of said shift. - """ - h_index = [] - for i, shift in enumerate(hsqc): - if abs(hshift - shift.H1) <= tolerance_h: - h_index.append(i) - - return h_index - -# match H1 and N15 shifts -def matchNH(hshift, nshift, hsqc, tolerance_h, tolerance_n): - """ - Loops through the HSQC list and matches the first proton and nitrogen shift of the NOE to a HSQC shift based on H and N cutoffs. - Returns the index of said shift. - """ - nh_index = [] - for i, shift in enumerate(hsqc): - if abs(hshift - shift.H1) <= tolerance_h and abs(nshift - shift.N15) <= tolerance_n: - nh_index.append(i) - - return nh_index - -def noe_combinations(noes, actual_shifts, tolerance_h=0.2, tolerance_n=0.2): - """ - Loops through the NOES to match the H1, N1 and H2 shifts to two HSQC shifts. - Lists the HSQC shifts separately (H1/N1 as contact1 and H2 as contact2) for each NOE. - Returns a lists of tuples of all possible shift index combinations for each NOE. - ex. [[(1,2),(1,3)],[(3,0)]] - 1 and 2 are a pair of potential shift indices for one NOE, 1 and 3 are another for the same NOE. - """ - contact1 = [] - contact2 = [] - for i, peak in enumerate(noes): - # matching H1 and N15 in NOE to shifts in HSQC - nh_index = matchNH(peak.H1, peak.N15, actual_shifts, tolerance_h, tolerance_n) - # matching H2 in NOE to H1 shifts in HSQC - h_index = matchH(peak.H2, actual_shifts, tolerance_h) - - contact1.append(nh_index) - contact2.append(h_index) - - # print(contact1) - # print(contact2) - - contacts = [] - for i, contact in enumerate(contact1): - prod = product(contact, contact2[i]) - # filter out same shift pairs and inverted duplicates (ex. (0,0) and [(0,1),(1,0)]) - pairs = set(tuple(sorted(combo)) for combo in prod if combo[0] != combo[1]) - contacts.append(list(pairs)) - - return contacts - if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('num_resid', type=int, help='Number of residues') @@ -158,5 +103,8 @@ def noe_combinations(noes, actual_shifts, tolerance_h=0.2, tolerance_n=0.2): num_resid = args.num_resid # These should be grouped to export into an environment coords, actual_shifts, predicted_shifts, noes = generate_data(num_resid) - combinations = noe_combinations(noes, actual_shifts) - print(combinations) + + # energy_obj = Energy(coords, actual_shifts, noes) + # test = energy_obj.setup_noe_restraints() + # test_energy = energy_obj.energy_loop(test) + # print(test_energy) \ No newline at end of file From 2b1d8b2cbb7d7f6c49cef4c10d0b31f90c624d03 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Sun, 24 Nov 2024 10:28:19 -0700 Subject: [PATCH 03/43] Updated energy function --- energy.py | 29 ++++++++++++++--------------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/energy.py b/energy.py index 6400674..4213032 100644 --- a/energy.py +++ b/energy.py @@ -88,35 +88,34 @@ def flat_bottom(self, x, tolerance): lambda x: x**2 - xmax*x]) return y - def energy_loop(self, restraints, assignments): + def get_energy(self, restraints, assignments): """ - Loops over the NOE restraints and then each possible shift option. - If these shift options have been assigned, the NOE is activated, if not the NOE restraint should be passed. + Loops over the NOE restraints and each possible shift option. + If these shift options have not been assigned the NOE is deactivated and passed, if not the NOE restraint stays activated. Activation means the coordinate indices are grabbed from the assignment dictionary and used in the flat_bottom function. - Energy is then summed across all activated NOEs. + The lowest energy restraint is then summed across all activated NOEs. """ total_energy = 0 - # assignments = {0: 0, 1 : 1, 2 : 2, 3 : 3} # this would be added to parameters as external dict for restraint in restraints: restraint_energy = math.inf - activated = False + activated = True for i, j in restraint: - if i in assignments.values() and j in assignments.values(): - activated = True + if i not in assignments.values() or j not in assignments.values(): + activated = False + break + elif activated: k, l = assignments.get(i), assignments.get(j) + # distance to energy evaluation dist = np.linalg.norm((np.array(self.coords[k]) - np.array(self.coords[l]))) energy_value = self.flat_bottom(dist, tolerance=0.5) + # keep the smallest energy (distance) restraint_energy = energy_value if energy_value < restraint_energy else restraint_energy - # add energy to the summation of the restraints - # problem with this adding the restraint energy each time in inner loop - total_energy += restraint_energy - - else: - # NOE not activated - break + + if restraint_energy != math.inf: + total_energy += restraint_energy return total_energy From b89eb16e0d13ffa612586fd6761900330801ff7b Mon Sep 17 00:00:00 2001 From: a-angel-s Date: Wed, 27 Nov 2024 10:28:56 -0700 Subject: [PATCH 04/43] added unittest for energy --- energy_unittest.py | 76 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 energy_unittest.py diff --git a/energy_unittest.py b/energy_unittest.py new file mode 100644 index 0000000..4cf1f20 --- /dev/null +++ b/energy_unittest.py @@ -0,0 +1,76 @@ +import unittest + +class TestEnergyMethods(unittest.TestCase): + + def setUp(self): + self.energy = Energy() + +#test for activated, if activated, the expected will be the energy calculation + def test_energy_activated(self): + + #random data + coords = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [0.91, 0.92, 0.93], [0.95,0.96, 0.97]] #x,y,z + restraints = [[(2, 1),(0, 3)]] + assignments = {2: 1, 0: 3, 1: 2, 3: 4} + + result = self.energy.get_energy(restraints, coords, assignments) + expected = 0.01019237886466845 + + self.assertEqual(result, expected) + + #test for not activated, if not activated, the expected will be 0 + def test_energy_not_activated(self): + + #random data + coords = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7,0.8,0.9], [0.91, 0.92, 0.93], [0.95,0.96, 0.97]] #x,y,z + restraints = [[(2, 1), (0, 3)]] + assignments = {0:3} + + result = self.energy.get_energy(restraints, coords, assignments) + expected = 0 + + self.assertEqual(result, expected) + + #test for sum of all energies for more then one activated NOE, the expected will be the total energy + #also works for NOE not activated, expected will be the same as test_energy_activated + def test_energy_sum(self): + coords = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [0.1, 0.92, 0.93], [0.95,0.96, 0.97]] #x,y,z + restraints = [(2, 1), (0, 3)], [(1, 2)] + assignments = {2: 1, 0: 3, 1: 2} + + result = self.energy.get_energy(restraints, coords, assignments) + expected = 0.5263195283063458 + + self.assertEqual(result, expected) + + #energy calc - pass + def test_flat_bottom(self): + tolerance = 1 + x = 5 #distance + result = self.energy.flat_bottom(x, tolerance) + expected = 20 # x^2 - tolerance*x + self.assertEqual(result, expected) + + ''' + calculations for the expected of each unittest + + coords = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [0.91, 0.92, 0.93], [0.95,0.96, 0.97]] #x,y,z + + #calculation for test_energy_activated + restraints = [[(2, 1),(0, 3)]] + + x = np.linalg.norm((coords[2])-np.array(coords[1])) + + expected_x = x**2 - 0.5*x + + print(expected_x) + + #calculation for test_energy_sum + + y = np.linalg.norm(np.array(coords[1])-np.array(coords[2])) + + expected_y = y**2 - 0.5*y + + print(expected_y + expected_x) + + ''' From 44217a0b57f1d6e96b05f59bde1871bc9e2ef2ce Mon Sep 17 00:00:00 2001 From: abbie-p Date: Thu, 28 Nov 2024 17:50:23 -0700 Subject: [PATCH 05/43] Added fractional activation --- assignment_order.py | 37 +++++++++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) create mode 100644 assignment_order.py diff --git a/assignment_order.py b/assignment_order.py new file mode 100644 index 0000000..57de31d --- /dev/null +++ b/assignment_order.py @@ -0,0 +1,37 @@ +import numpy as np + +class Assignment_order(): + + def __init__(self, num_resid, noes): + self.num_resid = num_resid # value would need to be the number of shifts/atoms + self.noes = noes + + def activation_loop(self, target, noe, assign_order): + s = set() + for combo in noe: + if combo[0] not in assign_order: # check if values have already been assigned as higher priority + s.add(combo[0]) + if combo[1] not in assign_order: + s.add(combo[1]) + + if target in s: + return 1/(len(s)) + else: + return 0 + + def fractional_activation(self): + assign_order = [] + for j in range(self.num_resid): # I don't know why I need this loop here it just makes it work + temp2 = [] + for i in range(self.num_resid): + temp = 0 + for noe in self.noes: + temp += self.activation_loop(i, noe, assign_order) # sum up the probabilities for 'i' (the given shift target) + temp2.append(temp) + + assign_order.append(np.argmax(temp2)) # append the max value's index (takes first occurence if equivalent) + + return assign_order + + + \ No newline at end of file From c30a450254934a8db85a855db832c7f642a11cea Mon Sep 17 00:00:00 2001 From: a-angel-s Date: Thu, 5 Dec 2024 07:39:56 -0700 Subject: [PATCH 06/43] Added unittest for fractional activation --- fractional_activations_unittest.py | 68 ++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) create mode 100644 fractional_activations_unittest.py diff --git a/fractional_activations_unittest.py b/fractional_activations_unittest.py new file mode 100644 index 0000000..b5e264e --- /dev/null +++ b/fractional_activations_unittest.py @@ -0,0 +1,68 @@ +#unittest for fractional activations +#has to be made into a class for it to work + +import unittest + +class TestFracAct(unittest.TestCase): + + def setUp(self): + self.frac_act = FracAct() + +#test for activation loop +#expected to return the noe peak divided by the total number of different/unique noe peaks + + def test_activation_loop(self): + assign_order = [] + noe = [(0, 1), (0, 2)] + target = 0 + + result = self.frac_act.activation_loop(target, noe, assign_order) + expected = 1/3 + + self.assertEqual(result, expected) + +#test if each noe peak will be evaluated to give the fractional activation + + def test_fractional_activation_of_all_noe_peaks(self): + num_resid = 3 + noes = [[(0, 1), (0, 2)]] + + result = self.frac_act.fractional_activation_summation(num_resid, noes) + expected = [1/3, 1/3, 1/3] + + self.assertEqual(result, expected) + +#for more then one noe, the probabilities are summed + + def test_fractional_activation_summation_of_noes(self): + num_resid = 4 + noes = [(0, 1), (0, 2)], [(2,1),(1,1)], [(1,0),(3,2)] + + result = self.frac_act.fractional_activation_summation(num_resid, noes) + expected = [7/12, 13/12, 13/12, 1/4] + + self.assertEqual(result, expected) + +#according to the fractional activations calculated, should give the correct order +#if it is a tie, then lowest value goes first + + def test_fractional_activation_for_assignment_order(self): + num_resid = 4 + noes = [(0, 1), (0, 2)], [(2,1),(1,1)], [(1,0),(3,2)] + + result = self.frac_act.fractional_activation(num_resid, noes) + expected = [1, 2, 0, 3] + + self.assertEqual(result, expected) + +#if number of residues is higher than the number of peaks, then it will loop through to give first/lowest value peak ex. will give 0 for every additional residue +#if number of residues is lower than the number of peaks, then it will return the first value again (0, in this case) + + def test_fractional_activation_for_num_resid(self): + num_resid = 5 + noes = [[(0, 1), (0, 2)]], [(2,1),(1,1)], [(1,0),(3,2)] + + result = self.frac_act.fractional_activation(num_resid, noes) + expected = [1, 2, 0, 3] + + self.assertEqual(result, expected) \ No newline at end of file From 0fbc7e4d3cac4dc25b465d0a977a3d04f30f565d Mon Sep 17 00:00:00 2001 From: a-angel-s Date: Thu, 5 Dec 2024 07:55:45 -0700 Subject: [PATCH 07/43] Added unittest for fractional activation --- fractional_activation_unittest.py | 68 +++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) create mode 100644 fractional_activation_unittest.py diff --git a/fractional_activation_unittest.py b/fractional_activation_unittest.py new file mode 100644 index 0000000..b5e264e --- /dev/null +++ b/fractional_activation_unittest.py @@ -0,0 +1,68 @@ +#unittest for fractional activations +#has to be made into a class for it to work + +import unittest + +class TestFracAct(unittest.TestCase): + + def setUp(self): + self.frac_act = FracAct() + +#test for activation loop +#expected to return the noe peak divided by the total number of different/unique noe peaks + + def test_activation_loop(self): + assign_order = [] + noe = [(0, 1), (0, 2)] + target = 0 + + result = self.frac_act.activation_loop(target, noe, assign_order) + expected = 1/3 + + self.assertEqual(result, expected) + +#test if each noe peak will be evaluated to give the fractional activation + + def test_fractional_activation_of_all_noe_peaks(self): + num_resid = 3 + noes = [[(0, 1), (0, 2)]] + + result = self.frac_act.fractional_activation_summation(num_resid, noes) + expected = [1/3, 1/3, 1/3] + + self.assertEqual(result, expected) + +#for more then one noe, the probabilities are summed + + def test_fractional_activation_summation_of_noes(self): + num_resid = 4 + noes = [(0, 1), (0, 2)], [(2,1),(1,1)], [(1,0),(3,2)] + + result = self.frac_act.fractional_activation_summation(num_resid, noes) + expected = [7/12, 13/12, 13/12, 1/4] + + self.assertEqual(result, expected) + +#according to the fractional activations calculated, should give the correct order +#if it is a tie, then lowest value goes first + + def test_fractional_activation_for_assignment_order(self): + num_resid = 4 + noes = [(0, 1), (0, 2)], [(2,1),(1,1)], [(1,0),(3,2)] + + result = self.frac_act.fractional_activation(num_resid, noes) + expected = [1, 2, 0, 3] + + self.assertEqual(result, expected) + +#if number of residues is higher than the number of peaks, then it will loop through to give first/lowest value peak ex. will give 0 for every additional residue +#if number of residues is lower than the number of peaks, then it will return the first value again (0, in this case) + + def test_fractional_activation_for_num_resid(self): + num_resid = 5 + noes = [[(0, 1), (0, 2)]], [(2,1),(1,1)], [(1,0),(3,2)] + + result = self.frac_act.fractional_activation(num_resid, noes) + expected = [1, 2, 0, 3] + + self.assertEqual(result, expected) \ No newline at end of file From a7a5a2a4572ab1fbe288e1b77f7c4f514aedc789 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Thu, 5 Dec 2024 08:42:55 -0700 Subject: [PATCH 08/43] Updated energy and fractional activation --- assignment_order.py | 19 ++++++++++-- assignment_unittest.py | 68 ++++++++++++++++++++++++++++++++++++++++++ energy.py | 2 +- 3 files changed, 86 insertions(+), 3 deletions(-) create mode 100644 assignment_unittest.py diff --git a/assignment_order.py b/assignment_order.py index 57de31d..7747f45 100644 --- a/assignment_order.py +++ b/assignment_order.py @@ -1,6 +1,6 @@ import numpy as np -class Assignment_order(): +class FracAct(): def __init__(self, num_resid, noes): self.num_resid = num_resid # value would need to be the number of shifts/atoms @@ -32,6 +32,21 @@ def fractional_activation(self): assign_order.append(np.argmax(temp2)) # append the max value's index (takes first occurence if equivalent) return assign_order + + # intermediate check for first set of probabilites assigned - not used as final calculation + def fractional_activation_summation(self): + + assign_order = [] + for j in range(self.num_resid): + temp2 = [] + for i in range(self.num_resid): + temp = 0 + for noe in self.noes: + temp += activation_loop(i, noe, assign_order) + temp2.append(temp) + + return temp2 + + - \ No newline at end of file diff --git a/assignment_unittest.py b/assignment_unittest.py new file mode 100644 index 0000000..ce63fc7 --- /dev/null +++ b/assignment_unittest.py @@ -0,0 +1,68 @@ +#unittest for fractional activations +#has to be made into a class for it to work + +import unittest + +class TestFracAct(unittest.TestCase): + + def setUp(self): + self.frac_act = FracAct() + + #test for activation loop + #expected to return the noe peak divided by the total number of different/unique noe peaks + + def test_activation_loop(self): + assign_order = [] + noe = [(0, 1), (0, 2)] + target = 0 + + result = self.frac_act.activation_loop(target, noe, assign_order) + expected = 1/3 + + self.assertEqual(result, expected) + + #test if each noe peak will be evaluated to give the fractional activation + + def test_fractional_activation_of_all_noe_peaks(self): + num_resid = 3 + noes = [[(0, 1), (0, 2)]] + + result = self.frac_act.fractional_activation_summation(num_resid, noes) + expected = [1/3, 1/3, 1/3] + + self.assertEqual(result, expected) + + #for more then one noe, the probabilities are summed + + def test_fractional_activation_summation_of_noes(self): + num_resid = 4 + noes = [(0, 1), (0, 2)], [(2,1),(1,1)], [(1,0),(3,2)] + + result = self.frac_act.fractional_activation_summation(num_resid, noes) + expected = [7/12, 13/12, 13/12, 1/4] + + self.assertEqual(result, expected) + + #according to the fractional activations calculated, should give the correct order + #if it is a tie, then lowest value goes first + + def test_fractional_activation_for_assignment_order(self): + num_resid = 4 + noes = [(0, 1), (0, 2)], [(2,1),(1,1)], [(1,0),(3,2)] + + result = self.frac_act.fractional_activation(num_resid, noes) + expected = [1, 2, 0, 3] + + self.assertEqual(result, expected) + + #if number of residues is higher than the number of peaks, then it will loop through to give first/lowest value peak ex. will give 0 for every additional residue + #if number of residues is lower than the number of peaks, then it will return the first value again (0, in this case) + + def test_fractional_activation_for_num_resid(self): + num_resid = 5 + noes = [[(0, 1), (0, 2)]], [(2,1),(1,1)], [(1,0),(3,2)] + + result = self.frac_act.fractional_activation(num_resid, noes) + expected = [1, 2, 0, 3] + + self.assertEqual(result, expected) \ No newline at end of file diff --git a/energy.py b/energy.py index 4213032..1657c92 100644 --- a/energy.py +++ b/energy.py @@ -101,7 +101,7 @@ def get_energy(self, restraints, assignments): restraint_energy = math.inf activated = True for i, j in restraint: - if i not in assignments.values() or j not in assignments.values(): + if i not in assignments.keys() or j not in assignments.keys(): activated = False break From 475408a757160f3be605c79ca6de9a94af7f337f Mon Sep 17 00:00:00 2001 From: abbie-p Date: Thu, 5 Dec 2024 13:18:47 -0700 Subject: [PATCH 09/43] Formatting --- energy.py | 210 ++++++++++++++++++++++----------------------- energy_unittest.py | 98 ++++++++++----------- fake_data.py | 158 +++++++++++++++++----------------- 3 files changed, 233 insertions(+), 233 deletions(-) diff --git a/energy.py b/energy.py index 1657c92..3d191be 100644 --- a/energy.py +++ b/energy.py @@ -4,118 +4,118 @@ # match H1 shifts def matchH(hshift, hsqc, tolerance_h): - """ - Loops through the HSQC list and matches the second proton shift of the NOE to a HSQC shift based on H cutoff. - Returns the index of said shift. - """ - h_index = [] - for i, shift in enumerate(hsqc): - if abs(hshift - shift.H1) <= tolerance_h: - h_index.append(i) + """ + Loops through the HSQC list and matches the second proton shift of the NOE to a HSQC shift based on H cutoff. + Returns the index of said shift. + """ + h_index = [] + for i, shift in enumerate(hsqc): + if abs(hshift - shift.H1) <= tolerance_h: + h_index.append(i) - return h_index + return h_index # match H1 and N15 shifts def matchNH(hshift, nshift, hsqc, tolerance_h, tolerance_n): - """ - Loops through the HSQC list and matches the first proton and nitrogen shift of the NOE to a HSQC shift based on H and N cutoffs. - Returns the index of said shift. - """ - nh_index = [] - for i, shift in enumerate(hsqc): - if abs(hshift - shift.H1) <= tolerance_h and abs(nshift - shift.N15) <= tolerance_n: - nh_index.append(i) + """ + Loops through the HSQC list and matches the first proton and nitrogen shift of the NOE to a HSQC shift based on H and N cutoffs. + Returns the index of said shift. + """ + nh_index = [] + for i, shift in enumerate(hsqc): + if abs(hshift - shift.H1) <= tolerance_h and abs(nshift - shift.N15) <= tolerance_n: + nh_index.append(i) - return nh_index + return nh_index def noe_combinations(noes, actual_shifts, tolerance_h=0.2, tolerance_n=0.2): - """ - Loops through the NOES to match the H1, N1 and H2 shifts to two HSQC shifts. - Lists the HSQC shifts separately (H1/N1 as contact1 and H2 as contact2) for each NOE. - Returns a lists of tuples of all possible shift index combinations for each NOE. - ex. [[(1,2),(1,3)],[(3,0)]] - 1 and 2 are a pair of potential shift indices for one NOE, 1 and 3 are another for the same NOE. - """ - contact1 = [] - contact2 = [] - for i, peak in enumerate(noes): - # matching H1 and N15 in NOE to shifts in HSQC - nh_index = matchNH(peak.H1, peak.N15, actual_shifts, tolerance_h, tolerance_n) - # matching H2 in NOE to H1 shifts in HSQC - h_index = matchH(peak.H2, actual_shifts, tolerance_h) - - contact1.append(nh_index) - contact2.append(h_index) - - # print(contact1) - # print(contact2) - - contacts = [] - for i, contact in enumerate(contact1): - prod = product(contact, contact2[i]) - # filter out same shift pairs and inverted duplicates (ex. (0,0) and [(0,1),(1,0)]) - pairs = set(tuple(sorted(combo)) for combo in prod if combo[0] != combo[1]) - contacts.append(list(pairs)) - - return contacts + """ + Loops through the NOES to match the H1, N1 and H2 shifts to two HSQC shifts. + Lists the HSQC shifts separately (H1/N1 as contact1 and H2 as contact2) for each NOE. + Returns a lists of tuples of all possible shift index combinations for each NOE. + ex. [[(1,2),(1,3)],[(3,0)]] + 1 and 2 are a pair of potential shift indices for one NOE, 1 and 3 are another for the same NOE. + """ + contact1 = [] + contact2 = [] + for i, peak in enumerate(noes): + # matching H1 and N15 in NOE to shifts in HSQC + nh_index = matchNH(peak.H1, peak.N15, actual_shifts, tolerance_h, tolerance_n) + # matching H2 in NOE to H1 shifts in HSQC + h_index = matchH(peak.H2, actual_shifts, tolerance_h) + + contact1.append(nh_index) + contact2.append(h_index) + + # print(contact1) + # print(contact2) + + contacts = [] + for i, contact in enumerate(contact1): + prod = product(contact, contact2[i]) + # filter out same shift pairs and inverted duplicates (ex. (0,0) and [(0,1),(1,0)]) + pairs = set(tuple(sorted(combo)) for combo in prod if combo[0] != combo[1]) + contacts.append(list(pairs)) + + return contacts class Energy(): - def __init__(self, coords, actual_shifts, noes): - self.coords = coords - self.actual_shifts = actual_shifts - self.noes = noes - - def setup_noe_restraints(self): - combos = noe_combinations(self.noes, self.actual_shifts) - return combos - - def flat_bottom(self, x, tolerance): - """ - Defines the flat bottom equation with a flat region from zero to the distance tolerance. - """ - assert x > 0 - - # ends of flat bottom region - xmax = tolerance - xmin = 0 - - # if x <= tolerance return energy of zero - # if x > tolerance return energy relating to quadratic - y = np.piecewise(x, - [(xmin <= x) & (x <= xmax), x > xmax], - [lambda x: 0, - lambda x: x**2 - xmax*x]) - return y - - def get_energy(self, restraints, assignments): - """ - Loops over the NOE restraints and each possible shift option. - If these shift options have not been assigned the NOE is deactivated and passed, if not the NOE restraint stays activated. - Activation means the coordinate indices are grabbed from the assignment dictionary and used in the flat_bottom function. - The lowest energy restraint is then summed across all activated NOEs. - """ - - total_energy = 0 - for restraint in restraints: - restraint_energy = math.inf - activated = True - for i, j in restraint: - if i not in assignments.keys() or j not in assignments.keys(): - activated = False - break - - elif activated: - k, l = assignments.get(i), assignments.get(j) - - # distance to energy evaluation - dist = np.linalg.norm((np.array(self.coords[k]) - np.array(self.coords[l]))) - energy_value = self.flat_bottom(dist, tolerance=0.5) - - # keep the smallest energy (distance) - restraint_energy = energy_value if energy_value < restraint_energy else restraint_energy - - if restraint_energy != math.inf: - total_energy += restraint_energy - - return total_energy + def __init__(self, coords, actual_shifts, noes): + self.coords = coords + self.actual_shifts = actual_shifts + self.noes = noes + + def setup_noe_restraints(self): + combos = noe_combinations(self.noes, self.actual_shifts) + return combos + + def flat_bottom(self, x, tolerance): + """ + Defines the flat bottom equation with a flat region from zero to the distance tolerance. + """ + assert x > 0 + + # ends of flat bottom region + xmax = tolerance + xmin = 0 + + # if x <= tolerance return energy of zero + # if x > tolerance return energy relating to quadratic + y = np.piecewise(x, + [(xmin <= x) & (x <= xmax), x > xmax], + [lambda x: 0, + lambda x: x**2 - xmax*x]) + return y + + def get_energy(self, restraints, assignments): + """ + Loops over the NOE restraints and each possible shift option. + If these shift options have not been assigned the NOE is deactivated and passed, if not the NOE restraint stays activated. + Activation means the coordinate indices are grabbed from the assignment dictionary and used in the flat_bottom function. + The lowest energy restraint is then summed across all activated NOEs. + """ + + total_energy = 0 + for restraint in restraints: + restraint_energy = math.inf + activated = True + for i, j in restraint: + if i not in assignments.keys() or j not in assignments.keys(): + activated = False + break + + elif activated: + k, l = assignments.get(i), assignments.get(j) + + # distance to energy evaluation + dist = np.linalg.norm((np.array(self.coords[k]) - np.array(self.coords[l]))) + energy_value = self.flat_bottom(dist, tolerance=0.5) + + # keep the smallest energy (distance) + restraint_energy = energy_value if energy_value < restraint_energy else restraint_energy + + if restraint_energy != math.inf: + total_energy += restraint_energy + + return total_energy diff --git a/energy_unittest.py b/energy_unittest.py index 4cf1f20..5b1375a 100644 --- a/energy_unittest.py +++ b/energy_unittest.py @@ -2,75 +2,75 @@ class TestEnergyMethods(unittest.TestCase): - def setUp(self): - self.energy = Energy() + def setUp(self): + self.energy = Energy() -#test for activated, if activated, the expected will be the energy calculation - def test_energy_activated(self): + #test for activated, if activated, the expected will be the energy calculation + def test_energy_activated(self): - #random data - coords = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [0.91, 0.92, 0.93], [0.95,0.96, 0.97]] #x,y,z - restraints = [[(2, 1),(0, 3)]] - assignments = {2: 1, 0: 3, 1: 2, 3: 4} + #random data + coords = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [0.91, 0.92, 0.93], [0.95,0.96, 0.97]] #x,y,z + restraints = [[(2, 1),(0, 3)]] + assignments = {2: 1, 0: 3, 1: 2, 3: 4} - result = self.energy.get_energy(restraints, coords, assignments) - expected = 0.01019237886466845 + result = self.energy.get_energy(restraints, coords, assignments) + expected = 0.01019237886466845 - self.assertEqual(result, expected) + self.assertEqual(result, expected) - #test for not activated, if not activated, the expected will be 0 - def test_energy_not_activated(self): + #test for not activated, if not activated, the expected will be 0 + def test_energy_not_activated(self): - #random data - coords = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7,0.8,0.9], [0.91, 0.92, 0.93], [0.95,0.96, 0.97]] #x,y,z - restraints = [[(2, 1), (0, 3)]] - assignments = {0:3} + #random data + coords = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7,0.8,0.9], [0.91, 0.92, 0.93], [0.95,0.96, 0.97]] #x,y,z + restraints = [[(2, 1), (0, 3)]] + assignments = {0:3} - result = self.energy.get_energy(restraints, coords, assignments) - expected = 0 + result = self.energy.get_energy(restraints, coords, assignments) + expected = 0 - self.assertEqual(result, expected) + self.assertEqual(result, expected) - #test for sum of all energies for more then one activated NOE, the expected will be the total energy - #also works for NOE not activated, expected will be the same as test_energy_activated - def test_energy_sum(self): - coords = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [0.1, 0.92, 0.93], [0.95,0.96, 0.97]] #x,y,z - restraints = [(2, 1), (0, 3)], [(1, 2)] - assignments = {2: 1, 0: 3, 1: 2} + #test for sum of all energies for more then one activated NOE, the expected will be the total energy + #also works for NOE not activated, expected will be the same as test_energy_activated + def test_energy_sum(self): + coords = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [0.1, 0.92, 0.93], [0.95,0.96, 0.97]] #x,y,z + restraints = [(2, 1), (0, 3)], [(1, 2)] + assignments = {2: 1, 0: 3, 1: 2} - result = self.energy.get_energy(restraints, coords, assignments) - expected = 0.5263195283063458 + result = self.energy.get_energy(restraints, coords, assignments) + expected = 0.5263195283063458 - self.assertEqual(result, expected) + self.assertEqual(result, expected) - #energy calc - pass - def test_flat_bottom(self): - tolerance = 1 - x = 5 #distance - result = self.energy.flat_bottom(x, tolerance) - expected = 20 # x^2 - tolerance*x - self.assertEqual(result, expected) + #energy calc - pass + def test_flat_bottom(self): + tolerance = 1 + x = 5 #distance + result = self.energy.flat_bottom(x, tolerance) + expected = 20 # x^2 - tolerance*x + self.assertEqual(result, expected) - ''' - calculations for the expected of each unittest + ''' + calculations for the expected of each unittest - coords = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [0.91, 0.92, 0.93], [0.95,0.96, 0.97]] #x,y,z + coords = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [0.91, 0.92, 0.93], [0.95,0.96, 0.97]] #x,y,z - #calculation for test_energy_activated - restraints = [[(2, 1),(0, 3)]] + #calculation for test_energy_activated + restraints = [[(2, 1),(0, 3)]] - x = np.linalg.norm((coords[2])-np.array(coords[1])) + x = np.linalg.norm((coords[2])-np.array(coords[1])) - expected_x = x**2 - 0.5*x + expected_x = x**2 - 0.5*x - print(expected_x) + print(expected_x) - #calculation for test_energy_sum + #calculation for test_energy_sum - y = np.linalg.norm(np.array(coords[1])-np.array(coords[2])) + y = np.linalg.norm(np.array(coords[1])-np.array(coords[2])) - expected_y = y**2 - 0.5*y + expected_y = y**2 - 0.5*y - print(expected_y + expected_x) + print(expected_y + expected_x) - ''' + ''' diff --git a/fake_data.py b/fake_data.py index bc2ad93..30cf7f5 100644 --- a/fake_data.py +++ b/fake_data.py @@ -7,104 +7,104 @@ from energy import Energy class HSQCPeak(NamedTuple): - H1: float - N15: float + H1: float + N15: float class NOEPeak(NamedTuple): - H1: float - N15: float - H2: float + H1: float + N15: float + H2: float class Protein(NamedTuple): - x: float - y: float - z: float + x: float + y: float + z: float # Sampled points from unit square/cube def sample_unit(n, num_sides, min=0, max=1): - return np.random.uniform(min, max, size=(n, num_sides)) + return np.random.uniform(min, max, size=(n, num_sides)) # Noise to sampled points def add_noise(point, scale=0.1, min=0, max=1): - """ - Random selection from normal (gaussian) distribution of 'scale' width from 0 (center). - 'size' makes sure that it's the same shape as the point we are adding noises to. - ex. point = [0,1,2], noise = [1,1,1], noisy_point = [1,2,3] - """ - noise = np.random.normal(0, scale, size=point.shape) - noisy_point = point + noise - - # fold over points outside of boundaries - for point in noisy_point: - for value in point: - if max < value: - value = (max-(value-max)) - elif value < min: - value = (min-value) - - return noisy_point + """ + Random selection from normal (gaussian) distribution of 'scale' width from 0 (center). + 'size' makes sure that it's the same shape as the point we are adding noises to. + ex. point = [0,1,2], noise = [1,1,1], noisy_point = [1,2,3] + """ + noise = np.random.normal(0, scale, size=point.shape) + noisy_point = point + noise + + # fold over points outside of boundaries + for point in noisy_point: + for value in point: + if max < value: + value = (max-(value-max)) + elif value < min: + value = (min-value) + + return noisy_point def calc_dist(p1, p2): - return np.linalg.norm((p2-p1)) + return np.linalg.norm((p2-p1)) def distance_noe(protein, shifts, cutoff): - """ - Grabs close coordinates and 'associated' HSQC shift peaks (by same index) to create NOES. - Adds gaussian noise to all points. + """ + Grabs close coordinates and 'associated' HSQC shift peaks (by same index) to create NOES. + Adds gaussian noise to all points. - **Can end up with no NOEs depending on the cutoff** - """ - noes = [] + **Can end up with no NOEs depending on the cutoff** + """ + noes = [] - for i, atom1 in enumerate(protein): - for j, atom2 in enumerate(protein): - if i != j: - dist = calc_dist(atom1, atom2) - if dist < cutoff: - # print(dist, shifts[i], shifts[j]) - noe = list(shifts[i][:]) # H1, N1 - noe.append(shifts[j][0]) # H2 - noes.append(noe) + for i, atom1 in enumerate(protein): + for j, atom2 in enumerate(protein): + if i != j: + dist = calc_dist(atom1, atom2) + if dist < cutoff: + # print(dist, shifts[i], shifts[j]) + noe = list(shifts[i][:]) # H1, N1 + noe.append(shifts[j][0]) # H2 + noes.append(noe) - noisy_noe = add_noise(np.array(noes), scale=0.01) + noisy_noe = add_noise(np.array(noes), scale=0.01) - return noisy_noe + return noisy_noe def generate_data(num_resid): - """ - Generates all fake data and orders it in lists of namedtuples. - """ - # 3D structure [x,y,z] - protein = sample_unit(num_resid, num_sides=3) - # "Actual" shifts [H1,N1] - actual_shifts = sample_unit(num_resid, num_sides=2) - # Predicted shifts [H1,N1] - predicted_shifts = add_noise(actual_shifts) - # NOES [H1,N1,H2] - noes = distance_noe(protein, actual_shifts, cutoff=0.5) # likely need to change dist cutoff and only accounting for actual_shifts right now - - # Lists of namedtuples (one object per residue) - coords = [Protein(x=resid[0], y=resid[1], z=resid[2]) for resid in protein] - actual_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in actual_shifts] - predicted_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in predicted_shifts] - noes = [NOEPeak(H1=shift[0], N15=shift[1], H2=shift[2]) for shift in noes] - - # print(coords) - # print(actual_shifts) - # print(noe) - - return coords, actual_shifts, predicted_shifts, noes + """ + Generates all fake data and orders it in lists of namedtuples. + """ + # 3D structure [x,y,z] + protein = sample_unit(num_resid, num_sides=3) + # "Actual" shifts [H1,N1] + actual_shifts = sample_unit(num_resid, num_sides=2) + # Predicted shifts [H1,N1] + predicted_shifts = add_noise(actual_shifts) + # NOES [H1,N1,H2] + noes = distance_noe(protein, actual_shifts, cutoff=0.5) # likely need to change dist cutoff and only accounting for actual_shifts right now + + # Lists of namedtuples (one object per residue) + coords = [Protein(x=resid[0], y=resid[1], z=resid[2]) for resid in protein] + actual_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in actual_shifts] + predicted_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in predicted_shifts] + noes = [NOEPeak(H1=shift[0], N15=shift[1], H2=shift[2]) for shift in noes] + + # print(coords) + # print(actual_shifts) + # print(noe) + + return coords, actual_shifts, predicted_shifts, noes if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('num_resid', type=int, help='Number of residues') - args = parser.parse_args() - - num_resid = args.num_resid - # These should be grouped to export into an environment - coords, actual_shifts, predicted_shifts, noes = generate_data(num_resid) - - # energy_obj = Energy(coords, actual_shifts, noes) - # test = energy_obj.setup_noe_restraints() - # test_energy = energy_obj.energy_loop(test) - # print(test_energy) \ No newline at end of file + parser = argparse.ArgumentParser() + parser.add_argument('num_resid', type=int, help='Number of residues') + args = parser.parse_args() + + num_resid = args.num_resid + # These should be grouped to export into an environment + coords, actual_shifts, predicted_shifts, noes = generate_data(num_resid) + + # energy_obj = Energy(coords, actual_shifts, noes) + # test = energy_obj.setup_noe_restraints() + # test_energy = energy_obj.energy_loop(test) + # print(test_energy) \ No newline at end of file From 89ab4a93dff38be1b722a3e5d6a1527402e6692d Mon Sep 17 00:00:00 2001 From: abbie-p Date: Tue, 10 Dec 2024 13:07:34 -0700 Subject: [PATCH 10/43] Fixed fractional activation when missing shifts --- assignment_order.py | 33 ++++++++++++++++++++++++++++++--- 1 file changed, 30 insertions(+), 3 deletions(-) diff --git a/assignment_order.py b/assignment_order.py index 7747f45..b8f646f 100644 --- a/assignment_order.py +++ b/assignment_order.py @@ -19,17 +19,44 @@ def activation_loop(self, target, noe, assign_order): else: return 0 + # def fractional_activation(self): + # assign_order = [] + # for j in range(self.num_resid): # I don't know why I need this loop here it just makes it work + # temp2 = [] + # for i in range(self.num_resid): + # temp = 0 + # for noe in self.noes: + # temp += self.activation_loop(i, noe, assign_order) # sum up the probabilities for 'i' (the given shift target) + # temp2.append(temp) + + # assign_order.append(np.argmax(temp2)) # append the max value's index (takes first occurence if equivalent) + + # return assign_order def fractional_activation(self): assign_order = [] - for j in range(self.num_resid): # I don't know why I need this loop here it just makes it work + missing_shifts = [] + + while len(assign_order) < self.num_resid: temp2 = [] for i in range(self.num_resid): - temp = 0 + temp = 0 for noe in self.noes: temp += self.activation_loop(i, noe, assign_order) # sum up the probabilities for 'i' (the given shift target) + + # identifies any missing shifts + if temp == 0 and len(assign_order) == 0: + missing_shifts.append(i) + + # appends sum to list (one sum per shift) temp2.append(temp) - assign_order.append(np.argmax(temp2)) # append the max value's index (takes first occurence if equivalent) + # makes sure extra zeros are not added during final iterations (if shifts are missing) + if any(temp2) == True: + # proper order based on max value + assign_order.append(np.argmax(temp2)) + else: + # add missing shifts to end of list + assign_order = assign_order + missing_shifts return assign_order From 5ff5bbebf35cbdf48e1db3a3964289a198c12fb7 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Thu, 12 Dec 2024 19:19:53 -0700 Subject: [PATCH 11/43] Added gym environment --- gym_env_unittest.py | 66 ++++++++++++++++++++++++++++++ nmr_gym_env.py | 94 +++++++++++++++++++++++++++++++++++++++++++ nmr_text_adventure.py | 29 +++++++++++++ 3 files changed, 189 insertions(+) create mode 100644 gym_env_unittest.py create mode 100644 nmr_gym_env.py create mode 100644 nmr_text_adventure.py diff --git a/gym_env_unittest.py b/gym_env_unittest.py new file mode 100644 index 0000000..89f73ff --- /dev/null +++ b/gym_env_unittest.py @@ -0,0 +1,66 @@ +import unittest +import numpy as np + +from nmr_gym_env import GymEnv + +class TestGymEnv(unittest.TestCase): + + def setUp(self): + self.num_resid = 4 + self.state = { + "coords": [[0.35, 0.83, 0.89], [0.69, 0.12, 0.92], [0.05, 0.94, 0.51], [0.45, 0.06, 0.70]], + "actual_shifts": [[0.17, 0.044], [0.41, 0.34], [0.18, 0.61], [0.64, 0.23]], + "noes": [[0.16, 0.040, 0.17], [0.40, 0.34, 0.63], [0.18, 0.61, 0.16], [0.64, 0.24, 0.39]], + "restraints": [[(0, 2)], [(1, 3)], [(0, 2)], [(1, 3)]], + "assignments": {}, + "total_energy": 0 + } + + self.gym_env = GymEnv(self.num_resid) + + def test_step_energy_zeros(self): + actions = [3, 1, 0, 2] + + expected_energy = [0, 0, 0, 0] + expected_reward = [0, 0, 0, 0] + + observation = self.gym_env.custom_state(self.state) + terminated = False + + for i, action in enumerate(actions): + observation, reward, terminated = self.gym_env.step(action) + + self.assertEqual(observation['total_energy'], expected_energy[i]) + self.assertEqual(reward, expected_reward[i]) + + def test_step_energy_goes_up(self): + actions = [2, 3, 0, 1] + + # dist1 = np.linalg.norm((np.array(self.state['coords'][2]) - np.array(self.state['coords'][3]))) + # dist2 = np.linalg.norm((np.array(self.state['coords'][1]) - np.array(self.state['coords'][0]))) + + # energy1 = dist1**2 - 0.5*dist1 # for restraint (0,2) + # energy2 = dist2**2 - 0.5*dist2 # for restraint (1,3) + # print(energy1*2 + energy2*2) + # print(energy1*2 - (energy1*2 + energy2*2)) # final reward + + """ + Step1: No NOEs activated --> energy 0 + Step2: Two NOEs activated --> energy 0.47/NOE + Step3: Two NOEs activated (no new additions) --> energy 0.47/NOE + Step4: All NOEs activated (two more NOEs activated) --> energy 0.47/NOE prior + 0.22/NOE new + + """ + expected_energy = [0, 0.9558604159815726, 0.9558604159815726, 1.4092787203323272] + expected_reward = [0, -0.9558604159815726, 0, -0.4534183043507546] + + observation = self.gym_env.custom_state(self.state) + terminated = False + + for i, action in enumerate(actions): + observation, reward, terminated = self.gym_env.step(action) + + self.assertEqual(observation['total_energy'], expected_energy[i]) + self.assertEqual(reward, expected_reward[i]) + +unittest.main(argv=[''], verbosity=2, exit=False) \ No newline at end of file diff --git a/nmr_gym_env.py b/nmr_gym_env.py new file mode 100644 index 0000000..90a09ba --- /dev/null +++ b/nmr_gym_env.py @@ -0,0 +1,94 @@ +import gym +from gym import error, spaces +import numpy as np + +from fake_data import generate_data +from energy import Energy, noe_combinations +from assignment_order import FracAct + +class GymEnv(gym.Env): + + def __init__(self, num_resid): + self.num_resid = num_resid + self.assignments = {} + self.assign_step = 0 + self.intermediate_energy = 0 + self.reward = 0 + + self.action_space = spaces.Discrete(self.num_resid) + + self.observation_space = spaces.Dict({ + "coords": spaces.Box(0, 1, shape=(self.num_resid, 3), dtype=np.float32), + "actual_shifts": spaces.Box(0, 1, shape=(self.num_resid, 2), dtype=np.float32), + "predicted_shifts": spaces.Box(0, 1, shape=(self.num_resid, 2), dtype=np.float32), + "noes": spaces.Box(0, 1, shape=(self.num_resid, 3), dtype=np.float32) + }) + + def custom_state(self, state): + self.state = state + frac_act = FracAct(self.num_resid, self.state['restraints']) + self.assign_order = frac_act.fractional_activation() + return self.state + + def reset(self): + # generate fake data + coords, actual_shifts, predicted_shifts, noes = generate_data(self.num_resid) + # get restraints + restraints = noe_combinations(noes, actual_shifts) + + # get assignment order + frac_act = FracAct(self.num_resid, restraints) + self.assign_order = frac_act.fractional_activation() + print(self.assign_order) + + self.assignments = {} + # self.assign_step = 0 + # self.intermediate_energy = 0 + self.total_energy = 0 + # self.reward = 0 + + self.state = { + "coords": coords, + "actual_shifts": actual_shifts, + "predicted_shifts": predicted_shifts, + "noes": noes, + "restraints": restraints, + "assignments": self.assignments, + "total_energy": self.total_energy + } + + return self.state + + def step(self, action): + + terminated = True if len(self.assignments) == self.num_resid else False + + if terminated: + # how to represent final reward? + return self.state, self.reward, terminated + + else: + energy = Energy(self.state['coords'], self.state['actual_shifts'], self.state['noes']) + + # add action to assignments + self.assignments[(self.assign_order[self.assign_step])] = action + print(self.assignments) + + # calc energy and store temporarily + temp_intermediate_energy = energy.get_energy(self.state['restraints'], self.assignments) + # print(f'this is the temp energy {temp_intermediate_energy}') + + # calc reward based on previous energy and temporary energy (+ve means energy went down, -ve means energy went up) + self.reward = (self.intermediate_energy - temp_intermediate_energy) #if self.intermediate_energy != 0 else 0 + # print(f'this is the reward {self.reward}') + + # store energy as intermediate + self.intermediate_energy = temp_intermediate_energy + + # assign intermediate energy as running total + self.state['total_energy'] = self.intermediate_energy + print(f"Running energy = {self.state['total_energy']}, Current reward = {self.reward}") + + self.assign_step += 1 + + return self.state, self.reward, terminated diff --git a/nmr_text_adventure.py b/nmr_text_adventure.py new file mode 100644 index 0000000..d6bf319 --- /dev/null +++ b/nmr_text_adventure.py @@ -0,0 +1,29 @@ +import gym +import argparse + +from nmr_gym_env import GymEnv + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('num_resid', type=int, help='Number of residues') + args = parser.parse_args() + + num_resid = args.num_resid + gym_env = GymEnv(num_resid) + episodes = 1 + + for episode in range(episodes): + observation = gym_env.reset() + terminated = False + + while not terminated: + action = gym_env.action_space.sample() + # action = int(input()) + # need to prevent sample from grabbing the same atom more than once + # but also need the loop to terminate (without last bit it won't go into else for last iteration) + if action in observation['assignments'].values() and len(observation['assignments']) != num_resid: + continue + else: + observation, reward, terminated = gym_env.step(action) + + print(f"Final assignments (shift:atom) = {observation['assignments']}\nFinal energy evaluation = {observation['total_energy']}") \ No newline at end of file From 5d71f2995d9f55d8b314a605a71b67e83c1fc554 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Fri, 20 Dec 2024 08:27:08 -0700 Subject: [PATCH 12/43] Added some text adventure examples --- fake_data.py | 27 ++++++++++-------- nmr_gym_env.py | 49 +++++++++++++++++---------------- nmr_text_adventure.py | 64 ++++++++++++++++++++++++++++++++----------- 3 files changed, 89 insertions(+), 51 deletions(-) diff --git a/fake_data.py b/fake_data.py index 30cf7f5..cc56881 100644 --- a/fake_data.py +++ b/fake_data.py @@ -95,16 +95,21 @@ def generate_data(num_resid): return coords, actual_shifts, predicted_shifts, noes -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('num_resid', type=int, help='Number of residues') - args = parser.parse_args() - - num_resid = args.num_resid - # These should be grouped to export into an environment - coords, actual_shifts, predicted_shifts, noes = generate_data(num_resid) - - # energy_obj = Energy(coords, actual_shifts, noes) - # test = energy_obj.setup_noe_restraints() +# if __name__ == '__main__': +# parser = argparse.ArgumentParser() +# parser.add_argument('num_resid', type=int, help='Number of residues') +# args = parser.parse_args() + +# num_resid = args.num_resid +# # These should be grouped to export into an environment +# coords, actual_shifts, predicted_shifts, noes = generate_data(num_resid) +# print(f'coords {(np.array(coords)).tolist()}\n shifts{(np.array(actual_shifts)).tolist()}\n noes{(np.array(noes)).tolist()}') +# # print(actual_shifts) +# # print(noes) + + +# energy_obj = Energy(coords, actual_shifts, noes) +# test = energy_obj.setup_noe_restraints() +# print(test) # test_energy = energy_obj.energy_loop(test) # print(test_energy) \ No newline at end of file diff --git a/nmr_gym_env.py b/nmr_gym_env.py index 90a09ba..8220414 100644 --- a/nmr_gym_env.py +++ b/nmr_gym_env.py @@ -28,6 +28,7 @@ def custom_state(self, state): self.state = state frac_act = FracAct(self.num_resid, self.state['restraints']) self.assign_order = frac_act.fractional_activation() + print(self.assign_order) return self.state def reset(self): @@ -61,34 +62,34 @@ def reset(self): def step(self, action): - terminated = True if len(self.assignments) == self.num_resid else False - - if terminated: - # how to represent final reward? - return self.state, self.reward, terminated - - else: - energy = Energy(self.state['coords'], self.state['actual_shifts'], self.state['noes']) + energy = Energy(self.state['coords'], self.state['actual_shifts'], self.state['noes']) - # add action to assignments - self.assignments[(self.assign_order[self.assign_step])] = action - print(self.assignments) + # add action to assignments + self.assignments[(self.assign_order[self.assign_step])] = action + print(self.assignments) - # calc energy and store temporarily - temp_intermediate_energy = energy.get_energy(self.state['restraints'], self.assignments) - # print(f'this is the temp energy {temp_intermediate_energy}') + # calc energy and store temporarily + temp_intermediate_energy = energy.get_energy(self.state['restraints'], self.assignments) + # print(f'this is the temp energy {temp_intermediate_energy}') - # calc reward based on previous energy and temporary energy (+ve means energy went down, -ve means energy went up) - self.reward = (self.intermediate_energy - temp_intermediate_energy) #if self.intermediate_energy != 0 else 0 - # print(f'this is the reward {self.reward}') + # calc reward based on previous energy and temporary energy (+ve means energy went down, -ve means energy went up) + self.reward = (self.intermediate_energy - temp_intermediate_energy) #if self.intermediate_energy != 0 else 0 + # print(f'this is the reward {self.reward}') - # store energy as intermediate - self.intermediate_energy = temp_intermediate_energy + # store energy as intermediate + self.intermediate_energy = temp_intermediate_energy - # assign intermediate energy as running total - self.state['total_energy'] = self.intermediate_energy - print(f"Running energy = {self.state['total_energy']}, Current reward = {self.reward}") - - self.assign_step += 1 + # assign intermediate energy as running total + self.state['total_energy'] = self.intermediate_energy + print(f"Running energy = {self.state['total_energy']}, Current reward = {self.reward}") + self.assign_step += 1 + terminated = True if len(self.assignments) == self.num_resid else False + + if terminated: + print(f"Final assignments (shift:atom) = {self.assignments}\nFinal energy evaluation = {self.state['total_energy']}") return self.state, self.reward, terminated + else: + return self.state, self.reward, terminated + + diff --git a/nmr_text_adventure.py b/nmr_text_adventure.py index d6bf319..e851a52 100644 --- a/nmr_text_adventure.py +++ b/nmr_text_adventure.py @@ -4,26 +4,58 @@ from nmr_gym_env import GymEnv if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('num_resid', type=int, help='Number of residues') - args = parser.parse_args() + # parser = argparse.ArgumentParser() + # parser.add_argument('num_resid', type=int, help='Number of residues') + # args = parser.parse_args() - num_resid = args.num_resid + # num_resid = args.num_resid + num_resid = 10 gym_env = GymEnv(num_resid) episodes = 1 for episode in range(episodes): - observation = gym_env.reset() + # observation = gym_env.reset() + + # random generated set of 4 + # state = { + # "coords": [[0.35, 0.83, 0.89], [0.69, 0.12, 0.92], [0.05, 0.94, 0.51], [0.45, 0.06, 0.70]], + # "actual_shifts": [[0.17, 0.044], [0.41, 0.34], [0.18, 0.61], [0.64, 0.23]], + # "noes": [[0.16, 0.040, 0.17], [0.40, 0.34, 0.63], [0.18, 0.61, 0.16], [0.64, 0.24, 0.39]], + # "restraints": [[(0, 2)], [(1, 3)], [(0, 2)], [(1, 3)]], + # "assignments": {}, + # "total_energy": 0 + # } + + # random generated set of 10 with many restraints + state = { + "coords": [[0.8136200069440616, 0.7789831346675045, 0.6474707035745629], [0.9458251909740947, 0.8744248260731231, 0.16939643090298961], [0.07941814153270477, 0.8569710876203085, 0.6679287503567067], [0.9295839450890829, 0.2591652267820922, 0.5767903211938266], [0.5686652108970609, 0.6600261693178074, 0.8353850759130969], [0.3049822956354953, 0.01597728824712763, 0.8007859163237383], [0.2707959067671988, 0.16592702782732538, 0.8530675522602625], [0.5906467171408669, 0.2275483180754293, 0.1179098748663463], [0.8089112838766771, 0.014346015197136075, 0.2651392664874339], [0.021085864510464902, 0.3582433470216929, 0.9002893595219581]], + "actual_shifts": [[0.7235029227185441, 0.8167276537744274], [0.44118846730192773, 0.8390592014969155], [0.3607330934400178, 0.813863583391123], [0.35720641978511525, 0.2931243342244849], [0.03585880729962143, 0.6045433843661965], [0.13284379187386453, 0.7429892824868469], [0.8157925033719939, 0.5225148209456146], [0.3240407057810011, 0.9258723658454763], [0.35659118663170697, 0.4579588797141084], [0.6229950482718816, 0.9870152102924102]], + "noes": [[0.7353272235348708, 0.7985692976248703, 0.024966821752980907], [0.3602327190429789, 0.27520736325235085, 0.3444804557069924], [0.030177646349364318, 0.6038110746526516, 0.7374671133868814], [0.1147855723335682, 0.7472335679962178, 0.8334688788656137], [0.14395802509472194, 0.7602008741457149, 0.6270020652720573], [0.8084393678158522, 0.49689451509103233, 0.1359092345145138], [0.8359026620374025, 0.5227806613419343, 0.6230942517135893], [0.3158886510580014, 0.9127784810952192, 0.36108155258485086], [0.35718811981030113, 0.46209771728576327, 0.36091347433509957], [0.3473018809331518, 0.45763566566297215, 0.3106537228148144], [0.63688033551314, 0.9916090993120613, 0.12731578248735478], [0.6196043419116839, 0.9903515010993935, 0.808008737938382]], + # "restraints": [[(4, 9), (5, 9), (0, 4), (0, 5)], [(3, 8), (3, 7), (1, 8), (2, 3), (1, 3), (7, 8), (2, 8)], [(0, 4), (4, 9), (4, 6), (5, 6), (0, 5), (5, 9)], [(4, 6), (0, 4), (5, 6), (0, 5)], [(0, 7), (0, 4), (1, 5), (4, 9), (4, 6), (1, 4), (7, 9), (6, 7), (1, 7), (5, 6), (0, 5), (5, 9)], [(6, 7), (4, 6), (5, 6)], [(1, 6), (6, 9), (0, 6)], [(1, 2), (2, 7), (1, 5), (5, 8), (3, 7), (1, 8), (5, 7), (2, 3), (1, 7), (7, 8), (2, 5), (1, 3), (3, 5), (2, 8)], [(3, 8), (3, 7), (1, 8), (2, 3), (1, 3), (7, 8), (2, 8)], [(3, 8), (5, 8), (3, 7), (1, 8), (2, 3), (7, 8), (1, 3), (3, 5), (2, 8)], [(0, 7), (0, 4), (1, 5), (4, 9), (1, 4), (7, 9), (1, 7), (0, 5), (5, 9)], [(0, 1), (0, 9), (0, 6), (1, 6), (6, 9), (1, 9)]], + "restraints": [[(3, 8), (5, 9)], [(5, 6), (0, 5), (5, 9)], [(0, 4), (4, 9), (4, 6), (5, 6)], [(3, 8), (3, 7), (1, 8)], [(1, 6), (6, 9)]], # there's something wrong with too many restraints, + "assignments": {}, + "total_energy": 0 + } + + # linear arrangement of 10 with easy restraints + # state = { + # "coords": [[0.37, 0.82, 0.00], [0.37, 0.82, 0.11], [0.37, 0.82, 0.24], [0.37, 0.82, 0.39], [0.37, 0.82, 0.52], [0.37, 0.82, 0.67], [0.37, 0.82, 0.75], [0.37, 0.82, 0.88], [0.37, 0.82, 0.93], [0.37, 0.82, 1.00]], + # "actual_shifts": [[0.6731658491622945, 0.657184525800008], [0.28748044870745615, 0.34032397854002194], [0.4921807498359013, 0.03828350226565602], [0.7763637896667895, 0.3617894550491416], [0.06840333800729603, 0.20738485015436503], [0.9913658149098207, 0.33931570183160287], [0.5459835943391216, 0.8469814942099086], [0.09701544891895453, 0.7197743401411141], [0.5107870945621455, 0.7086010383972177], [0.5738658802959612, 0.31748338342430493]], + # "noes": [[0.7028015287380062, 0.5021893114424876], [0.33894896143919023, 0.45179229762303535], [0.46784325494458656, -0.0019242124744091177], [0.7161042434070474, 0.3465503455737962], [0.00801328834772183, 0.3164743274291713], [0.9337074775227354, 0.3689512869911665], [0.47077422692743703, 0.8353558577228544], [0.2038502444268092, 0.5189774171248352], [0.50122164179375, 0.5801271184891338], [0.5307875874176243, 0.2934724693638126]], + # "restraints": [[(0, 1)], [(2, 3)], [(4, 5)], [(6, 7)], [(8, 9)]], + # "assignments": {}, + # "total_energy": 0 + # } + + observation = gym_env.custom_state(state) + terminated = False - + action_assigned = [] + while not terminated: - action = gym_env.action_space.sample() - # action = int(input()) - # need to prevent sample from grabbing the same atom more than once - # but also need the loop to terminate (without last bit it won't go into else for last iteration) - if action in observation['assignments'].values() and len(observation['assignments']) != num_resid: - continue - else: - observation, reward, terminated = gym_env.step(action) - - print(f"Final assignments (shift:atom) = {observation['assignments']}\nFinal energy evaluation = {observation['total_energy']}") \ No newline at end of file + # action = gym_env.action_space.sample() + action = int(input()) + if action < num_resid: + if action not in action_assigned: + action_assigned.append(action) + observation, reward, terminated = gym_env.step(action) From 2a883357537241ef5dfdf35f8fa360fa8e84cc83 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Tue, 14 Jan 2025 08:30:53 -0700 Subject: [PATCH 13/43] Energy checks for activation --- energy.py | 59 ++++++++++++++++++++++++++++++++----------------------- 1 file changed, 34 insertions(+), 25 deletions(-) diff --git a/energy.py b/energy.py index 3d191be..29cca26 100644 --- a/energy.py +++ b/energy.py @@ -87,35 +87,44 @@ def flat_bottom(self, x, tolerance): [lambda x: 0, lambda x: x**2 - xmax*x]) return y - - def get_energy(self, restraints, assignments): + + def calc_restraint_energy(self, assignments, restraint): + """ + Goes over shift possibilities in the restraint, grabs associated coordinates from assignment, and calculates energy according to distance. + Only the smallest energy is returned. + """ + restraint_energy = math.inf + for i, j in restraint: + k, l = assignments.get(i), assignments.get(j) + dist = np.linalg.norm((np.array(self.coords[k]) - np.array(self.coords[l]))) + energy_value = self.flat_bottom(dist, tolerance=0.5) + restraint_energy = energy_value if energy_value < restraint_energy else restraint_energy + return restraint_energy + + def noe_activation(self, assignments, restraint): + """ + Checks if all shift possibilities in the restraint have been assigned (aka activated). + Calculates energy if they have, returns zero if they have not. + """ + x = set(assignments.keys()) + y = set() + for i, j in restraint: + y.add(i) + y.add(j) + d = y - x + if d: + return 0 + else: + return self.calc_restraint_energy(assignments, restraint) + + def get_total_energy(self, restraints, assignments): """ - Loops over the NOE restraints and each possible shift option. - If these shift options have not been assigned the NOE is deactivated and passed, if not the NOE restraint stays activated. - Activation means the coordinate indices are grabbed from the assignment dictionary and used in the flat_bottom function. - The lowest energy restraint is then summed across all activated NOEs. + Loops over the NOE restraints to sum up calculated energies. """ total_energy = 0 for restraint in restraints: - restraint_energy = math.inf - activated = True - for i, j in restraint: - if i not in assignments.keys() or j not in assignments.keys(): - activated = False - break - - elif activated: - k, l = assignments.get(i), assignments.get(j) - - # distance to energy evaluation - dist = np.linalg.norm((np.array(self.coords[k]) - np.array(self.coords[l]))) - energy_value = self.flat_bottom(dist, tolerance=0.5) - - # keep the smallest energy (distance) - restraint_energy = energy_value if energy_value < restraint_energy else restraint_energy - - if restraint_energy != math.inf: - total_energy += restraint_energy + energy_value = self.noe_activation(assignments, restraint) + total_energy += energy_value return total_energy From ade15950eee3f4f0155c7461adeb1b0c7c78db3f Mon Sep 17 00:00:00 2001 From: abbie-p Date: Tue, 28 Jan 2025 14:11:18 -0700 Subject: [PATCH 14/43] Added visualization components --- .gitignore | 4 +++ visualize_data.py | 90 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+) create mode 100644 visualize_data.py diff --git a/.gitignore b/.gitignore index 82f9275..2f45271 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,10 @@ __pycache__/ *.py[cod] *$py.class +# Pickled data and visualization +*.pkl +*.png + # C extensions *.so diff --git a/visualize_data.py b/visualize_data.py new file mode 100644 index 0000000..c7a669e --- /dev/null +++ b/visualize_data.py @@ -0,0 +1,90 @@ +from typing import NamedTuple +import matplotlib.pyplot as plt +import numpy as np +from mpl_toolkits.mplot3d import Axes3D + +def get_data(actual_shifts, coords, named_tuple_used=True): + if named_tuple_used: + # Actual + actual_x = [shift.H1 for shift in actual_shifts] + actual_y = [shift.N15 for shift in actual_shifts] + # Predicted + predict_x = [shift.H1 for shift in coords] + predict_y = [shift.N15 for shift in coords] + + else: + # Actual + actual_x = [shift[0] for shift in actual_shifts] + actual_y = [shift[1] for shift in actual_shifts] + # Predicted + predict_x = [shift[3] for shift in coords] + predict_y = [shift[4] for shift in coords] + + return actual_x, actual_y, predict_x, predict_y + +# def new_plot(coords): + +# fig = plt.figure(figsize=(5,5), layout='tight') +# ax = fig.add_subplot(111, projection='3d') +# x = [] +# y = [] +# z = [] +# for i, coord in enumerate(coords): +# # x.append(coord.x) +# # y.append(coord.y) +# # z.append(coord.z) +# x = coord.x +# y = coord.y +# z = coord.z + +# ax.scatter(x,y,z, s=40, label=f'shift {i}') + +# ax.legend() + +# plt.show() +# #plt.savefig("coord_plot.png", dpi = 150) + +def close_contacts(coords, cutoff=0.37): + """ + Calculates close contacts based on coordinates. + """ + contacts = [] + + for i, atom1 in enumerate(coords): + for j, atom2 in enumerate(coords): + if i != j: + dist = np.linalg.norm((np.array(atom1[:3]) - np.array(atom2[:3]))) + if dist < cutoff: + contacts.append((i,j)) + + return contacts + +def plot_shifts(actual_shifts, restraints, coords, named_tuple_used=True): + + actual_x, actual_y, predict_x, predict_y = get_data(actual_shifts, coords, named_tuple_used=named_tuple_used) + + plt.scatter(predict_x, predict_y, color='blue', label='Predicted shifts') + + plt.scatter(actual_x, actual_y, color='red', label='Measured shifts') + + # Label measured shifts by index and predicted shifts by associated coordinate + for i in range(len(actual_shifts)): + plt.annotate(i, (actual_x[i]+0.01, actual_y[i]+0.01)) + plt.annotate(i, (predict_x[i]+0.01, predict_y[i]+0.01)) + + # Draw restraints + for restraint in restraints: + for j,k in restraint: + plt.plot((actual_x[j], actual_x[k]), (actual_y[j], actual_y[k]), color='red') + + # Draw close contacts + contacts = close_contacts(coords) + for i,j in contacts: + plt.plot((predict_x[i], predict_x[j]), (predict_y[i], predict_y[j]), color='blue', ls='dotted', alpha=0.5) + + plt.xlabel('H1') + plt.ylabel('N15') + plt.legend() + + plt.savefig("shifts_plot.png", dpi = 150) + plt.close() \ No newline at end of file From c35b655402aa877504271554809000cbb4094ed1 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Tue, 28 Jan 2025 14:12:17 -0700 Subject: [PATCH 15/43] Changes to shift/atom correlation and data pickling --- energy.py | 11 +++-- fake_data.py | 94 +++++++++++++++++++++++++++---------------- nmr_gym_env.py | 70 +++++++++++++++++++++++--------- nmr_text_adventure.py | 82 +++++++++++++++++++------------------ 4 files changed, 159 insertions(+), 98 deletions(-) diff --git a/energy.py b/energy.py index 29cca26..d26f0e5 100644 --- a/energy.py +++ b/energy.py @@ -28,7 +28,7 @@ def matchNH(hshift, nshift, hsqc, tolerance_h, tolerance_n): return nh_index -def noe_combinations(noes, actual_shifts, tolerance_h=0.2, tolerance_n=0.2): +def noe_combinations(noes, actual_shifts, tolerance_h=0.02, tolerance_n=0.02): """ Loops through the NOES to match the H1, N1 and H2 shifts to two HSQC shifts. Lists the HSQC shifts separately (H1/N1 as contact1 and H2 as contact2) for each NOE. @@ -96,7 +96,7 @@ def calc_restraint_energy(self, assignments, restraint): restraint_energy = math.inf for i, j in restraint: k, l = assignments.get(i), assignments.get(j) - dist = np.linalg.norm((np.array(self.coords[k]) - np.array(self.coords[l]))) + dist = np.linalg.norm((np.array(self.coords[k][:3]) - np.array(self.coords[l][:3]))) energy_value = self.flat_bottom(dist, tolerance=0.5) restraint_energy = energy_value if energy_value < restraint_energy else restraint_energy return restraint_energy @@ -124,7 +124,10 @@ def get_total_energy(self, restraints, assignments): total_energy = 0 for restraint in restraints: - energy_value = self.noe_activation(assignments, restraint) - total_energy += energy_value + if restraint == []: + pass + else: + energy_value = self.noe_activation(assignments, restraint) + total_energy += energy_value return total_energy diff --git a/fake_data.py b/fake_data.py index cc56881..0fb32b1 100644 --- a/fake_data.py +++ b/fake_data.py @@ -3,8 +3,7 @@ from typing import NamedTuple import random import argparse - -from energy import Energy +import pickle class HSQCPeak(NamedTuple): H1: float @@ -19,6 +18,8 @@ class Protein(NamedTuple): x: float y: float z: float + H1: float + N15: float # Sampled points from unit square/cube def sample_unit(n, num_sides, min=0, max=1): @@ -36,11 +37,11 @@ def add_noise(point, scale=0.1, min=0, max=1): # fold over points outside of boundaries for point in noisy_point: - for value in point: + for i, value in enumerate(point): if max < value: - value = (max-(value-max)) - elif value < min: - value = (min-value) + point[i] = (max-(value-max)) + if value < min: + point[i] = (min-value) return noisy_point @@ -49,11 +50,14 @@ def calc_dist(p1, p2): def distance_noe(protein, shifts, cutoff): """ - Grabs close coordinates and 'associated' HSQC shift peaks (by same index) to create NOES. - Adds gaussian noise to all points. + Grabs close coordinates and 'associated' HSQC shift peaks (randomized index) to create NOES. + Predicted shifts assigned as the new randomized shift order - this order is associated with the coordinate order (ie. the answer key). + Adds gaussian noise to all points at the end. **Can end up with no NOEs depending on the cutoff** """ + randomized_shifts = np.random.permutation(shifts) + noes = [] for i, atom1 in enumerate(protein): @@ -62,15 +66,30 @@ def distance_noe(protein, shifts, cutoff): dist = calc_dist(atom1, atom2) if dist < cutoff: # print(dist, shifts[i], shifts[j]) - noe = list(shifts[i][:]) # H1, N1 - noe.append(shifts[j][0]) # H2 + # noe = list(shifts[i][:]) # H1, N1 + # noe.append(shifts[j][0]) # H2 + noe = list(randomized_shifts[i][:]) + noe.append(randomized_shifts[j][0]) noes.append(noe) noisy_noe = add_noise(np.array(noes), scale=0.01) + predicted_shifts = add_noise(np.array(randomized_shifts), scale=0.1) + + return noisy_noe, predicted_shifts + +def close_contacts(protein, cutoff): + contacts = [] + + for i, atom1 in enumerate(protein): + for j, atom2 in enumerate(protein): + if i != j: + dist = calc_dist(atom1, atom2) + if dist < cutoff: + contacts.append((i,j)) - return noisy_noe + return contacts -def generate_data(num_resid): +def generate_data(num_resid, pickle_data=True, example=True): """ Generates all fake data and orders it in lists of namedtuples. """ @@ -78,38 +97,43 @@ def generate_data(num_resid): protein = sample_unit(num_resid, num_sides=3) # "Actual" shifts [H1,N1] actual_shifts = sample_unit(num_resid, num_sides=2) - # Predicted shifts [H1,N1] - predicted_shifts = add_noise(actual_shifts) - # NOES [H1,N1,H2] - noes = distance_noe(protein, actual_shifts, cutoff=0.5) # likely need to change dist cutoff and only accounting for actual_shifts right now + # NOES [H1,N1,H2] and predicted shifts [H1,N1] + noes, predicted_shifts = distance_noe(protein, actual_shifts, cutoff=0.5) # Lists of namedtuples (one object per residue) - coords = [Protein(x=resid[0], y=resid[1], z=resid[2]) for resid in protein] + coords = [Protein(x=resid[0], y=resid[1], z=resid[2], H1=shift[0], N15=shift[1]) for resid, shift in zip(protein, predicted_shifts)] actual_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in actual_shifts] - predicted_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in predicted_shifts] noes = [NOEPeak(H1=shift[0], N15=shift[1], H2=shift[2]) for shift in noes] - # print(coords) - # print(actual_shifts) - # print(noe) + if pickle_data: + name = f'fakedata_r{num_resid}.pkl' if example else f'current_run.pkl' - return coords, actual_shifts, predicted_shifts, noes + with open(name, 'wb') as f: + pickle.dump(coords, f) + pickle.dump(actual_shifts, f) + pickle.dump(noes, f) + + else: + return coords, actual_shifts, noes # if __name__ == '__main__': -# parser = argparse.ArgumentParser() -# parser.add_argument('num_resid', type=int, help='Number of residues') -# args = parser.parse_args() +# parser = argparse.ArgumentParser() +# parser.add_argument('num_resid', type=int, help='Number of residues') +# args = parser.parse_args() + +# num_resid = args.num_resid +# generate_data(num_resid) -# num_resid = args.num_resid # # These should be grouped to export into an environment # coords, actual_shifts, predicted_shifts, noes = generate_data(num_resid) -# print(f'coords {(np.array(coords)).tolist()}\n shifts{(np.array(actual_shifts)).tolist()}\n noes{(np.array(noes)).tolist()}') -# # print(actual_shifts) -# # print(noes) - - -# energy_obj = Energy(coords, actual_shifts, noes) -# test = energy_obj.setup_noe_restraints() -# print(test) - # test_energy = energy_obj.energy_loop(test) + # print(f'coords {(np.array(coords)).tolist()}')#\n shifts{(np.array(actual_shifts)).tolist()}\n noes{(np.array(noes)).tolist()}') + # # print(actual_shifts) + # # print(noes) + + # energy_obj = Energy(coords, actual_shifts, noes) + # test = energy_obj.setup_noe_restraints() + # # print(test) + # # print(coords) + + # test_energy = energy_obj.get_total_energy(test, {0:2, 1:1, 2:0}) # print(test_energy) \ No newline at end of file diff --git a/nmr_gym_env.py b/nmr_gym_env.py index 8220414..906c89f 100644 --- a/nmr_gym_env.py +++ b/nmr_gym_env.py @@ -1,26 +1,32 @@ import gym from gym import error, spaces import numpy as np +import pickle +import glob as glob from fake_data import generate_data from energy import Energy, noe_combinations from assignment_order import FracAct +from visualize_data import plot_shifts + class GymEnv(gym.Env): def __init__(self, num_resid): + self.state = None self.num_resid = num_resid + self.assignments = {} self.assign_step = 0 - self.intermediate_energy = 0 - self.reward = 0 + self.total_energy = 0.0 + self.intermediate_energy = 0.0 + self.reward = 0.0 self.action_space = spaces.Discrete(self.num_resid) self.observation_space = spaces.Dict({ - "coords": spaces.Box(0, 1, shape=(self.num_resid, 3), dtype=np.float32), + "coords": spaces.Box(0, 1, shape=(self.num_resid, 5), dtype=np.float32), "actual_shifts": spaces.Box(0, 1, shape=(self.num_resid, 2), dtype=np.float32), - "predicted_shifts": spaces.Box(0, 1, shape=(self.num_resid, 2), dtype=np.float32), "noes": spaces.Box(0, 1, shape=(self.num_resid, 3), dtype=np.float32) }) @@ -29,39 +35,60 @@ def custom_state(self, state): frac_act = FracAct(self.num_resid, self.state['restraints']) self.assign_order = frac_act.fractional_activation() print(self.assign_order) + + # get plot of shifts and restraints + plot_shifts(self.state['actual_shifts'], self.state['restraints'], self.state['coords'], named_tuple_used=False) + return self.state - def reset(self): - # generate fake data - coords, actual_shifts, predicted_shifts, noes = generate_data(self.num_resid) + def reset(self, pickled=False, pickle_data=True, example=True): + + name = f"./*{self.num_resid}.pkl" if example else "./current_run.pkl" + + if self.intermediate_energy > 0 or pickled: + pickle_file = glob.glob(name) + with open(pickle_file[0], 'rb') as f: + coords = pickle.load(f) + actual_shifts = pickle.load(f) + noes = pickle.load(f) + else: + generate_data(self.num_resid, pickle_data=pickle_data, example=example) + pickle_file = glob.glob(name) + with open(pickle_file[0], 'rb') as f: + coords = pickle.load(f) + actual_shifts = pickle.load(f) + noes = pickle.load(f) + # get restraints restraints = noe_combinations(noes, actual_shifts) + # get plot of shifts and restraints + plot_shifts(actual_shifts, restraints, coords, named_tuple_used=True) + # get assignment order frac_act = FracAct(self.num_resid, restraints) self.assign_order = frac_act.fractional_activation() print(self.assign_order) self.assignments = {} - # self.assign_step = 0 - # self.intermediate_energy = 0 - self.total_energy = 0 - # self.reward = 0 + self.assign_step = 0 + self.total_energy = 0.0 + self.intermediate_energy = 0.0 + self.reward = 0.0 self.state = { "coords": coords, "actual_shifts": actual_shifts, - "predicted_shifts": predicted_shifts, "noes": noes, "restraints": restraints, "assignments": self.assignments, - "total_energy": self.total_energy + "total_energy": self.total_energy, + "reward": self.reward } return self.state def step(self, action): - energy = Energy(self.state['coords'], self.state['actual_shifts'], self.state['noes']) # add action to assignments @@ -69,11 +96,14 @@ def step(self, action): print(self.assignments) # calc energy and store temporarily - temp_intermediate_energy = energy.get_energy(self.state['restraints'], self.assignments) + temp_intermediate_energy = energy.get_total_energy(self.state['restraints'], self.assignments) # print(f'this is the temp energy {temp_intermediate_energy}') - # calc reward based on previous energy and temporary energy (+ve means energy went down, -ve means energy went up) - self.reward = (self.intermediate_energy - temp_intermediate_energy) #if self.intermediate_energy != 0 else 0 + # calc reward based on previous energy and temporary energy + # before correction, +ve means energy went down - we DO NOT want this, -ve means energy went up + self.state['reward'] = (self.intermediate_energy - temp_intermediate_energy) + assert self.state['reward'] <= 0 + self.state['reward'] = abs(self.state['reward']) # print(f'this is the reward {self.reward}') # store energy as intermediate @@ -81,15 +111,15 @@ def step(self, action): # assign intermediate energy as running total self.state['total_energy'] = self.intermediate_energy - print(f"Running energy = {self.state['total_energy']}, Current reward = {self.reward}") + print(f"Running energy = {self.state['total_energy']}, Current reward = {self.state['reward']}") self.assign_step += 1 terminated = True if len(self.assignments) == self.num_resid else False if terminated: print(f"Final assignments (shift:atom) = {self.assignments}\nFinal energy evaluation = {self.state['total_energy']}") - return self.state, self.reward, terminated + return self.state, self.state['reward'], terminated, self.state['total_energy'] else: - return self.state, self.reward, terminated + return self.state, self.state['reward'], terminated, self.state['total_energy'] diff --git a/nmr_text_adventure.py b/nmr_text_adventure.py index e851a52..718ac44 100644 --- a/nmr_text_adventure.py +++ b/nmr_text_adventure.py @@ -1,24 +1,51 @@ import gym +import math import argparse from nmr_gym_env import GymEnv if __name__ == '__main__': - # parser = argparse.ArgumentParser() - # parser.add_argument('num_resid', type=int, help='Number of residues') - # args = parser.parse_args() + + parser = argparse.ArgumentParser() + parser.add_argument('num_resid', type=int, help='Number of residues') + args = parser.parse_args() - # num_resid = args.num_resid - num_resid = 10 + num_resid = args.num_resid + gym_env = GymEnv(num_resid) - episodes = 1 + total_energy = math.inf - for episode in range(episodes): - # observation = gym_env.reset() + while total_energy > 0: - # random generated set of 4 + ###### Available options to work with right now ###### + + # 1. Need to generate and save an example? Saved in all instances as fakedata_r{resid_num}.pkl + # observation = gym_env.reset(pickled=False, pickle_data=True) + + # 2. Grab one of these previous example? Make sure example exists, with desired resid number. + # observation = gym_env.reset(pickled=True, pickle_data=False) + + # 3. Run a new example? Saved in all instances as current_run.pkl + observation = gym_env.reset(pickled=False, pickle_data=True, example=False) + + # 4. Use a custom hardcoded state? May be some issues with data types and visualization. + # observation = gym_env.custom_state(state) + + ###################################################### + + terminated = False + action_assigned = [] + + while not terminated: + action = int(input()) + if action < num_resid: + if action not in action_assigned: + action_assigned.append(action) + observation, reward, terminated, total_energy = gym_env.step(action) + +# random generated set of 4 # state = { - # "coords": [[0.35, 0.83, 0.89], [0.69, 0.12, 0.92], [0.05, 0.94, 0.51], [0.45, 0.06, 0.70]], + # "coords": [[0.35, 0.83, 0.89, 0.16, 0.034], [0.69, 0.12, 0.92, 0.41, 0.44], [0.05, 0.94, 0.51, 0.189, 0.62], [0.45, 0.06, 0.70, 0.63, 0.15]], # "actual_shifts": [[0.17, 0.044], [0.41, 0.34], [0.18, 0.61], [0.64, 0.23]], # "noes": [[0.16, 0.040, 0.17], [0.40, 0.34, 0.63], [0.18, 0.61, 0.16], [0.64, 0.24, 0.39]], # "restraints": [[(0, 2)], [(1, 3)], [(0, 2)], [(1, 3)]], @@ -27,35 +54,12 @@ # } # random generated set of 10 with many restraints - state = { - "coords": [[0.8136200069440616, 0.7789831346675045, 0.6474707035745629], [0.9458251909740947, 0.8744248260731231, 0.16939643090298961], [0.07941814153270477, 0.8569710876203085, 0.6679287503567067], [0.9295839450890829, 0.2591652267820922, 0.5767903211938266], [0.5686652108970609, 0.6600261693178074, 0.8353850759130969], [0.3049822956354953, 0.01597728824712763, 0.8007859163237383], [0.2707959067671988, 0.16592702782732538, 0.8530675522602625], [0.5906467171408669, 0.2275483180754293, 0.1179098748663463], [0.8089112838766771, 0.014346015197136075, 0.2651392664874339], [0.021085864510464902, 0.3582433470216929, 0.9002893595219581]], - "actual_shifts": [[0.7235029227185441, 0.8167276537744274], [0.44118846730192773, 0.8390592014969155], [0.3607330934400178, 0.813863583391123], [0.35720641978511525, 0.2931243342244849], [0.03585880729962143, 0.6045433843661965], [0.13284379187386453, 0.7429892824868469], [0.8157925033719939, 0.5225148209456146], [0.3240407057810011, 0.9258723658454763], [0.35659118663170697, 0.4579588797141084], [0.6229950482718816, 0.9870152102924102]], - "noes": [[0.7353272235348708, 0.7985692976248703, 0.024966821752980907], [0.3602327190429789, 0.27520736325235085, 0.3444804557069924], [0.030177646349364318, 0.6038110746526516, 0.7374671133868814], [0.1147855723335682, 0.7472335679962178, 0.8334688788656137], [0.14395802509472194, 0.7602008741457149, 0.6270020652720573], [0.8084393678158522, 0.49689451509103233, 0.1359092345145138], [0.8359026620374025, 0.5227806613419343, 0.6230942517135893], [0.3158886510580014, 0.9127784810952192, 0.36108155258485086], [0.35718811981030113, 0.46209771728576327, 0.36091347433509957], [0.3473018809331518, 0.45763566566297215, 0.3106537228148144], [0.63688033551314, 0.9916090993120613, 0.12731578248735478], [0.6196043419116839, 0.9903515010993935, 0.808008737938382]], - # "restraints": [[(4, 9), (5, 9), (0, 4), (0, 5)], [(3, 8), (3, 7), (1, 8), (2, 3), (1, 3), (7, 8), (2, 8)], [(0, 4), (4, 9), (4, 6), (5, 6), (0, 5), (5, 9)], [(4, 6), (0, 4), (5, 6), (0, 5)], [(0, 7), (0, 4), (1, 5), (4, 9), (4, 6), (1, 4), (7, 9), (6, 7), (1, 7), (5, 6), (0, 5), (5, 9)], [(6, 7), (4, 6), (5, 6)], [(1, 6), (6, 9), (0, 6)], [(1, 2), (2, 7), (1, 5), (5, 8), (3, 7), (1, 8), (5, 7), (2, 3), (1, 7), (7, 8), (2, 5), (1, 3), (3, 5), (2, 8)], [(3, 8), (3, 7), (1, 8), (2, 3), (1, 3), (7, 8), (2, 8)], [(3, 8), (5, 8), (3, 7), (1, 8), (2, 3), (7, 8), (1, 3), (3, 5), (2, 8)], [(0, 7), (0, 4), (1, 5), (4, 9), (1, 4), (7, 9), (1, 7), (0, 5), (5, 9)], [(0, 1), (0, 9), (0, 6), (1, 6), (6, 9), (1, 9)]], - "restraints": [[(3, 8), (5, 9)], [(5, 6), (0, 5), (5, 9)], [(0, 4), (4, 9), (4, 6), (5, 6)], [(3, 8), (3, 7), (1, 8)], [(1, 6), (6, 9)]], # there's something wrong with too many restraints, - "assignments": {}, - "total_energy": 0 - } - - # linear arrangement of 10 with easy restraints # state = { - # "coords": [[0.37, 0.82, 0.00], [0.37, 0.82, 0.11], [0.37, 0.82, 0.24], [0.37, 0.82, 0.39], [0.37, 0.82, 0.52], [0.37, 0.82, 0.67], [0.37, 0.82, 0.75], [0.37, 0.82, 0.88], [0.37, 0.82, 0.93], [0.37, 0.82, 1.00]], - # "actual_shifts": [[0.6731658491622945, 0.657184525800008], [0.28748044870745615, 0.34032397854002194], [0.4921807498359013, 0.03828350226565602], [0.7763637896667895, 0.3617894550491416], [0.06840333800729603, 0.20738485015436503], [0.9913658149098207, 0.33931570183160287], [0.5459835943391216, 0.8469814942099086], [0.09701544891895453, 0.7197743401411141], [0.5107870945621455, 0.7086010383972177], [0.5738658802959612, 0.31748338342430493]], - # "noes": [[0.7028015287380062, 0.5021893114424876], [0.33894896143919023, 0.45179229762303535], [0.46784325494458656, -0.0019242124744091177], [0.7161042434070474, 0.3465503455737962], [0.00801328834772183, 0.3164743274291713], [0.9337074775227354, 0.3689512869911665], [0.47077422692743703, 0.8353558577228544], [0.2038502444268092, 0.5189774171248352], [0.50122164179375, 0.5801271184891338], [0.5307875874176243, 0.2934724693638126]], - # "restraints": [[(0, 1)], [(2, 3)], [(4, 5)], [(6, 7)], [(8, 9)]], + # "coords": [[0.8136200069440616, 0.7789831346675045, 0.6474707035745629, 0.8235029227185441, 0.6167276537744274], [0.9458251909740947, 0.8744248260731231, 0.16939643090298961, 0.54118846730192773, 0.7390592014969155], [0.07941814153270477, 0.8569710876203085, 0.6679287503567067, 0.3807330934400178, 0.713863583391123], [0.9295839450890829, 0.2591652267820922, 0.5767903211938266, 0.43720641978511525, 0.3031243342244849], [0.5686652108970609, 0.6600261693178074, 0.8353850759130969, 0.01585880729962143, 0.6545433843661965], [0.3049822956354953, 0.01597728824712763, 0.8007859163237383, 0.15284379187386453, 0.7329892824868469], [0.2707959067671988, 0.16592702782732538, 0.8530675522602625, 0.8257925033719939, 0.6225148209456146], [0.5906467171408669, 0.2275483180754293, 0.1179098748663463, 0.4240407057810011, 0.9558723658454763], [0.8089112838766771, 0.014346015197136075, 0.2651392664874339, 0.45659118663170697, 0.6579588797141084], [0.021085864510464902, 0.3582433470216929, 0.9002893595219581, 0.5229950482718816, 0.9970152102924102]], + # "actual_shifts": [[0.7235029227185441, 0.8167276537744274], [0.44118846730192773, 0.8390592014969155], [0.3607330934400178, 0.813863583391123], [0.35720641978511525, 0.2931243342244849], [0.03585880729962143, 0.6045433843661965], [0.13284379187386453, 0.7429892824868469], [0.8157925033719939, 0.5225148209456146], [0.3240407057810011, 0.9258723658454763], [0.35659118663170697, 0.4579588797141084], [0.6229950482718816, 0.9870152102924102]], + # "noes": [[0.7353272235348708, 0.7985692976248703, 0.024966821752980907], [0.3602327190429789, 0.27520736325235085, 0.3444804557069924], [0.030177646349364318, 0.6038110746526516, 0.7374671133868814], [0.1147855723335682, 0.7472335679962178, 0.8334688788656137], [0.14395802509472194, 0.7602008741457149, 0.6270020652720573], [0.8084393678158522, 0.49689451509103233, 0.1359092345145138], [0.8359026620374025, 0.5227806613419343, 0.6230942517135893], [0.3158886510580014, 0.9127784810952192, 0.36108155258485086], [0.35718811981030113, 0.46209771728576327, 0.36091347433509957], [0.3473018809331518, 0.45763566566297215, 0.3106537228148144], [0.63688033551314, 0.9916090993120613, 0.12731578248735478], [0.6196043419116839, 0.9903515010993935, 0.808008737938382]], + # # "restraints": [[(4, 9), (5, 9), (0, 4), (0, 5)], [(3, 8), (3, 7), (1, 8), (2, 3), (1, 3), (7, 8), (2, 8)], [(0, 4), (4, 9), (4, 6), (5, 6), (0, 5), (5, 9)], [(4, 6), (0, 4), (5, 6), (0, 5)], [(0, 7), (0, 4), (1, 5), (4, 9), (4, 6), (1, 4), (7, 9), (6, 7), (1, 7), (5, 6), (0, 5), (5, 9)], [(6, 7), (4, 6), (5, 6)], [(1, 6), (6, 9), (0, 6)], [(1, 2), (2, 7), (1, 5), (5, 8), (3, 7), (1, 8), (5, 7), (2, 3), (1, 7), (7, 8), (2, 5), (1, 3), (3, 5), (2, 8)], [(3, 8), (3, 7), (1, 8), (2, 3), (1, 3), (7, 8), (2, 8)], [(3, 8), (5, 8), (3, 7), (1, 8), (2, 3), (7, 8), (1, 3), (3, 5), (2, 8)], [(0, 7), (0, 4), (1, 5), (4, 9), (1, 4), (7, 9), (1, 7), (0, 5), (5, 9)], [(0, 1), (0, 9), (0, 6), (1, 6), (6, 9), (1, 9)]], + # "restraints": [[(3, 8), (5, 9)], [(5, 6), (0, 5), (5, 9)], [(0, 4), (4, 9), (4, 6), (5, 6)], [(3, 8), (3, 7), (1, 8)], [(1, 6), (6, 9)]], # there's something wrong with too many restraints, # "assignments": {}, # "total_energy": 0 - # } - - observation = gym_env.custom_state(state) - - terminated = False - action_assigned = [] - - while not terminated: - # action = gym_env.action_space.sample() - action = int(input()) - if action < num_resid: - if action not in action_assigned: - action_assigned.append(action) - observation, reward, terminated = gym_env.step(action) + # } From 7f25ea45242c379235a9dcc90c057e93569ed2e9 Mon Sep 17 00:00:00 2001 From: a-angel-s Date: Fri, 14 Feb 2025 14:50:30 -0700 Subject: [PATCH 16/43] added connectivity and updated generate_data --- fake_data.py | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/fake_data.py b/fake_data.py index 0fb32b1..187c8f1 100644 --- a/fake_data.py +++ b/fake_data.py @@ -21,6 +21,11 @@ class Protein(NamedTuple): H1: float N15: float +class Connectivity(NamedTuple): + atom1: float + atom2: float + distance: float + # Sampled points from unit square/cube def sample_unit(n, num_sides, min=0, max=1): return np.random.uniform(min, max, size=(n, num_sides)) @@ -89,6 +94,18 @@ def close_contacts(protein, cutoff): return contacts +def connectivity_data(protein, cutoff): + connectivity = [] + + for i, atom1 in enumerate(protein): + for j, atom2 in enumerate(protein): + if i != j: + dist = calc_dist(atom1, atom2) #dist = np.linalg.norm(atom1 - atom2) + if dist < cutoff: + connectivity.append((atom1, atom2, dist)) + + return connectivity + def generate_data(num_resid, pickle_data=True, example=True): """ Generates all fake data and orders it in lists of namedtuples. @@ -99,11 +116,14 @@ def generate_data(num_resid, pickle_data=True, example=True): actual_shifts = sample_unit(num_resid, num_sides=2) # NOES [H1,N1,H2] and predicted shifts [H1,N1] noes, predicted_shifts = distance_noe(protein, actual_shifts, cutoff=0.5) + # connectivity [atom1,atom2,dist] + connectivity = connectivity_data(protein, cutoff=0.5) # Lists of namedtuples (one object per residue) coords = [Protein(x=resid[0], y=resid[1], z=resid[2], H1=shift[0], N15=shift[1]) for resid, shift in zip(protein, predicted_shifts)] actual_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in actual_shifts] noes = [NOEPeak(H1=shift[0], N15=shift[1], H2=shift[2]) for shift in noes] + connectivity = [Connectivity(atom1=connect[0], atom2=connect[1], distance=connect[2]) for connect in connectivity] if pickle_data: name = f'fakedata_r{num_resid}.pkl' if example else f'current_run.pkl' @@ -112,9 +132,10 @@ def generate_data(num_resid, pickle_data=True, example=True): pickle.dump(coords, f) pickle.dump(actual_shifts, f) pickle.dump(noes, f) + pickle.dump(connectivity, f) else: - return coords, actual_shifts, noes + return coords, actual_shifts, noes, connectivity # if __name__ == '__main__': # parser = argparse.ArgumentParser() From d2c318094710e99cb620f98855a17a22b62352d6 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Thu, 20 Feb 2025 21:45:21 -0700 Subject: [PATCH 17/43] Added fake histories and distributed connectivity param --- fake_data.py | 38 +++++++---------- fake_histories.py | 101 ++++++++++++++++++++++++++++++++++++++++++++++ nmr_gym_env.py | 26 ++++++++---- visualize_data.py | 36 ++++++++--------- 4 files changed, 151 insertions(+), 50 deletions(-) create mode 100644 fake_histories.py diff --git a/fake_data.py b/fake_data.py index 187c8f1..3106997 100644 --- a/fake_data.py +++ b/fake_data.py @@ -61,7 +61,8 @@ def distance_noe(protein, shifts, cutoff): **Can end up with no NOEs depending on the cutoff** """ - randomized_shifts = np.random.permutation(shifts) + # Can use this to randomize - not required + # shifts = np.random.permutation(shifts) noes = [] @@ -71,38 +72,27 @@ def distance_noe(protein, shifts, cutoff): dist = calc_dist(atom1, atom2) if dist < cutoff: # print(dist, shifts[i], shifts[j]) - # noe = list(shifts[i][:]) # H1, N1 - # noe.append(shifts[j][0]) # H2 - noe = list(randomized_shifts[i][:]) - noe.append(randomized_shifts[j][0]) + noe = list(shifts[i][:]) # H1, N1 + noe.append(shifts[j][0]) # H2 noes.append(noe) noisy_noe = add_noise(np.array(noes), scale=0.01) - predicted_shifts = add_noise(np.array(randomized_shifts), scale=0.1) + predicted_shifts = add_noise(np.array(shifts), scale=0.1) return noisy_noe, predicted_shifts -def close_contacts(protein, cutoff): - contacts = [] - - for i, atom1 in enumerate(protein): - for j, atom2 in enumerate(protein): - if i != j: - dist = calc_dist(atom1, atom2) - if dist < cutoff: - contacts.append((i,j)) - - return contacts - def connectivity_data(protein, cutoff): + """ + Calculates close contacts based on coordinates. + """ connectivity = [] for i, atom1 in enumerate(protein): - for j, atom2 in enumerate(protein): - if i != j: - dist = calc_dist(atom1, atom2) #dist = np.linalg.norm(atom1 - atom2) - if dist < cutoff: - connectivity.append((atom1, atom2, dist)) + for j, atom2 in enumerate(protein): + if i != j: + dist = calc_dist(atom1, atom2) #dist = np.linalg.norm(atom1 - atom2) + if dist < cutoff: + connectivity.append((i, j, dist)) return connectivity @@ -117,7 +107,7 @@ def generate_data(num_resid, pickle_data=True, example=True): # NOES [H1,N1,H2] and predicted shifts [H1,N1] noes, predicted_shifts = distance_noe(protein, actual_shifts, cutoff=0.5) # connectivity [atom1,atom2,dist] - connectivity = connectivity_data(protein, cutoff=0.5) + connectivity = connectivity_data(protein, cutoff=0.37) # Lists of namedtuples (one object per residue) coords = [Protein(x=resid[0], y=resid[1], z=resid[2], H1=shift[0], N15=shift[1]) for resid, shift in zip(protein, predicted_shifts)] diff --git a/fake_histories.py b/fake_histories.py new file mode 100644 index 0000000..a691ec2 --- /dev/null +++ b/fake_histories.py @@ -0,0 +1,101 @@ +import numpy as np +from fake_data import * +from energy import noe_combinations + +def perturb_data(coords): + """ + Adds noise to original fake data and generates new 'actual' shifts based on predicted shifts. + """ + protein = add_noise(np.array(coords)[:, :3], scale=0.1) + predicted_shifts = add_noise(np.array(coords)[:, 3:], scale=0.1) + actual_fake_shifts = add_noise(np.array(coords)[:, 3:], scale=0.1) + + return protein, predicted_shifts, actual_fake_shifts + +def recalculate_noe(protein, shifts, shift_order, cutoff): + """ + Grabs close coordinates and 'associated' HSQC shift peaks to create new NOES. + Adds gaussian noise to all points at the end. + """ + # Can use this to randomize if it was used in the original + # shifts = [shifts[i] for i in shift_order.values()] + + noes = [] + + for i, atom1 in enumerate(protein): + for j, atom2 in enumerate(protein): + if i != j: + dist = calc_dist(atom1, atom2) + if dist < cutoff: + # print(dist, shifts[i], shifts[j]) + noe = list(shifts[i][:]) # H1, N1 + noe.append(shifts[j][0]) # H2 + noes.append(noe) + + noisy_noe = add_noise(np.array(noes), scale=0.01) + + return noisy_noe + +def weighted_error(num_resid, scale=1): + """ + Generates weights for number of errors added based on exponential distribution. + Chooses error number to add based on weights. + """ + # Weight each point + weights = np.exp(-np.arange(num_resid)/scale) + # Normalized + weights = weights/np.sum(weights) + # Number of errors based on weight + error_num = np.random.choice(num_resid, 1, p=weights) + + # Can't make a singular error - needs to be replaced by another choice + return error_num if error_num > 1 else error_num*2 + +def add_errors(answers): + """ + If errors are present chooses that number of random shifts and moves answers one over from true. + """ + error_num = weighted_error(len(answers), scale=1) + + if error_num == 0: + return answers + else: + selections = np.random.choice(list(answers.values()), error_num, replace=False) + mutations = np.roll(selections, 1) + for i, j in enumerate(selections): + answers[j] = mutations[i] + + return answers + +def generate_history(original, history_len=5): + """ + Loops over chosen history length, adding noise to original data, reorganizes, and edits answer key if errors are introduced. + """ + answers = {} + history = [] + + for i in range(history_len): + protein, predicted_shifts, actual_shifts = perturb_data(original['coords']) + noes = recalculate_noe(protein, actual_shifts, original['assignments'], cutoff=0.5) + connectivity = connectivity_data(protein, cutoff=0.37) + + # Lists of namedtuples (one object per residue) + coords = [Protein(x=resid[0], y=resid[1], z=resid[2], H1=shift[0], N15=shift[1]) for resid, shift in zip(protein, predicted_shifts)] + actual_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in actual_shifts] + noes = [NOEPeak(H1=shift[0], N15=shift[1], H2=shift[2]) for shift in noes] + connectivity = [Connectivity(atom1=connect[0], atom2=connect[1], distance=connect[2]) for connect in connectivity] + + restraints = noe_combinations(noes, actual_shifts) + + history.append(coords) + history.append(actual_shifts) + history.append(noes) + history.append(restraints) + history.append(connectivity) + + answers[i] = dict(original['assignments']) + answers[i] = add_errors(answers[i]) + + return history, answers + + \ No newline at end of file diff --git a/nmr_gym_env.py b/nmr_gym_env.py index 906c89f..9e687d5 100644 --- a/nmr_gym_env.py +++ b/nmr_gym_env.py @@ -18,6 +18,7 @@ def __init__(self, num_resid): self.assignments = {} self.assign_step = 0 + self.assign_shift = 0 self.total_energy = 0.0 self.intermediate_energy = 0.0 self.reward = 0.0 @@ -51,6 +52,7 @@ def reset(self, pickled=False, pickle_data=True, example=True): coords = pickle.load(f) actual_shifts = pickle.load(f) noes = pickle.load(f) + connectivity = pickle.load(f) else: generate_data(self.num_resid, pickle_data=pickle_data, example=example) pickle_file = glob.glob(name) @@ -58,24 +60,28 @@ def reset(self, pickled=False, pickle_data=True, example=True): coords = pickle.load(f) actual_shifts = pickle.load(f) noes = pickle.load(f) + connectivity = pickle.load(f) # get restraints restraints = noe_combinations(noes, actual_shifts) # get plot of shifts and restraints - plot_shifts(actual_shifts, restraints, coords, named_tuple_used=True) - - # get assignment order - frac_act = FracAct(self.num_resid, restraints) - self.assign_order = frac_act.fractional_activation() - print(self.assign_order) + plot_shifts(actual_shifts, restraints, coords, connectivity, named_tuple_used=True) self.assignments = {} self.assign_step = 0 + self.assign_shift = 0 self.total_energy = 0.0 self.intermediate_energy = 0.0 self.reward = 0.0 + # get assignment order + frac_act = FracAct(self.num_resid, restraints) + self.assign_order = frac_act.fractional_activation() + self.assign_shift = self.assign_order[self.assign_step] + print(self.assign_order) + print(f'Shift to be assigned: {self.assign_shift}') + self.state = { "coords": coords, "actual_shifts": actual_shifts, @@ -100,8 +106,8 @@ def step(self, action): # print(f'this is the temp energy {temp_intermediate_energy}') # calc reward based on previous energy and temporary energy - # before correction, +ve means energy went down - we DO NOT want this, -ve means energy went up self.state['reward'] = (self.intermediate_energy - temp_intermediate_energy) + # before correction, +ve means energy went down - we DO NOT want this, -ve means energy went up assert self.state['reward'] <= 0 self.state['reward'] = abs(self.state['reward']) # print(f'this is the reward {self.reward}') @@ -114,12 +120,16 @@ def step(self, action): print(f"Running energy = {self.state['total_energy']}, Current reward = {self.state['reward']}") self.assign_step += 1 + # self.assign_shift = self.assign_order[self.assign_step] terminated = True if len(self.assignments) == self.num_resid else False if terminated: print(f"Final assignments (shift:atom) = {self.assignments}\nFinal energy evaluation = {self.state['total_energy']}") return self.state, self.state['reward'], terminated, self.state['total_energy'] - else: + + self.assign_shift = self.assign_order[self.assign_step] + if not terminated: + print(f'Shift to be assigned: {self.assign_shift}') return self.state, self.state['reward'], terminated, self.state['total_energy'] diff --git a/visualize_data.py b/visualize_data.py index c7a669e..fbb0f1a 100644 --- a/visualize_data.py +++ b/visualize_data.py @@ -3,7 +3,7 @@ import numpy as np from mpl_toolkits.mplot3d import Axes3D -def get_data(actual_shifts, coords, named_tuple_used=True): +def get_data(actual_shifts, coords, connectivity, named_tuple_used=True): if named_tuple_used: # Actual actual_x = [shift.H1 for shift in actual_shifts] @@ -11,6 +11,8 @@ def get_data(actual_shifts, coords, named_tuple_used=True): # Predicted predict_x = [shift.H1 for shift in coords] predict_y = [shift.N15 for shift in coords] + # Connectivity + contacts = [(contact.atom1, contact.atom2) for contact in connectivity] else: # Actual @@ -20,7 +22,7 @@ def get_data(actual_shifts, coords, named_tuple_used=True): predict_x = [shift[3] for shift in coords] predict_y = [shift[4] for shift in coords] - return actual_x, actual_y, predict_x, predict_y + return actual_x, actual_y, predict_x, predict_y, contacts # def new_plot(coords): @@ -44,24 +46,23 @@ def get_data(actual_shifts, coords, named_tuple_used=True): # plt.show() # #plt.savefig("coord_plot.png", dpi = 150) -def close_contacts(coords, cutoff=0.37): - """ - Calculates close contacts based on coordinates. - """ - contacts = [] +# def close_contacts(coords, cutoff=0.37): +# """ +# Calculates close contacts based on coordinates. +# """ +# contacts = [] - for i, atom1 in enumerate(coords): - for j, atom2 in enumerate(coords): - if i != j: - dist = np.linalg.norm((np.array(atom1[:3]) - np.array(atom2[:3]))) - if dist < cutoff: - contacts.append((i,j)) +# for i, atom1 in enumerate(coords): +# for j, atom2 in enumerate(coords): +# if i != j: +# dist = np.linalg.norm((np.array(atom1[:3]) - np.array(atom2[:3]))) +# if dist < cutoff: +# contacts.append((i,j)) +# return contacts - return contacts +def plot_shifts(actual_shifts, restraints, coords, connectivity, named_tuple_used=True): -def plot_shifts(actual_shifts, restraints, coords, named_tuple_used=True): - - actual_x, actual_y, predict_x, predict_y = get_data(actual_shifts, coords, named_tuple_used=named_tuple_used) + actual_x, actual_y, predict_x, predict_y, contacts = get_data(actual_shifts, coords, connectivity, named_tuple_used=named_tuple_used) plt.scatter(predict_x, predict_y, color='blue', label='Predicted shifts') @@ -78,7 +79,6 @@ def plot_shifts(actual_shifts, restraints, coords, named_tuple_used=True): plt.plot((actual_x[j], actual_x[k]), (actual_y[j], actual_y[k]), color='red') # Draw close contacts - contacts = close_contacts(coords) for i,j in contacts: plt.plot((predict_x[i], predict_x[j]), (predict_y[i], predict_y[j]), color='blue', ls='dotted', alpha=0.5) From d5d02c08d24b5a8319290533ba9947a51d26d470 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Wed, 5 Mar 2025 12:23:58 -0700 Subject: [PATCH 18/43] Added agent to generate synthetic histories, and distributed some functions --- fake_data.py | 49 +++++++++++++++++++++++--------- fake_histories.py | 63 +++++++++++------------------------------ fake_histories_agent.py | 44 ++++++++++++++++++++++++++++ nmr_gym_env.py | 26 ++++++++++++----- nmr_text_adventure.py | 9 +++--- 5 files changed, 119 insertions(+), 72 deletions(-) create mode 100644 fake_histories_agent.py diff --git a/fake_data.py b/fake_data.py index 3106997..0c55766 100644 --- a/fake_data.py +++ b/fake_data.py @@ -77,9 +77,10 @@ def distance_noe(protein, shifts, cutoff): noes.append(noe) noisy_noe = add_noise(np.array(noes), scale=0.01) - predicted_shifts = add_noise(np.array(shifts), scale=0.1) + # predicted_shifts = add_noise(np.array(shifts), scale=0.1) - return noisy_noe, predicted_shifts + # return noisy_noe, predicted_shifts + return noisy_noe def connectivity_data(protein, cutoff): """ @@ -96,33 +97,55 @@ def connectivity_data(protein, cutoff): return connectivity -def generate_data(num_resid, pickle_data=True, example=True): +def generate_raw_data(num_resid): """ - Generates all fake data and orders it in lists of namedtuples. + Generates all qualities of the system as individual arrays. """ # 3D structure [x,y,z] protein = sample_unit(num_resid, num_sides=3) # "Actual" shifts [H1,N1] actual_shifts = sample_unit(num_resid, num_sides=2) + # Predicted shifts [H1,N1] + predicted_shifts = add_noise(actual_shifts, scale=0.1) # NOES [H1,N1,H2] and predicted shifts [H1,N1] - noes, predicted_shifts = distance_noe(protein, actual_shifts, cutoff=0.5) + noes = distance_noe(protein, actual_shifts, cutoff=0.5) # connectivity [atom1,atom2,dist] connectivity = connectivity_data(protein, cutoff=0.37) - # Lists of namedtuples (one object per residue) + return protein, actual_shifts, predicted_shifts, noes, connectivity + +def order_data(protein, actual_shifts, predicted_shifts, noes, connectivity): + """ + Compiles data in a list of named tuples with one object per residue. + """ coords = [Protein(x=resid[0], y=resid[1], z=resid[2], H1=shift[0], N15=shift[1]) for resid, shift in zip(protein, predicted_shifts)] actual_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in actual_shifts] noes = [NOEPeak(H1=shift[0], N15=shift[1], H2=shift[2]) for shift in noes] connectivity = [Connectivity(atom1=connect[0], atom2=connect[1], distance=connect[2]) for connect in connectivity] - if pickle_data: - name = f'fakedata_r{num_resid}.pkl' if example else f'current_run.pkl' + return coords, actual_shifts, noes, connectivity + +def save_as_pickle(coords, actual_shifts, noes, connectivity, num_resid, example=True): + """ + Saves data to disk. Run is given a proper name if used as a saved example. + """ + name = f'fakedata_r{num_resid}.pkl' if example else f'current_run.pkl' + + with open(name, 'wb') as f: + pickle.dump(coords, f) + pickle.dump(actual_shifts, f) + pickle.dump(noes, f) + pickle.dump(connectivity, f) - with open(name, 'wb') as f: - pickle.dump(coords, f) - pickle.dump(actual_shifts, f) - pickle.dump(noes, f) - pickle.dump(connectivity, f) +def generate_data(num_resid, pickle_data=True, example=True): + """ + Generates all fake data, orders it in lists of namedtuples, and pickles if required. + """ + protein, actual_shifts, predicted_shifts, noes, connectivity = generate_raw_data(num_resid) + coords, actual_shifts, noes, connectivity = order_data(protein, actual_shifts, predicted_shifts, noes, connectivity) + + if pickle_data: + save_as_pickle(coords, actual_shifts, noes, connectivity, num_resid, example=example) else: return coords, actual_shifts, noes, connectivity diff --git a/fake_histories.py b/fake_histories.py index a691ec2..1d59450 100644 --- a/fake_histories.py +++ b/fake_histories.py @@ -12,30 +12,6 @@ def perturb_data(coords): return protein, predicted_shifts, actual_fake_shifts -def recalculate_noe(protein, shifts, shift_order, cutoff): - """ - Grabs close coordinates and 'associated' HSQC shift peaks to create new NOES. - Adds gaussian noise to all points at the end. - """ - # Can use this to randomize if it was used in the original - # shifts = [shifts[i] for i in shift_order.values()] - - noes = [] - - for i, atom1 in enumerate(protein): - for j, atom2 in enumerate(protein): - if i != j: - dist = calc_dist(atom1, atom2) - if dist < cutoff: - # print(dist, shifts[i], shifts[j]) - noe = list(shifts[i][:]) # H1, N1 - noe.append(shifts[j][0]) # H2 - noes.append(noe) - - noisy_noe = add_noise(np.array(noes), scale=0.01) - - return noisy_noe - def weighted_error(num_resid, scale=1): """ Generates weights for number of errors added based on exponential distribution. @@ -51,51 +27,44 @@ def weighted_error(num_resid, scale=1): # Can't make a singular error - needs to be replaced by another choice return error_num if error_num > 1 else error_num*2 -def add_errors(answers): +def add_errors(answer): """ - If errors are present chooses that number of random shifts and moves answers one over from true. + If errors are present chooses that number of random shifts and move answer one over from true. """ - error_num = weighted_error(len(answers), scale=1) + error_num = weighted_error(len(answer), scale=1) if error_num == 0: - return answers + return answer else: - selections = np.random.choice(list(answers.values()), error_num, replace=False) + selections = np.random.choice(list(answer.values()), error_num, replace=False) mutations = np.roll(selections, 1) for i, j in enumerate(selections): - answers[j] = mutations[i] + answer[j] = mutations[i] - return answers + return answer -def generate_history(original, history_len=5): +def generate_history(original, history_length=5): """ Loops over chosen history length, adding noise to original data, reorganizes, and edits answer key if errors are introduced. """ - answers = {} + answer = {} history = [] - for i in range(history_len): + for i in range(history_length): protein, predicted_shifts, actual_shifts = perturb_data(original['coords']) - noes = recalculate_noe(protein, actual_shifts, original['assignments'], cutoff=0.5) + noes = distance_noe(protein, actual_shifts, cutoff=0.5) connectivity = connectivity_data(protein, cutoff=0.37) - # Lists of namedtuples (one object per residue) - coords = [Protein(x=resid[0], y=resid[1], z=resid[2], H1=shift[0], N15=shift[1]) for resid, shift in zip(protein, predicted_shifts)] - actual_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in actual_shifts] - noes = [NOEPeak(H1=shift[0], N15=shift[1], H2=shift[2]) for shift in noes] - connectivity = [Connectivity(atom1=connect[0], atom2=connect[1], distance=connect[2]) for connect in connectivity] + coords, actual_shifts, noes, connectivity = order_data(protein, actual_shifts, predicted_shifts, noes, connectivity) restraints = noe_combinations(noes, actual_shifts) history.append(coords) - history.append(actual_shifts) - history.append(noes) - history.append(restraints) - history.append(connectivity) + history.extend((actual_shifts, noes, restraints, connectivity)) - answers[i] = dict(original['assignments']) - answers[i] = add_errors(answers[i]) + answer[i] = dict(original['assignments']) + answer[i] = add_errors(answer[i]) - return history, answers + return history, answer \ No newline at end of file diff --git a/fake_histories_agent.py b/fake_histories_agent.py new file mode 100644 index 0000000..5952901 --- /dev/null +++ b/fake_histories_agent.py @@ -0,0 +1,44 @@ +import gym +import math +import argparse +import pickle + +from nmr_gym_env import GymEnv +from fake_histories import generate_history +from visualize_data import plot_shifts + +if __name__ == '__main__': + """ + Runs through protein examples of various sizes (min to max range) following the 1:1 shift to atom assignment. + Once example is assigned, generates histories and answers of original and moves on to next protein size. + All histories and answers are then saved to disk. + """ + parser = argparse.ArgumentParser() + parser.add_argument('min_protein', type=int, help='Minimum protein size') + parser.add_argument('max_protein', type=int, help='Max protein size') + parser.add_argument('history_length', type=int, help='Number of synthetic examples to generate') + args = parser.parse_args() + + min_protein = args.min_protein + max_protein = args.max_protein + history_length = args.history_length + + histories = [] + answers = [] + + for protein_length in range(min_protein, max_protein, 10): + gym_env = GymEnv(protein_length) + observation = gym_env.reset(pickled=False, pickle_data=False) + + terminated = False + while not terminated: + for action in observation['assign_order']: + observation, reward, terminated, total_energy = gym_env.step(action) + + histories, answer = generate_history(observation, history_length=history_length) + histories.append(histories) + answers.append(answer) + + with open('fake_histories.pkl', 'wb') as f: + pickle.dump(histories, f) + pickle.dump(answers, f) \ No newline at end of file diff --git a/nmr_gym_env.py b/nmr_gym_env.py index 9e687d5..de9ef23 100644 --- a/nmr_gym_env.py +++ b/nmr_gym_env.py @@ -2,6 +2,7 @@ from gym import error, spaces import numpy as np import pickle +import math import glob as glob from fake_data import generate_data @@ -19,6 +20,7 @@ def __init__(self, num_resid): self.assignments = {} self.assign_step = 0 self.assign_shift = 0 + self.restraints = 0 self.total_energy = 0.0 self.intermediate_energy = 0.0 self.reward = 0.0 @@ -28,7 +30,11 @@ def __init__(self, num_resid): self.observation_space = spaces.Dict({ "coords": spaces.Box(0, 1, shape=(self.num_resid, 5), dtype=np.float32), "actual_shifts": spaces.Box(0, 1, shape=(self.num_resid, 2), dtype=np.float32), - "noes": spaces.Box(0, 1, shape=(self.num_resid, 3), dtype=np.float32) + "noes": spaces.Box(0, 1, shape=(self.num_resid, 3), dtype=np.float32), + "assignments": spaces.Box(0, num_resid, shape=(self.num_resid, 3), dtype=np.float32), + "assign_order": spaces.Box(0, num_resid, shape=(self.num_resid, 1), dtype=np.float32), + "total_energy": spaces.Box(0, math.inf, shape=(1, 1), dtype=np.float32), + "reward": spaces.Box(0, math.inf, shape=(1, 1), dtype=np.float32) }) def custom_state(self, state): @@ -53,7 +59,8 @@ def reset(self, pickled=False, pickle_data=True, example=True): actual_shifts = pickle.load(f) noes = pickle.load(f) connectivity = pickle.load(f) - else: + + elif pickle_data: generate_data(self.num_resid, pickle_data=pickle_data, example=example) pickle_file = glob.glob(name) with open(pickle_file[0], 'rb') as f: @@ -62,11 +69,14 @@ def reset(self, pickled=False, pickle_data=True, example=True): noes = pickle.load(f) connectivity = pickle.load(f) + else: + coords, actual_shifts, noes, connectivity = generate_data(self.num_resid, pickle_data=pickle_data, example=example) + # get restraints - restraints = noe_combinations(noes, actual_shifts) + self.restraints = noe_combinations(noes, actual_shifts) # get plot of shifts and restraints - plot_shifts(actual_shifts, restraints, coords, connectivity, named_tuple_used=True) + plot_shifts(actual_shifts, self.restraints, coords, connectivity, named_tuple_used=True) self.assignments = {} self.assign_step = 0 @@ -76,7 +86,7 @@ def reset(self, pickled=False, pickle_data=True, example=True): self.reward = 0.0 # get assignment order - frac_act = FracAct(self.num_resid, restraints) + frac_act = FracAct(self.num_resid, self.restraints) self.assign_order = frac_act.fractional_activation() self.assign_shift = self.assign_order[self.assign_step] print(self.assign_order) @@ -86,8 +96,8 @@ def reset(self, pickled=False, pickle_data=True, example=True): "coords": coords, "actual_shifts": actual_shifts, "noes": noes, - "restraints": restraints, "assignments": self.assignments, + "assign_order": self.assign_order, "total_energy": self.total_energy, "reward": self.reward } @@ -95,6 +105,8 @@ def reset(self, pickled=False, pickle_data=True, example=True): return self.state def step(self, action): + assert action < self.num_resid and action not in self.state['assignments'].values() + energy = Energy(self.state['coords'], self.state['actual_shifts'], self.state['noes']) # add action to assignments @@ -102,7 +114,7 @@ def step(self, action): print(self.assignments) # calc energy and store temporarily - temp_intermediate_energy = energy.get_total_energy(self.state['restraints'], self.assignments) + temp_intermediate_energy = energy.get_total_energy(self.restraints, self.assignments) # print(f'this is the temp energy {temp_intermediate_energy}') # calc reward based on previous energy and temporary energy diff --git a/nmr_text_adventure.py b/nmr_text_adventure.py index 718ac44..16e430c 100644 --- a/nmr_text_adventure.py +++ b/nmr_text_adventure.py @@ -3,6 +3,7 @@ import argparse from nmr_gym_env import GymEnv +from visualize_data import plot_shifts if __name__ == '__main__': @@ -19,10 +20,10 @@ ###### Available options to work with right now ###### - # 1. Need to generate and save an example? Saved in all instances as fakedata_r{resid_num}.pkl + # 1. Need to generate and save an example? Saved in all instances as fakedata_r{num_resid}.pkl # observation = gym_env.reset(pickled=False, pickle_data=True) - # 2. Grab one of these previous example? Make sure example exists, with desired resid number. + # 2. Grab one of these previous examples? Make sure example exists with desired resid number. # observation = gym_env.reset(pickled=True, pickle_data=False) # 3. Run a new example? Saved in all instances as current_run.pkl @@ -34,13 +35,11 @@ ###################################################### terminated = False - action_assigned = [] while not terminated: action = int(input()) if action < num_resid: - if action not in action_assigned: - action_assigned.append(action) + if action not in observation['assignments'].values(): observation, reward, terminated, total_energy = gym_env.step(action) # random generated set of 4 From e8ba0d8a459c012cb1457cb6553a518ab4bc6838 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Thu, 6 Mar 2025 19:04:26 -0700 Subject: [PATCH 19/43] Value to go and energy transformation --- energy.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/energy.py b/energy.py index d26f0e5..324a2ed 100644 --- a/energy.py +++ b/energy.py @@ -131,3 +131,16 @@ def get_total_energy(self, restraints, assignments): total_energy += energy_value return total_energy + + def value_to_go(self, restraints, answer, energy): + """ + Calculates energy difference from current to final energy state. + """ + total_energy = self.get_total_energy(restraints, answer) + + return total_energy - energy + + def transform_energy(self, energy): + return np.log((energy + 1)) + + From c31b0b72dc47fb2998677161a38460c2c47c46a0 Mon Sep 17 00:00:00 2001 From: a-angel-s Date: Mon, 10 Mar 2025 11:17:46 -0600 Subject: [PATCH 20/43] added nmr_transformer --- nmr_transformer.py | 324 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 324 insertions(+) create mode 100644 nmr_transformer.py diff --git a/nmr_transformer.py b/nmr_transformer.py new file mode 100644 index 0000000..f02ba23 --- /dev/null +++ b/nmr_transformer.py @@ -0,0 +1,324 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import NamedTuple, List, Optional + +class NMRInput(NamedTuple): + # input array of observed chemical shifts + # 2-dimensions: N_shift, H_shift + # (n_peaks, 2) + obs_chemical_shifts: torch.Tensor + + # input array of observed chemical shifts + # 2-dimensions: N_shift, H_shift + # (n_peaks, 2) + pred_chemical_shifts: torch.Tensor + + # input array of observed noes + # 3-dimensions: N_shift1, H_shift1, H_shift2 + # (n_noe, 3) + obs_noes: Optional[torch.Tensor] + + # input array of close distances + # 4-dimensions: N_shift1, H_shift1, H_shift2, H_shift2 + # note: each close distance should be listed twice, with the residue order swapped + # (n_close, 4) + close_distances: Optional[torch.Tensor] + + # input array of peaks assigned so far + # 4-dimensions: N_obs_shift, H_obs_shift, N_pred_shift, H_pred_shift + # (n_assigned, 4) + # note: can set to None if no peaks have been assigned + assigned_peaks: Optional[torch.Tensor] + + # which peak should be assigned? + peak_to_assign: int + + +class NMRTransformer(torch.nn.Module): + """ + A transformer model for predicting the next peak to assign in an NMR spectrum. + """ + def __init__(self, n_hidden: int=128, n_heads: int=4, n_layers: int=4): + super(NMRTransformer, self).__init__() + self.embedder = NMRInitialEmbedding(n_hidden) + encoder_modules = [TransformerEncoderLayer(d_model=n_hidden, nhead=n_heads) for _ in range(n_layers)] + self.encoder_layers = nn.Sequential(*encoder_modules) + self.policy_linear1 = nn.Linear(2*n_hidden, 2*n_hidden) + self.policy_linear2 = nn.Linear(2*n_hidden, 1) + self.value_linear1 = nn.Linear(n_hidden, n_hidden) + self.value_linear2 = nn.Linear(n_hidden, 1) + + def forward(self, nmr_inputs: List[NMRInput]) -> torch.Tensor: + embedded = self.embedder(nmr_inputs) + x = self.encoder_layers(embedded) + + # compute the policy for each input in the batch + policies = [] + for i in range(len(nmr_inputs)): + n_res = len(nmr_inputs[i].pred_chemical_shifts) + residue_embeddings = x[i][:n_res] + peak_to_assign = nmr_inputs[i].peak_to_assign + assert peak_to_assign >=0 and peak_to_assign < len(nmr_inputs[i].obs_chemical_shifts) + peak_embedding = x[i][n_res + peak_to_assign].expand(n_res, -1) + embeddings = torch.cat([residue_embeddings, peak_embedding], dim=1) + policy = self.policy_linear2(F.relu(self.policy_linear1(embeddings))).reshape(-1) + policy = F.softmax(policy) + policies.append(policy) + + # compute the value for each input in the batch + embeddings = [] + for i in range(len(nmr_inputs)): + # the global token is always the last one + embedding = x[i][-1] + embeddings.append(embedding) + embeddings = torch.stack(embeddings) + values = self.value_linear2(F.relu(self.value_linear1(embeddings))).reshape(-1) + + return policies, values + + +class NMRInitialEmbedding(nn.Module): + """ + Embeds the initial input for the NMR transformer. + """ + def __init__(self, n_hidden: int=128): + super(NMRInitialEmbedding, self).__init__() + self.linear_obs_shifts = nn.Linear(2, n_hidden) + self.linear_pred_shifts = nn.Linear(2, n_hidden) + self.linear_obs_noes = nn.Linear(3, n_hidden) + self.linear_close_distances = nn.Linear(4, n_hidden) + self.linear_assigned_peaks = nn.Linear(4, n_hidden) + self.embeddings = nn.Embedding(6, n_hidden) + + def forward(self, nmr_inputs: List[NMRInput]) -> torch.Tensor: + # embed each of the inputs + embedded_obs_shifts = self._handle_input( + [nmr.obs_chemical_shifts for nmr in nmr_inputs], + self.linear_obs_shifts, + 0 + ) + embedded_pred_shifts = self._handle_input( + [nmr.pred_chemical_shifts for nmr in nmr_inputs], + self.linear_pred_shifts, + 1 + ) + embedded_obs_noes = self._handle_input( + [nmr.obs_noes for nmr in nmr_inputs], + self.linear_obs_noes, + 2 + ) + embedded_close_distances = self._handle_input( + [nmr.close_distances for nmr in nmr_inputs], + self.linear_close_distances, + 3 + ) + embedded_assigned_peaks = self._handle_input( + [nmr.assigned_peaks for nmr in nmr_inputs], + self.linear_assigned_peaks, + 4 + ) + + # collect the "global" token for each input + embedded_global = [] + embed = self.embeddings(torch.tensor([5])) + for _ in nmr_inputs: + embedded_global.append(embed) + + # concatenate the two sets of embeddings + concatendated = [] + variables_to_zip = [ + embedded_pred_shifts, + embedded_obs_shifts, + embedded_obs_noes, + embedded_close_distances, + embedded_assigned_peaks, + embedded_global + ] + for x1, x2, x3, x4, x5, x6 in zip(*variables_to_zip): + concatendated.append(torch.cat([x1, x2, x3, x4, x5, x6], dim=0)) + + return torch.nested.nested_tensor(concatendated, layout=torch.jagged) + + def _handle_input(self, xs, linear, embed_index): + embedded = [] + for x in xs: + if x is None: + embedded.append(torch.Tensor()) + else: + x = linear(x) + x += self.embeddings(torch.tensor([embed_index])) + embedded.append(x) + return embedded + + +class TransformerEncoderLayer(nn.Module): + """ + TransformerEncoderLayer is made up of self-attn and feedforward network + with residual connections. + """ + def __init__( + self, + d_model, + nhead, + dim_feedforward=2048, + dropout=0.1, + activation : nn.Module = torch.nn.functional.relu, + layer_norm_eps=1e-5, + norm_first=True, + bias=True, + device=None, + dtype=None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.self_attn = MultiHeadAttention( + d_model, + d_model, + d_model, + d_model, + nhead, + dropout=dropout, + bias=bias, + **factory_kwargs, + ) + self.linear1 = nn.Linear(d_model, dim_feedforward, bias=bias, **factory_kwargs) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model, bias=bias, **factory_kwargs) + + self.norm_first = norm_first + self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs) + self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs) + + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + self.activation = activation + + + def _sa_block(self, x): + x = self.self_attn(x, x, x) + return self.dropout1(x) + + def _ff_block(self, x): + x = self.linear2(self.dropout(self.activation(self.linear1(x)))) + return self.dropout2(x) + + def forward(self, src): + ''' + Arguments: + src: (batch_size, seq_len, d_model) + ''' + x = src + if self.norm_first: + x = x + self._sa_block(self.norm1(x)) + x = x + self._ff_block(self.norm2(x)) + else: + x = self.norm1(x + self._sa_block(x)) + x = self.norm2(x + self._ff_block(x)) + return x + + +class MultiHeadAttention(nn.Module): + """ + Computes multi-head attention. Supports nested or padded tensors. + + Args: + E_q (int): Size of embedding dim for query + E_k (int): Size of embedding dim for key + E_v (int): Size of embedding dim for value + E_total (int): Total embedding dim of combined heads post input projection. Each head + has dim E_total // nheads + nheads (int): Number of heads + dropout (float, optional): Dropout probability. Default: 0.0 + bias (bool, optional): Whether to add bias to input projection. Default: True + """ + def __init__( + self, + E_q: int, + E_k: int, + E_v: int, + E_total: int, + nheads: int, + dropout: float = 0.0, + bias=True, + device=None, + dtype=None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.nheads = nheads + self.dropout = dropout + self.q_proj = nn.Linear(E_q, E_total, bias=bias, **factory_kwargs) + self.k_proj = nn.Linear(E_k, E_total, bias=bias, **factory_kwargs) + self.v_proj = nn.Linear(E_v, E_total, bias=bias, **factory_kwargs) + E_out = E_q + self.out_proj = nn.Linear(E_total, E_out, bias=bias, **factory_kwargs) + assert E_total % nheads == 0, "Embedding dim is not divisible by nheads" + self.E_head = E_total // nheads + self.bias = bias + + def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + """ + Forward pass; runs the following process: + 1. Apply input projection + 2. Split heads and prepare for SDPA + 3. Run SDPA + 4. Apply output projection + + Args: + query (torch.Tensor): query of shape (N, L_q, E_qk) + key (torch.Tensor): key of shape (N, L_kv, E_qk) + value (torch.Tensor): value of shape (N, L_kv, E_v) + + Returns: + attn_output (torch.Tensor): output of shape (N, L_t, E_q) + """ + # Step 1. Apply input projection + query = self.q_proj(query) + key = self.k_proj(key) + value = self.v_proj(value) + + # Step 2. Split heads and prepare for SDPA + # reshape query, key, value to separate by head + # (N, L_t, E_total) -> (N, L_t, nheads, E_head) -> (N, nheads, L_t, E_head) + query = query.unflatten(-1, [self.nheads, self.E_head]).transpose(1, 2) + # (N, L_s, E_total) -> (N, L_s, nheads, E_head) -> (N, nheads, L_s, E_head) + key = key.unflatten(-1, [self.nheads, self.E_head]).transpose(1, 2) + # (N, L_s, E_total) -> (N, L_s, nheads, E_head) -> (N, nheads, L_s, E_head) + value = value.unflatten(-1, [self.nheads, self.E_head]).transpose(1, 2) + + # Step 3. Run SDPA + # (N, nheads, L_t, E_head) + attn_output = F.scaled_dot_product_attention( + query, key, value, dropout_p=self.dropout) + # (N, nheads, L_t, E_head) -> (N, L_t, nheads, E_head) -> (N, L_t, E_total) + attn_output = attn_output.transpose(1, 2).flatten(-2) + + # Step 4. Apply output projection + # (N, L_t, E_total) -> (N, L_t, E_out) + attn_output = self.out_proj(attn_output) + + return attn_output + + +if __name__ == "__main__": + # stupid code to find simple errors + nmr_inputs = [ + NMRInput( + obs_chemical_shifts=torch.randn(2, 2), + pred_chemical_shifts=torch.randn(2, 2), + obs_noes=torch.randn(3, 3), + close_distances=torch.randn(2, 4), + assigned_peaks=torch.randn(2, 4), + peak_to_assign=0), + NMRInput( + obs_chemical_shifts=torch.randn(5, 2), + pred_chemical_shifts=torch.randn(5, 2), + obs_noes=torch.randn(4, 3), + close_distances=torch.randn(7, 4), + assigned_peaks=None, + peak_to_assign=0) + ] + trans = NMRTransformer() + output = trans(nmr_inputs) + print(output) \ No newline at end of file From 42906d1d2dcd42cbeb33f92d820a4c4e9c739dc4 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Wed, 12 Mar 2025 14:10:37 -0600 Subject: [PATCH 21/43] Attempt to scale NOEs and restraints --- fake_data.py | 32 +++++++++++++++----------------- fake_histories.py | 7 +++---- nmr_gym_env.py | 14 ++++++++------ nmr_text_adventure.py | 6 +++--- 4 files changed, 29 insertions(+), 30 deletions(-) diff --git a/fake_data.py b/fake_data.py index 0c55766..7ed9c9e 100644 --- a/fake_data.py +++ b/fake_data.py @@ -53,16 +53,16 @@ def add_noise(point, scale=0.1, min=0, max=1): def calc_dist(p1, p2): return np.linalg.norm((p2-p1)) -def distance_noe(protein, shifts, cutoff): +def distance_noe(protein, shifts, cutoff=0.5, random_key=False): """ - Grabs close coordinates and 'associated' HSQC shift peaks (randomized index) to create NOES. - Predicted shifts assigned as the new randomized shift order - this order is associated with the coordinate order (ie. the answer key). + Grabs close coordinates and 'associated' HSQC shift peaks (randomized index or 1:1 correlation) to create NOES. + Predicted shifts assigned as shift order - this order is associated with the coordinate order if randomized (ie. the answer key). Adds gaussian noise to all points at the end. **Can end up with no NOEs depending on the cutoff** """ - # Can use this to randomize - not required - # shifts = np.random.permutation(shifts) + if random_key: + shifts = np.random.permutation(shifts) noes = [] @@ -71,16 +71,16 @@ def distance_noe(protein, shifts, cutoff): if i != j: dist = calc_dist(atom1, atom2) if dist < cutoff: - # print(dist, shifts[i], shifts[j]) + #print(dist, shifts[i], shifts[j]) noe = list(shifts[i][:]) # H1, N1 noe.append(shifts[j][0]) # H2 noes.append(noe) + noisy_noe = add_noise(np.array(noes), scale=0.01) - # predicted_shifts = add_noise(np.array(shifts), scale=0.1) - - # return noisy_noe, predicted_shifts - return noisy_noe + predicted_shifts = add_noise(np.array(shifts), scale=0.1) + + return noisy_noe, predicted_shifts def connectivity_data(protein, cutoff): """ @@ -97,7 +97,7 @@ def connectivity_data(protein, cutoff): return connectivity -def generate_raw_data(num_resid): +def generate_raw_data(num_resid, random_key): """ Generates all qualities of the system as individual arrays. """ @@ -105,12 +105,10 @@ def generate_raw_data(num_resid): protein = sample_unit(num_resid, num_sides=3) # "Actual" shifts [H1,N1] actual_shifts = sample_unit(num_resid, num_sides=2) - # Predicted shifts [H1,N1] - predicted_shifts = add_noise(actual_shifts, scale=0.1) # NOES [H1,N1,H2] and predicted shifts [H1,N1] - noes = distance_noe(protein, actual_shifts, cutoff=0.5) + noes, predicted_shifts = distance_noe(protein, actual_shifts, cutoff=float(1/np.cbrt(num_resid)), random_key=random_key) # connectivity [atom1,atom2,dist] - connectivity = connectivity_data(protein, cutoff=0.37) + connectivity = connectivity_data(protein, cutoff=float(0.8/np.cbrt(num_resid))) #0.37 return protein, actual_shifts, predicted_shifts, noes, connectivity @@ -137,11 +135,11 @@ def save_as_pickle(coords, actual_shifts, noes, connectivity, num_resid, example pickle.dump(noes, f) pickle.dump(connectivity, f) -def generate_data(num_resid, pickle_data=True, example=True): +def generate_data(num_resid, pickle_data=True, example=True, random_key=False): """ Generates all fake data, orders it in lists of namedtuples, and pickles if required. """ - protein, actual_shifts, predicted_shifts, noes, connectivity = generate_raw_data(num_resid) + protein, actual_shifts, predicted_shifts, noes, connectivity = generate_raw_data(num_resid, random_key=random_key) coords, actual_shifts, noes, connectivity = order_data(protein, actual_shifts, predicted_shifts, noes, connectivity) if pickle_data: diff --git a/fake_histories.py b/fake_histories.py index 1d59450..fa3de10 100644 --- a/fake_histories.py +++ b/fake_histories.py @@ -7,10 +7,9 @@ def perturb_data(coords): Adds noise to original fake data and generates new 'actual' shifts based on predicted shifts. """ protein = add_noise(np.array(coords)[:, :3], scale=0.1) - predicted_shifts = add_noise(np.array(coords)[:, 3:], scale=0.1) actual_fake_shifts = add_noise(np.array(coords)[:, 3:], scale=0.1) - return protein, predicted_shifts, actual_fake_shifts + return protein, actual_fake_shifts def weighted_error(num_resid, scale=1): """ @@ -51,8 +50,8 @@ def generate_history(original, history_length=5): history = [] for i in range(history_length): - protein, predicted_shifts, actual_shifts = perturb_data(original['coords']) - noes = distance_noe(protein, actual_shifts, cutoff=0.5) + protein, actual_shifts = perturb_data(original['coords']) + noes, predicted_shifts = distance_noe(protein, actual_shifts, cutoff=0.5) connectivity = connectivity_data(protein, cutoff=0.37) coords, actual_shifts, noes, connectivity = order_data(protein, actual_shifts, predicted_shifts, noes, connectivity) diff --git a/nmr_gym_env.py b/nmr_gym_env.py index de9ef23..a2bb1bd 100644 --- a/nmr_gym_env.py +++ b/nmr_gym_env.py @@ -31,8 +31,9 @@ def __init__(self, num_resid): "coords": spaces.Box(0, 1, shape=(self.num_resid, 5), dtype=np.float32), "actual_shifts": spaces.Box(0, 1, shape=(self.num_resid, 2), dtype=np.float32), "noes": spaces.Box(0, 1, shape=(self.num_resid, 3), dtype=np.float32), - "assignments": spaces.Box(0, num_resid, shape=(self.num_resid, 3), dtype=np.float32), - "assign_order": spaces.Box(0, num_resid, shape=(self.num_resid, 1), dtype=np.float32), + "assignments": spaces.Box(0, self.num_resid, shape=(self.num_resid, 3), dtype=np.float32), + "assign_order": spaces.Box(0, self.num_resid, shape=(self.num_resid, 1), dtype=np.float32), + "assign_shift": spaces.Discrete(self.num_resid), "total_energy": spaces.Box(0, math.inf, shape=(1, 1), dtype=np.float32), "reward": spaces.Box(0, math.inf, shape=(1, 1), dtype=np.float32) }) @@ -48,7 +49,7 @@ def custom_state(self, state): return self.state - def reset(self, pickled=False, pickle_data=True, example=True): + def reset(self, pickled=False, pickle_data=True, example=True, random_key=False): name = f"./*{self.num_resid}.pkl" if example else "./current_run.pkl" @@ -61,7 +62,7 @@ def reset(self, pickled=False, pickle_data=True, example=True): connectivity = pickle.load(f) elif pickle_data: - generate_data(self.num_resid, pickle_data=pickle_data, example=example) + generate_data(self.num_resid, pickle_data=pickle_data, example=example, random_key=random_key) pickle_file = glob.glob(name) with open(pickle_file[0], 'rb') as f: coords = pickle.load(f) @@ -70,10 +71,10 @@ def reset(self, pickled=False, pickle_data=True, example=True): connectivity = pickle.load(f) else: - coords, actual_shifts, noes, connectivity = generate_data(self.num_resid, pickle_data=pickle_data, example=example) + coords, actual_shifts, noes, connectivity = generate_data(self.num_resid, pickle_data=pickle_data, example=example, random_key=random_key) # get restraints - self.restraints = noe_combinations(noes, actual_shifts) + self.restraints = noe_combinations(noes, actual_shifts, tolerance_h=float(0.02/np.cbrt(self.num_resid)), tolerance_n=float(0.02/np.cbrt(self.num_resid))) # get plot of shifts and restraints plot_shifts(actual_shifts, self.restraints, coords, connectivity, named_tuple_used=True) @@ -98,6 +99,7 @@ def reset(self, pickled=False, pickle_data=True, example=True): "noes": noes, "assignments": self.assignments, "assign_order": self.assign_order, + "assign_shift": self.assign_shift, "total_energy": self.total_energy, "reward": self.reward } diff --git a/nmr_text_adventure.py b/nmr_text_adventure.py index 16e430c..694f742 100644 --- a/nmr_text_adventure.py +++ b/nmr_text_adventure.py @@ -21,13 +21,13 @@ ###### Available options to work with right now ###### # 1. Need to generate and save an example? Saved in all instances as fakedata_r{num_resid}.pkl - # observation = gym_env.reset(pickled=False, pickle_data=True) + # observation = gym_env.reset(pickled=False, pickle_data=True, random_key=True) # 2. Grab one of these previous examples? Make sure example exists with desired resid number. - # observation = gym_env.reset(pickled=True, pickle_data=False) + # observation = gym_env.reset(pickled=True, pickle_data=False, random_key=True) # 3. Run a new example? Saved in all instances as current_run.pkl - observation = gym_env.reset(pickled=False, pickle_data=True, example=False) + observation = gym_env.reset(pickled=False, pickle_data=True, example=False, random_key=True) # 4. Use a custom hardcoded state? May be some issues with data types and visualization. # observation = gym_env.custom_state(state) From 9832f4f50927ace137520b97b841417d48432c88 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Thu, 13 Mar 2025 11:00:34 -0600 Subject: [PATCH 22/43] Connected energy to the NOE distance scale --- energy.py | 14 +++++++------- nmr_gym_env.py | 13 ++++++++----- 2 files changed, 15 insertions(+), 12 deletions(-) diff --git a/energy.py b/energy.py index 324a2ed..17b3c00 100644 --- a/energy.py +++ b/energy.py @@ -67,7 +67,7 @@ def __init__(self, coords, actual_shifts, noes): self.noes = noes def setup_noe_restraints(self): - combos = noe_combinations(self.noes, self.actual_shifts) + combos = noe_combinations(self.noes, self.actual_shifts, tolerance_h=0.02, tolerance_n=0.02) return combos def flat_bottom(self, x, tolerance): @@ -88,7 +88,7 @@ def flat_bottom(self, x, tolerance): lambda x: x**2 - xmax*x]) return y - def calc_restraint_energy(self, assignments, restraint): + def calc_restraint_energy(self, assignments, restraint, tolerance): """ Goes over shift possibilities in the restraint, grabs associated coordinates from assignment, and calculates energy according to distance. Only the smallest energy is returned. @@ -97,11 +97,11 @@ def calc_restraint_energy(self, assignments, restraint): for i, j in restraint: k, l = assignments.get(i), assignments.get(j) dist = np.linalg.norm((np.array(self.coords[k][:3]) - np.array(self.coords[l][:3]))) - energy_value = self.flat_bottom(dist, tolerance=0.5) + energy_value = self.flat_bottom(dist, tolerance=tolerance) restraint_energy = energy_value if energy_value < restraint_energy else restraint_energy return restraint_energy - def noe_activation(self, assignments, restraint): + def noe_activation(self, assignments, restraint, tolerance): """ Checks if all shift possibilities in the restraint have been assigned (aka activated). Calculates energy if they have, returns zero if they have not. @@ -115,9 +115,9 @@ def noe_activation(self, assignments, restraint): if d: return 0 else: - return self.calc_restraint_energy(assignments, restraint) + return self.calc_restraint_energy(assignments, restraint, tolerance) - def get_total_energy(self, restraints, assignments): + def get_total_energy(self, restraints, assignments, tolerance): """ Loops over the NOE restraints to sum up calculated energies. """ @@ -127,7 +127,7 @@ def get_total_energy(self, restraints, assignments): if restraint == []: pass else: - energy_value = self.noe_activation(assignments, restraint) + energy_value = self.noe_activation(assignments, restraint, tolerance=tolerance) total_energy += energy_value return total_energy diff --git a/nmr_gym_env.py b/nmr_gym_env.py index a2bb1bd..5f653e2 100644 --- a/nmr_gym_env.py +++ b/nmr_gym_env.py @@ -51,7 +51,7 @@ def custom_state(self, state): def reset(self, pickled=False, pickle_data=True, example=True, random_key=False): - name = f"./*{self.num_resid}.pkl" if example else "./current_run.pkl" + name = f"./*r{self.num_resid}.pkl" if example else "./current_run.pkl" if self.intermediate_energy > 0 or pickled: pickle_file = glob.glob(name) @@ -60,7 +60,7 @@ def reset(self, pickled=False, pickle_data=True, example=True, random_key=False) actual_shifts = pickle.load(f) noes = pickle.load(f) connectivity = pickle.load(f) - + elif pickle_data: generate_data(self.num_resid, pickle_data=pickle_data, example=example, random_key=random_key) pickle_file = glob.glob(name) @@ -74,8 +74,11 @@ def reset(self, pickled=False, pickle_data=True, example=True, random_key=False) coords, actual_shifts, noes, connectivity = generate_data(self.num_resid, pickle_data=pickle_data, example=example, random_key=random_key) # get restraints - self.restraints = noe_combinations(noes, actual_shifts, tolerance_h=float(0.02/np.cbrt(self.num_resid)), tolerance_n=float(0.02/np.cbrt(self.num_resid))) - + """ + float(float(0.001*(self.num_resid)+0.01)) shift tolerance needs to increase with protein size somehow or else it can miss the correct answer, + but this will increase the complexity of the problem + """ + self.restraints = noe_combinations(noes, actual_shifts, tolerance_h=0.02, tolerance_n=0.02) # get plot of shifts and restraints plot_shifts(actual_shifts, self.restraints, coords, connectivity, named_tuple_used=True) @@ -116,7 +119,7 @@ def step(self, action): print(self.assignments) # calc energy and store temporarily - temp_intermediate_energy = energy.get_total_energy(self.restraints, self.assignments) + temp_intermediate_energy = energy.get_total_energy(self.restraints, self.assignments, tolerance=float(1/np.cbrt(self.num_resid))) # print(f'this is the temp energy {temp_intermediate_energy}') # calc reward based on previous energy and temporary energy From 469994f33a387fa79cba767bda4a9c1be78d3a7a Mon Sep 17 00:00:00 2001 From: abbie-p Date: Sat, 15 Mar 2025 21:48:27 -0600 Subject: [PATCH 23/43] Data processing --- fake_histories.py | 23 ++++++++------ fake_histories_agent.py | 3 +- training_loop.py | 67 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 82 insertions(+), 11 deletions(-) create mode 100644 training_loop.py diff --git a/fake_histories.py b/fake_histories.py index fa3de10..bc2e0ea 100644 --- a/fake_histories.py +++ b/fake_histories.py @@ -46,24 +46,29 @@ def generate_history(original, history_length=5): """ Loops over chosen history length, adding noise to original data, reorganizes, and edits answer key if errors are introduced. """ - answer = {} + answer = [] history = [] for i in range(history_length): protein, actual_shifts = perturb_data(original['coords']) - noes, predicted_shifts = distance_noe(protein, actual_shifts, cutoff=0.5) - connectivity = connectivity_data(protein, cutoff=0.37) + noes, predicted_shifts = distance_noe(protein, actual_shifts, cutoff=float(1/np.cbrt(len(actual_shifts)))) + connectivity = connectivity_data(protein, cutoff=float(0.8/np.cbrt(len(actual_shifts)))) coords, actual_shifts, noes, connectivity = order_data(protein, actual_shifts, predicted_shifts, noes, connectivity) restraints = noe_combinations(noes, actual_shifts) - history.append(coords) - history.extend((actual_shifts, noes, restraints, connectivity)) + history.append([coords]) + history[i].extend((actual_shifts, noes, restraints, connectivity)) + # history.append(actual_shifts) + # history.append(noes) + # history.append(restraints) + # history.append(connectivity) - answer[i] = dict(original['assignments']) - answer[i] = add_errors(answer[i]) + # answer[i] = dict(original['assignments']) + # answer[i] = add_errors(answer[i]) - return history, answer + answer.append([dict(original['assignments'])]) + answer[i] = add_errors(answer[i]) - \ No newline at end of file + return history, answer \ No newline at end of file diff --git a/fake_histories_agent.py b/fake_histories_agent.py index 5952901..b12dee4 100644 --- a/fake_histories_agent.py +++ b/fake_histories_agent.py @@ -1,5 +1,4 @@ import gym -import math import argparse import pickle @@ -37,7 +36,7 @@ histories, answer = generate_history(observation, history_length=history_length) histories.append(histories) - answers.append(answer) + answers.extend(answer) with open('fake_histories.pkl', 'wb') as f: pickle.dump(histories, f) diff --git a/training_loop.py b/training_loop.py new file mode 100644 index 0000000..6495bb5 --- /dev/null +++ b/training_loop.py @@ -0,0 +1,67 @@ +import torch +import pickle +import numpy as np +from nmr_transformer import * + +def extract_data(pickle_file='fake_histories.pkl'): + with open(pickle_file, 'rb') as f: + histories = pickle.load(f) + answers = pickle.load(f) + return histories, answers + +def preprocess_data(histories, answers, batch_size): + """ + Loops over chosen histories and answers, splits up components, and adds shifts if indices were used. + *should be reorganized for clarity after fake histories is fixed* + """ + predicted_shifts = [] + obs_chemical_shifts = [] + noes = [] + close_dist = [] + assignments = [] + peak_to_assign = [] + + for history, answer in zip(histories[:batch_size], answers[:batch_size]): + + predicted_shifts_temp = ([(i.H1, i.N15) for i in history[0]]) + predicted_shifts.append(predicted_shifts_temp) + + obs_chemical_shifts_temp = ([(i.H1, i.N15) for i in history[1]]) + obs_chemical_shifts.append(obs_chemical_shifts_temp) + + noes.append(history[2]) + + # these have parentheses issues from list comprehension + close_dist.append([(predicted_shifts_temp[i.atom1], predicted_shifts_temp[i.atom2]) for i in history[4]]) + assignments.append([(obs_chemical_shifts_temp[i], predicted_shifts_temp[j]) for i,j in list(answer[0].items())]) #shift:atom assignment, therefore obs_shift:predicted_shift + + # peak_to_assign += [] # need to fix the fake histories for this + + return predicted_shifts, obs_chemical_shifts, noes, close_dist, assignments, peak_to_assign + +def organize_nmr_inputs(histories, answers, batch_size): + """ + Sets up NMR inputs as list of named tuples. + """ + predicted_shifts, actual_shifts, noes, close_dist, assignments, assign_peak = preprocess_data(histories, answers, batch_size=batch_size) + # print(actual_shifts) + # print(assignments) + nmr_inputs = [ + NMRInput( + obs_chemical_shifts=torch.tensor(actual_shifts[i]), + pred_chemical_shifts=torch.tensor(predicted_shifts[i]), + obs_noes=torch.tensor(noes[i]), + close_distances=torch.tensor(close_dist[i]), + assigned_peaks=torch.tensor(assignments[i]), + peak_to_assign=0) # need to fix + for i in range(batch_size) + ] + + print(nmr_inputs) + +if __name__ == "__main__": + # need to loop over here + histories, answers = extract_data() + organize_nmr_inputs(histories, answers, batch_size=2) + + \ No newline at end of file From 019a7b6408fe93643df70b021468ded524469823 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Thu, 20 Mar 2025 10:06:24 -0600 Subject: [PATCH 24/43] Updated fake histories and nmr inputs --- fake_histories.py | 51 ++++++++++++++++++++++------------------- fake_histories_agent.py | 36 ++++++++++++++++++----------- nmr_gym_env.py | 39 +++++++++++++++++++++---------- training_loop.py | 49 +++++++++++++++++++-------------------- 4 files changed, 102 insertions(+), 73 deletions(-) diff --git a/fake_histories.py b/fake_histories.py index bc2e0ea..87c57ec 100644 --- a/fake_histories.py +++ b/fake_histories.py @@ -36,39 +36,44 @@ def add_errors(answer): return answer else: selections = np.random.choice(list(answer.values()), error_num, replace=False) + # print(selections) mutations = np.roll(selections, 1) + # print(mutations) for i, j in enumerate(selections): + # print(answer, answer[j], mutations[i], error_num) answer[j] = mutations[i] - + return answer -def generate_history(original, history_length=5): +def generate_history(original, i, history_length=5): """ Loops over chosen history length, adding noise to original data, reorganizes, and edits answer key if errors are introduced. """ - answer = [] - history = [] - - for i in range(history_length): - protein, actual_shifts = perturb_data(original['coords']) - noes, predicted_shifts = distance_noe(protein, actual_shifts, cutoff=float(1/np.cbrt(len(actual_shifts)))) - connectivity = connectivity_data(protein, cutoff=float(0.8/np.cbrt(len(actual_shifts)))) - - coords, actual_shifts, noes, connectivity = order_data(protein, actual_shifts, predicted_shifts, noes, connectivity) - restraints = noe_combinations(noes, actual_shifts) + protein, actual_shifts = perturb_data(original['coords']) + noes, predicted_shifts = distance_noe(protein, actual_shifts, cutoff=float(1/np.cbrt(len(actual_shifts)))) + + connectivity = connectivity_data(protein, cutoff=float(0.8/np.cbrt(len(actual_shifts)))) - history.append([coords]) - history[i].extend((actual_shifts, noes, restraints, connectivity)) - # history.append(actual_shifts) - # history.append(noes) - # history.append(restraints) - # history.append(connectivity) + coords, actual_shifts, noes, connectivity = order_data(protein, actual_shifts, predicted_shifts, noes, connectivity) - # answer[i] = dict(original['assignments']) - # answer[i] = add_errors(answer[i]) + # answer = original['assignments'] + # answer, error_num = add_errors(answer) + + # issue with permanence - need to recreate full assignment for now + answer = {i:i for i in original['assign_order']} + answer = add_errors(answer) - answer.append([dict(original['assignments'])]) - answer[i] = add_errors(answer[i]) + state = { + "coords": coords, + "actual_shifts": actual_shifts, + "noes": noes, + "connectivity": connectivity, + "assignments": {}, + "assign_order": list(answer.values()), + "assign_shift": 0, + "total_energy": 0.0, + "reward": 0.0 + } - return history, answer \ No newline at end of file + return state \ No newline at end of file diff --git a/fake_histories_agent.py b/fake_histories_agent.py index b12dee4..1e3905f 100644 --- a/fake_histories_agent.py +++ b/fake_histories_agent.py @@ -1,6 +1,7 @@ import gym import argparse import pickle +import copy from nmr_gym_env import GymEnv from fake_histories import generate_history @@ -8,36 +9,45 @@ if __name__ == '__main__': """ - Runs through protein examples of various sizes (min to max range) following the 1:1 shift to atom assignment. - Once example is assigned, generates histories and answers of original and moves on to next protein size. - All histories and answers are then saved to disk. + Runs through protein examples of one size following the 1:1 shift to atom assignment. + Once example is assigned, generates history and repeats assignment process. + All histories are then saved to disk. """ parser = argparse.ArgumentParser() - parser.add_argument('min_protein', type=int, help='Minimum protein size') + # parser.add_argument('min_protein', type=int, help='Minimum protein size') parser.add_argument('max_protein', type=int, help='Max protein size') parser.add_argument('history_length', type=int, help='Number of synthetic examples to generate') args = parser.parse_args() - min_protein = args.min_protein + # min_protein = args.min_protein max_protein = args.max_protein history_length = args.history_length + gym_env = GymEnv(max_protein) + observation = gym_env.reset(pickled=False, pickle_data=False) + histories = [] - answers = [] - for protein_length in range(min_protein, max_protein, 10): - gym_env = GymEnv(protein_length) - observation = gym_env.reset(pickled=False, pickle_data=False) + for i in range(history_length): + if i > 0: + observation = gym_env.custom_state(fake_observation) + + histories.append(copy.deepcopy(observation)) terminated = False while not terminated: for action in observation['assign_order']: observation, reward, terminated, total_energy = gym_env.step(action) + histories.append(copy.deepcopy(observation)) + + # don't know if this is needed - issues with permanence of observation + if i <= 0: + original_observation = copy.deepcopy(observation) - histories, answer = generate_history(observation, history_length=history_length) - histories.append(histories) - answers.extend(answer) + fake_observation = generate_history(original_observation, i, history_length=history_length) + # print(histories) with open('fake_histories.pkl', 'wb') as f: pickle.dump(histories, f) - pickle.dump(answers, f) \ No newline at end of file + + print(f'{history_length} fake histories saved') \ No newline at end of file diff --git a/nmr_gym_env.py b/nmr_gym_env.py index 5f653e2..5a7c8ff 100644 --- a/nmr_gym_env.py +++ b/nmr_gym_env.py @@ -40,12 +40,25 @@ def __init__(self, num_resid): def custom_state(self, state): self.state = state - frac_act = FracAct(self.num_resid, self.state['restraints']) - self.assign_order = frac_act.fractional_activation() - print(self.assign_order) + self.restraints = noe_combinations(self.state['noes'], self.state['actual_shifts'], tolerance_h=0.02, tolerance_n=0.02) - # get plot of shifts and restraints - plot_shifts(self.state['actual_shifts'], self.state['restraints'], self.state['coords'], named_tuple_used=False) + self.assignments = {} + self.assign_step = 0 + self.assign_shift = 0 + self.total_energy = 0.0 + self.intermediate_energy = 0.0 + self.reward = 0.0 + + # get assignment order + # frac_act = FracAct(self.num_resid, self.restraints) + # self.assign_order = frac_act.fractional_activation() + self.assign_shift = self.state['assign_order'][self.assign_step] + # print(self.assign_order) + # print(f'Shift to be assigned: {self.assign_shift}') + + # self.state['assign_order'] = self.assign_order + self.state['assign_shift'] = self.assign_shift + self.state['assignments'] = self.assignments return self.state @@ -93,13 +106,14 @@ def reset(self, pickled=False, pickle_data=True, example=True, random_key=False) frac_act = FracAct(self.num_resid, self.restraints) self.assign_order = frac_act.fractional_activation() self.assign_shift = self.assign_order[self.assign_step] - print(self.assign_order) - print(f'Shift to be assigned: {self.assign_shift}') + # print(self.assign_order) + # print(f'Shift to be assigned: {self.assign_shift}') self.state = { "coords": coords, "actual_shifts": actual_shifts, "noes": noes, + "connectivity": connectivity, "assignments": self.assignments, "assign_order": self.assign_order, "assign_shift": self.assign_shift, @@ -110,13 +124,14 @@ def reset(self, pickled=False, pickle_data=True, example=True, random_key=False) return self.state def step(self, action): + # print(action) assert action < self.num_resid and action not in self.state['assignments'].values() energy = Energy(self.state['coords'], self.state['actual_shifts'], self.state['noes']) # add action to assignments self.assignments[(self.assign_order[self.assign_step])] = action - print(self.assignments) + # print(self.assignments) # calc energy and store temporarily temp_intermediate_energy = energy.get_total_energy(self.restraints, self.assignments, tolerance=float(1/np.cbrt(self.num_resid))) @@ -134,19 +149,19 @@ def step(self, action): # assign intermediate energy as running total self.state['total_energy'] = self.intermediate_energy - print(f"Running energy = {self.state['total_energy']}, Current reward = {self.state['reward']}") + # print(f"Running energy = {self.state['total_energy']}, Current reward = {self.state['reward']}") self.assign_step += 1 # self.assign_shift = self.assign_order[self.assign_step] terminated = True if len(self.assignments) == self.num_resid else False if terminated: - print(f"Final assignments (shift:atom) = {self.assignments}\nFinal energy evaluation = {self.state['total_energy']}") + # print(f"Final assignments (shift:atom) = {self.assignments}\nFinal energy evaluation = {self.state['total_energy']}") return self.state, self.state['reward'], terminated, self.state['total_energy'] - self.assign_shift = self.assign_order[self.assign_step] + self.state['assign_shift'] = self.assign_order[self.assign_step] if not terminated: - print(f'Shift to be assigned: {self.assign_shift}') + # print(f'Shift to be assigned: {self.state["assign_shift"]}') return self.state, self.state['reward'], terminated, self.state['total_energy'] diff --git a/training_loop.py b/training_loop.py index 6495bb5..d76babf 100644 --- a/training_loop.py +++ b/training_loop.py @@ -6,13 +6,11 @@ def extract_data(pickle_file='fake_histories.pkl'): with open(pickle_file, 'rb') as f: histories = pickle.load(f) - answers = pickle.load(f) - return histories, answers + return histories -def preprocess_data(histories, answers, batch_size): +def preprocess_data(histories, batch_size): """ Loops over chosen histories and answers, splits up components, and adds shifts if indices were used. - *should be reorganized for clarity after fake histories is fixed* """ predicted_shifts = [] obs_chemical_shifts = [] @@ -21,31 +19,32 @@ def preprocess_data(histories, answers, batch_size): assignments = [] peak_to_assign = [] - for history, answer in zip(histories[:batch_size], answers[:batch_size]): - - predicted_shifts_temp = ([(i.H1, i.N15) for i in history[0]]) + # batch size needs to be reworked in here (not randomized - selects from start up until 'batch_size' index) + for history in histories[:batch_size]: + # Predicted shifts + predicted_shifts_temp = ([(i.H1, i.N15) for i in history['coords']]) predicted_shifts.append(predicted_shifts_temp) - - obs_chemical_shifts_temp = ([(i.H1, i.N15) for i in history[1]]) + # Observed shifts + obs_chemical_shifts_temp = ([(i.H1, i.N15) for i in history['actual_shifts']]) obs_chemical_shifts.append(obs_chemical_shifts_temp) - - noes.append(history[2]) - - # these have parentheses issues from list comprehension - close_dist.append([(predicted_shifts_temp[i.atom1], predicted_shifts_temp[i.atom2]) for i in history[4]]) - assignments.append([(obs_chemical_shifts_temp[i], predicted_shifts_temp[j]) for i,j in list(answer[0].items())]) #shift:atom assignment, therefore obs_shift:predicted_shift - - # peak_to_assign += [] # need to fix the fake histories for this + # NOEs + noes.append(history['noes']) + # Close distances (duplicated for reverse pair order) + close_dist.append([(predicted_shifts_temp[i.atom1], predicted_shifts_temp[i.atom2]) for i in history['connectivity']]) + close_dist.append([(predicted_shifts_temp[i.atom2], predicted_shifts_temp[i.atom1]) for i in history['connectivity']]) + # Assignments + assignments.append([(obs_chemical_shifts_temp[i], predicted_shifts_temp[j]) for i,j in list(history['assignments'].items())]) #shift:atom assignment, therefore obs_shift:predicted_shift + # Peak to assign + peak_to_assign.append(obs_chemical_shifts_temp[history['assign_shift']]) return predicted_shifts, obs_chemical_shifts, noes, close_dist, assignments, peak_to_assign -def organize_nmr_inputs(histories, answers, batch_size): +def organize_nmr_inputs(histories, batch_size): """ Sets up NMR inputs as list of named tuples. """ - predicted_shifts, actual_shifts, noes, close_dist, assignments, assign_peak = preprocess_data(histories, answers, batch_size=batch_size) - # print(actual_shifts) - # print(assignments) + predicted_shifts, actual_shifts, noes, close_dist, assignments, assign_peak = preprocess_data(histories, batch_size=batch_size) + nmr_inputs = [ NMRInput( obs_chemical_shifts=torch.tensor(actual_shifts[i]), @@ -53,15 +52,15 @@ def organize_nmr_inputs(histories, answers, batch_size): obs_noes=torch.tensor(noes[i]), close_distances=torch.tensor(close_dist[i]), assigned_peaks=torch.tensor(assignments[i]), - peak_to_assign=0) # need to fix + peak_to_assign=assign_peak[i]) # int or tensor for H,N peak? for i in range(batch_size) ] - print(nmr_inputs) + return nmr_inputs if __name__ == "__main__": # need to loop over here - histories, answers = extract_data() - organize_nmr_inputs(histories, answers, batch_size=2) + histories = extract_data() + nmr_inputs = organize_nmr_inputs(histories, batch_size=3) \ No newline at end of file From 9149a76165ae75ece8e8a94fc616efa0df1cdc47 Mon Sep 17 00:00:00 2001 From: Justin MacCallum Date: Sat, 29 Mar 2025 10:05:28 -0600 Subject: [PATCH 25/43] Initial implementation of training --- fake_histories_agent.py | 5 +- nmr_gym_env.py | 4 +- nmr_transformer.py | 93 +++++++++++++++------------ training_loop.py | 138 +++++++++++++++++++++++++++++++--------- 4 files changed, 168 insertions(+), 72 deletions(-) diff --git a/fake_histories_agent.py b/fake_histories_agent.py index 1e3905f..4b90fb0 100644 --- a/fake_histories_agent.py +++ b/fake_histories_agent.py @@ -1,4 +1,4 @@ -import gym +import gymnasium as gym import argparse import pickle import copy @@ -38,7 +38,8 @@ while not terminated: for action in observation['assign_order']: observation, reward, terminated, total_energy = gym_env.step(action) - histories.append(copy.deepcopy(observation)) + if not terminated: + histories.append(copy.deepcopy(observation)) # don't know if this is needed - issues with permanence of observation if i <= 0: diff --git a/nmr_gym_env.py b/nmr_gym_env.py index 5a7c8ff..7fa796a 100644 --- a/nmr_gym_env.py +++ b/nmr_gym_env.py @@ -1,5 +1,5 @@ -import gym -from gym import error, spaces +import gymnasium as gym +from gymnasium import error, spaces import numpy as np import pickle import math diff --git a/nmr_transformer.py b/nmr_transformer.py index f02ba23..839bad6 100644 --- a/nmr_transformer.py +++ b/nmr_transformer.py @@ -3,6 +3,7 @@ import torch.nn.functional as F from typing import NamedTuple, List, Optional + class NMRInput(NamedTuple): # input array of observed chemical shifts # 2-dimensions: N_shift, H_shift @@ -33,19 +34,29 @@ class NMRInput(NamedTuple): # which peak should be assigned? peak_to_assign: int - + class NMRTransformer(torch.nn.Module): """ A transformer model for predicting the next peak to assign in an NMR spectrum. """ - def __init__(self, n_hidden: int=128, n_heads: int=4, n_layers: int=4): + + def __init__( + self, + n_hidden: int = 128, + n_heads: int = 4, + n_layers: int = 4, + dropout: float = 0.1, + ): super(NMRTransformer, self).__init__() self.embedder = NMRInitialEmbedding(n_hidden) - encoder_modules = [TransformerEncoderLayer(d_model=n_hidden, nhead=n_heads) for _ in range(n_layers)] + encoder_modules = [ + TransformerEncoderLayer(d_model=n_hidden, nhead=n_heads, dropout=dropout) + for _ in range(n_layers) + ] self.encoder_layers = nn.Sequential(*encoder_modules) - self.policy_linear1 = nn.Linear(2*n_hidden, 2*n_hidden) - self.policy_linear2 = nn.Linear(2*n_hidden, 1) + self.policy_linear1 = nn.Linear(2 * n_hidden, 2 * n_hidden) + self.policy_linear2 = nn.Linear(2 * n_hidden, 1) self.value_linear1 = nn.Linear(n_hidden, n_hidden) self.value_linear2 = nn.Linear(n_hidden, 1) @@ -59,11 +70,14 @@ def forward(self, nmr_inputs: List[NMRInput]) -> torch.Tensor: n_res = len(nmr_inputs[i].pred_chemical_shifts) residue_embeddings = x[i][:n_res] peak_to_assign = nmr_inputs[i].peak_to_assign - assert peak_to_assign >=0 and peak_to_assign < len(nmr_inputs[i].obs_chemical_shifts) + assert peak_to_assign >= 0 and peak_to_assign < len( + nmr_inputs[i].obs_chemical_shifts + ) peak_embedding = x[i][n_res + peak_to_assign].expand(n_res, -1) embeddings = torch.cat([residue_embeddings, peak_embedding], dim=1) - policy = self.policy_linear2(F.relu(self.policy_linear1(embeddings))).reshape(-1) - policy = F.softmax(policy) + policy = self.policy_linear2( + F.relu(self.policy_linear1(embeddings)) + ).reshape(-1) policies.append(policy) # compute the value for each input in the batch @@ -82,7 +96,8 @@ class NMRInitialEmbedding(nn.Module): """ Embeds the initial input for the NMR transformer. """ - def __init__(self, n_hidden: int=128): + + def __init__(self, n_hidden: int = 128): super(NMRInitialEmbedding, self).__init__() self.linear_obs_shifts = nn.Linear(2, n_hidden) self.linear_pred_shifts = nn.Linear(2, n_hidden) @@ -90,33 +105,23 @@ def __init__(self, n_hidden: int=128): self.linear_close_distances = nn.Linear(4, n_hidden) self.linear_assigned_peaks = nn.Linear(4, n_hidden) self.embeddings = nn.Embedding(6, n_hidden) - + def forward(self, nmr_inputs: List[NMRInput]) -> torch.Tensor: # embed each of the inputs embedded_obs_shifts = self._handle_input( - [nmr.obs_chemical_shifts for nmr in nmr_inputs], - self.linear_obs_shifts, - 0 + [nmr.obs_chemical_shifts for nmr in nmr_inputs], self.linear_obs_shifts, 0 ) embedded_pred_shifts = self._handle_input( - [nmr.pred_chemical_shifts for nmr in nmr_inputs], - self.linear_pred_shifts, - 1 + [nmr.pred_chemical_shifts for nmr in nmr_inputs], self.linear_pred_shifts, 1 ) embedded_obs_noes = self._handle_input( - [nmr.obs_noes for nmr in nmr_inputs], - self.linear_obs_noes, - 2 + [nmr.obs_noes for nmr in nmr_inputs], self.linear_obs_noes, 2 ) embedded_close_distances = self._handle_input( - [nmr.close_distances for nmr in nmr_inputs], - self.linear_close_distances, - 3 + [nmr.close_distances for nmr in nmr_inputs], self.linear_close_distances, 3 ) embedded_assigned_peaks = self._handle_input( - [nmr.assigned_peaks for nmr in nmr_inputs], - self.linear_assigned_peaks, - 4 + [nmr.assigned_peaks for nmr in nmr_inputs], self.linear_assigned_peaks, 4 ) # collect the "global" token for each input @@ -133,7 +138,7 @@ def forward(self, nmr_inputs: List[NMRInput]) -> torch.Tensor: embedded_obs_noes, embedded_close_distances, embedded_assigned_peaks, - embedded_global + embedded_global, ] for x1, x2, x3, x4, x5, x6 in zip(*variables_to_zip): concatendated.append(torch.cat([x1, x2, x3, x4, x5, x6], dim=0)) @@ -157,13 +162,14 @@ class TransformerEncoderLayer(nn.Module): TransformerEncoderLayer is made up of self-attn and feedforward network with residual connections. """ + def __init__( self, d_model, nhead, dim_feedforward=2048, dropout=0.1, - activation : nn.Module = torch.nn.functional.relu, + activation: nn.Module = torch.nn.functional.relu, layer_norm_eps=1e-5, norm_first=True, bias=True, @@ -187,13 +193,16 @@ def __init__( self.linear2 = nn.Linear(dim_feedforward, d_model, bias=bias, **factory_kwargs) self.norm_first = norm_first - self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs) - self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs) + self.norm1 = nn.LayerNorm( + d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs + ) + self.norm2 = nn.LayerNorm( + d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs + ) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.activation = activation - def _sa_block(self, x): x = self.self_attn(x, x, x) @@ -204,10 +213,10 @@ def _ff_block(self, x): return self.dropout2(x) def forward(self, src): - ''' + """ Arguments: src: (batch_size, seq_len, d_model) - ''' + """ x = src if self.norm_first: x = x + self._sa_block(self.norm1(x)) @@ -216,7 +225,7 @@ def forward(self, src): x = self.norm1(x + self._sa_block(x)) x = self.norm2(x + self._ff_block(x)) return x - + class MultiHeadAttention(nn.Module): """ @@ -232,6 +241,7 @@ class MultiHeadAttention(nn.Module): dropout (float, optional): Dropout probability. Default: 0.0 bias (bool, optional): Whether to add bias to input projection. Default: True """ + def __init__( self, E_q: int, @@ -257,7 +267,9 @@ def __init__( self.E_head = E_total // nheads self.bias = bias - def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + def forward( + self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor + ) -> torch.Tensor: """ Forward pass; runs the following process: 1. Apply input projection @@ -290,7 +302,8 @@ def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) - # Step 3. Run SDPA # (N, nheads, L_t, E_head) attn_output = F.scaled_dot_product_attention( - query, key, value, dropout_p=self.dropout) + query, key, value, dropout_p=self.dropout + ) # (N, nheads, L_t, E_head) -> (N, L_t, nheads, E_head) -> (N, L_t, E_total) attn_output = attn_output.transpose(1, 2).flatten(-2) @@ -299,7 +312,7 @@ def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) - attn_output = self.out_proj(attn_output) return attn_output - + if __name__ == "__main__": # stupid code to find simple errors @@ -310,15 +323,17 @@ def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) - obs_noes=torch.randn(3, 3), close_distances=torch.randn(2, 4), assigned_peaks=torch.randn(2, 4), - peak_to_assign=0), + peak_to_assign=0, + ), NMRInput( obs_chemical_shifts=torch.randn(5, 2), pred_chemical_shifts=torch.randn(5, 2), obs_noes=torch.randn(4, 3), close_distances=torch.randn(7, 4), assigned_peaks=None, - peak_to_assign=0) + peak_to_assign=0, + ), ] trans = NMRTransformer() output = trans(nmr_inputs) - print(output) \ No newline at end of file + print(output) diff --git a/training_loop.py b/training_loop.py index d76babf..e7c18f5 100644 --- a/training_loop.py +++ b/training_loop.py @@ -1,14 +1,18 @@ import torch +import torcheval import pickle import numpy as np from nmr_transformer import * +from torch.utils.tensorboard import SummaryWriter -def extract_data(pickle_file='fake_histories.pkl'): - with open(pickle_file, 'rb') as f: + +def extract_data(pickle_file="fake_histories.pkl"): + with open(pickle_file, "rb") as f: histories = pickle.load(f) return histories -def preprocess_data(histories, batch_size): + +def preprocess_data(histories): """ Loops over chosen histories and answers, splits up components, and adds shifts if indices were used. """ @@ -18,49 +22,125 @@ def preprocess_data(histories, batch_size): close_dist = [] assignments = [] peak_to_assign = [] + correct_answer = [] # batch size needs to be reworked in here (not randomized - selects from start up until 'batch_size' index) - for history in histories[:batch_size]: + for history in histories: # Predicted shifts - predicted_shifts_temp = ([(i.H1, i.N15) for i in history['coords']]) + predicted_shifts_temp = [(i.H1, i.N15) for i in history["coords"]] predicted_shifts.append(predicted_shifts_temp) # Observed shifts - obs_chemical_shifts_temp = ([(i.H1, i.N15) for i in history['actual_shifts']]) + obs_chemical_shifts_temp = [(i.H1, i.N15) for i in history["actual_shifts"]] obs_chemical_shifts.append(obs_chemical_shifts_temp) # NOEs - noes.append(history['noes']) + noes.append(history["noes"]) + # Close distances (duplicated for reverse pair order) - close_dist.append([(predicted_shifts_temp[i.atom1], predicted_shifts_temp[i.atom2]) for i in history['connectivity']]) - close_dist.append([(predicted_shifts_temp[i.atom2], predicted_shifts_temp[i.atom1]) for i in history['connectivity']]) + close = [] + for i in history["connectivity"]: + atom1 = predicted_shifts_temp[i.atom1] + atom2 = predicted_shifts_temp[i.atom2] + close.append((atom1[0], atom1[1], atom2[0], atom2[1])) + close.append((atom2[0], atom2[1], atom1[0], atom1[1])) + close_dist.append(close) + # Assignments - assignments.append([(obs_chemical_shifts_temp[i], predicted_shifts_temp[j]) for i,j in list(history['assignments'].items())]) #shift:atom assignment, therefore obs_shift:predicted_shift + assign = [] + for i, j in list(history["assignments"].items()): + shift1 = obs_chemical_shifts_temp[i] + shift2 = predicted_shifts_temp[j] + assign.append((shift1[0], shift1[1], shift2[0], shift2[1])) + if assign: + assignments.append(assign) + else: + assignments.append(None) + # Peak to assign - peak_to_assign.append(obs_chemical_shifts_temp[history['assign_shift']]) - - return predicted_shifts, obs_chemical_shifts, noes, close_dist, assignments, peak_to_assign + peak_to_assign.append(history["assign_shift"]) + correct_answer.append(history["assign_shift"]) -def organize_nmr_inputs(histories, batch_size): + return ( + predicted_shifts, + obs_chemical_shifts, + noes, + close_dist, + assignments, + peak_to_assign, + correct_answer, + ) + + +def organize_nmr_inputs(histories): """ Sets up NMR inputs as list of named tuples. """ - predicted_shifts, actual_shifts, noes, close_dist, assignments, assign_peak = preprocess_data(histories, batch_size=batch_size) - + ( + predicted_shifts, + actual_shifts, + noes, + close_dist, + assignments, + assign_peak, + correct_answer, + ) = preprocess_data(histories) + nmr_inputs = [ - NMRInput( - obs_chemical_shifts=torch.tensor(actual_shifts[i]), - pred_chemical_shifts=torch.tensor(predicted_shifts[i]), - obs_noes=torch.tensor(noes[i]), - close_distances=torch.tensor(close_dist[i]), - assigned_peaks=torch.tensor(assignments[i]), - peak_to_assign=assign_peak[i]) # int or tensor for H,N peak? - for i in range(batch_size) + ( + NMRInput( + obs_chemical_shifts=torch.tensor(actual_shifts[i], dtype=torch.float), + pred_chemical_shifts=torch.tensor(predicted_shifts[i], dtype=torch.float), + obs_noes=torch.tensor(noes[i], dtype=torch.float), + close_distances=torch.tensor(close_dist[i], dtype=torch.float), + assigned_peaks=torch.tensor(assignments[i], dtype=torch.float) if assignments[i] else None, + peak_to_assign=assign_peak[i], + ), + correct_answer[i], + ) + for i in range(len(histories)) ] - print(nmr_inputs) return nmr_inputs + if __name__ == "__main__": - # need to loop over here - histories = extract_data() - nmr_inputs = organize_nmr_inputs(histories, batch_size=3) + # set up tensor board + writer = SummaryWriter() + + # load data + nmr_inputs = organize_nmr_inputs(extract_data()) + + # create our network and optimizer + net = NMRTransformer(dropout=0.0) + opt = torch.optim.Adam(net.parameters()) + policy_loss = torch.nn.CrossEntropyLoss() + + # training loop + iteration = 0 + while True: + iteration += 1 + xs = [inp[0] for inp in nmr_inputs] + ys = [inp[1] for inp in nmr_inputs] + + y_hats, _ = net(xs) + + loss = 0 + for y_hat, y in zip(y_hats, ys): + delta = policy_loss(y_hat, torch.tensor(y)) + loss += delta + + print(loss.item()) + + # correct = 0 + # trials = 0 + # for y_hat, y in zip(y_hats, ys): + # y_pred = torch.argmax(y_hat) + # if y_pred == y: + # correct += 1 + # trials += 1 + # accuracy = correct / trials + + writer.add_scalar("Loss/train", loss.item(), iteration) + # writer.add_scalar("Loss/accuracy", accuracy, iteration) + opt.zero_grad() + loss.backward() + opt.step() - \ No newline at end of file From 7e1fe64d56efe98061dc3edb32c7610de90984ff Mon Sep 17 00:00:00 2001 From: Justin MacCallum Date: Sat, 29 Mar 2025 11:22:58 -0600 Subject: [PATCH 26/43] padded version of transformer --- nmr_transformer.py | 376 +++++++++++++-------------------------------- training_loop.py | 53 ++++--- 2 files changed, 136 insertions(+), 293 deletions(-) diff --git a/nmr_transformer.py b/nmr_transformer.py index 839bad6..081e529 100644 --- a/nmr_transformer.py +++ b/nmr_transformer.py @@ -1,7 +1,7 @@ import torch import torch.nn as nn import torch.nn.functional as F -from typing import NamedTuple, List, Optional +from typing import NamedTuple, List, Optional, Tuple class NMRInput(NamedTuple): @@ -36,7 +36,7 @@ class NMRInput(NamedTuple): peak_to_assign: int -class NMRTransformer(torch.nn.Module): +class NMRTransformer(nn.Module): """ A transformer model for predicting the next peak to assign in an NMR spectrum. """ @@ -47,293 +47,125 @@ def __init__( n_heads: int = 4, n_layers: int = 4, dropout: float = 0.1, + device = None ): super(NMRTransformer, self).__init__() - self.embedder = NMRInitialEmbedding(n_hidden) - encoder_modules = [ - TransformerEncoderLayer(d_model=n_hidden, nhead=n_heads, dropout=dropout) - for _ in range(n_layers) - ] - self.encoder_layers = nn.Sequential(*encoder_modules) - self.policy_linear1 = nn.Linear(2 * n_hidden, 2 * n_hidden) - self.policy_linear2 = nn.Linear(2 * n_hidden, 1) - self.value_linear1 = nn.Linear(n_hidden, n_hidden) - self.value_linear2 = nn.Linear(n_hidden, 1) - - def forward(self, nmr_inputs: List[NMRInput]) -> torch.Tensor: - embedded = self.embedder(nmr_inputs) - x = self.encoder_layers(embedded) - - # compute the policy for each input in the batch + self.embedder = NMRInitialEmbedding(n_hidden, device=device) + encoder_layers = nn.TransformerEncoderLayer( + d_model=n_hidden, nhead=n_heads, dropout=dropout, device=device + ) + self.encoder = nn.TransformerEncoder(encoder_layers, num_layers=n_layers) + self.policy_linear1 = nn.Linear(2 * n_hidden, 2 * n_hidden, device=device) + self.policy_linear2 = nn.Linear(2 * n_hidden, 1, device=device) + self.value_linear1 = nn.Linear(n_hidden, n_hidden, device=device) + self.value_linear2 = nn.Linear(n_hidden, 1, device=device) + + def forward(self, nmr_inputs: List[NMRInput]) -> Tuple[torch.Tensor, torch.Tensor]: + embedded, mask = self.embedder(nmr_inputs) # (N, S, D), (N, S) + embedded = embedded.transpose( + 0, 1 + ) # (S, N, D) for nn.Transformer compatibility + + x = self.encoder(embedded, src_key_padding_mask=mask, is_causal=False) # (S, N, D) + x = x.transpose(0, 1) # (N, S, D) + policies = [] - for i in range(len(nmr_inputs)): - n_res = len(nmr_inputs[i].pred_chemical_shifts) - residue_embeddings = x[i][:n_res] - peak_to_assign = nmr_inputs[i].peak_to_assign - assert peak_to_assign >= 0 and peak_to_assign < len( - nmr_inputs[i].obs_chemical_shifts - ) - peak_embedding = x[i][n_res + peak_to_assign].expand(n_res, -1) + for i, nmr_input in enumerate(nmr_inputs): + n_res = len(nmr_input.pred_chemical_shifts) + residue_embeddings = x[i, 1 : (n_res + 1), :] + peak_to_assign = nmr_input.peak_to_assign + peak_embedding = x[i, n_res + peak_to_assign + 1, :].expand(n_res, -1) embeddings = torch.cat([residue_embeddings, peak_embedding], dim=1) policy = self.policy_linear2( F.relu(self.policy_linear1(embeddings)) ).reshape(-1) policies.append(policy) - # compute the value for each input in the batch - embeddings = [] - for i in range(len(nmr_inputs)): - # the global token is always the last one - embedding = x[i][-1] - embeddings.append(embedding) - embeddings = torch.stack(embeddings) + embeddings = x[:, 0, :] # Global tokens for each input in the batch values = self.value_linear2(F.relu(self.value_linear1(embeddings))).reshape(-1) return policies, values class NMRInitialEmbedding(nn.Module): - """ - Embeds the initial input for the NMR transformer. - """ - - def __init__(self, n_hidden: int = 128): + def __init__(self, n_hidden: int = 128, device=None): super(NMRInitialEmbedding, self).__init__() - self.linear_obs_shifts = nn.Linear(2, n_hidden) - self.linear_pred_shifts = nn.Linear(2, n_hidden) - self.linear_obs_noes = nn.Linear(3, n_hidden) - self.linear_close_distances = nn.Linear(4, n_hidden) - self.linear_assigned_peaks = nn.Linear(4, n_hidden) - self.embeddings = nn.Embedding(6, n_hidden) - - def forward(self, nmr_inputs: List[NMRInput]) -> torch.Tensor: - # embed each of the inputs - embedded_obs_shifts = self._handle_input( - [nmr.obs_chemical_shifts for nmr in nmr_inputs], self.linear_obs_shifts, 0 - ) - embedded_pred_shifts = self._handle_input( - [nmr.pred_chemical_shifts for nmr in nmr_inputs], self.linear_pred_shifts, 1 - ) - embedded_obs_noes = self._handle_input( - [nmr.obs_noes for nmr in nmr_inputs], self.linear_obs_noes, 2 - ) - embedded_close_distances = self._handle_input( - [nmr.close_distances for nmr in nmr_inputs], self.linear_close_distances, 3 - ) - embedded_assigned_peaks = self._handle_input( - [nmr.assigned_peaks for nmr in nmr_inputs], self.linear_assigned_peaks, 4 - ) - - # collect the "global" token for each input - embedded_global = [] - embed = self.embeddings(torch.tensor([5])) - for _ in nmr_inputs: - embedded_global.append(embed) - - # concatenate the two sets of embeddings - concatendated = [] - variables_to_zip = [ - embedded_pred_shifts, - embedded_obs_shifts, - embedded_obs_noes, - embedded_close_distances, - embedded_assigned_peaks, - embedded_global, - ] - for x1, x2, x3, x4, x5, x6 in zip(*variables_to_zip): - concatendated.append(torch.cat([x1, x2, x3, x4, x5, x6], dim=0)) - - return torch.nested.nested_tensor(concatendated, layout=torch.jagged) - - def _handle_input(self, xs, linear, embed_index): - embedded = [] - for x in xs: - if x is None: - embedded.append(torch.Tensor()) + self.device = device + self.n_hidden = n_hidden + self.linear_obs_shifts = nn.Linear(2, n_hidden, device=device) + self.linear_pred_shifts = nn.Linear(2, n_hidden, device=device) + self.linear_obs_noes = nn.Linear(3, n_hidden, device=device) + self.linear_close_distances = nn.Linear(4, n_hidden, device=device) + self.linear_assigned_peaks = nn.Linear(4, n_hidden, device=device) + self.embeddings = nn.Embedding(6, n_hidden, device=device) + + def forward(self, nmr_inputs: List[NMRInput]): + batch_embeddings = [] + lengths = [] + + for nmr in nmr_inputs: + parts = [] + + global_token = self.embeddings( + torch.tensor([5], device=nmr.pred_chemical_shifts.device) + ).expand(1, -1) + parts.append(global_token) + + x = self.linear_pred_shifts(nmr.pred_chemical_shifts) + x += self.embeddings(torch.tensor([0], device=x.device)) + parts.append(x) + + x = self.linear_obs_shifts(nmr.obs_chemical_shifts) + x += self.embeddings(torch.tensor([1], device=x.device)) + parts.append(x) + + if nmr.obs_noes is not None: + x = self.linear_obs_noes(nmr.obs_noes) + x += self.embeddings(torch.tensor([2], device=x.device)) + parts.append(x) + + if nmr.close_distances is not None: + x = self.linear_close_distances(nmr.close_distances) + x += self.embeddings(torch.tensor([3], device=x.device)) + parts.append(x) + + if nmr.assigned_peaks is not None: + x = self.linear_assigned_peaks(nmr.assigned_peaks) + x += self.embeddings(torch.tensor([4], device=x.device)) + parts.append(x) + + full_seq = torch.cat(parts, dim=0) + batch_embeddings.append(full_seq) + lengths.append(full_seq.size(0)) + + max_len = max(lengths) + padded = [] + padding_mask = [] + + for emb in batch_embeddings: + emb_len = emb.size(0) + + if emb_len < max_len: + # Create a padding tensor to append + pad_size = max_len - emb_len + pad_tensor = torch.zeros(pad_size, self.n_hidden, device=emb.device) + padded_emb = torch.cat([emb, pad_tensor], dim=0) + + # Create padding mask + mask = torch.cat( + [ + torch.zeros(emb_len, dtype=torch.bool, device=emb.device), + torch.ones(pad_size, dtype=torch.bool, device=emb.device), + ] + ) else: - x = linear(x) - x += self.embeddings(torch.tensor([embed_index])) - embedded.append(x) - return embedded - - -class TransformerEncoderLayer(nn.Module): - """ - TransformerEncoderLayer is made up of self-attn and feedforward network - with residual connections. - """ + padded_emb = emb + mask = torch.zeros(emb_len, dtype=torch.bool, device=emb.device) - def __init__( - self, - d_model, - nhead, - dim_feedforward=2048, - dropout=0.1, - activation: nn.Module = torch.nn.functional.relu, - layer_norm_eps=1e-5, - norm_first=True, - bias=True, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.self_attn = MultiHeadAttention( - d_model, - d_model, - d_model, - d_model, - nhead, - dropout=dropout, - bias=bias, - **factory_kwargs, - ) - self.linear1 = nn.Linear(d_model, dim_feedforward, bias=bias, **factory_kwargs) - self.dropout = nn.Dropout(dropout) - self.linear2 = nn.Linear(dim_feedforward, d_model, bias=bias, **factory_kwargs) + padded.append(padded_emb) + padding_mask.append(mask) - self.norm_first = norm_first - self.norm1 = nn.LayerNorm( - d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs - ) - self.norm2 = nn.LayerNorm( - d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs - ) + padded = torch.stack(padded) + padding_mask = torch.stack(padding_mask) - self.dropout1 = nn.Dropout(dropout) - self.dropout2 = nn.Dropout(dropout) - self.activation = activation - - def _sa_block(self, x): - x = self.self_attn(x, x, x) - return self.dropout1(x) - - def _ff_block(self, x): - x = self.linear2(self.dropout(self.activation(self.linear1(x)))) - return self.dropout2(x) - - def forward(self, src): - """ - Arguments: - src: (batch_size, seq_len, d_model) - """ - x = src - if self.norm_first: - x = x + self._sa_block(self.norm1(x)) - x = x + self._ff_block(self.norm2(x)) - else: - x = self.norm1(x + self._sa_block(x)) - x = self.norm2(x + self._ff_block(x)) - return x - - -class MultiHeadAttention(nn.Module): - """ - Computes multi-head attention. Supports nested or padded tensors. - - Args: - E_q (int): Size of embedding dim for query - E_k (int): Size of embedding dim for key - E_v (int): Size of embedding dim for value - E_total (int): Total embedding dim of combined heads post input projection. Each head - has dim E_total // nheads - nheads (int): Number of heads - dropout (float, optional): Dropout probability. Default: 0.0 - bias (bool, optional): Whether to add bias to input projection. Default: True - """ - - def __init__( - self, - E_q: int, - E_k: int, - E_v: int, - E_total: int, - nheads: int, - dropout: float = 0.0, - bias=True, - device=None, - dtype=None, - ): - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.nheads = nheads - self.dropout = dropout - self.q_proj = nn.Linear(E_q, E_total, bias=bias, **factory_kwargs) - self.k_proj = nn.Linear(E_k, E_total, bias=bias, **factory_kwargs) - self.v_proj = nn.Linear(E_v, E_total, bias=bias, **factory_kwargs) - E_out = E_q - self.out_proj = nn.Linear(E_total, E_out, bias=bias, **factory_kwargs) - assert E_total % nheads == 0, "Embedding dim is not divisible by nheads" - self.E_head = E_total // nheads - self.bias = bias - - def forward( - self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor - ) -> torch.Tensor: - """ - Forward pass; runs the following process: - 1. Apply input projection - 2. Split heads and prepare for SDPA - 3. Run SDPA - 4. Apply output projection - - Args: - query (torch.Tensor): query of shape (N, L_q, E_qk) - key (torch.Tensor): key of shape (N, L_kv, E_qk) - value (torch.Tensor): value of shape (N, L_kv, E_v) - - Returns: - attn_output (torch.Tensor): output of shape (N, L_t, E_q) - """ - # Step 1. Apply input projection - query = self.q_proj(query) - key = self.k_proj(key) - value = self.v_proj(value) - - # Step 2. Split heads and prepare for SDPA - # reshape query, key, value to separate by head - # (N, L_t, E_total) -> (N, L_t, nheads, E_head) -> (N, nheads, L_t, E_head) - query = query.unflatten(-1, [self.nheads, self.E_head]).transpose(1, 2) - # (N, L_s, E_total) -> (N, L_s, nheads, E_head) -> (N, nheads, L_s, E_head) - key = key.unflatten(-1, [self.nheads, self.E_head]).transpose(1, 2) - # (N, L_s, E_total) -> (N, L_s, nheads, E_head) -> (N, nheads, L_s, E_head) - value = value.unflatten(-1, [self.nheads, self.E_head]).transpose(1, 2) - - # Step 3. Run SDPA - # (N, nheads, L_t, E_head) - attn_output = F.scaled_dot_product_attention( - query, key, value, dropout_p=self.dropout - ) - # (N, nheads, L_t, E_head) -> (N, L_t, nheads, E_head) -> (N, L_t, E_total) - attn_output = attn_output.transpose(1, 2).flatten(-2) - - # Step 4. Apply output projection - # (N, L_t, E_total) -> (N, L_t, E_out) - attn_output = self.out_proj(attn_output) - - return attn_output - - -if __name__ == "__main__": - # stupid code to find simple errors - nmr_inputs = [ - NMRInput( - obs_chemical_shifts=torch.randn(2, 2), - pred_chemical_shifts=torch.randn(2, 2), - obs_noes=torch.randn(3, 3), - close_distances=torch.randn(2, 4), - assigned_peaks=torch.randn(2, 4), - peak_to_assign=0, - ), - NMRInput( - obs_chemical_shifts=torch.randn(5, 2), - pred_chemical_shifts=torch.randn(5, 2), - obs_noes=torch.randn(4, 3), - close_distances=torch.randn(7, 4), - assigned_peaks=None, - peak_to_assign=0, - ), - ] - trans = NMRTransformer() - output = trans(nmr_inputs) - print(output) + return padded, padding_mask diff --git a/training_loop.py b/training_loop.py index e7c18f5..64fb54d 100644 --- a/training_loop.py +++ b/training_loop.py @@ -70,7 +70,7 @@ def preprocess_data(histories): ) -def organize_nmr_inputs(histories): +def organize_nmr_inputs(histories, device): """ Sets up NMR inputs as list of named tuples. """ @@ -87,11 +87,11 @@ def organize_nmr_inputs(histories): nmr_inputs = [ ( NMRInput( - obs_chemical_shifts=torch.tensor(actual_shifts[i], dtype=torch.float), - pred_chemical_shifts=torch.tensor(predicted_shifts[i], dtype=torch.float), - obs_noes=torch.tensor(noes[i], dtype=torch.float), - close_distances=torch.tensor(close_dist[i], dtype=torch.float), - assigned_peaks=torch.tensor(assignments[i], dtype=torch.float) if assignments[i] else None, + obs_chemical_shifts=torch.tensor(actual_shifts[i], dtype=torch.float, device=device), + pred_chemical_shifts=torch.tensor(predicted_shifts[i], dtype=torch.float, device=device), + obs_noes=torch.tensor(noes[i], dtype=torch.float, device=device), + close_distances=torch.tensor(close_dist[i], dtype=torch.float, device=device), + assigned_peaks=torch.tensor(assignments[i], dtype=torch.float, device=device) if assignments[i] else None, peak_to_assign=assign_peak[i], ), correct_answer[i], @@ -102,15 +102,24 @@ def organize_nmr_inputs(histories): if __name__ == "__main__": + # set device + device = "cpu" + # set up tensor board writer = SummaryWriter() # load data - nmr_inputs = organize_nmr_inputs(extract_data()) + nmr_inputs = organize_nmr_inputs(extract_data(), device=device) # create our network and optimizer - net = NMRTransformer(dropout=0.0) - opt = torch.optim.Adam(net.parameters()) + net = NMRTransformer(dropout=0.0, device=device) + for layer in net.modules(): + if isinstance(layer, (nn.Conv2d, nn.Linear)): + torch.nn.init.kaiming_normal_(layer.weight, mode='fan_out', nonlinearity='relu') + if layer.bias is not None: + torch.nn.init.zeros_(layer.bias) + + opt = torch.optim.Adam(net.parameters(), lr=1e-5) policy_loss = torch.nn.CrossEntropyLoss() # training loop @@ -123,23 +132,25 @@ def organize_nmr_inputs(histories): y_hats, _ = net(xs) loss = 0 + count = 0 for y_hat, y in zip(y_hats, ys): - delta = policy_loss(y_hat, torch.tensor(y)) + delta = policy_loss(y_hat, torch.tensor(y, device=device)) loss += delta + count += 1 + loss = loss / count - print(loss.item()) - - # correct = 0 - # trials = 0 - # for y_hat, y in zip(y_hats, ys): - # y_pred = torch.argmax(y_hat) - # if y_pred == y: - # correct += 1 - # trials += 1 - # accuracy = correct / trials + correct = 0 + trials = 0 + for y_hat, y in zip(y_hats, ys): + y_pred = torch.argmax(y_hat) + if y_pred == y: + correct += 1 + trials += 1 + accuracy = correct / trials + print(loss.item(), accuracy) writer.add_scalar("Loss/train", loss.item(), iteration) - # writer.add_scalar("Loss/accuracy", accuracy, iteration) + writer.add_scalar("Loss/accuracy", accuracy, iteration) opt.zero_grad() loss.backward() opt.step() From 201fc1aa42a8208ffe4dfea31c316974c228ff6e Mon Sep 17 00:00:00 2001 From: Justin MacCallum Date: Sat, 29 Mar 2025 11:31:07 -0600 Subject: [PATCH 27/43] Update training --- training_loop.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/training_loop.py b/training_loop.py index e7c18f5..c904d87 100644 --- a/training_loop.py +++ b/training_loop.py @@ -110,6 +110,7 @@ def organize_nmr_inputs(histories): # create our network and optimizer net = NMRTransformer(dropout=0.0) + net.compile() opt = torch.optim.Adam(net.parameters()) policy_loss = torch.nn.CrossEntropyLoss() @@ -123,23 +124,27 @@ def organize_nmr_inputs(histories): y_hats, _ = net(xs) loss = 0 + count = 0 for y_hat, y in zip(y_hats, ys): delta = policy_loss(y_hat, torch.tensor(y)) loss += delta + count += 1 + + loss = loss / count - print(loss.item()) + correct = 0 + trials = 0 + for y_hat, y in zip(y_hats, ys): + y_pred = torch.argmax(y_hat) + if y_pred == y: + correct += 1 + trials += 1 + accuracy = correct / trials - # correct = 0 - # trials = 0 - # for y_hat, y in zip(y_hats, ys): - # y_pred = torch.argmax(y_hat) - # if y_pred == y: - # correct += 1 - # trials += 1 - # accuracy = correct / trials + print(loss.item(), accuracy) writer.add_scalar("Loss/train", loss.item(), iteration) - # writer.add_scalar("Loss/accuracy", accuracy, iteration) + writer.add_scalar("Loss/accuracy", accuracy, iteration) opt.zero_grad() loss.backward() opt.step() From 538ac1636f1cf44f2d1fcdf5b760af3d99282d58 Mon Sep 17 00:00:00 2001 From: Justin MacCallum Date: Sat, 29 Mar 2025 12:24:59 -0600 Subject: [PATCH 28/43] Implement train-test split --- training_loop.py | 132 ++++++++++++++++++++++++++++++++--------------- 1 file changed, 89 insertions(+), 43 deletions(-) diff --git a/training_loop.py b/training_loop.py index 64fb54d..eb47ade 100644 --- a/training_loop.py +++ b/training_loop.py @@ -1,9 +1,9 @@ import torch -import torcheval import pickle import numpy as np from nmr_transformer import * from torch.utils.tensorboard import SummaryWriter +import random def extract_data(pickle_file="fake_histories.pkl"): @@ -101,57 +101,103 @@ def organize_nmr_inputs(histories, device): return nmr_inputs + if __name__ == "__main__": - # set device + # Set device device = "cpu" - # set up tensor board + # Set up TensorBoard writer = SummaryWriter() - # load data + # Load data nmr_inputs = organize_nmr_inputs(extract_data(), device=device) - # create our network and optimizer - net = NMRTransformer(dropout=0.0, device=device) - for layer in net.modules(): - if isinstance(layer, (nn.Conv2d, nn.Linear)): - torch.nn.init.kaiming_normal_(layer.weight, mode='fan_out', nonlinearity='relu') - if layer.bias is not None: - torch.nn.init.zeros_(layer.bias) + # Shuffle the data + random.shuffle(nmr_inputs) + + # Split dataset manually into training and testing + train_size = int(0.8 * len(nmr_inputs)) + train_data = nmr_inputs[:train_size] + test_data = nmr_inputs[train_size:] + # Create our network and optimizer + net = NMRTransformer(dropout=0.1, device=device) opt = torch.optim.Adam(net.parameters(), lr=1e-5) policy_loss = torch.nn.CrossEntropyLoss() - # training loop - iteration = 0 - while True: - iteration += 1 - xs = [inp[0] for inp in nmr_inputs] - ys = [inp[1] for inp in nmr_inputs] - - y_hats, _ = net(xs) - - loss = 0 - count = 0 - for y_hat, y in zip(y_hats, ys): - delta = policy_loss(y_hat, torch.tensor(y, device=device)) - loss += delta - count += 1 - loss = loss / count - - correct = 0 - trials = 0 - for y_hat, y in zip(y_hats, ys): - y_pred = torch.argmax(y_hat) - if y_pred == y: - correct += 1 - trials += 1 - accuracy = correct / trials - print(loss.item(), accuracy) - - writer.add_scalar("Loss/train", loss.item(), iteration) - writer.add_scalar("Loss/accuracy", accuracy, iteration) - opt.zero_grad() - loss.backward() - opt.step() + # Training loop + batch_size = 8 + epochs = 10_000 # Define the number of epochs + eval_interval = 100 # Evaluate the model every 100 iterations + iteration = 0 + for epoch in range(epochs): + net.train() + random.shuffle(train_data) # Shuffle training data each epoch + + # Break training data into batches manually + for i in range(0, len(train_data), batch_size): + batch = train_data[i:i + batch_size] + xs = [inp[0] for inp in batch] + ys = [inp[1] for inp in batch] + + iteration += 1 + + y_hats, _ = net(xs) + + loss = 0 + correct = 0 + count = 0 + + for y_hat, y in zip(y_hats, ys): + y = torch.tensor(y, device=device) + loss += policy_loss(y_hat, y) + + y_pred = torch.argmax(y_hat) + if y_pred == y: + correct += 1 + count += 1 + + loss = loss / count + accuracy = correct / count + + print(f"Iteration {iteration} - Loss: {loss.item()} - Accuracy: {accuracy}") + + writer.add_scalar("Loss/train", loss.item(), iteration) + writer.add_scalar("Accuracy/train", accuracy, iteration) + + opt.zero_grad() + loss.backward() + opt.step() + + # Evaluate on test set periodically + if iteration % eval_interval == 0: + net.eval() + test_loss = 0 + test_correct = 0 + test_trials = 0 + + with torch.no_grad(): + for j in range(0, len(test_data), batch_size): + test_batch = test_data[j:j + batch_size] + xs = [inp[0] for inp in test_batch] + ys = [inp[1] for inp in test_batch] + + test_y_hats, _ = net(xs) + + batch_loss = 0 + for y_hat, y in zip(test_y_hats, ys): + y = torch.tensor(y, device=device) + batch_loss += policy_loss(y_hat, y).item() + if torch.argmax(y_hat) == y: + test_correct += 1 + test_trials += 1 + + test_loss += batch_loss / len(test_batch) + + test_loss /= (len(test_data) / batch_size) + test_accuracy = test_correct / test_trials + + print(f"Test Loss: {test_loss} - Test Accuracy: {test_accuracy}") + writer.add_scalar("Loss/test", test_loss, iteration) + writer.add_scalar("Accuracy/test", test_accuracy, iteration) From caef94de686eaaea35f429834b28e5811e46bc99 Mon Sep 17 00:00:00 2001 From: Justin MacCallum Date: Mon, 7 Apr 2025 16:09:11 -0600 Subject: [PATCH 29/43] Add device support --- nmr_transformer.py | 19 +++++++++++-------- training_loop.py | 12 ++++++------ 2 files changed, 17 insertions(+), 14 deletions(-) diff --git a/nmr_transformer.py b/nmr_transformer.py index 081e529..23e00af 100644 --- a/nmr_transformer.py +++ b/nmr_transformer.py @@ -47,12 +47,16 @@ def __init__( n_heads: int = 4, n_layers: int = 4, dropout: float = 0.1, - device = None + device=None, ): super(NMRTransformer, self).__init__() self.embedder = NMRInitialEmbedding(n_hidden, device=device) encoder_layers = nn.TransformerEncoderLayer( - d_model=n_hidden, nhead=n_heads, dropout=dropout, device=device + d_model=n_hidden, + dim_feedforward=n_hidden * n_heads, + nhead=n_heads, + dropout=dropout, + device=device, ) self.encoder = nn.TransformerEncoder(encoder_layers, num_layers=n_layers) self.policy_linear1 = nn.Linear(2 * n_hidden, 2 * n_hidden, device=device) @@ -66,7 +70,9 @@ def forward(self, nmr_inputs: List[NMRInput]) -> Tuple[torch.Tensor, torch.Tenso 0, 1 ) # (S, N, D) for nn.Transformer compatibility - x = self.encoder(embedded, src_key_padding_mask=mask, is_causal=False) # (S, N, D) + x = self.encoder( + embedded, src_key_padding_mask=mask, is_causal=False + ) # (S, N, D) x = x.transpose(0, 1) # (N, S, D) policies = [] @@ -74,11 +80,8 @@ def forward(self, nmr_inputs: List[NMRInput]) -> Tuple[torch.Tensor, torch.Tenso n_res = len(nmr_input.pred_chemical_shifts) residue_embeddings = x[i, 1 : (n_res + 1), :] peak_to_assign = nmr_input.peak_to_assign - peak_embedding = x[i, n_res + peak_to_assign + 1, :].expand(n_res, -1) - embeddings = torch.cat([residue_embeddings, peak_embedding], dim=1) - policy = self.policy_linear2( - F.relu(self.policy_linear1(embeddings)) - ).reshape(-1) + peak_embedding = x[i, n_res + peak_to_assign + 1, :] + policy = torch.matmul(residue_embeddings, peak_embedding) policies.append(policy) embeddings = x[:, 0, :] # Global tokens for each input in the batch diff --git a/training_loop.py b/training_loop.py index eb47ade..b76baa8 100644 --- a/training_loop.py +++ b/training_loop.py @@ -104,7 +104,7 @@ def organize_nmr_inputs(histories, device): if __name__ == "__main__": # Set device - device = "cpu" + device = "mps" # Set up TensorBoard writer = SummaryWriter() @@ -121,22 +121,22 @@ def organize_nmr_inputs(histories, device): test_data = nmr_inputs[train_size:] # Create our network and optimizer - net = NMRTransformer(dropout=0.1, device=device) - opt = torch.optim.Adam(net.parameters(), lr=1e-5) + net = NMRTransformer(dropout=0.1, n_layers=4, device=device) + opt = torch.optim.AdamW(net.parameters(), lr=1e-4, weight_decay=0.01) policy_loss = torch.nn.CrossEntropyLoss() # Training loop - batch_size = 8 - epochs = 10_000 # Define the number of epochs + batch_size = 10 + epochs = 500_000 # Define the number of epochs eval_interval = 100 # Evaluate the model every 100 iterations iteration = 0 for epoch in range(epochs): - net.train() random.shuffle(train_data) # Shuffle training data each epoch # Break training data into batches manually for i in range(0, len(train_data), batch_size): + net.train() batch = train_data[i:i + batch_size] xs = [inp[0] for inp in batch] ys = [inp[1] for inp in batch] From 9fc7eda99a5c4e4a65c0991270b5f730eaf99f41 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Fri, 23 May 2025 14:40:53 -0600 Subject: [PATCH 30/43] Linear mapping --- nmr_transformer.py | 88 ++++++++++++++++++++++++++++++++++++++++++++-- training_loop.py | 18 +++++----- 2 files changed, 95 insertions(+), 11 deletions(-) diff --git a/nmr_transformer.py b/nmr_transformer.py index 23e00af..016f368 100644 --- a/nmr_transformer.py +++ b/nmr_transformer.py @@ -51,6 +51,7 @@ def __init__( ): super(NMRTransformer, self).__init__() self.embedder = NMRInitialEmbedding(n_hidden, device=device) + self.linearmapping = NMRInputMapping() encoder_layers = nn.TransformerEncoderLayer( d_model=n_hidden, dim_feedforward=n_hidden * n_heads, @@ -59,29 +60,37 @@ def __init__( device=device, ) self.encoder = nn.TransformerEncoder(encoder_layers, num_layers=n_layers) + self.layer_norm = nn.LayerNorm(n_hidden, device=device)############### + self.policy_linear1 = nn.Linear(2 * n_hidden, 2 * n_hidden, device=device) self.policy_linear2 = nn.Linear(2 * n_hidden, 1, device=device) self.value_linear1 = nn.Linear(n_hidden, n_hidden, device=device) self.value_linear2 = nn.Linear(n_hidden, 1, device=device) def forward(self, nmr_inputs: List[NMRInput]) -> Tuple[torch.Tensor, torch.Tensor]: - embedded, mask = self.embedder(nmr_inputs) # (N, S, D), (N, S) + + linear_values = self.linearmapping(nmr_inputs) + embedded, mask = self.embedder(linear_values) + + # embedded, mask = self.embedder(nmr_inputs) # (N, S, D), (N, S) embedded = embedded.transpose( 0, 1 ) # (S, N, D) for nn.Transformer compatibility + # embedded_norm = self.layer_norm(embedded) + x = self.encoder( embedded, src_key_padding_mask=mask, is_causal=False ) # (S, N, D) x = x.transpose(0, 1) # (N, S, D) policies = [] - for i, nmr_input in enumerate(nmr_inputs): + for i, nmr_input in enumerate(linear_values): n_res = len(nmr_input.pred_chemical_shifts) residue_embeddings = x[i, 1 : (n_res + 1), :] peak_to_assign = nmr_input.peak_to_assign peak_embedding = x[i, n_res + peak_to_assign + 1, :] - policy = torch.matmul(residue_embeddings, peak_embedding) + policy = (torch.matmul(residue_embeddings, peak_embedding)) policies.append(policy) embeddings = x[:, 0, :] # Global tokens for each input in the batch @@ -90,6 +99,79 @@ def forward(self, nmr_inputs: List[NMRInput]) -> Tuple[torch.Tensor, torch.Tenso return policies, values +class NMRInputMapping(nn.Module): + def __init__( + self, + xmin_h: int = 6, + xmax_h: int = 10, + xmin_n: int = 100, + xmax_n: int = 135 + ): + super(NMRInputMapping, self).__init__() + self.xmin_h = xmin_h + self.xmax_h = xmax_h + self.xmin_n = xmin_n + self.xmax_n = xmax_n + + def linear_transform(self, x, xmin, xmax): + value = 2*((x - xmin)/(xmax - xmin)) - 1 + return value + + def forward(self, nmr_inputs: List[NMRInput]): #-> Tuple[torch.Tensor, torch.Tensor]: + + linear_values = [] + for i, nmr_input in enumerate(nmr_inputs): + pred_shift1 = self.linear_transform(nmr_input.pred_chemical_shifts[:,0], self.xmin_h, self.xmax_h) + pred_shift2 = self.linear_transform(nmr_input.pred_chemical_shifts[:,1], self.xmin_n, self.xmax_n) + pred_chemical_shifts = torch.stack((pred_shift1, pred_shift2), dim=-1) + # print(pred_chemical_shifts.shape) + + obs_shift1 = self.linear_transform(nmr_input.obs_chemical_shifts[:,0], self.xmin_h, self.xmax_h) + obs_shift2 = self.linear_transform(nmr_input.obs_chemical_shifts[:,1], self.xmin_n, self.xmax_n) + obs_chemical_shifts = torch.stack((obs_shift1, obs_shift2), dim=-1) + # print(obs_chemical_shifts.shape) + + if nmr_input.obs_noes is not None: + noe1 = self.linear_transform(nmr_input.obs_noes[:,0], self.xmin_h, self.xmax_h) + noe2 = self.linear_transform(nmr_input.obs_noes[:,1], self.xmin_n, self.xmax_n) + noe3 = self.linear_transform(nmr_input.obs_noes[:,2], self.xmin_h, self.xmax_h) + obs_noes = torch.stack((noe1, noe2, noe3), dim=-1) + else: + obs_noes = nmr_input.obs_noes + # print(obs_noes.shape) + + if nmr_input.close_distances is not None: + dist1 = self.linear_transform(nmr_input.close_distances[:,0], self.xmin_h, self.xmax_h) + dist2 = self.linear_transform(nmr_input.close_distances[:,1], self.xmin_n, self.xmax_n) + dist3 = self.linear_transform(nmr_input.close_distances[:,2], self.xmin_h, self.xmax_h) + dist4 = self.linear_transform(nmr_input.close_distances[:,3], self.xmin_n, self.xmax_n) + close_distances = torch.stack((dist1, dist2, dist3, dist4), dim=-1) + else: + close_distances = nmr_input.close_distances + # print(close_distances.shape) + + if nmr_input.assigned_peaks is not None: + peak1 = self.linear_transform(nmr_input.assigned_peaks[:,0], self.xmin_h, self.xmax_h) + peak2 = self.linear_transform(nmr_input.assigned_peaks[:,1], self.xmin_n, self.xmax_n) + peak3 = self.linear_transform(nmr_input.assigned_peaks[:,2], self.xmin_h, self.xmax_h) + peak4 = self.linear_transform(nmr_input.assigned_peaks[:,3], self.xmin_n, self.xmax_n) + assigned_peaks = torch.stack((peak1, peak2, peak3, peak4), dim=-1) + else: + assigned_peaks = nmr_input.assigned_peaks + + peak_to_assign = nmr_input.peak_to_assign + + linear_values.append(NMRInput(obs_chemical_shifts=obs_chemical_shifts, + pred_chemical_shifts=pred_chemical_shifts, + obs_noes=obs_noes, + close_distances=close_distances, + assigned_peaks=assigned_peaks, + peak_to_assign=peak_to_assign)) + + return linear_values + + + class NMRInitialEmbedding(nn.Module): def __init__(self, n_hidden: int = 128, device=None): super(NMRInitialEmbedding, self).__init__() diff --git a/training_loop.py b/training_loop.py index b76baa8..0381600 100644 --- a/training_loop.py +++ b/training_loop.py @@ -6,7 +6,8 @@ import random -def extract_data(pickle_file="fake_histories.pkl"): + +def extract_data(pickle_file="fake_histories_r20_1.pkl"): with open(pickle_file, "rb") as f: histories = pickle.load(f) return histories @@ -24,13 +25,12 @@ def preprocess_data(histories): peak_to_assign = [] correct_answer = [] - # batch size needs to be reworked in here (not randomized - selects from start up until 'batch_size' index) for history in histories: # Predicted shifts - predicted_shifts_temp = [(i.H1, i.N15) for i in history["coords"]] + predicted_shifts_temp = [(i.H1, i.N15) for i in history["coordinates"]] predicted_shifts.append(predicted_shifts_temp) # Observed shifts - obs_chemical_shifts_temp = [(i.H1, i.N15) for i in history["actual_shifts"]] + obs_chemical_shifts_temp = [(i.H1, i.N15) for i in history["obs_chemical_shifts"]] obs_chemical_shifts.append(obs_chemical_shifts_temp) # NOEs noes.append(history["noes"]) @@ -56,8 +56,8 @@ def preprocess_data(histories): assignments.append(None) # Peak to assign - peak_to_assign.append(history["assign_shift"]) - correct_answer.append(history["assign_shift"]) + peak_to_assign.append(history["shift_to_assign"]) + correct_answer.append(history["shift_to_assign"]) return ( predicted_shifts, @@ -66,7 +66,7 @@ def preprocess_data(histories): close_dist, assignments, peak_to_assign, - correct_answer, + correct_answer ) @@ -103,8 +103,10 @@ def organize_nmr_inputs(histories, device): if __name__ == "__main__": + # Set device - device = "mps" + # device = "mps" + device ='cuda' # Set up TensorBoard writer = SummaryWriter() From 140e71381c031ebe5f467bbe8436641a6717376e Mon Sep 17 00:00:00 2001 From: abbie-p Date: Fri, 23 May 2025 14:41:27 -0600 Subject: [PATCH 31/43] Scaling visualization --- visualize_data.py | 217 +++++++++++++++++++++++++++++----------------- write_pdb.py | 63 ++++++++++++++ 2 files changed, 201 insertions(+), 79 deletions(-) create mode 100644 write_pdb.py diff --git a/visualize_data.py b/visualize_data.py index fbb0f1a..9c8bde4 100644 --- a/visualize_data.py +++ b/visualize_data.py @@ -2,89 +2,148 @@ import matplotlib.pyplot as plt import numpy as np from mpl_toolkits.mplot3d import Axes3D - -def get_data(actual_shifts, coords, connectivity, named_tuple_used=True): - if named_tuple_used: - # Actual - actual_x = [shift.H1 for shift in actual_shifts] - actual_y = [shift.N15 for shift in actual_shifts] - # Predicted - predict_x = [shift.H1 for shift in coords] - predict_y = [shift.N15 for shift in coords] - # Connectivity - contacts = [(contact.atom1, contact.atom2) for contact in connectivity] - - else: - # Actual - actual_x = [shift[0] for shift in actual_shifts] - actual_y = [shift[1] for shift in actual_shifts] - # Predicted - predict_x = [shift[3] for shift in coords] - predict_y = [shift[4] for shift in coords] - - return actual_x, actual_y, predict_x, predict_y, contacts - -# def new_plot(coords): - -# fig = plt.figure(figsize=(5,5), layout='tight') -# ax = fig.add_subplot(111, projection='3d') -# x = [] -# y = [] -# z = [] -# for i, coord in enumerate(coords): -# # x.append(coord.x) -# # y.append(coord.y) -# # z.append(coord.z) -# x = coord.x -# y = coord.y -# z = coord.z +from itertools import chain +import math + + + +class Visualization(): + + def __init__(self, coordinates, obs_chemical_shifts, connectivity, restraints): + self.coordinates = coordinates + self.obs_chemical_shifts = obs_chemical_shifts + self.connectivity = connectivity + self.restraints = restraints + + def get_data(self, named_tuple_used=True): + if named_tuple_used: + # Actual + actual_x = [shift.H1 for shift in self.obs_chemical_shifts] + actual_y = [shift.N15 for shift in self.obs_chemical_shifts] + # Predicted + predict_x = [shift.H1 for shift in self.coordinates] + predict_y = [shift.N15 for shift in self.coordinates] + # Connectivity + contacts = [(contact.atom1, contact.atom2) for contact in self.connectivity] + + else: + # Actual + actual_x = [shift[0] for shift in self.obs_chemical_shifts] + actual_y = [shift[1] for shift in self.obs_chemical_shifts] + # Predicted + predict_x = [shift[3] for shift in self.coordinates] + predict_y = [shift[4] for shift in self.coordinates] -# ax.scatter(x,y,z, s=40, label=f'shift {i}') - -# ax.legend() - -# plt.show() -# #plt.savefig("coord_plot.png", dpi = 150) - -# def close_contacts(coords, cutoff=0.37): -# """ -# Calculates close contacts based on coordinates. -# """ -# contacts = [] + return actual_x, actual_y, predict_x, predict_y, contacts -# for i, atom1 in enumerate(coords): -# for j, atom2 in enumerate(coords): -# if i != j: -# dist = np.linalg.norm((np.array(atom1[:3]) - np.array(atom2[:3]))) -# if dist < cutoff: -# contacts.append((i,j)) -# return contacts + def plot_shifts(self, named_tuple_used=True): -def plot_shifts(actual_shifts, restraints, coords, connectivity, named_tuple_used=True): + actual_x, actual_y, predict_x, predict_y, contacts = self.get_data(named_tuple_used=named_tuple_used) - actual_x, actual_y, predict_x, predict_y, contacts = get_data(actual_shifts, coords, connectivity, named_tuple_used=named_tuple_used) + plt.scatter(predict_x, predict_y, color='blue', label='Predicted shifts') - plt.scatter(predict_x, predict_y, color='blue', label='Predicted shifts') + plt.scatter(actual_x, actual_y, color='red', label='Measured shifts') - plt.scatter(actual_x, actual_y, color='red', label='Measured shifts') + # Label measured shifts by index and predicted shifts by associated coordinate + for i in range(len(self.obs_chemical_shifts)): + plt.annotate(i, (actual_x[i]+0.01, actual_y[i]+0.01)) + plt.annotate(i, (predict_x[i]+0.01, predict_y[i]+0.01)) - # Label measured shifts by index and predicted shifts by associated coordinate - for i in range(len(actual_shifts)): - plt.annotate(i, (actual_x[i]+0.01, actual_y[i]+0.01)) - plt.annotate(i, (predict_x[i]+0.01, predict_y[i]+0.01)) - - # Draw restraints - for restraint in restraints: - for j,k in restraint: - plt.plot((actual_x[j], actual_x[k]), (actual_y[j], actual_y[k]), color='red') - - # Draw close contacts - for i,j in contacts: - plt.plot((predict_x[i], predict_x[j]), (predict_y[i], predict_y[j]), color='blue', ls='dotted', alpha=0.5) - - plt.xlabel('H1') - plt.ylabel('N15') - plt.legend() - - plt.savefig("shifts_plot.png", dpi = 150) - plt.close() \ No newline at end of file + # Draw restraints + for restraint in self.restraints: + for j,k in restraint: + plt.plot((actual_x[j], actual_x[k]), (actual_y[j], actual_y[k]), color='red') + + # Draw close contacts + for i,j in contacts: + plt.plot((predict_x[i], predict_x[j]), (predict_y[i], predict_y[j]), color='blue', ls='dotted', alpha=0.5) + + plt.xlim([6, 10]) + plt.ylim([100, 135]) + + plt.xlabel('H1') + plt.ylabel('N15') + plt.legend() + + plt.savefig("shifts_plot.png", dpi = 150) + plt.close() + + def plot_step(self, shift_to_assign, assign_order): + + plt.subplot(1, 2, 1) + + self.plot_shifts() + + plt.subplot(1, 2, 2) + + actual_x, actual_y, predict_x, predict_y, contacts = self.get_data(named_tuple_used=True) + + plt.scatter(actual_x[shift_to_assign], actual_y[shift_to_assign], color='green', label='Shift to Assign') + plt.annotate(shift_to_assign, (actual_x[shift_to_assign]+0.01, actual_y[shift_to_assign]+0.01)) + + # to be plotted as restraints + test_values = [restraint for restraint in chain.from_iterable(self.restraints) if shift_to_assign in restraint] + test_values = set(chain.from_iterable(test_values)).difference({shift_to_assign}) + + if len(test_values) > 0: + for j in test_values: + plt.scatter(actual_x[j], actual_y[j], color='red') + plt.plot((actual_x[j], actual_x[shift_to_assign]), (actual_y[j], actual_y[shift_to_assign]), color='red') + plt.annotate(j, (actual_x[j]+0.01, actual_y[j]+0.01)) + + for i in range(len(assign_order)): + if abs(predict_x[i] - actual_x[shift_to_assign]) <= 0.15: + plt.scatter(predict_x[i], predict_y[i], color='blue') + [plt.plot((predict_x[i], predict_x[j]), (predict_y[i], predict_y[j]), color='blue', ls='dotted', alpha=0.5) for i,j in contacts if i == shift_to_assign or j == shift_to_assign] + plt.annotate(i, (predict_x[i]+0.02, predict_y[i]+0.02)) + + if abs(predict_y[i] - actual_y[shift_to_assign]) <= 0.15: + plt.scatter(predict_x[i], predict_y[i], color='blue') + [plt.plot((predict_x[i], predict_x[j]), (predict_y[i], predict_y[j]), color='blue', ls='dotted', alpha=0.5) for i,j in contacts if i == shift_to_assign or j == shift_to_assign] + plt.annotate(i, (predict_x[i]+0.02, predict_y[i]+0.02)) + + plt.xlim([6, 10]) + plt.ylim([100, 135]) + + plt.xlabel('H1') + plt.ylabel('N15') + plt.legend() + + plt.savefig("shifts_plot.png", dpi = 150) + plt.close() + + # def new_plot(coordinates): + + # fig = plt.figure(figsize=(5,5), layout='tight') + # ax = fig.add_subplot(111, projection='3d') + # x = [] + # y = [] + # z = [] + # for i, coord in enumerate(coordinates): + # # x.append(coord.x) + # # y.append(coord.y) + # # z.append(coord.z) + # x = coord.x + # y = coord.y + # z = coord.z + + # ax.scatter(x,y,z, s=40, label=f'shift {i}') + + # ax.legend() + + # plt.show() + # #plt.savefig("coord_plot.png", dpi = 150) + + # def close_contacts(coordinates, cutoff=0.37): + # """ + # Calculates close contacts based on coordinates. + # """ + # contacts = [] + + # for i, atom1 in enumerate(coordinates): + # for j, atom2 in enumerate(coordinates): + # if i != j: + # dist = np.linalg.norm((np.array(atom1[:3]) - np.array(atom2[:3]))) + # if dist < cutoff: + # contacts.append((i,j)) + # return contacts \ No newline at end of file diff --git a/write_pdb.py b/write_pdb.py new file mode 100644 index 0000000..0b291e2 --- /dev/null +++ b/write_pdb.py @@ -0,0 +1,63 @@ +from fake_data import FakeDataGenerator +from typing import NamedTuple +import numpy as np +import argparse + + + +def write_to_file(num_resid, coordinates, noes, protein=False): + if protein: + name = f"protein_r{num_resid}_noe{len(noes)}.pdb" + else: + name = f"coordinates_r{num_resid}_noe{len(noes)}.pdb" + + with open(name, "w") as f: + for (index, coord) in zip(range(num_resid), coordinates): + + x = f"{coord[0]*10:.3f}" + y = f"{coord[1]*10:.3f}" + z = f"{coord[2]*10:.3f}" + + f.write(f"ATOM{index:>7} N A A{index:>4} {x:>7} {y:>7} {z:>7} 1.00 0.00 N\n") + + +def parse_pdb(PDB): + coords = [] + with open(f"{PDB}.pdb", "r") as f: + for line in f: + if line.startswith("TER"): + break + if line.startswith("ATOM"): + parts = line.split() + if parts[2] == "CA": + x = float(parts[6])/10 + y = float(parts[7])/10 + z = float(parts[8])/10 + coords.append((x,y,z)) + + return np.array(coords) + + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument('num_resid', type=int, help='Number of residues') + parser.add_argument('pdb', type=int, help='PDB being used (any value) or not (zero)') + args = parser.parse_args() + + num_resid = args.num_resid + pdb = args.pdb + + fake_data = FakeDataGenerator(num_resid) + + coordinates, obs_chemical_shifts, pred_chemical_shifts, noes, connectivity = fake_data.generate_data_arrays(random_key=False) + + protein = False + + if pdb: + protein = True + coordinates = parse_pdb("") + noes, pred_chemical_shifts = fake_data.create_noes(coordinates, obs_chemical_shifts) + + print(len(noes)) + write_to_file(num_resid, coordinates, noes, protein=protein) From f227b21314b1c9ff7d45bb61f8c7f2196c0346da Mon Sep 17 00:00:00 2001 From: abbie-p Date: Fri, 23 May 2025 14:43:00 -0600 Subject: [PATCH 32/43] Coordinate scaling and optimization --- assignment_order.py | 137 +++++++++------- energy.py | 150 +++++++++-------- fake_data.py | 349 ++++++++++++++++++++++++---------------- fake_histories.py | 164 ++++++++++--------- fake_histories_agent.py | 75 +++++++-- nmr_gym_env.py | 215 ++++++++++++------------- nmr_text_adventure.py | 104 ++++++++++-- 7 files changed, 712 insertions(+), 482 deletions(-) diff --git a/assignment_order.py b/assignment_order.py index b8f646f..a9ff933 100644 --- a/assignment_order.py +++ b/assignment_order.py @@ -1,78 +1,97 @@ import numpy as np +from itertools import chain -class FracAct(): - def __init__(self, num_resid, noes): - self.num_resid = num_resid # value would need to be the number of shifts/atoms - self.noes = noes + +class FractionalActivation(): + + def __init__(self, num_resid, restraints): + self.num_resid = num_resid + self.restraints = restraints def activation_loop(self, target, noe, assign_order): - s = set() - for combo in noe: - if combo[0] not in assign_order: # check if values have already been assigned as higher priority - s.add(combo[0]) - if combo[1] not in assign_order: - s.add(combo[1]) - - if target in s: - return 1/(len(s)) + """ + Checks if target is in NOE, if true, checks if it has already been assigned as higher priority. + """ + if target in noe and target not in assign_order: + return 1/(len(noe)) else: return 0 + + def fractional_activation(self): + """ + Sets up residue assignment order based on fractional prevelence across NOEs. + + """ + assign_order = [] + + # Get all possible values from NOEs + noes_set = set(chain.from_iterable(chain.from_iterable(self.restraints))) + + for i in range(len(noes_set)): + running_sums = np.zeros(self.num_resid) + for noe in self.restraints: + # values from current NOE, excluding assigned values + noe_set = (set(chain.from_iterable(noe))).difference(set(assign_order)) + running_sum = ([self.activation_loop(j, noe_set, assign_order) for j in range(self.num_resid)]) + # print(running_sum) + running_sums += running_sum + assign_order.append(np.argmax(running_sums)) + # print(i, assign_order) + + missing_shifts = list(sorted(set(range(self.num_resid)).difference(set(assign_order)))) # need to preserve sorted index order for manual check + assign_order.extend(missing_shifts) + + return assign_order + + # def activation_loop(self, target, noe, assign_order): + # s = set() + # for combo in noe: + # if combo[0] not in assign_order: # check if values have already been assigned as higher priority + # s.add(combo[0]) + # if combo[1] not in assign_order: + # s.add(combo[1]) + + # if target in s: + # return 1/(len(s)) + # else: + # return 0 # def fractional_activation(self): + # assign_order = [] - # for j in range(self.num_resid): # I don't know why I need this loop here it just makes it work + # missing_shifts = [] + + # while len(assign_order) < self.num_resid: # temp2 = [] # for i in range(self.num_resid): - # temp = 0 - # for noe in self.noes: + # temp = 0 + # for noe in self.restraints: + # # print(i) + # # print([(self.activation_loop(j, noe, assign_order)) for j, noe in zip(self.restraints)]) # temp += self.activation_loop(i, noe, assign_order) # sum up the probabilities for 'i' (the given shift target) - # temp2.append(temp) + # # print(assign_order, temp) + # # print(i, temp2) - # assign_order.append(np.argmax(temp2)) # append the max value's index (takes first occurence if equivalent) + # # identifies any missing shifts + # if temp == 0 and len(assign_order) == 0: + # missing_shifts.append(i) - # return assign_order - def fractional_activation(self): - assign_order = [] - missing_shifts = [] - - while len(assign_order) < self.num_resid: - temp2 = [] - for i in range(self.num_resid): - temp = 0 - for noe in self.noes: - temp += self.activation_loop(i, noe, assign_order) # sum up the probabilities for 'i' (the given shift target) - - # identifies any missing shifts - if temp == 0 and len(assign_order) == 0: - missing_shifts.append(i) - - # appends sum to list (one sum per shift) - temp2.append(temp) - - # makes sure extra zeros are not added during final iterations (if shifts are missing) - if any(temp2) == True: - # proper order based on max value - assign_order.append(np.argmax(temp2)) - else: - # add missing shifts to end of list - assign_order = assign_order + missing_shifts + # # appends sum to list (one sum per shift) + # temp2.append(temp) + # # print(temp2) + # # makes sure extra zeros are not added during final iterations (if shifts are missing) + # if any(temp2) == True: + # # for x in temp2: + # # proper order based on max value + # assign_order.append(np.argmax(temp2)) + # # temp2[np.argmax(temp2)] = 0 + # # print(temp2) + # else: + # # add missing shifts to end of list + # assign_order = assign_order + missing_shifts - return assign_order - - # intermediate check for first set of probabilites assigned - not used as final calculation - def fractional_activation_summation(self): - - assign_order = [] - for j in range(self.num_resid): - temp2 = [] - for i in range(self.num_resid): - temp = 0 - for noe in self.noes: - temp += activation_loop(i, noe, assign_order) - temp2.append(temp) - - return temp2 + # return assign_order diff --git a/energy.py b/energy.py index 17b3c00..2e99608 100644 --- a/energy.py +++ b/energy.py @@ -1,74 +1,78 @@ import numpy as np +from typing import NamedTuple from itertools import product import math +from scipy.spatial.distance import pdist, squareform + -# match H1 shifts -def matchH(hshift, hsqc, tolerance_h): - """ - Loops through the HSQC list and matches the second proton shift of the NOE to a HSQC shift based on H cutoff. - Returns the index of said shift. - """ - h_index = [] - for i, shift in enumerate(hsqc): - if abs(hshift - shift.H1) <= tolerance_h: - h_index.append(i) - - return h_index - -# match H1 and N15 shifts -def matchNH(hshift, nshift, hsqc, tolerance_h, tolerance_n): - """ - Loops through the HSQC list and matches the first proton and nitrogen shift of the NOE to a HSQC shift based on H and N cutoffs. - Returns the index of said shift. - """ - nh_index = [] - for i, shift in enumerate(hsqc): - if abs(hshift - shift.H1) <= tolerance_h and abs(nshift - shift.N15) <= tolerance_n: - nh_index.append(i) - - return nh_index - -def noe_combinations(noes, actual_shifts, tolerance_h=0.02, tolerance_n=0.02): - """ - Loops through the NOES to match the H1, N1 and H2 shifts to two HSQC shifts. - Lists the HSQC shifts separately (H1/N1 as contact1 and H2 as contact2) for each NOE. - Returns a lists of tuples of all possible shift index combinations for each NOE. - ex. [[(1,2),(1,3)],[(3,0)]] - 1 and 2 are a pair of potential shift indices for one NOE, 1 and 3 are another for the same NOE. - """ - contact1 = [] - contact2 = [] - for i, peak in enumerate(noes): - # matching H1 and N15 in NOE to shifts in HSQC - nh_index = matchNH(peak.H1, peak.N15, actual_shifts, tolerance_h, tolerance_n) - # matching H2 in NOE to H1 shifts in HSQC - h_index = matchH(peak.H2, actual_shifts, tolerance_h) - - contact1.append(nh_index) - contact2.append(h_index) - - # print(contact1) - # print(contact2) - - contacts = [] - for i, contact in enumerate(contact1): - prod = product(contact, contact2[i]) - # filter out same shift pairs and inverted duplicates (ex. (0,0) and [(0,1),(1,0)]) - pairs = set(tuple(sorted(combo)) for combo in prod if combo[0] != combo[1]) - contacts.append(list(pairs)) - - return contacts class Energy(): - def __init__(self, coords, actual_shifts, noes): - self.coords = coords - self.actual_shifts = actual_shifts + def __init__(self, coordinates, obs_chemical_shifts, noes): + self.tolerance_h = 0.02 + self.tolerance_n = 0.2 + self.tolerance_dist = 0.5 + self.coordinates = coordinates + self.obs_chemical_shifts = obs_chemical_shifts self.noes = noes - def setup_noe_restraints(self): - combos = noe_combinations(self.noes, self.actual_shifts, tolerance_h=0.02, tolerance_n=0.02) - return combos + # Match H1 shifts + def matchH(self, hshift, hsqc): + """ + Loops through the HSQC list and matches the second proton shift of the NOE to a HSQC shift based on H cutoff. + Returns the index of said shift. + """ + h_index = [] + for i, shift in enumerate(hsqc): + if abs(hshift - shift.H1) <= np.float64(self.tolerance_h): + h_index.append(i) + + return h_index + + # Match H1 and N15 shifts + def matchNH(self, hshift, nshift, hsqc): + """ + Loops through the HSQC list and matches the first proton and nitrogen shift of the NOE to a HSQC shift based on H and N cutoffs. + Returns the index of said shift. + """ + nh_index = [] + for i, shift in enumerate(hsqc): + if (abs((hshift) - (shift.H1)) <= self.tolerance_h) and (abs((nshift) - (shift.N15)) <= self.tolerance_n): + nh_index.append(i) + + return nh_index + + def noe_combinations(self): + """ + Loops through the NOES to match the H1, N1 and H2 shifts to two HSQC shifts. + Lists the HSQC shifts separately (H1/N1 as contact1 and H2 as contact2) for each NOE. + Returns a lists of tuples of all possible shift index combinations for each NOE. + ex. [[(1,2),(1,3)],[(3,0)]] + 1 and 2 are a pair of potential shift indices for one NOE, 1 and 3 are another for the same NOE. + """ + contact1 = [] + contact2 = [] + + for peak in self.noes: + # matching H1 and N15 in NOE to shifts in HSQC + nh_index = self.matchNH(peak.H1, peak.N15, self.obs_chemical_shifts) + # matching H2 in NOE to H1 shifts in HSQC + h_index = self.matchH(peak.H2, self.obs_chemical_shifts) + + contact1.append(nh_index) + contact2.append(h_index) + + # print(contact1) + # print(contact2) + + contacts = [] + for i, contact in enumerate(contact1): + prod = product(contact, contact2[i]) + # filter out same shift pairs and inverted duplicates (ex. (0,0) and [(0,1),(1,0)]) + pairs = set(tuple(sorted(combo)) for combo in prod if combo[0] != combo[1]) + contacts.append(list(pairs)) + + return contacts def flat_bottom(self, x, tolerance): """ @@ -88,7 +92,14 @@ def flat_bottom(self, x, tolerance): lambda x: x**2 - xmax*x]) return y - def calc_restraint_energy(self, assignments, restraint, tolerance): + def calc_pdist(self): + coords = [coord[:3] for coord in self.coordinates] # can be removed if predicted shifts and coordinates are not combined + pairwise_dists = pdist(coords) + square_dists = squareform(pairwise_dists) + + return square_dists + + def calc_restraint_energy(self, assignments, restraint, dist_grid, tolerance): """ Goes over shift possibilities in the restraint, grabs associated coordinates from assignment, and calculates energy according to distance. Only the smallest energy is returned. @@ -96,12 +107,14 @@ def calc_restraint_energy(self, assignments, restraint, tolerance): restraint_energy = math.inf for i, j in restraint: k, l = assignments.get(i), assignments.get(j) - dist = np.linalg.norm((np.array(self.coords[k][:3]) - np.array(self.coords[l][:3]))) + # dist = np.linalg.norm((np.array(self.coordinates[k][:3]) - np.array(self.coordinates[l][:3]))) + dist = dist_grid[k][l] energy_value = self.flat_bottom(dist, tolerance=tolerance) restraint_energy = energy_value if energy_value < restraint_energy else restraint_energy + return restraint_energy - def noe_activation(self, assignments, restraint, tolerance): + def noe_activation(self, assignments, restraint, dist_grid, tolerance): """ Checks if all shift possibilities in the restraint have been assigned (aka activated). Calculates energy if they have, returns zero if they have not. @@ -115,19 +128,18 @@ def noe_activation(self, assignments, restraint, tolerance): if d: return 0 else: - return self.calc_restraint_energy(assignments, restraint, tolerance) + return self.calc_restraint_energy(assignments, restraint, dist_grid, tolerance) - def get_total_energy(self, restraints, assignments, tolerance): + def get_total_energy(self, restraints, assignments, dist_grid): """ Loops over the NOE restraints to sum up calculated energies. """ - total_energy = 0 for restraint in restraints: if restraint == []: pass else: - energy_value = self.noe_activation(assignments, restraint, tolerance=tolerance) + energy_value = self.noe_activation(assignments, restraint, dist_grid, tolerance=self.tolerance_dist) total_energy += energy_value return total_energy diff --git a/fake_data.py b/fake_data.py index 7ed9c9e..38ae6f8 100644 --- a/fake_data.py +++ b/fake_data.py @@ -4,16 +4,20 @@ import random import argparse import pickle +import glob as glob + class HSQCPeak(NamedTuple): H1: float N15: float + class NOEPeak(NamedTuple): H1: float N15: float H2: float + class Protein(NamedTuple): x: float y: float @@ -21,151 +25,220 @@ class Protein(NamedTuple): H1: float N15: float + class Connectivity(NamedTuple): atom1: float atom2: float distance: float -# Sampled points from unit square/cube -def sample_unit(n, num_sides, min=0, max=1): - return np.random.uniform(min, max, size=(n, num_sides)) - -# Noise to sampled points -def add_noise(point, scale=0.1, min=0, max=1): - """ - Random selection from normal (gaussian) distribution of 'scale' width from 0 (center). - 'size' makes sure that it's the same shape as the point we are adding noises to. - ex. point = [0,1,2], noise = [1,1,1], noisy_point = [1,2,3] - """ - noise = np.random.normal(0, scale, size=point.shape) - noisy_point = point + noise - - # fold over points outside of boundaries - for point in noisy_point: - for i, value in enumerate(point): - if max < value: - point[i] = (max-(value-max)) - if value < min: - point[i] = (min-value) - - return noisy_point - -def calc_dist(p1, p2): - return np.linalg.norm((p2-p1)) - -def distance_noe(protein, shifts, cutoff=0.5, random_key=False): - """ - Grabs close coordinates and 'associated' HSQC shift peaks (randomized index or 1:1 correlation) to create NOES. - Predicted shifts assigned as shift order - this order is associated with the coordinate order if randomized (ie. the answer key). - Adds gaussian noise to all points at the end. - - **Can end up with no NOEs depending on the cutoff** - """ - if random_key: - shifts = np.random.permutation(shifts) - - noes = [] - - for i, atom1 in enumerate(protein): - for j, atom2 in enumerate(protein): - if i != j: - dist = calc_dist(atom1, atom2) - if dist < cutoff: - #print(dist, shifts[i], shifts[j]) - noe = list(shifts[i][:]) # H1, N1 - noe.append(shifts[j][0]) # H2 - noes.append(noe) +class FakeDataGenerator(): + + def __init__(self, num_resid): + self.num_resid: int = num_resid + self.nshift_min = 100 + self.nshift_max = 135 + self.hshift_min = 6 + self.hshift_max = 10 + self.cutoff = 0.5 + + def sample_unit(self, num_points, num_sides, min_len=0, max_len=1): + """ + Samples points from unit square or cube. + """ + return np.random.uniform(low=min_len, high=max_len, size=(num_points, num_sides)) + + def scale_unit(self, coordinates): + """ + Scaling box size to reflect a globular protein. + PDB:1CRC (~100 resid, globular, ~3.2 nm diameter) + + Rg = RN**v + Scaling factor (v) of 0.4, instead of 0.3 (cube-root) + R of 0.2 nm used in paper for fit + DOI: 10.1142/S021972002050050X + """ + radius_gyration = 0.2*(self.num_resid**0.4) + + com = np.mean(coordinates, axis=0) # should be ~0.5 + distances_sampled = np.linalg.norm(coordinates - com, axis=1) # distances around origin + radius_sampled = np.sqrt(np.mean((distances_sampled**2))) # radius of gyration based on these distances + + # Expected is radius_gyration, current is radius_sampled --> scale to fit + # Will overlap occur? + scale = radius_gyration / radius_sampled + coordinates_scaled = com + ((coordinates - com) * scale) + + return coordinates_scaled + + def create_hsqc(self): + """ + Sample points within NMR window for H1 and N15. Stack values to create HSQC peaks. + """ + # H shifts + h_shifts = self.sample_unit(self.num_resid, num_sides=1, min_len=self.hshift_min, max_len=self.hshift_max) + # N shifts + n_shifts = self.sample_unit(self.num_resid, num_sides=1, min_len=self.nshift_min, max_len=self.nshift_max) + + return np.hstack((h_shifts, n_shifts)) + + def add_noise(self, points, scale=0.1): + """ + Random selection from normal (gaussian) distribution of 'scale' width from 0 (center). + 'size' makes sure that it's the same shape as the point we are adding noises to. + ex. point = [0,1,2], noise = [1,1,1], noisy_point = [1,2,3] + """ + noise = np.random.normal(loc=0, scale=scale, size=points.shape) + noisy_point = points + noise + + # fold over points outside of boundaries + # for point in noisy_point: + # for i, axis in enumerate(point): + # if max_len < axis: + # point[i] = (max_len - (axis - max_len)) + # if axis < min_len: + # point[i] = (min_len - axis) + # else: + # pass + + return noisy_point + + def calculate_dist(self, p1, p2): + """ + Euclidian distance between two points. + """ + return np.linalg.norm((p2 - p1)) + + def create_noes(self, coordinates, shifts, random_key=False): + """ + Grabs close coordinates and 'associated' HSQC shift peaks (randomized index or 1:1 correlation) to create noes. + Predicted shifts assigned as shift order - this order is associated with the coordinate order if randomized (ie. the answer key). + Adds gaussian noise to all points at the end. + + **Can end up with no noes depending on the cutoff** + """ + if random_key: + shifts = np.random.permutation(shifts) + + noe_list = [] + for i, atom1 in enumerate(coordinates): + for j, atom2 in enumerate(coordinates): + if i != j: + dist = self.calculate_dist(atom1, atom2) + # dist = dist_grid[i][j] + if dist < self.cutoff: + #print(dist, shifts[i], shifts[j]) + noe = list(shifts[i][:]) # H1, N15 + noe.append(shifts[j][0]) # H2 + noe_list.append(noe) + + noisy_noes = self.add_noise(np.array(noe_list), scale=0.01) + pred_chemical_shifts = self.add_noise(np.array(shifts), scale=0.1) + + return noisy_noes, pred_chemical_shifts + + def calculate_connectivity(self, coordinates): + """ + Calculates close contacts based on protein coordinates. + """ + contacts = [] + + for i, atom1 in enumerate(coordinates): + for j, atom2 in enumerate(coordinates): + if i != j: + dist = self.calculate_dist(atom1, atom2) + # dist = dist_grid[i][j] + if dist < self.cutoff: + contacts.append((i, j, dist)) + + return contacts + + # def pdist_test(self, coordinates): + # pairwise_dists = pdist(coordinates) + # square_dists = squareform(pairwise_dists) + # return square_dists + + def generate_data_arrays(self, random_key): + """ + Generates all qualities of the system as individual arrays. + """ + # 3D structure [x,y,z] + pred_coordinates = self.sample_unit(self.num_resid, num_sides=3) + pred_coordinates = self.scale_unit(pred_coordinates) - noisy_noe = add_noise(np.array(noes), scale=0.01) - predicted_shifts = add_noise(np.array(shifts), scale=0.1) + # "Actual" shifts [H1,N15] + obs_chemical_shifts = self.create_hsqc() + + # Distance grid + # dist_grid = self.pdist_test(pred_coordinates) + + # NOEs [H1,N15,H2] and predicted shifts [H1,N15] + noes, pred_chemical_shifts = self.create_noes(pred_coordinates, obs_chemical_shifts, random_key=random_key) + + # Connectivity [atom1,atom2,dist] + connectivity = self.calculate_connectivity(pred_coordinates) + + return pred_coordinates, obs_chemical_shifts, pred_chemical_shifts, noes, connectivity + + def order_data(self, pred_coordinates, obs_chemical_shifts, pred_chemical_shifts, noes, connectivity): + """ + Compiles data in a list of named tuples with one object per residue. + """ + pred_coordinates = [Protein(x=resid[0], y=resid[1], z=resid[2], H1=shift[0], N15=shift[1]) for resid, shift in zip(pred_coordinates, pred_chemical_shifts)] + # pred_chemical_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in pred_chemical_shifts] + obs_chemical_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in obs_chemical_shifts] + noes = [NOEPeak(H1=shift[0], N15=shift[1], H2=shift[2]) for shift in noes] + connectivity = [Connectivity(atom1=contact[0], atom2=contact[1], distance=contact[2]) for contact in connectivity] + + return pred_coordinates, obs_chemical_shifts, noes, connectivity + + def dump_pickle(self, pred_coordinates, obs_chemical_shifts, noes, connectivity, example=True): + """ + Saves data to disk. Run is given a proper name if used as a saved example. + """ + name = f'fakedata_r{self.num_resid}.pkl' if example else f'current_run.pkl' + + with open(name, 'wb') as f: + pickle.dump(pred_coordinates, f) + pickle.dump(obs_chemical_shifts, f) + pickle.dump(noes, f) + pickle.dump(connectivity, f) - return noisy_noe, predicted_shifts + def load_pickle(self, example): -def connectivity_data(protein, cutoff): - """ - Calculates close contacts based on coordinates. - """ - connectivity = [] - - for i, atom1 in enumerate(protein): - for j, atom2 in enumerate(protein): - if i != j: - dist = calc_dist(atom1, atom2) #dist = np.linalg.norm(atom1 - atom2) - if dist < cutoff: - connectivity.append((i, j, dist)) - - return connectivity - -def generate_raw_data(num_resid, random_key): - """ - Generates all qualities of the system as individual arrays. - """ - # 3D structure [x,y,z] - protein = sample_unit(num_resid, num_sides=3) - # "Actual" shifts [H1,N1] - actual_shifts = sample_unit(num_resid, num_sides=2) - # NOES [H1,N1,H2] and predicted shifts [H1,N1] - noes, predicted_shifts = distance_noe(protein, actual_shifts, cutoff=float(1/np.cbrt(num_resid)), random_key=random_key) - # connectivity [atom1,atom2,dist] - connectivity = connectivity_data(protein, cutoff=float(0.8/np.cbrt(num_resid))) #0.37 - - return protein, actual_shifts, predicted_shifts, noes, connectivity - -def order_data(protein, actual_shifts, predicted_shifts, noes, connectivity): - """ - Compiles data in a list of named tuples with one object per residue. - """ - coords = [Protein(x=resid[0], y=resid[1], z=resid[2], H1=shift[0], N15=shift[1]) for resid, shift in zip(protein, predicted_shifts)] - actual_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in actual_shifts] - noes = [NOEPeak(H1=shift[0], N15=shift[1], H2=shift[2]) for shift in noes] - connectivity = [Connectivity(atom1=connect[0], atom2=connect[1], distance=connect[2]) for connect in connectivity] - - return coords, actual_shifts, noes, connectivity - -def save_as_pickle(coords, actual_shifts, noes, connectivity, num_resid, example=True): - """ - Saves data to disk. Run is given a proper name if used as a saved example. - """ - name = f'fakedata_r{num_resid}.pkl' if example else f'current_run.pkl' - - with open(name, 'wb') as f: - pickle.dump(coords, f) - pickle.dump(actual_shifts, f) - pickle.dump(noes, f) - pickle.dump(connectivity, f) - -def generate_data(num_resid, pickle_data=True, example=True, random_key=False): - """ - Generates all fake data, orders it in lists of namedtuples, and pickles if required. - """ - protein, actual_shifts, predicted_shifts, noes, connectivity = generate_raw_data(num_resid, random_key=random_key) - coords, actual_shifts, noes, connectivity = order_data(protein, actual_shifts, predicted_shifts, noes, connectivity) - - if pickle_data: - save_as_pickle(coords, actual_shifts, noes, connectivity, num_resid, example=example) - - else: - return coords, actual_shifts, noes, connectivity - -# if __name__ == '__main__': -# parser = argparse.ArgumentParser() -# parser.add_argument('num_resid', type=int, help='Number of residues') -# args = parser.parse_args() - -# num_resid = args.num_resid -# generate_data(num_resid) - -# # These should be grouped to export into an environment -# coords, actual_shifts, predicted_shifts, noes = generate_data(num_resid) - # print(f'coords {(np.array(coords)).tolist()}')#\n shifts{(np.array(actual_shifts)).tolist()}\n noes{(np.array(noes)).tolist()}') - # # print(actual_shifts) - # # print(noes) - - # energy_obj = Energy(coords, actual_shifts, noes) - # test = energy_obj.setup_noe_restraints() - # # print(test) - # # print(coords) + name = f"./*r{self.num_resid}.pkl" if example else "./current_run.pkl" + + pickle_file = glob.glob(name) + with open(pickle_file[0], 'rb') as f: + coordinates = pickle.load(f) + obs_chemical_shifts = pickle.load(f) + noes = pickle.load(f) + connectivity = pickle.load(f) - # test_energy = energy_obj.get_total_energy(test, {0:2, 1:1, 2:0}) - # print(test_energy) \ No newline at end of file + return coordinates, obs_chemical_shifts, noes, connectivity + + def generate_data(self, pickle_data=True, example=True, random_key=False): + """ + Generates all fake data, orders it in lists of namedtuples, and pickles if required. + """ + pred_coordinates, obs_chemical_shifts, pred_chemical_shifts, noes, connectivity = self.generate_data_arrays(random_key=random_key) + pred_coordinates, obs_chemical_shifts, noes, connectivity = self.order_data(pred_coordinates, obs_chemical_shifts, pred_chemical_shifts, noes, connectivity) + + if pickle_data: + self.dump_pickle(pred_coordinates, obs_chemical_shifts, noes, connectivity, example=example) + + else: + return pred_coordinates, obs_chemical_shifts, noes, connectivity + + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('num_resid', type=int, help='Number of residues') + args = parser.parse_args() + + num_resid = args.num_resid + + fakedata = FakeDataGenerator(num_resid) + pred_coordinates, obs_chemical_shifts, noes, connectivity = fakedata.generate_data(pickle_data=False, example=False) + \ No newline at end of file diff --git a/fake_histories.py b/fake_histories.py index 87c57ec..0e19108 100644 --- a/fake_histories.py +++ b/fake_histories.py @@ -1,79 +1,89 @@ import numpy as np -from fake_data import * -from energy import noe_combinations - -def perturb_data(coords): - """ - Adds noise to original fake data and generates new 'actual' shifts based on predicted shifts. - """ - protein = add_noise(np.array(coords)[:, :3], scale=0.1) - actual_fake_shifts = add_noise(np.array(coords)[:, 3:], scale=0.1) - - return protein, actual_fake_shifts - -def weighted_error(num_resid, scale=1): - """ - Generates weights for number of errors added based on exponential distribution. - Chooses error number to add based on weights. - """ - # Weight each point - weights = np.exp(-np.arange(num_resid)/scale) - # Normalized - weights = weights/np.sum(weights) - # Number of errors based on weight - error_num = np.random.choice(num_resid, 1, p=weights) - - # Can't make a singular error - needs to be replaced by another choice - return error_num if error_num > 1 else error_num*2 - -def add_errors(answer): - """ - If errors are present chooses that number of random shifts and move answer one over from true. - """ - error_num = weighted_error(len(answer), scale=1) - - if error_num == 0: + +from fake_data import FakeDataGenerator, HSQCPeak, Protein +from energy import Energy +from assignment_order import FractionalActivation + + + +class FakeHistoryGenerator(): + + def __init__(self, num_resid): + self.num_resid = num_resid + self.fakedata = FakeDataGenerator(num_resid) + self.cutoff = 0.5 + + def perturb_data(self, coordinates): + """ + Adds noise to original fake data and generates new 'actual' shifts based on predicted shifts. + """ + coords = self.fakedata.add_noise(np.array(coordinates)[:, :3], scale=0.1) + synthetic_shifts = self.fakedata.add_noise(np.array(coordinates)[:, 3:], scale=0.1) + + return coords, synthetic_shifts + + def weighted_error(self, scale=1): + """ + Generates weights for number of errors added based on exponential distribution. + Chooses error number to add based on weights. + """ + # Weight each point + weights = np.exp(-np.arange(self.num_resid)/scale) + # Normalized + weights = weights/np.sum(weights) + # Number of errors based on weight + error_num = np.random.choice(self.num_resid, 1, p=weights) + + # Can't make a singular error - needs to be replaced by another choice + return error_num if error_num > 1 else error_num*2 + + def add_errors(self, answer): + """ + If errors are present chooses that number of random shifts and move answer one over from true. + """ + error_num = self.weighted_error(scale=1) + + if error_num == 0: + return answer + else: + selections = np.random.choice(list(answer.values()), error_num, replace=False) + # print(selections) + mutations = np.roll(selections, 1) + # print(mutations) + for i, j in enumerate(selections): + # print(answer, answer[j], mutations[i], error_num) + answer[j] = mutations[i] + return answer - else: - selections = np.random.choice(list(answer.values()), error_num, replace=False) - # print(selections) - mutations = np.roll(selections, 1) - # print(mutations) - for i, j in enumerate(selections): - # print(answer, answer[j], mutations[i], error_num) - answer[j] = mutations[i] - - return answer - -def generate_history(original, i, history_length=5): - """ - Loops over chosen history length, adding noise to original data, reorganizes, and edits answer key if errors are introduced. - """ - - protein, actual_shifts = perturb_data(original['coords']) - noes, predicted_shifts = distance_noe(protein, actual_shifts, cutoff=float(1/np.cbrt(len(actual_shifts)))) - - connectivity = connectivity_data(protein, cutoff=float(0.8/np.cbrt(len(actual_shifts)))) - - coords, actual_shifts, noes, connectivity = order_data(protein, actual_shifts, predicted_shifts, noes, connectivity) - - # answer = original['assignments'] - # answer, error_num = add_errors(answer) - - # issue with permanence - need to recreate full assignment for now - answer = {i:i for i in original['assign_order']} - answer = add_errors(answer) - - state = { - "coords": coords, - "actual_shifts": actual_shifts, - "noes": noes, - "connectivity": connectivity, - "assignments": {}, - "assign_order": list(answer.values()), - "assign_shift": 0, - "total_energy": 0.0, - "reward": 0.0 - } - - return state \ No newline at end of file + + def generate_history(self, original): + """ + Adds noise to original data, reorganizes, and edits answer key if errors are introduced. + """ + + coordinates, obs_chemical_shifts = self.perturb_data(original['coordinates']) + + noes, pred_chemical_shifts = self.fakedata.create_noes(coordinates, obs_chemical_shifts) + + connectivity = self.fakedata.calculate_connectivity(coordinates) + + coordinates, obs_chemical_shifts, noes, connectivity = self.fakedata.order_data(coordinates, obs_chemical_shifts, pred_chemical_shifts, noes, connectivity) + + # Errors not included currently + # answer = original['assignments'] + # answer = add_errors(answer) + + + state = { + "coordinates": coordinates, + "obs_chemical_shifts": obs_chemical_shifts, + "noes": noes, + "connectivity": connectivity, + "assignments": {}, + "assign_order": [], + "shift_to_assign": 0, + "total_energy": 0.0, + "reward": 0.0 + } + + return state diff --git a/fake_histories_agent.py b/fake_histories_agent.py index 4b90fb0..0289c0f 100644 --- a/fake_histories_agent.py +++ b/fake_histories_agent.py @@ -2,36 +2,77 @@ import argparse import pickle import copy +import numpy as np from nmr_gym_env import GymEnv -from fake_histories import generate_history -from visualize_data import plot_shifts +from fake_data import FakeDataGenerator +from fake_histories import FakeHistoryGenerator +from visualize_data import Visualization + +from nmr_text_adventure import get_data + +import time + + if __name__ == '__main__': + start = time.process_time() """ Runs through protein examples of one size following the 1:1 shift to atom assignment. Once example is assigned, generates history and repeats assignment process. All histories are then saved to disk. """ parser = argparse.ArgumentParser() - # parser.add_argument('min_protein', type=int, help='Minimum protein size') - parser.add_argument('max_protein', type=int, help='Max protein size') + parser.add_argument('num_resid', type=int, help='Max protein size') parser.add_argument('history_length', type=int, help='Number of synthetic examples to generate') args = parser.parse_args() - # min_protein = args.min_protein - max_protein = args.max_protein + num_resid = args.num_resid history_length = args.history_length - gym_env = GymEnv(max_protein) - observation = gym_env.reset(pickled=False, pickle_data=False) + gym_env = GymEnv(num_resid) + + fakedata = FakeDataGenerator(num_resid) + fakehistory = FakeHistoryGenerator(num_resid) + + ###################################################### + + print("\nAvailable options to work with:\n1. Generate and save an example\n2. Grab a previous saved example\n3. Run a new example (not saved)") + + selection = int(input("Option Selection: ")) + + if selection == 1: + # 1. Need to generate and save an example? Saved in all instances as fakedata_r{num_resid}.pkl + coordinates, obs_chemical_shifts, noes, connectivity = get_data(num_resid, pickled=False, pickle_data=True, example=True, random_key=True) + if selection == 2: + # 2. Grab one of the previous examples? Make sure example exists with desired resid number + coordinates, obs_chemical_shifts, noes, connectivity = get_data(num_resid, pickled=True, pickle_data=False, random_key=True) + if selection == 3: + # 3. Run a new example? Saved in all instances as current_run.pkl + coordinates, obs_chemical_shifts, noes, connectivity = get_data(num_resid, pickled=False, pickle_data=True, example=False, random_key=True) + + ###################################################### + + print(f'\nCumulative Time:') + print(f'{time.process_time() - start}s --> Data Generated') + + observation = gym_env.reset(coordinates, obs_chemical_shifts, noes, connectivity) + + print(f'{time.process_time() - start}s --> Observation Reset') histories = [] - for i in range(history_length): + for i in range(history_length+1): + + # If solving for fake histories off the original if i > 0: observation = gym_env.custom_state(fake_observation) + print(f'{time.process_time() - start}s --> Custom Reset ') + if i <= 0: + original_observation = copy.deepcopy(observation) + + # Initial observation (nothing assigned) histories.append(copy.deepcopy(observation)) terminated = False @@ -41,14 +82,16 @@ if not terminated: histories.append(copy.deepcopy(observation)) - # don't know if this is needed - issues with permanence of observation - if i <= 0: - original_observation = copy.deepcopy(observation) + # Generate new history + if i != history_length + 1: + fake_observation = fakehistory.generate_history(original_observation) + + # Final observation (everything assigned) + histories.append(copy.deepcopy(observation)) - fake_observation = generate_history(original_observation, i, history_length=history_length) + print(f"{time.process_time() - start} --> {i} Histories Complete") - # print(histories) - with open('fake_histories.pkl', 'wb') as f: + with open(f'fake_histories_r{num_resid}_{history_length}.pkl', 'wb') as f: pickle.dump(histories, f) - print(f'{history_length} fake histories saved') \ No newline at end of file + print(f'{time.process_time() - start}s --> {history_length} Histories Saved\n') \ No newline at end of file diff --git a/nmr_gym_env.py b/nmr_gym_env.py index 7fa796a..049b259 100644 --- a/nmr_gym_env.py +++ b/nmr_gym_env.py @@ -3,13 +3,16 @@ import numpy as np import pickle import math -import glob as glob -from fake_data import generate_data -from energy import Energy, noe_combinations -from assignment_order import FracAct +from energy import Energy +from assignment_order import FractionalActivation +from visualize_data import Visualization + +import cProfile +import pstats +import time + -from visualize_data import plot_shifts class GymEnv(gym.Env): @@ -17,151 +20,145 @@ def __init__(self, num_resid): self.state = None self.num_resid = num_resid - self.assignments = {} - self.assign_step = 0 - self.assign_shift = 0 - self.restraints = 0 - self.total_energy = 0.0 - self.intermediate_energy = 0.0 - self.reward = 0.0 + self.assignments: Dict[int, int] = None + self.assign_step: int = None + self.shift_to_assign: int = None + self.restraints = None + self.intermediate_energy: float = None + self.dist_grid = [] + + self.energy: Energy = None + self.activation: FractionalActivation = None + self.visualize: Visualization = None self.action_space = spaces.Discrete(self.num_resid) self.observation_space = spaces.Dict({ - "coords": spaces.Box(0, 1, shape=(self.num_resid, 5), dtype=np.float32), - "actual_shifts": spaces.Box(0, 1, shape=(self.num_resid, 2), dtype=np.float32), + "coordinates": spaces.Box(0, 1, shape=(self.num_resid, 5), dtype=np.float32), + "obs_chemical_shifts": spaces.Box(0, 1, shape=(self.num_resid, 2), dtype=np.float32), "noes": spaces.Box(0, 1, shape=(self.num_resid, 3), dtype=np.float32), "assignments": spaces.Box(0, self.num_resid, shape=(self.num_resid, 3), dtype=np.float32), "assign_order": spaces.Box(0, self.num_resid, shape=(self.num_resid, 1), dtype=np.float32), - "assign_shift": spaces.Discrete(self.num_resid), + "shift_to_assign": spaces.Discrete(self.num_resid), "total_energy": spaces.Box(0, math.inf, shape=(1, 1), dtype=np.float32), "reward": spaces.Box(0, math.inf, shape=(1, 1), dtype=np.float32) }) + def step(self, action): + assert action < self.num_resid and action not in self.state['assignments'].values() + + # Add action to assignments + self.assignments[(self.assign_order[self.assign_step])] = action + # print(self.assignments) + + # Calculate energy and store temporarily + temp_intermediate_energy = self.energy.get_total_energy(self.restraints, self.assignments, self.dist_grid) + + ### TEST WHEN REMOVING ENERGY FUNCTION ### + # temp_intermediate_energy = 0 + ########################################## + + # Calculate reward based on previous energy and temporary energy + self.state['reward'] = (self.intermediate_energy - temp_intermediate_energy) + + # Before correction, +ve means energy went down - we DO NOT want this, -ve means energy went up + assert self.state['reward'] <= 0 + self.state['reward'] = abs(self.state['reward']) + + # Store energy as intermediate + self.intermediate_energy = temp_intermediate_energy + + # Assign intermediate energy as running total + self.state['total_energy'] = self.intermediate_energy + # print(f"Running energy = {self.state['total_energy']}, Current reward = {self.state['reward']}") + + self.assign_step += 1 + terminated = True if len(self.assignments) == self.num_resid else False + + if terminated: + # print(f"Final assignments (shift:atom) = {self.assignments}\nFinal energy evaluation = {self.state['total_energy']}") + self.state['shift_to_assign'] = None + return self.state, self.state['reward'], terminated, self.state['total_energy'] + + self.state['shift_to_assign'] = self.assign_order[self.assign_step] + if not terminated: + # print(f'Shift to be assigned: {self.state["shift_to_assign"]}') + return self.state, self.state['reward'], terminated, self.state['total_energy'] + def custom_state(self, state): self.state = state - self.restraints = noe_combinations(self.state['noes'], self.state['actual_shifts'], tolerance_h=0.02, tolerance_n=0.02) + + # Get restraints + self.energy = Energy(self.state['coordinates'], self.state['obs_chemical_shifts'], self.state['noes']) + self.restraints = self.energy.noe_combinations() + + self.dist_grid = self.energy.calc_pdist() self.assignments = {} self.assign_step = 0 - self.assign_shift = 0 - self.total_energy = 0.0 self.intermediate_energy = 0.0 - self.reward = 0.0 - # get assignment order - # frac_act = FracAct(self.num_resid, self.restraints) - # self.assign_order = frac_act.fractional_activation() - self.assign_shift = self.state['assign_order'][self.assign_step] - # print(self.assign_order) - # print(f'Shift to be assigned: {self.assign_shift}') + # Get assignment order + self.activation = FractionalActivation(self.num_resid, self.restraints) + self.assign_order = self.activation.fractional_activation() - # self.state['assign_order'] = self.assign_order - self.state['assign_shift'] = self.assign_shift + ### TEST WHEN REMOVING FRACTIONAL ACTIVATION ### + # test_range = np.arange(0, self.num_resid) + # self.assign_order = test_range + ################################################ + + self.state["assign_order"] = self.assign_order + self.shift_to_assign = self.assign_order[self.assign_step] + + self.state['shift_to_assign'] = self.shift_to_assign self.state['assignments'] = self.assignments return self.state - def reset(self, pickled=False, pickle_data=True, example=True, random_key=False): + def reset(self, coordinates, obs_chemical_shifts, noes, connectivity, seed=None, options=None): #pickled=False, pickle_data=True, example=True, random_key=False, seed=None, options=None): + super().reset(seed=seed) - name = f"./*r{self.num_resid}.pkl" if example else "./current_run.pkl" + # Get restraints + self.energy = Energy(coordinates, obs_chemical_shifts, noes) + self.restraints = self.energy.noe_combinations() - if self.intermediate_energy > 0 or pickled: - pickle_file = glob.glob(name) - with open(pickle_file[0], 'rb') as f: - coords = pickle.load(f) - actual_shifts = pickle.load(f) - noes = pickle.load(f) - connectivity = pickle.load(f) - - elif pickle_data: - generate_data(self.num_resid, pickle_data=pickle_data, example=example, random_key=random_key) - pickle_file = glob.glob(name) - with open(pickle_file[0], 'rb') as f: - coords = pickle.load(f) - actual_shifts = pickle.load(f) - noes = pickle.load(f) - connectivity = pickle.load(f) - - else: - coords, actual_shifts, noes, connectivity = generate_data(self.num_resid, pickle_data=pickle_data, example=example, random_key=random_key) - - # get restraints - """ - float(float(0.001*(self.num_resid)+0.01)) shift tolerance needs to increase with protein size somehow or else it can miss the correct answer, - but this will increase the complexity of the problem - """ - self.restraints = noe_combinations(noes, actual_shifts, tolerance_h=0.02, tolerance_n=0.02) - # get plot of shifts and restraints - plot_shifts(actual_shifts, self.restraints, coords, connectivity, named_tuple_used=True) + self.dist_grid = self.energy.calc_pdist() self.assignments = {} self.assign_step = 0 - self.assign_shift = 0 - self.total_energy = 0.0 self.intermediate_energy = 0.0 - self.reward = 0.0 - # get assignment order - frac_act = FracAct(self.num_resid, self.restraints) - self.assign_order = frac_act.fractional_activation() - self.assign_shift = self.assign_order[self.assign_step] + # Get assignment order + self.activation = FractionalActivation(self.num_resid, self.restraints) + self.assign_order = self.activation.fractional_activation() + + ### TEST WHEN REMOVING FRACTIONAL ACTIVATION ### + # test_range = np.arange(0, self.num_resid) + # self.assign_order = test_range + ################################################ + + self.shift_to_assign = self.assign_order[self.assign_step] # print(self.assign_order) - # print(f'Shift to be assigned: {self.assign_shift}') + # print(f'Shift to be assigned: {self.shift_to_assign}') self.state = { - "coords": coords, - "actual_shifts": actual_shifts, + "coordinates": coordinates, + "obs_chemical_shifts": obs_chemical_shifts, "noes": noes, "connectivity": connectivity, "assignments": self.assignments, "assign_order": self.assign_order, - "assign_shift": self.assign_shift, - "total_energy": self.total_energy, - "reward": self.reward + "shift_to_assign": self.shift_to_assign, + "total_energy": 0.0, + "reward": 0.0 } return self.state - def step(self, action): - # print(action) - assert action < self.num_resid and action not in self.state['assignments'].values() - - energy = Energy(self.state['coords'], self.state['actual_shifts'], self.state['noes']) - - # add action to assignments - self.assignments[(self.assign_order[self.assign_step])] = action - # print(self.assignments) - - # calc energy and store temporarily - temp_intermediate_energy = energy.get_total_energy(self.restraints, self.assignments, tolerance=float(1/np.cbrt(self.num_resid))) - # print(f'this is the temp energy {temp_intermediate_energy}') + def render(self): + self.visualize = Visualization(self.state['coordinates'], self.state['obs_chemical_shifts'], self.state['connectivity'], self.restraints) + self.visualize.plot_shifts(named_tuple_used=True) + # self.visualize.plot_step(self.state['shift_to_assign'], self.assign_order) - # calc reward based on previous energy and temporary energy - self.state['reward'] = (self.intermediate_energy - temp_intermediate_energy) - # before correction, +ve means energy went down - we DO NOT want this, -ve means energy went up - assert self.state['reward'] <= 0 - self.state['reward'] = abs(self.state['reward']) - # print(f'this is the reward {self.reward}') - - # store energy as intermediate - self.intermediate_energy = temp_intermediate_energy - - # assign intermediate energy as running total - self.state['total_energy'] = self.intermediate_energy - # print(f"Running energy = {self.state['total_energy']}, Current reward = {self.state['reward']}") - - self.assign_step += 1 - # self.assign_shift = self.assign_order[self.assign_step] - terminated = True if len(self.assignments) == self.num_resid else False - - if terminated: - # print(f"Final assignments (shift:atom) = {self.assignments}\nFinal energy evaluation = {self.state['total_energy']}") - return self.state, self.state['reward'], terminated, self.state['total_energy'] - - self.state['assign_shift'] = self.assign_order[self.assign_step] - if not terminated: - # print(f'Shift to be assigned: {self.state["assign_shift"]}') - return self.state, self.state['reward'], terminated, self.state['total_energy'] diff --git a/nmr_text_adventure.py b/nmr_text_adventure.py index 694f742..cab1f37 100644 --- a/nmr_text_adventure.py +++ b/nmr_text_adventure.py @@ -1,12 +1,37 @@ -import gym +import gymnasium as gym import math import argparse from nmr_gym_env import GymEnv -from visualize_data import plot_shifts +from fake_data import FakeDataGenerator +from visualize_data import Visualization -if __name__ == '__main__': + + +def get_data(num_resid, pickled=False, pickle_data=True, example=True, random_key=False): + assert not pickled or not pickle_data + + fakedata = FakeDataGenerator(num_resid) + name = f"./*r{num_resid}.pkl" if example else "./current_run.pkl" + + # Grab a previous example (already pickled) + if pickled: + coordinates, obs_chemical_shifts, noes, connectivity = fakedata.load_pickle(example=example) + + # Create a new example and pickle + elif pickle_data: + fakedata.generate_data(pickle_data=True, example=example, random_key=random_key) + coordinates, obs_chemical_shifts, noes, connectivity = fakedata.load_pickle(example=example) + + # Run a new example and don't pickle + else: + coordinates, obs_chemical_shifts, noes, connectivity = fakedata.generate_data(pickle_data=False, example=example, random_key=random_key) + return coordinates, obs_chemical_shifts, noes, connectivity + + +def text_adventure(): + parser = argparse.ArgumentParser() parser.add_argument('num_resid', type=int, help='Number of residues') args = parser.parse_args() @@ -16,32 +41,83 @@ gym_env = GymEnv(num_resid) total_energy = math.inf - while total_energy > 0: - - ###### Available options to work with right now ###### + ###################################################### - # 1. Need to generate and save an example? Saved in all instances as fakedata_r{num_resid}.pkl - # observation = gym_env.reset(pickled=False, pickle_data=True, random_key=True) + print("\nAvailable options to work with:\n1. Generate and save an example\n2. Grab a previous saved example\n3. Run a new example (not saved)") - # 2. Grab one of these previous examples? Make sure example exists with desired resid number. - # observation = gym_env.reset(pickled=True, pickle_data=False, random_key=True) + selection = int(input("Option Selection: ")) + if selection == 1: + # 1. Need to generate and save an example? Saved in all instances as fakedata_r{num_resid}.pkl + coordinates, obs_chemical_shifts, noes, connectivity = get_data(num_resid, pickled=False, pickle_data=True, example=True, random_key=True) + if selection == 2: + # 2. Grab one of the previous examples? Make sure example exists with desired resid number + coordinates, obs_chemical_shifts, noes, connectivity = get_data(num_resid, pickled=True, pickle_data=False, random_key=True) + if selection == 3: # 3. Run a new example? Saved in all instances as current_run.pkl - observation = gym_env.reset(pickled=False, pickle_data=True, example=False, random_key=True) + coordinates, obs_chemical_shifts, noes, connectivity = get_data(num_resid, pickled=False, pickle_data=True, example=False, random_key=True) - # 4. Use a custom hardcoded state? May be some issues with data types and visualization. - # observation = gym_env.custom_state(state) + ###################################################### - ###################################################### + # Game repeats until energy of zero is reached + while total_energy > 0: + observation = gym_env.reset(coordinates, obs_chemical_shifts, noes, connectivity) terminated = False while not terminated: + gym_env.render() action = int(input()) if action < num_resid: if action not in observation['assignments'].values(): observation, reward, terminated, total_energy = gym_env.step(action) + + +if __name__ == '__main__': + text_adventure() + +# parser = argparse.ArgumentParser() +# parser.add_argument('num_resid', type=int, help='Number of residues') +# args = parser.parse_args() + +# num_resid = args.num_resid + +# gym_env = GymEnv(num_resid) +# total_energy = math.inf + +# while total_energy > 0: + +# ###### Available options to work with right now ###### + +# # 1. Need to generate and save an example? Saved in all instances as fakedata_r{num_resid}.pkl +# # observation = gym_env.reset(pickled=False, pickle_data=True, random_key=False) + +# # 2. Grab one of these previous examples? Make sure example exists with desired resid number. +# # observation = gym_env.reset(pickled=True, pickle_data=False, random_key=False) + +# # 3. Run a new example? Saved in all instances as current_run.pkl +# # with cProfile.Profile() as pr: +# observation = gym_env.reset(pickled=False, pickle_data=True, example=False, random_key=False) +# # stats = pstats.Stats(pr) +# # stats.sort_stats('time').print_stats() + +# # 4. Use a custom hardcoded state? May be some issues with data types and visualization. +# # observation = gym_env.custom_state(state) + +# ###################################################### + +# terminated = False + +# while not terminated: +# action = int(input()) +# if action < num_resid: +# if action not in observation['assignments'].values(): +# observation, reward, terminated, total_energy = gym_env.step(action) + +# if __name__ == '__main__': + # cProfile.run('noe_combinations()') + # test_run() # random generated set of 4 # state = { # "coords": [[0.35, 0.83, 0.89, 0.16, 0.034], [0.69, 0.12, 0.92, 0.41, 0.44], [0.05, 0.94, 0.51, 0.189, 0.62], [0.45, 0.06, 0.70, 0.63, 0.15]], From 6d3866679310e6626feb74e1c18db732b66141f5 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Sat, 24 May 2025 23:18:21 -0600 Subject: [PATCH 33/43] Unittesting for updates --- energy_unittest.py | 66 ++++++++++++--------- fractional_activation_unittest.py | 95 +++++++++++++------------------ 2 files changed, 78 insertions(+), 83 deletions(-) diff --git a/energy_unittest.py b/energy_unittest.py index 5b1375a..6edd4bd 100644 --- a/energy_unittest.py +++ b/energy_unittest.py @@ -1,56 +1,70 @@ import unittest +import numpy as np + +from energy import Energy + + class TestEnergyMethods(unittest.TestCase): def setUp(self): - self.energy = Energy() + self.coords = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [0.91, 0.92, 0.93], [0.95, 0.96, 0.97]] # x,y,z + self.energy = Energy(self.coords, 0, 0) + self.dist_grid = self.energy.calc_pdist() - #test for activated, if activated, the expected will be the energy calculation + # Test for NOE activation - if activated, the expected value will be the lowest energy calculated def test_energy_activated(self): - - #random data - coords = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [0.91, 0.92, 0.93], [0.95,0.96, 0.97]] #x,y,z - restraints = [[(2, 1),(0, 3)]] + restraint = [(2, 1), (0, 3)] assignments = {2: 1, 0: 3, 1: 2, 3: 4} - - result = self.energy.get_energy(restraints, coords, assignments) - expected = 0.01019237886466845 + + result = self.energy.calc_restraint_energy(assignments, restraint, self.dist_grid, 0.5) + expected = 0.0 self.assertEqual(result, expected) - #test for not activated, if not activated, the expected will be 0 + # Test if NOE not activated - if not activated, the expected value will be 0 (no energy calculated) def test_energy_not_activated(self): - - #random data - coords = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7,0.8,0.9], [0.91, 0.92, 0.93], [0.95,0.96, 0.97]] #x,y,z restraints = [[(2, 1), (0, 3)]] - assignments = {0:3} + assignments = {0: 3} - result = self.energy.get_energy(restraints, coords, assignments) + result = self.energy.get_total_energy(restraints, assignments, self.dist_grid) expected = 0 self.assertEqual(result, expected) - #test for sum of all energies for more then one activated NOE, the expected will be the total energy - #also works for NOE not activated, expected will be the same as test_energy_activated - def test_energy_sum(self): - coords = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [0.1, 0.92, 0.93], [0.95,0.96, 0.97]] #x,y,z - restraints = [(2, 1), (0, 3)], [(1, 2)] + # Test for sum of all energies for one activated NOE, the expected will be single restraint calculation + def test_energy_sum_not_activated(self): + restraints = [[(2, 1), (0, 3)], [(1, 2)]] assignments = {2: 1, 0: 3, 1: 2} - result = self.energy.get_energy(restraints, coords, assignments) - expected = 0.5263195283063458 + result = self.energy.get_total_energy(restraints, assignments, self.dist_grid) + expected = 0.01019237886466845 self.assertEqual(result, expected) - #energy calc - pass + # Test for sum of all energies for more then one activated NOE, the expected will be the sum of the two restraints + def test_energy_sum_activated(self): + restraints = [[(2, 1), (0, 3)], [(1, 2)]] + assignments = {2: 1, 0: 3, 1: 2, 3: 0} + + result = self.energy.get_total_energy(restraints, assignments, self.dist_grid) + expected = 0.01019237886466845*2 + + self.assertEqual(result, expected) + + # Test for flat bottom restraint, distance above tolerance def test_flat_bottom(self): - tolerance = 1 - x = 5 #distance + tolerance = 0.5 + x = 5 # distance result = self.energy.flat_bottom(x, tolerance) - expected = 20 # x^2 - tolerance*x + expected = 22 # x^2 - tolerance*x self.assertEqual(result, expected) + + +if __name__ == '__main__': + unittest.main() + ''' calculations for the expected of each unittest diff --git a/fractional_activation_unittest.py b/fractional_activation_unittest.py index b5e264e..e1ceb04 100644 --- a/fractional_activation_unittest.py +++ b/fractional_activation_unittest.py @@ -1,68 +1,49 @@ -#unittest for fractional activations -#has to be made into a class for it to work - import unittest +from itertools import chain -class TestFracAct(unittest.TestCase): - - def setUp(self): - self.frac_act = FracAct() - -#test for activation loop -#expected to return the noe peak divided by the total number of different/unique noe peaks - - def test_activation_loop(self): - assign_order = [] - noe = [(0, 1), (0, 2)] - target = 0 - - result = self.frac_act.activation_loop(target, noe, assign_order) - expected = 1/3 - - self.assertEqual(result, expected) - -#test if each noe peak will be evaluated to give the fractional activation +from assignment_order import FractionalActivation - def test_fractional_activation_of_all_noe_peaks(self): - num_resid = 3 - noes = [[(0, 1), (0, 2)]] - result = self.frac_act.fractional_activation_summation(num_resid, noes) - expected = [1/3, 1/3, 1/3] - self.assertEqual(result, expected) - -#for more then one noe, the probabilities are summed - - def test_fractional_activation_summation_of_noes(self): - num_resid = 4 - noes = [(0, 1), (0, 2)], [(2,1),(1,1)], [(1,0),(3,2)] - - result = self.frac_act.fractional_activation_summation(num_resid, noes) - expected = [7/12, 13/12, 13/12, 1/4] - - self.assertEqual(result, expected) - -#according to the fractional activations calculated, should give the correct order -#if it is a tie, then lowest value goes first +class TestFracAct(unittest.TestCase): + # Test for activation loop + # Expected to return the target noe peak divided by the total number of different/unique noe peak options + def test_activation_loop(self): + + cases = [ + {'noe': [(0, 1), (0, 2)], 'noe_set': {0, 1, 2}, 'assign_order': [], 'target': 0, 'expected': 1/3, 'message': 'CASE 1: First assignment'}, + {'noe': [(0, 1), (0, 2)], 'noe_set': {1, 2}, 'assign_order': [0], 'target': 0, 'expected': 0, 'message': 'CASE 2: Target already assigned'}, + {'noe': [(0, 1), (0, 2)], 'noe_set': {2}, 'assign_order': [0, 1],'target': 2, 'expected': 1.0, 'message': 'CASE 3: Final assignment'}, + {'noe': [(0, 1), (0, 2)], 'noe_set': {}, 'assign_order': [0, 1, 2], 'target': 0, 'expected': 0, 'message': 'CASE 4: Everything already assigned'}, + {'noe': [(0, 1), (0, 2)], 'noe_set': {0, 1, 2}, 'assign_order': [], 'target': 3, 'expected': 0, 'message': 'CASE 5: Target not present in NOE'} + ] + + for case in cases: + with self.subTest(message=case['message']): + frac_act = FractionalActivation(3, case['noe']) + result = frac_act.activation_loop(case['target'], case['noe_set'], case['assign_order']) + self.assertEqual(result, case['expected']) + + # According to the fractional activations calculated, should give the correct order + # If it is a tie, then lowest index goes first def test_fractional_activation_for_assignment_order(self): - num_resid = 4 - noes = [(0, 1), (0, 2)], [(2,1),(1,1)], [(1,0),(3,2)] - - result = self.frac_act.fractional_activation(num_resid, noes) - expected = [1, 2, 0, 3] - - self.assertEqual(result, expected) + + cases = [ + {'restraints': ([[(0, 1), (0, 2)], [(2, 1), (1, 1)], [(1, 0), (3, 2)]]), 'num_resid': 4, 'expected': [1, 2, 0, 3], 'message': 'CASE 1: Acceptable order'}, + {'restraints': ([[(0, 1), (0, 2)], [(2, 1), (1, 1)], [(1, 0), (3, 2)]]), 'num_resid': 5, 'expected': [1, 2, 0, 3, 4], 'message': 'CASE 2: Higher resid count --> Missing residues added last'}, + # {'restraints': ([[(2, 3)], [(2, 1)], [(1, 0)], [(1, 3)], [(1, 0)], [(2, 3)], [(1, 3)], [(2, 0)]]), 'num_resid': 3, 'expected': [1, 3, 2, 0], 'message': 'CASE 2: Lower resid count --> Expected Failure'}, + {'restraints': ([[(0, 6)], [(1, 4)], [(1, 6)], [(2, 3)], [(2, 3)], [(1, 4)], [], [(0, 6)], [], [(4, 8)]]), 'num_resid': 10, 'expected': [1, 4, 6, 0, 2, 3, 8, 5, 7, 9], 'message': 'CASE 3: Handling blank NOES (true case)'}, + {'restraints': ([[(0, 3)], [(0, 1)], [(1, 2)], [(1, 3)], [(1, 2)], [(0, 3)], [(1, 3)], [(0, 2)]]), 'num_resid': 4, 'expected': [1, 3, 0, 2], 'message': 'CASE 4: Hand solved, single restraints'} + ] -#if number of residues is higher than the number of peaks, then it will loop through to give first/lowest value peak ex. will give 0 for every additional residue -#if number of residues is lower than the number of peaks, then it will return the first value again (0, in this case) + for case in cases: + with self.subTest(message=case['message']): + frac_act = FractionalActivation(case['num_resid'], case['restraints']) + result = frac_act.fractional_activation() + self.assertEqual(result, case['expected']) - def test_fractional_activation_for_num_resid(self): - num_resid = 5 - noes = [[(0, 1), (0, 2)]], [(2,1),(1,1)], [(1,0),(3,2)] - result = self.frac_act.fractional_activation(num_resid, noes) - expected = [1, 2, 0, 3] - self.assertEqual(result, expected) \ No newline at end of file +if __name__ == '__main__': + unittest.main() \ No newline at end of file From 44df9901a8e118c8873364b69ae5a1f4664074b3 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Fri, 11 Jul 2025 08:18:45 -0600 Subject: [PATCH 34/43] Sample PyG work for triples --- pyg_triple_update.py | 163 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 163 insertions(+) create mode 100644 pyg_triple_update.py diff --git a/pyg_triple_update.py b/pyg_triple_update.py new file mode 100644 index 0000000..ee64b4a --- /dev/null +++ b/pyg_triple_update.py @@ -0,0 +1,163 @@ +import torch +import torch_geometric +from torch_geometric.data import Data + +from torch_geometric.loader import DataLoader +from torch_geometric.nn import MessagePassing +from torch_geometric.data import HeteroData +import torch_geometric.transforms as T +import torch.nn as nn + +# Set up data (random to test) +data = HeteroData() + +# Spatial and non-spatial +data['NOE'].x = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32) +data['NOE'].f = torch.zeros((2, 1), dtype=torch.float32) +data['RES'].x = torch.tensor([[7, 8, 9, 1, 2], [10, 11, 12, 7, 3]], dtype=torch.float32) +data['RES'].f = torch.zeros((2, 1), dtype=torch.float32) +data['SHIFT'].x = torch.tensor([[4, 5], [1, 2]], dtype=torch.float32) +data['SHIFT'].f = torch.zeros((2, 1), dtype=torch.float32) + +# Triple +data['TRIPLE'].x = torch.tensor(()) + +# Edges +data['NOE', 'NOE_extract', 'TRIPLE'].edge_index = torch.tensor([[0], [0]]) +data['RES', 'NH1_extract', 'TRIPLE'].edge_index = torch.tensor([[0], [0]]) +data['RES', 'NH2_extract', 'TRIPLE'].edge_index = torch.tensor([[1], [0]]) +data['SHIFT', 'NH1_extract', 'TRIPLE'].edge_index = torch.tensor([[0], [0]]) +data['SHIFT', 'NH2_extract', 'TRIPLE'].edge_index = torch.tensor([[1], [0]]) +data['TRIPLE', 'update', 'TRIPLE'].edge_index = torch.tensor([[0], [0]]) # self loop + +# Generate reverse for message passing +data = T.ToUndirected()(data) + +def triple_in(data, node, edge): + """ + Source node to triple via indexing. + """ + # Spatial features + xi = data[node].x[data[edge].edge_index[0]] + # Non-spatial features + fi = data[node].f[data[edge].edge_index[0]] + + return xi, fi + + +class TripleUpdateResResNoe(MessagePassing): + def __init__(self): + super().__init__(aggr='add') + + self.NOE_N1 = 0 + self.NOE_H1 = 1 + self.NOE_H2 = 2 + + self.RES_X = 0 + self.RES_Y = 1 + self.RES_Z = 2 + self.RES_N = 3 + self.RES_H = 4 + + self.SHIFT_N = 0 + self.SHIFT_H = 1 + + self.hidden = 64 + + self.mlp1 = nn.Sequential( + nn.Linear(9, self.hidden), + nn.ReLU(), + nn.Linear(self.hidden, 9)) + self.mlp2 = nn.Sequential( + nn.Linear(3, self.hidden), + nn.ReLU(), + nn.Linear(self.hidden, 3)) + + def calc_noe_difference(self, x1, x2, noe, N1, H1, H2): + diff_N = (noe[:, self.NOE_N1] - x1[:, N1]).unsqueeze(1) # N + diff_H1 = (noe[:, self.NOE_H1] - x1[:, H1]).unsqueeze(1) # H' + diff_H2 = (noe[:, self.NOE_H2] - x2[:, H2]).unsqueeze(1) # H" + return diff_N, diff_H1, diff_H2 + + def calc_res_distance(self, x1, x2): + rel_positions = ((x1[:, self.RES_X:self.RES_Z+1]) - (x2[:, self.RES_X:self.RES_Z+1])).squeeze(dim=-1) + distances = torch.norm(rel_positions, dim=-1, keepdim=True)**2 + return rel_positions, distances + + def forward(self, x1, x2, x3, f1, f2, f3, edge_index): + out = self.propagate(edge_index, x1=x1, x2=x2, x3=x3, f1=f1, f2=f2, f3=f3) + return out + + def message(self, x1_i, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j): + # Residue distances + rel_positions, distances = self.calc_res_distance(x1_j, x2_j) + + # Differences relative to NOE shifts (N, H', H") + # SHIFT 1 + diff1, diff2, diff3 = self.calc_noe_difference(x1_j, x1_j, x3_j, self.RES_N, self.RES_H, self.RES_H) + # SHIFT 2 + diff4, diff5, diff6 = self.calc_noe_difference(x2_j, x2_j, x3_j, self.RES_N, self.RES_H, self.RES_H) + + # Inputs for MLP + # (N, H', H", features) + mlp_input1 = (torch.cat((diff1, diff2, diff3, diff4, diff5, diff6, f1_j, f2_j, f3_j), dim=-1)) + # (distances, features) + mlp_input2 = (torch.cat((distances, f1_j, f2_j), dim=-1)) + + # MLP + mlp_out1 = self.mlp1(mlp_input1) + mlp_out2 = self.mlp2(mlp_input2) + + # NOE deltas + delta1x = (diff1 * mlp_out1[:, 0].unsqueeze(1)) + (diff4 * mlp_out1[:, 3].unsqueeze(1)) # N + delta2x = (diff2 * mlp_out1[:, 1].unsqueeze(1)) + (diff5 * mlp_out1[:, 4].unsqueeze(1)) # H' + delta3x = (diff3 * mlp_out1[:, 2].unsqueeze(1)) + (diff6 * mlp_out1[:, 5].unsqueeze(1)) # H" + deltaf3 = mlp_out1[:, 8].unsqueeze(1) # features + + # SHIFT deltas + # SHIFT1 + delta4x = -diff1 * mlp_out1[:, 0].unsqueeze(1) # N1 + delta5x = (-diff2 * mlp_out1[:, 1].unsqueeze(1)) + (-diff3 * mlp_out1[:, 2].unsqueeze(1)) # H1 + deltaf1 = mlp_out1[:, 6] + # SHIFT2 + delta6x = -diff4 * mlp_out1[:, 3].unsqueeze(1) # N2 + delta7x = (-diff5 * mlp_out1[:, 4].unsqueeze(1)) + (-diff6 * mlp_out1[:, 5].unsqueeze(1)) # H2 + deltaf2 = mlp_out1[:, 7] + + # DISTANCE deltas? + rel_positions = rel_positions + delta8x = (rel_positions * (mlp_out2[:, 0]).unsqueeze(1)) + deltaf1 = (deltaf1 + mlp_out2[:, 1]).unsqueeze(1) + deltaf2 = (deltaf2 + mlp_out2[:, 2]).unsqueeze(1) + + return (torch.cat((delta1x, delta2x, delta3x, delta4x, delta5x, delta6x, delta7x, delta8x, deltaf1, deltaf2, deltaf3), dim=-1)) + + def update(self, aggr_out, x1, x2, x3, f1, f2, f3): + update_NOE, update_x1, update_x2, update_f1, update_f2, update_f3 = aggr_out[:, 0:3], torch.cat((aggr_out[:, 7:10], aggr_out[:, 3:5]), dim=-1), torch.cat((aggr_out[:, 7:10]*-1, aggr_out[:, 5:7]), dim=-1), aggr_out[:, 10], aggr_out[:, 11], aggr_out[:, 12] + return (x3 + update_NOE, + x1 + update_x1, + x2 + update_x2, + f1 + update_f1, + f2 + update_f2, + f3 + update_f3) + + # return aggr_out + + +if __name__ == "__main__": + + # Can grab these from data - would need some form of string input for each combination (node and edge) + # data.node_types + # data.edge_types + node1, node2, node3 = 'RES', 'RES', 'NOE' + edge1, edge2, edge3 = ('RES', 'NH1_extract', 'TRIPLE'), ('RES', 'NH2_extract', 'TRIPLE'), ('NOE', 'NOE_extract', 'TRIPLE') + + # Data --> Triple node + x1, f1 = triple_in(data, node=node1, edge=edge1) + x2, f2 = triple_in(data, node=node2, edge=edge2) + x3, f3 = triple_in(data, node=node3, edge=edge3) + + gnn = TripleUpdateResResNoe() + + out = gnn(x1, x2, x3, f1, f2, f3, data['TRIPLE', 'update', 'TRIPLE'].edge_index) + print(f'Update: {out}') From 13682a7a81ef979d5bfed5276a7a129ebd5bbd12 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Wed, 30 Jul 2025 09:18:42 -0600 Subject: [PATCH 35/43] Modified triple out message passing and overall in/self/out update calls --- pyg_triple_update.py | 310 +++++++++++++++++++++++++++++-------------- 1 file changed, 211 insertions(+), 99 deletions(-) diff --git a/pyg_triple_update.py b/pyg_triple_update.py index ee64b4a..2d94bf4 100644 --- a/pyg_triple_update.py +++ b/pyg_triple_update.py @@ -23,141 +23,253 @@ data['TRIPLE'].x = torch.tensor(()) # Edges -data['NOE', 'NOE_extract', 'TRIPLE'].edge_index = torch.tensor([[0], [0]]) -data['RES', 'NH1_extract', 'TRIPLE'].edge_index = torch.tensor([[0], [0]]) -data['RES', 'NH2_extract', 'TRIPLE'].edge_index = torch.tensor([[1], [0]]) -data['SHIFT', 'NH1_extract', 'TRIPLE'].edge_index = torch.tensor([[0], [0]]) -data['SHIFT', 'NH2_extract', 'TRIPLE'].edge_index = torch.tensor([[1], [0]]) -data['TRIPLE', 'update', 'TRIPLE'].edge_index = torch.tensor([[0], [0]]) # self loop +data['NOE', 'NOE_extract', 'TRIPLE'].edge_index = torch.tensor([[0, 1], [0, 1]]) +data['RES', 'NH1_extract', 'TRIPLE'].edge_index = torch.tensor([[0, 1], [0, 1]]) +data['RES', 'NH2_extract', 'TRIPLE'].edge_index = torch.tensor([[1, 0], [0, 1]]) +data['SHIFT', 'NH1_extract', 'TRIPLE'].edge_index = torch.tensor([[0, 1], [0, 1]]) +data['SHIFT', 'NH2_extract', 'TRIPLE'].edge_index = torch.tensor([[1, 0], [0, 1]]) -# Generate reverse for message passing -data = T.ToUndirected()(data) +data['TRIPLE', 'update', 'TRIPLE'].edge_index = torch.tensor([[0, 1], [0, 1]]) # self loop -def triple_in(data, node, edge): - """ - Source node to triple via indexing. - """ - # Spatial features - xi = data[node].x[data[edge].edge_index[0]] - # Non-spatial features - fi = data[node].f[data[edge].edge_index[0]] +data['TRIPLE', 'NOE_add', 'NOE'].edge_index = torch.tensor([[0, 1], [0, 1]]) +data['TRIPLE', 'NH1_add', 'RES'].edge_index = torch.tensor([[0, 1], [0, 1]]) +data['TRIPLE', 'NH2_add', 'RES'].edge_index = torch.tensor([[0, 1], [1, 0]]) +data['TRIPLE', 'res1_add', 'RES'].edge_index = torch.tensor([[0, 1], [0, 1]]) +data['TRIPLE', 'res2_add', 'RES'].edge_index = torch.tensor([[0, 1], [1, 0]]) +data['TRIPLE', 'NH1_add', 'SHIFT'].edge_index = torch.tensor([[0, 1], [0, 1]]) +data['TRIPLE', 'NH2_add', 'SHIFT'].edge_index = torch.tensor([[0, 1], [1, 0]]) - return xi, fi +class TripleIn(): + def __init__(self, data): + self.data = data + + def grab_node(self, node_type, edge_type): + """ + Source node to triple via indexing. + """ + # Spatial features + xi = self.data[node_type].x[self.data[edge_type].edge_index[0]] + # Non-spatial features + fi = self.data[node_type].f[self.data[edge_type].edge_index[0]] + + return xi, fi + + def construct_triple(self, edge_type1, edge_type2, edge_type3): + """ + Constructs triple based on type: + - residue, residue, NOE + - residue, shift, NOE + - shift, residue, NOE + - shift, shift, NOE + """ + x1, f1 = self.grab_node(node_type=edge_type1[0], edge_type=edge_type1) + x2, f2 = self.grab_node(node_type=edge_type2[0], edge_type=edge_type2) + x3, f3 = self.grab_node(node_type=edge_type3[0], edge_type=edge_type3) + + return x1, x2, x3, f1, f2, f3 class TripleUpdateResResNoe(MessagePassing): + """ + Message passing class for triple self updates (only set up for residue, residue, NOE triple type). + Residue, residue, NOE triple type is the only one that deviates in calculations (coordinate calculation used), + other three would have same base set up with some differences in indexing. + Pairwise calculations would use the same functions but would have a different message passing scheme. + + Options --> if/else statements for triples because of similarity, different class for pairwise? + """ def __init__(self): super().__init__(aggr='add') - self.NOE_N1 = 0 - self.NOE_H1 = 1 - self.NOE_H2 = 2 + # Input index organization - slices to keep proper tensor dimension + self.NOE_N1 = slice(0,1) + self.NOE_H1 = slice(1,2) + self.NOE_H2 = slice(2,3) - self.RES_X = 0 - self.RES_Y = 1 - self.RES_Z = 2 - self.RES_N = 3 - self.RES_H = 4 + self.RES_XYZ = slice(0,3) + self.RES_N = slice(3,4) + self.RES_H = slice(4,5) - self.SHIFT_N = 0 - self.SHIFT_H = 1 + self.SHIFT_N = slice(0,1) + self.SHIFT_H = slice(1,2) self.hidden = 64 self.mlp1 = nn.Sequential( nn.Linear(9, self.hidden), nn.ReLU(), - nn.Linear(self.hidden, 9)) - self.mlp2 = nn.Sequential( - nn.Linear(3, self.hidden), - nn.ReLU(), - nn.Linear(self.hidden, 3)) + nn.Linear(self.hidden, 16)) def calc_noe_difference(self, x1, x2, noe, N1, H1, H2): - diff_N = (noe[:, self.NOE_N1] - x1[:, N1]).unsqueeze(1) # N - diff_H1 = (noe[:, self.NOE_H1] - x1[:, H1]).unsqueeze(1) # H' - diff_H2 = (noe[:, self.NOE_H2] - x2[:, H2]).unsqueeze(1) # H" + """ + Calculates shift difference between NOE and residue/measured shifts (direct/indirect only reverse option is available by index). + """ + diff_N = (noe[:, self.NOE_N1] - x1[:, N1]) # N [n, 1] + diff_H1 = (noe[:, self.NOE_H1] - x1[:, H1]) # H' [n, 1] + diff_H2 = (noe[:, self.NOE_H2] - x2[:, H2]) # H" [n, 1] return diff_N, diff_H1, diff_H2 - + + def calc_shift_difference(self, x1, x2, N1, H1): + """ + Calculates shift difference between residue/measured shifts. + """ + diff_N = (x1[:, N1] - x2[:, N1]) # N [n, 1] + diff_H = (x1[:, H1] - x2[:, H1]) # H [n, 1] + return diff_N, diff_H + def calc_res_distance(self, x1, x2): - rel_positions = ((x1[:, self.RES_X:self.RES_Z+1]) - (x2[:, self.RES_X:self.RES_Z+1])).squeeze(dim=-1) - distances = torch.norm(rel_positions, dim=-1, keepdim=True)**2 - return rel_positions, distances + """ + Calculates relative distance between residues and this value squared for equivariant calculations. + """ + rel_dist = ((x1[:, self.RES_XYZ]) - (x2[:, self.RES_XYZ])) # [n, 3] + dist2 = torch.norm(rel_dist, dim=-1, keepdim=True)**2 # [n, 1] + return rel_dist, dist2 def forward(self, x1, x2, x3, f1, f2, f3, edge_index): out = self.propagate(edge_index, x1=x1, x2=x2, x3=x3, f1=f1, f2=f2, f3=f3) return out - def message(self, x1_i, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j): + def message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j): # Residue distances - rel_positions, distances = self.calc_res_distance(x1_j, x2_j) + rel_dist, dist2 = self.calc_res_distance(x1_j, x2_j) # Differences relative to NOE shifts (N, H', H") - # SHIFT 1 - diff1, diff2, diff3 = self.calc_noe_difference(x1_j, x1_j, x3_j, self.RES_N, self.RES_H, self.RES_H) - # SHIFT 2 - diff4, diff5, diff6 = self.calc_noe_difference(x2_j, x2_j, x3_j, self.RES_N, self.RES_H, self.RES_H) + diff1, diff2, diff3 = self.calc_noe_difference(x1_j, x2_j, x3_j, self.RES_N, self.RES_H, self.RES_H) - # Inputs for MLP - # (N, H', H", features) - mlp_input1 = (torch.cat((diff1, diff2, diff3, diff4, diff5, diff6, f1_j, f2_j, f3_j), dim=-1)) - # (distances, features) - mlp_input2 = (torch.cat((distances, f1_j, f2_j), dim=-1)) + # Shift differences + diff4, diff5 = self.calc_shift_difference(x1_j, x2_j, self.RES_N, self.RES_H) - # MLP - mlp_out1 = self.mlp1(mlp_input1) - mlp_out2 = self.mlp2(mlp_input2) + # Input for MLP + # (N, H', H", N, H, dist2, features) + mlp_in = (torch.cat((diff1, diff2, diff3, diff4, diff5, dist2, f1_j, f2_j, f3_j), dim=-1)) # [n, 9] + + # Output from MLP + # Out should include values for each 'change' wanting to make + # (N, H', H", N1, H1, N2, H2, dist1, dist2, features) + mlp_out = self.mlp1(mlp_in) # [n, 16] # NOE deltas - delta1x = (diff1 * mlp_out1[:, 0].unsqueeze(1)) + (diff4 * mlp_out1[:, 3].unsqueeze(1)) # N - delta2x = (diff2 * mlp_out1[:, 1].unsqueeze(1)) + (diff5 * mlp_out1[:, 4].unsqueeze(1)) # H' - delta3x = (diff3 * mlp_out1[:, 2].unsqueeze(1)) + (diff6 * mlp_out1[:, 5].unsqueeze(1)) # H" - deltaf3 = mlp_out1[:, 8].unsqueeze(1) # features - - # SHIFT deltas - # SHIFT1 - delta4x = -diff1 * mlp_out1[:, 0].unsqueeze(1) # N1 - delta5x = (-diff2 * mlp_out1[:, 1].unsqueeze(1)) + (-diff3 * mlp_out1[:, 2].unsqueeze(1)) # H1 - deltaf1 = mlp_out1[:, 6] - # SHIFT2 - delta6x = -diff4 * mlp_out1[:, 3].unsqueeze(1) # N2 - delta7x = (-diff5 * mlp_out1[:, 4].unsqueeze(1)) + (-diff6 * mlp_out1[:, 5].unsqueeze(1)) # H2 - deltaf2 = mlp_out1[:, 7] - - # DISTANCE deltas? - rel_positions = rel_positions - delta8x = (rel_positions * (mlp_out2[:, 0]).unsqueeze(1)) - deltaf1 = (deltaf1 + mlp_out2[:, 1]).unsqueeze(1) - deltaf2 = (deltaf2 + mlp_out2[:, 2]).unsqueeze(1) - - return (torch.cat((delta1x, delta2x, delta3x, delta4x, delta5x, delta6x, delta7x, delta8x, deltaf1, deltaf2, deltaf3), dim=-1)) - - def update(self, aggr_out, x1, x2, x3, f1, f2, f3): - update_NOE, update_x1, update_x2, update_f1, update_f2, update_f3 = aggr_out[:, 0:3], torch.cat((aggr_out[:, 7:10], aggr_out[:, 3:5]), dim=-1), torch.cat((aggr_out[:, 7:10]*-1, aggr_out[:, 5:7]), dim=-1), aggr_out[:, 10], aggr_out[:, 11], aggr_out[:, 12] - return (x3 + update_NOE, - x1 + update_x1, - x2 + update_x2, - f1 + update_f1, - f2 + update_f2, - f3 + update_f3) - - # return aggr_out + delta1x = (diff1 * mlp_out[:, 0:1]) # N [n, 1] + delta2x = (diff2 * mlp_out[:, 1:2]) # H' [n, 1] + delta3x = (diff3 * mlp_out[:, 2:3]) # H" [n, 1] + # SHIFT deltas residue + delta4x = (diff4 * mlp_out[:, 3:4]) # N1 [n, 1] + delta5x = (diff5 * mlp_out[:, 4:5]) # H1 [n, 1] + delta6x = (diff4 * mlp_out[:, 5:6]) # N2 [n, 1] + delta7x = (diff5 * mlp_out[:, 6:7]) # H2 [n, 1] -if __name__ == "__main__": + # DISTANCE deltas + delta12x = (rel_dist * (mlp_out[:, 7:10])) # [n, 3] + delta13x = (rel_dist * (mlp_out[:, 10:13])) # [n, 3] + + # FEATURE deltas + delta1f = mlp_out[:, 13:14] # [n, 1] + delta2f = mlp_out[:, 14:15] # [n, 1] + delta3f = mlp_out[:, 15:] # [n, 1] + + return torch.cat((delta1x, delta2x, delta3x, delta4x, delta5x, delta6x, delta7x, delta12x, delta13x, delta1f, delta2f, delta3f), dim=-1) # [n, 16] + + def update(self, aggr_out): + # res involved in multiple different triples --> need to add across these (in different message passing class) + # output shapes: noe[n, 3], shift1[n, 2], shift2[n, 2], dist1[n, 3], dist2[n, 3], f1[n, 1], f2[n, 1], f3[n, 1] + delta_noe, delta_shift1, delta_shift2, delta_dist1, delta_dist2, delta_f1, delta_f2, delta_f3 = aggr_out[:, 0:3], aggr_out[:, 3:5], aggr_out[:, 5:7], aggr_out[:, 7:10], aggr_out[:, 10:13], aggr_out[:, 13:14], aggr_out[:, 14:15], aggr_out[:, 15:] + return delta_shift1, delta_shift2, delta_noe, delta_dist1, delta_dist2, delta_f1, delta_f2, delta_f3 + + + +class TripleMessagePass(MessagePassing): + """ + Standard message passing class for outgoing triple messages. + """ + def __init__(self): + super().__init__(aggr='add') - # Can grab these from data - would need some form of string input for each combination (node and edge) - # data.node_types - # data.edge_types - node1, node2, node3 = 'RES', 'RES', 'NOE' - edge1, edge2, edge3 = ('RES', 'NH1_extract', 'TRIPLE'), ('RES', 'NH2_extract', 'TRIPLE'), ('NOE', 'NOE_extract', 'TRIPLE') + def forward(self, x_source, x_target, edge_index): + # SIZE (n, m) (source, target) + return self.propagate(edge_index=edge_index, x=(x_source, x_target), size=(x_source.size(0), x_target.size(0))) - # Data --> Triple node - x1, f1 = triple_in(data, node=node1, edge=edge1) - x2, f2 = triple_in(data, node=node2, edge=edge2) - x3, f3 = triple_in(data, node=node3, edge=edge3) + def message(self, x_j): + return x_j - gnn = TripleUpdateResResNoe() + def update(self, aggr_out, x): + # non-spatial features can also be included in same call but requires more work if they're not being updated (don't want to do residue feature updates 2x - only call once on shifts or coordinates) + # x_val, f_val = aggr_out[:, :len(x[1][1])], aggr_out[:, len(x[1][1]):] + # outf = f_val + f[1] + # outx = x_val + x[1] + + out = aggr_out + x[1] + return out + + + +class TripleOut(): + def __init__(self, data): + self.data = data + self.triple_messgage = TripleMessagePass() + + def update_data(self, x_source, f_source, edge_type, x_index, update_f=True): + """ + Calls message passing and directly updates the heterodata object based on target. + """ + source, edge, target = edge_type + edge_index = self.data[edge_type].edge_index + + # Message passing/update for spatial features + self.data[target].x[:, x_index] = self.triple_messgage(x_source, self.data[target].x[:, x_index], edge_index) + + # Message passing/update for non-spatial features if needed (don't want a duplicate update for residue features) + if update_f: + self.data[target].f = self.triple_messgage(f_source, self.data[target].f, edge_index) + + return self.data + + + +class ProteinGNN(): + """ + Test class that: + - construct the triple based on edge types (only res, res, NOE as example) + - self updates + - sends message back to original nodes (note number of reverse edges depends on triple type) + + Can make separate classes to construct each triple type and do following updates or generalize this as a base class? + """ + def __init__(self, data): + self.data = data + self.triple_in = TripleIn(data) + self.triple_self = TripleUpdateResResNoe() # update would need to match data coming in + self.triple_out = TripleOut(data) + + # Original data object indexing + self.NOE = slice(0,3) + self.RES_XYZ = slice(0,3) + self.RES_NH = slice(3,5) + self.SHIFT_NH = slice(0,2) + + def forward(self): + # Triple in + x1, x2, x3, f1, f2, f3 = self.triple_in.construct_triple(('RES', 'NH1_extract', 'TRIPLE'), ('RES', 'NH2_extract', 'TRIPLE'), ('NOE', 'NOE_extract', 'TRIPLE')) + + # Self update + delta_shift1, delta_shift2, delta_noe, delta_dist1, delta_dist2, deltaf1, deltaf2, deltaf3 = self.triple_self(x1, x2, x3, f1, f2, f3, self.data['TRIPLE', 'update', 'TRIPLE'].edge_index) + + # Triple out (typically 3 reverse edges - 5 if coordinates are present) + self.data = self.triple_out.update_data(delta_shift1, deltaf1, ('TRIPLE', 'NH1_add', 'RES'), self.RES_NH) + self.data = self.triple_out.update_data(delta_shift2, deltaf2, ('TRIPLE', 'NH2_add', 'RES'), self.RES_NH) + self.data = self.triple_out.update_data(delta_noe, deltaf3, ('TRIPLE', 'NOE_add', 'NOE'), self.NOE) + self.data = self.triple_out.update_data(delta_dist1, deltaf1, ('TRIPLE', 'res1_add', 'RES'), self.RES_XYZ, update_f=False) + self.data = self.triple_out.update_data(delta_dist2, deltaf2, ('TRIPLE', 'res2_add', 'RES'), self.RES_XYZ, update_f=False) + + return self.data + + + +if __name__ == "__main__": + print(f"original RES: {data['RES']}") + print(f"original NOE: {data['NOE']}") + testgnn = ProteinGNN(data) + data = testgnn.forward() + print(f"updated RES: {data['RES']}") + print(f"updated NOE: {data['NOE']}") - out = gnn(x1, x2, x3, f1, f2, f3, data['TRIPLE', 'update', 'TRIPLE'].edge_index) - print(f'Update: {out}') From c55fa56777b954416a0b2ee98d9b461619d99197 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Thu, 21 Aug 2025 13:34:23 -0600 Subject: [PATCH 36/43] Split triple node types and modified setup to work with batching --- pyg_triple_update.py | 457 ++++++++++++++++++++++++++++++------------- 1 file changed, 317 insertions(+), 140 deletions(-) diff --git a/pyg_triple_update.py b/pyg_triple_update.py index 2d94bf4..4bc1d0f 100644 --- a/pyg_triple_update.py +++ b/pyg_triple_update.py @@ -8,137 +8,234 @@ import torch_geometric.transforms as T import torch.nn as nn -# Set up data (random to test) +from itertools import product + +# Set up data data = HeteroData() # Spatial and non-spatial -data['NOE'].x = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32) +data['NOE'].x = torch.tensor([[1, 2, 3]], dtype=torch.float32) data['NOE'].f = torch.zeros((2, 1), dtype=torch.float32) data['RES'].x = torch.tensor([[7, 8, 9, 1, 2], [10, 11, 12, 7, 3]], dtype=torch.float32) data['RES'].f = torch.zeros((2, 1), dtype=torch.float32) data['SHIFT'].x = torch.tensor([[4, 5], [1, 2]], dtype=torch.float32) data['SHIFT'].f = torch.zeros((2, 1), dtype=torch.float32) -# Triple -data['TRIPLE'].x = torch.tensor(()) +# Triples (nothing is ever stored in these - they are placeholders to be used in edge types) +data['TRIPLE'].x = torch.zeros((1, 1), dtype=torch.float32) # easier just to use separate triple for self updates +data['TRIPLE0'].x = torch.zeros((1, 1), dtype=torch.float32) +data['TRIPLE1'].x = torch.zeros((1, 1), dtype=torch.float32) +data['TRIPLE2'].x = torch.zeros((1, 1), dtype=torch.float32) +data['TRIPLE3'].x = torch.zeros((1, 1), dtype=torch.float32) + +# # Edges +# data['NOE', 'NOE_extract', 'TRIPLE'].edge_index = torch.tensor([[0, 0], [0, 0]]) +# data['RES', 'NH1_extract', 'TRIPLE'].edge_index = torch.tensor([[0, 1], [0, 0]]) +# data['RES', 'NH2_extract', 'TRIPLE'].edge_index = torch.tensor([[1, 0], [0, 0]]) #[[1, 0], [0, 1] +# data['SHIFT', 'NH1_extract', 'TRIPLE'].edge_index = torch.tensor([[0, 1], [0, 0]]) +# data['SHIFT', 'NH2_extract', 'TRIPLE'].edge_index = torch.tensor([[1, 0], [0, 0]]) + +# data['TRIPLE', 'update', 'TRIPLE'].edge_index = torch.tensor([[0], [0]]) # self loop + +# data['TRIPLE', 'NOE_add', 'NOE'].edge_index = torch.tensor([[0, 0], [0, 0]]) +# data['TRIPLE', 'NH1_add', 'RES'].edge_index = torch.tensor([[0, 0], [0, 1]]) +# data['TRIPLE', 'NH2_add', 'RES'].edge_index = torch.tensor([[0, 0], [1, 0]]) +# data['TRIPLE', 'res1_add', 'RES'].edge_index = torch.tensor([[0, 0], [0, 1]]) +# data['TRIPLE', 'res2_add', 'RES'].edge_index = torch.tensor([[0, 0], [1, 0]]) +# data['TRIPLE', 'NH1_add', 'SHIFT'].edge_index = torch.tensor([[0, 0], [0, 1]]) +# data['TRIPLE', 'NH2_add', 'SHIFT'].edge_index = torch.tensor([[0, 0], [1, 0]]) + +class ConstructNodes(): + def __init__(self, data_history): + self.data_history = data_history + + def construct_data_nodes(self, data): + # will need to extract each component of data from data_history + # data['NOE'].x + # data['NOE'].f + # data['SHIFT'].x + # data['SHIFT'].f + # data['RES'].x + # data['RES'].f + return data + + def construct_triple_nodes(self, data): + # Triples (nothing is ever stored in these - they are placeholders to be used in edge types) + data['TRIPLE'].x = torch.zeros((1, 1), dtype=torch.float32) # easier just to use separate triple for self updates + data['TRIPLE0'].x = torch.zeros((1, 1), dtype=torch.float32) + data['TRIPLE1'].x = torch.zeros((1, 1), dtype=torch.float32) + data['TRIPLE2'].x = torch.zeros((1, 1), dtype=torch.float32) + data['TRIPLE3'].x = torch.zeros((1, 1), dtype=torch.float32) + return data + + def construct_data(self): + data = HeteroData() + self.construct_data_nodes(data) + self.construct_triple_nodes(data) + return data -# Edges -data['NOE', 'NOE_extract', 'TRIPLE'].edge_index = torch.tensor([[0, 1], [0, 1]]) -data['RES', 'NH1_extract', 'TRIPLE'].edge_index = torch.tensor([[0, 1], [0, 1]]) -data['RES', 'NH2_extract', 'TRIPLE'].edge_index = torch.tensor([[1, 0], [0, 1]]) -data['SHIFT', 'NH1_extract', 'TRIPLE'].edge_index = torch.tensor([[0, 1], [0, 1]]) -data['SHIFT', 'NH2_extract', 'TRIPLE'].edge_index = torch.tensor([[1, 0], [0, 1]]) +class ConstructEdges(): + def __init__(self, data): + self.data = data + self.num_noe = len(data['NOE'].x) + self.num_shift = len(data['SHIFT'].x) + self.num_res = len(data['RES'].x) + + def get_triple_edges(self): + """ + Grabs combinations from the three ranges. + Orders these into source and target indices. + """ + combo = list(product(range(self.num_noe), range(self.num_shift), range( self.num_res))) + combo_tensor = torch.tensor(combo) + # switches tensor dimension (columns set up for each edge) + source_nodes = torch.transpose(combo_tensor, 0, 1) + # repeats target edges for number of occurences in source (3 for the triple in this instance) + target_nodes = torch.tensor(range(len(source_nodes[0]))).repeat(len(source_nodes), 1) + return torch.stack([source_nodes, target_nodes], dim=0) + + def get_edges(self, triple_num, triple_type): + source1, source2, source3 = triple_type + + edges_in = self.get_triple_edges() + self.data['NOE', 'NOE_extract', f'TRIPLE{triple_num}'].edge_index = edges_in[:, 0] + self.data[f'{source1}', 'NH1_extract', f'TRIPLE{triple_num}'].edge_index = edges_in[:, 1] + self.data[f'{source2}', 'NH2_extract', f'TRIPLE{triple_num}'].edge_index = edges_in[:, 2] + + # reverse ordering for nodes going out + edges_out = torch.stack([edges_in[1], edges_in[0]], dim=0) + self.data[f'TRIPLE{triple_num}', 'NOE_add', 'NOE'].edge_index = edges_out[:, 0] + self.data[f'TRIPLE{triple_num}', 'NH1_add', f'{source1}'].edge_index = edges_out[:, 1] + self.data[f'TRIPLE{triple_num}', 'NH2_add', f'{source2}'].edge_index = edges_out[:, 2] + + # only residue nodes will have coordinate edges + if source1 == 'RES': + self.data[f'TRIPLE{triple_num}', 'res1_add', 'RES'].edge_index = edges_out[:, 1] + if source2 == 'RES': + self.data[f'TRIPLE{triple_num}', 'res2_add', 'RES'].edge_index = edges_out[:, 2] + + # self updates are the same across triple types - only need one + if triple_num == 0: + self.data['TRIPLE', 'update', 'TRIPLE'].edge_index = torch.stack([edges_in[1, 0], edges_in[1, 0]], dim=0) # self loop + return self.data + + def generate_edge_indices(self): + self.data = self.get_edges(0, ('RES', 'RES', 'NOE')) + self.data = self.get_edges(1, ('RES', 'SHIFT', 'NOE')) + self.data = self.get_edges(2, ('SHIFT', 'RES', 'NOE')) + self.data = self.get_edges(3, ('SHIFT', 'SHIFT', 'NOE')) + return self.data -data['TRIPLE', 'update', 'TRIPLE'].edge_index = torch.tensor([[0, 1], [0, 1]]) # self loop -data['TRIPLE', 'NOE_add', 'NOE'].edge_index = torch.tensor([[0, 1], [0, 1]]) -data['TRIPLE', 'NH1_add', 'RES'].edge_index = torch.tensor([[0, 1], [0, 1]]) -data['TRIPLE', 'NH2_add', 'RES'].edge_index = torch.tensor([[0, 1], [1, 0]]) -data['TRIPLE', 'res1_add', 'RES'].edge_index = torch.tensor([[0, 1], [0, 1]]) -data['TRIPLE', 'res2_add', 'RES'].edge_index = torch.tensor([[0, 1], [1, 0]]) -data['TRIPLE', 'NH1_add', 'SHIFT'].edge_index = torch.tensor([[0, 1], [0, 1]]) -data['TRIPLE', 'NH2_add', 'SHIFT'].edge_index = torch.tensor([[0, 1], [1, 0]]) class TripleIn(): - def __init__(self, data): - self.data = data + def __init__(self): + # Input index organization - slices to keep proper tensor dimension + self.RES_XYZ = slice(0,3) + self.RES_NH = slice(3,5) - def grab_node(self, node_type, edge_type): + def grab_node(self, data, node_type, edge_type): """ Source node to triple via indexing. + For residue data, extracts shift and coordinate information to process separately. """ # Spatial features - xi = self.data[node_type].x[self.data[edge_type].edge_index[0]] + xi = data[node_type].x[data[edge_type].edge_index[0]] # Non-spatial features - fi = self.data[node_type].f[self.data[edge_type].edge_index[0]] + fi = data[node_type].f[data[edge_type].edge_index[0]] - return xi, fi + # Splits coordinate and shift features if combined + xj = None + if edge_type[0] == 'RES': + xi, xj = xi[:, self.RES_NH], xi[:, self.RES_XYZ] + return xi, xj, fi - def construct_triple(self, edge_type1, edge_type2, edge_type3): + def construct_triple(self, data, edge_type1, edge_type2, edge_type3): """ - Constructs triple based on type: - - residue, residue, NOE - - residue, shift, NOE - - shift, residue, NOE - - shift, shift, NOE + Constructs triple based on type. """ - x1, f1 = self.grab_node(node_type=edge_type1[0], edge_type=edge_type1) - x2, f2 = self.grab_node(node_type=edge_type2[0], edge_type=edge_type2) - x3, f3 = self.grab_node(node_type=edge_type3[0], edge_type=edge_type3) + # Measured shift or residue (shifts will return nonetype value) + x1, x12, f1 = self.grab_node(data, node_type=edge_type1[0], edge_type=edge_type1) + x2, x22, f2 = self.grab_node(data, node_type=edge_type2[0], edge_type=edge_type2) + # NOE (does not require any value split) + x3, _, f3 = self.grab_node(data, node_type=edge_type3[0], edge_type=edge_type3) + return x1, x12, x2, x22, x3, f1, f2, f3 - return x1, x2, x3, f1, f2, f3 -class TripleUpdateResResNoe(MessagePassing): +class CalculationManager(): """ - Message passing class for triple self updates (only set up for residue, residue, NOE triple type). - Residue, residue, NOE triple type is the only one that deviates in calculations (coordinate calculation used), - other three would have same base set up with some differences in indexing. - Pairwise calculations would use the same functions but would have a different message passing scheme. - - Options --> if/else statements for triples because of similarity, different class for pairwise? + Common calculations used across triple and pairwise comparisons. """ def __init__(self): - super().__init__(aggr='add') - # Input index organization - slices to keep proper tensor dimension self.NOE_N1 = slice(0,1) self.NOE_H1 = slice(1,2) self.NOE_H2 = slice(2,3) - self.RES_XYZ = slice(0,3) - self.RES_N = slice(3,4) - self.RES_H = slice(4,5) - self.SHIFT_N = slice(0,1) self.SHIFT_H = slice(1,2) - - self.hidden = 64 - - self.mlp1 = nn.Sequential( - nn.Linear(9, self.hidden), - nn.ReLU(), - nn.Linear(self.hidden, 16)) - - def calc_noe_difference(self, x1, x2, noe, N1, H1, H2): + + def calc_noe_difference(self, x1, x2, noe): """ Calculates shift difference between NOE and residue/measured shifts (direct/indirect only reverse option is available by index). """ - diff_N = (noe[:, self.NOE_N1] - x1[:, N1]) # N [n, 1] - diff_H1 = (noe[:, self.NOE_H1] - x1[:, H1]) # H' [n, 1] - diff_H2 = (noe[:, self.NOE_H2] - x2[:, H2]) # H" [n, 1] + diff_N = (noe[:, self.NOE_N1] - x1[:, self.SHIFT_N]) # N [n, 1] + diff_H1 = (noe[:, self.NOE_H1] - x1[:, self.SHIFT_H]) # H' [n, 1] + diff_H2 = (noe[:, self.NOE_H2] - x2[:, self.SHIFT_H]) # H" [n, 1] return diff_N, diff_H1, diff_H2 - def calc_shift_difference(self, x1, x2, N1, H1): + def calc_shift_difference(self, x1, x2): """ Calculates shift difference between residue/measured shifts. """ - diff_N = (x1[:, N1] - x2[:, N1]) # N [n, 1] - diff_H = (x1[:, H1] - x2[:, H1]) # H [n, 1] + diff_N = (x1[:, self.SHIFT_N] - x2[:, self.SHIFT_N]) # N [n, 1] + diff_H = (x1[:, self.SHIFT_H] - x2[:, self.SHIFT_H]) # H [n, 1] return diff_N, diff_H def calc_res_distance(self, x1, x2): """ Calculates relative distance between residues and this value squared for equivariant calculations. """ - rel_dist = ((x1[:, self.RES_XYZ]) - (x2[:, self.RES_XYZ])) # [n, 3] + rel_dist = (x1 - x2) # [n, 3] dist2 = torch.norm(rel_dist, dim=-1, keepdim=True)**2 # [n, 1] - return rel_dist, dist2 - - def forward(self, x1, x2, x3, f1, f2, f3, edge_index): - out = self.propagate(edge_index, x1=x1, x2=x2, x3=x3, f1=f1, f2=f2, f3=f3) - return out + return rel_dist, dist2 - def message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j): + + +class TripleUpdate(MessagePassing): + """ + Message passing class for triple self updates. + """ + def __init__(self): + super().__init__(aggr='add') + self.calc_manager = CalculationManager() + + self.hidden = 64 + + self.mlp1 = nn.Sequential( + nn.Linear(9, self.hidden), + nn.ReLU(), + nn.Linear(self.hidden, 16)) + + self.mlp2 = nn.Sequential( + nn.Linear(8, self.hidden), + nn.ReLU(), + nn.Linear(self.hidden, 10)) + + def resresnoe_message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j, x22_j): + """ + Constructs message for (res, res, NOE) triple type. + """ # Residue distances - rel_dist, dist2 = self.calc_res_distance(x1_j, x2_j) + rel_dist, dist2 = self.calc_manager.calc_res_distance(x12_j, x22_j) # Differences relative to NOE shifts (N, H', H") - diff1, diff2, diff3 = self.calc_noe_difference(x1_j, x2_j, x3_j, self.RES_N, self.RES_H, self.RES_H) + diff1, diff2, diff3 = self.calc_manager.calc_noe_difference(x1_j, x2_j, x3_j) # Shift differences - diff4, diff5 = self.calc_shift_difference(x1_j, x2_j, self.RES_N, self.RES_H) - + diff4, diff5 = self.calc_manager.calc_shift_difference(x1_j, x2_j) + # Input for MLP # (N, H', H", N, H, dist2, features) mlp_in = (torch.cat((diff1, diff2, diff3, diff4, diff5, dist2, f1_j, f2_j, f3_j), dim=-1)) # [n, 9] @@ -148,33 +245,83 @@ def message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j): # (N, H', H", N1, H1, N2, H2, dist1, dist2, features) mlp_out = self.mlp1(mlp_in) # [n, 16] - # NOE deltas - delta1x = (diff1 * mlp_out[:, 0:1]) # N [n, 1] - delta2x = (diff2 * mlp_out[:, 1:2]) # H' [n, 1] - delta3x = (diff3 * mlp_out[:, 2:3]) # H" [n, 1] - - # SHIFT deltas residue - delta4x = (diff4 * mlp_out[:, 3:4]) # N1 [n, 1] - delta5x = (diff5 * mlp_out[:, 4:5]) # H1 [n, 1] - delta6x = (diff4 * mlp_out[:, 5:6]) # N2 [n, 1] - delta7x = (diff5 * mlp_out[:, 6:7]) # H2 [n, 1] + # NOE deltas [n, 1] + delta1x = (diff1 * mlp_out[:, 0:1]) # N + delta2x = (diff2 * mlp_out[:, 1:2]) # H' + delta3x = (diff3 * mlp_out[:, 2:3]) # H" + + # SHIFT deltas residue [n, 1] + delta4x = (diff4 * mlp_out[:, 3:4]) # N1 + delta5x = (diff5 * mlp_out[:, 4:5]) # H1 + delta6x = (diff4 * mlp_out[:, 5:6]) # N2 + delta7x = (diff5 * mlp_out[:, 6:7]) # H2 + + # DISTANCE deltas [n, 3] + delta12x = (rel_dist * (mlp_out[:, 7:10])) + delta13x = (rel_dist * (mlp_out[:, 10:13])) + + # FEATURE deltas [n, 1] + delta1f = mlp_out[:, 13:14] + delta2f = mlp_out[:, 14:15] + delta3f = mlp_out[:, 15:] + return torch.cat((delta1x, delta2x, delta3x, delta4x, delta5x, delta6x, delta7x, delta12x, delta13x, delta1f, delta2f, delta3f), dim=-1) # [n, 16] + + def shiftshiftnoe_message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j): + """ + Constructs message for (shift, shift, NOE), (shift, res, NOE), (res, shift, NOE) triple types. + """ + # Differences relative to NOE shifts (N, H', H") + diff1, diff2, diff3 = self.calc_manager.calc_noe_difference(x1_j, x2_j, x3_j) - # DISTANCE deltas - delta12x = (rel_dist * (mlp_out[:, 7:10])) # [n, 3] - delta13x = (rel_dist * (mlp_out[:, 10:13])) # [n, 3] + # Shift differences + diff4, diff5 = self.calc_manager.calc_shift_difference(x1_j, x2_j) + + # Input for MLP + # (N, H', H", N, H, dist2, features) + mlp_in = (torch.cat((diff1, diff2, diff3, diff4, diff5, f1_j, f2_j, f3_j), dim=-1)) # [n, 8] - # FEATURE deltas - delta1f = mlp_out[:, 13:14] # [n, 1] - delta2f = mlp_out[:, 14:15] # [n, 1] - delta3f = mlp_out[:, 15:] # [n, 1] + # Output from MLP + # Out should include values for each 'change' wanting to make + # (N, H', H", N1, H1, N2, H2, features) + mlp_out = self.mlp2(mlp_in) # [n, 10] + + # NOE deltas [n, 1] + delta1x = (diff1 * mlp_out[:, 0:1]) # N + delta2x = (diff2 * mlp_out[:, 1:2]) # H' + delta3x = (diff3 * mlp_out[:, 2:3]) # H" + + # SHIFT deltas residue [n, 1] + delta4x = (diff4 * mlp_out[:, 3:4]) # N1 + delta5x = (diff5 * mlp_out[:, 4:5]) # H1 + delta6x = (diff4 * mlp_out[:, 5:6]) # N2 + delta7x = (diff5 * mlp_out[:, 6:7]) # H2 + + # FEATURE deltas [n, 1] + delta1f = mlp_out[:, 7:8] + delta2f = mlp_out[:, 8:9] + delta3f = mlp_out[:, 9:] + return torch.cat((delta1x, delta2x, delta3x, delta4x, delta5x, delta6x, delta7x, delta1f, delta2f, delta3f), dim=-1) # [n, 10] + + def forward(self, x1, x12, x2, x22, x3, f1, f2, f3, edge_index): + out = self.propagate(edge_index, x1=x1, x2=x2, x3=x3, f1=f1, f2=f2, f3=f3, x12=x12, x22=x22) # not sure how to deal with size here + return out - return torch.cat((delta1x, delta2x, delta3x, delta4x, delta5x, delta6x, delta7x, delta12x, delta13x, delta1f, delta2f, delta3f), dim=-1) # [n, 16] + def message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j=None, x22_j=None): + # residue based triple type will have two sets of coordinates (shifts will have nonetype) + if x12_j != None and x22_j != None: + return self.resresnoe_message(x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j, x22_j) + else: + return self.shiftshiftnoe_message(x1_j, x2_j, x3_j, f1_j, f2_j, f3_j) - def update(self, aggr_out): - # res involved in multiple different triples --> need to add across these (in different message passing class) - # output shapes: noe[n, 3], shift1[n, 2], shift2[n, 2], dist1[n, 3], dist2[n, 3], f1[n, 1], f2[n, 1], f3[n, 1] - delta_noe, delta_shift1, delta_shift2, delta_dist1, delta_dist2, delta_f1, delta_f2, delta_f3 = aggr_out[:, 0:3], aggr_out[:, 3:5], aggr_out[:, 5:7], aggr_out[:, 7:10], aggr_out[:, 10:13], aggr_out[:, 13:14], aggr_out[:, 14:15], aggr_out[:, 15:] - return delta_shift1, delta_shift2, delta_noe, delta_dist1, delta_dist2, delta_f1, delta_f2, delta_f3 + def update(self, aggr_out, x12, x22): + if x12 != None and x22 != None: + # output shapes: noe[n, 3], shift1[n, 2], shift2[n, 2], dist1[n, 3], dist2[n, 3], f1[n, 1], f2[n, 1], f3[n, 1] + delta_noe, delta_shift1, delta_shift2, delta_dist1, delta_dist2, delta_f1, delta_f2, delta_f3 = aggr_out[:, 0:3], aggr_out[:, 3:5], aggr_out[:, 5:7], aggr_out[:, 7:10], aggr_out[:, 10:13], aggr_out[:, 13:14], aggr_out[:, 14:15], aggr_out[:, 15:] + return delta_shift1, delta_shift2, delta_noe, delta_dist1, delta_dist2, delta_f1, delta_f2, delta_f3 + else: + # output shapes: noe[n, 3], shift1[n, 2], shift2[n, 2], f1[n, 1], f2[n, 1], f3[n, 1] + delta_noe, delta_shift1, delta_shift2, delta_f1, delta_f2, delta_f3 = aggr_out[:, 0:3], aggr_out[:, 3:5], aggr_out[:, 5:7], aggr_out[:, 7:8], aggr_out[:, 8:9], aggr_out[:, 9:] + return delta_shift1, delta_shift2, delta_noe, delta_f1, delta_f2, delta_f3 @@ -193,83 +340,113 @@ def message(self, x_j): return x_j def update(self, aggr_out, x): - # non-spatial features can also be included in same call but requires more work if they're not being updated (don't want to do residue feature updates 2x - only call once on shifts or coordinates) + # non-spatial features can also be included in same call but requires more work if they're not being updated + # (don't want to do residue feature updates 2x - only call once on shifts or coordinates) # x_val, f_val = aggr_out[:, :len(x[1][1])], aggr_out[:, len(x[1][1]):] # outf = f_val + f[1] # outx = x_val + x[1] - out = aggr_out + x[1] return out class TripleOut(): - def __init__(self, data): - self.data = data + def __init__(self): self.triple_messgage = TripleMessagePass() - def update_data(self, x_source, f_source, edge_type, x_index, update_f=True): + def update_data(self, data, x_source, f_source, edge_type, x_index, update_f=True): """ Calls message passing and directly updates the heterodata object based on target. + Set to update non-spatial features in all cases, but can be turned off case by case to prevent duplicate updates. """ source, edge, target = edge_type - edge_index = self.data[edge_type].edge_index + edge_index = data[edge_type].edge_index # Message passing/update for spatial features - self.data[target].x[:, x_index] = self.triple_messgage(x_source, self.data[target].x[:, x_index], edge_index) + data[target].x[:, x_index] = self.triple_messgage(x_source, data[target].x[:, x_index], edge_index) # Message passing/update for non-spatial features if needed (don't want a duplicate update for residue features) if update_f: - self.data[target].f = self.triple_messgage(f_source, self.data[target].f, edge_index) - - return self.data + data[target].f = self.triple_messgage(f_source, data[target].f, edge_index) + return data class ProteinGNN(): - """ - Test class that: - - construct the triple based on edge types (only res, res, NOE as example) - - self updates - - sends message back to original nodes (note number of reverse edges depends on triple type) - - Can make separate classes to construct each triple type and do following updates or generalize this as a base class? - """ - def __init__(self, data): - self.data = data - self.triple_in = TripleIn(data) - self.triple_self = TripleUpdateResResNoe() # update would need to match data coming in - self.triple_out = TripleOut(data) + def __init__(self): + self.triple_in = TripleIn() + self.triple_self = TripleUpdate() + self.triple_out = TripleOut() - # Original data object indexing self.NOE = slice(0,3) self.RES_XYZ = slice(0,3) self.RES_NH = slice(3,5) self.SHIFT_NH = slice(0,2) - - def forward(self): - # Triple in - x1, x2, x3, f1, f2, f3 = self.triple_in.construct_triple(('RES', 'NH1_extract', 'TRIPLE'), ('RES', 'NH2_extract', 'TRIPLE'), ('NOE', 'NOE_extract', 'TRIPLE')) - # Self update - delta_shift1, delta_shift2, delta_noe, delta_dist1, delta_dist2, deltaf1, deltaf2, deltaf3 = self.triple_self(x1, x2, x3, f1, f2, f3, self.data['TRIPLE', 'update', 'TRIPLE'].edge_index) + def update_noes(self, data, noe_delta, feature, i): + data = self.triple_out.update_data(data, noe_delta, feature, (f'TRIPLE{i}', 'NOE_add', 'NOE'), self.NOE) + return data + + def update_coordinates(self, data, dist_delta1, dist_delta2, feature1, feature2, i): + data = self.triple_out.update_data(data, dist_delta1, feature1, (f'TRIPLE{i}', 'res1_add', 'RES'), self.RES_XYZ, update_f=False) + data = self.triple_out.update_data(data, dist_delta2, feature2, (f'TRIPLE{i}', 'res2_add', 'RES'), self.RES_XYZ, update_f=False) + return data + + def update_shifts(self, data, shift_delta1, shift_delta2, feature1, feature2, target1, target2, i): + # selects target indexing for residue or shift node + target_range1 = self.RES_NH if target1 == 'RES' else self.SHIFT_NH + target_range2 = self.RES_NH if target2 == 'RES' else self.SHIFT_NH + + data = self.triple_out.update_data(data, shift_delta1, feature1, (f'TRIPLE{i}', 'NH1_add', target1), target_range1) + data = self.triple_out.update_data(data, shift_delta2, feature2, (f'TRIPLE{i}', 'NH2_add', target2), target_range2) + return data + + def do_updates(self, data, triple_type, i): + # Triple in + shift_x1, coord_x1, shift_x2, coord_x2, noe_x, shift_f1, shift_f2, noe_f = self.triple_in.construct_triple(data, (triple_type[0], 'NH1_extract', f'TRIPLE{i}'), (triple_type[1], 'NH2_extract', f'TRIPLE{i}'), ('NOE', 'NOE_extract', f'TRIPLE{i}')) - # Triple out (typically 3 reverse edges - 5 if coordinates are present) - self.data = self.triple_out.update_data(delta_shift1, deltaf1, ('TRIPLE', 'NH1_add', 'RES'), self.RES_NH) - self.data = self.triple_out.update_data(delta_shift2, deltaf2, ('TRIPLE', 'NH2_add', 'RES'), self.RES_NH) - self.data = self.triple_out.update_data(delta_noe, deltaf3, ('TRIPLE', 'NOE_add', 'NOE'), self.NOE) - self.data = self.triple_out.update_data(delta_dist1, deltaf1, ('TRIPLE', 'res1_add', 'RES'), self.RES_XYZ, update_f=False) - self.data = self.triple_out.update_data(delta_dist2, deltaf2, ('TRIPLE', 'res2_add', 'RES'), self.RES_XYZ, update_f=False) + # Self update and output + if triple_type == ('RES', 'RES', 'NOE'): + delta_shift1, delta_shift2, delta_noe, delta_dist1, delta_dist2, deltaf1, deltaf2, deltaf3 = self.triple_self(shift_x1, coord_x1, shift_x2, coord_x2, noe_x, shift_f1, shift_f2, noe_f, data['TRIPLE', 'update', 'TRIPLE'].edge_index) + data = self.update_noes(data, delta_noe, deltaf3, i) + data = self.update_coordinates(data, delta_dist1, delta_dist2, deltaf1, deltaf2, i) + data = self.update_shifts(data, delta_shift1, delta_shift2, deltaf1, deltaf2, triple_type[0], triple_type[1], i) + + if triple_type in (('SHIFT', 'RES', 'NOE'), ('RES', 'SHIFT', 'NOE'), ('SHIFT', 'SHIFT', 'NOE')): + delta_shift1, delta_shift2, delta_noe, deltaf1, deltaf2, deltaf3 = self.triple_self(shift_x1, coord_x1, shift_x2, coord_x2, noe_x, shift_f1, shift_f2, noe_f, data['TRIPLE', 'update', 'TRIPLE'].edge_index) + data = self.update_noes(data, delta_noe, deltaf3, i) + data = self.update_shifts(data, delta_shift1, delta_shift2, deltaf1, deltaf2, triple_type[0], triple_type[1], i) + return data + + def forward(self, data): + data = self.do_updates(data, ('RES', 'RES', 'NOE'), 0) + data = self.do_updates(data, ('RES', 'SHIFT', 'NOE'), 1) + data = self.do_updates(data, ('SHIFT', 'RES', 'NOE'), 2) + data = self.do_updates(data, ('SHIFT', 'SHIFT', 'NOE'), 3) + return data - return self.data +if __name__ == "__main__": + edges = ConstructEdges(data) + data = edges.generate_edge_indices() + # testgnn = ProteinGNN() + # data = testgnn.forward(data) + # print(f"updated RES: {data['RES']}") + # print(f"updated NOE: {data['NOE']}") + # print(f"updated SHIFT: {data['SHIFT']}") -if __name__ == "__main__": - print(f"original RES: {data['RES']}") - print(f"original NOE: {data['NOE']}") - testgnn = ProteinGNN(data) - data = testgnn.forward() - print(f"updated RES: {data['RES']}") - print(f"updated NOE: {data['NOE']}") + graph_list = [data, data] + data_loader = DataLoader(graph_list, batch_size=1, shuffle=False) + + unbatched_data_list = [] + testgnn = ProteinGNN() + for batch in data_loader: + test = testgnn.forward(batch) + unbatched_data_list += test.to_data_list() + + # results = [] + # for subgraph in unbatched_data_list: + # results += torch.matmul(subgraph['RES'].x[:, 3:5], subgraph['SHIFT'].x.T) + # print(results) \ No newline at end of file From eb04c953bef7142318e0d65be1faa7acc1b11f44 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Tue, 30 Sep 2025 23:45:04 -0600 Subject: [PATCH 37/43] Organization into composable units with value and policy --- pyg_triple_update.py | 323 ++++++++++++++++++++++++++++++++----------- 1 file changed, 246 insertions(+), 77 deletions(-) diff --git a/pyg_triple_update.py b/pyg_triple_update.py index 4bc1d0f..8315044 100644 --- a/pyg_triple_update.py +++ b/pyg_triple_update.py @@ -5,75 +5,88 @@ from torch_geometric.loader import DataLoader from torch_geometric.nn import MessagePassing from torch_geometric.data import HeteroData -import torch_geometric.transforms as T import torch.nn as nn from itertools import product +import pickle # Set up data data = HeteroData() # Spatial and non-spatial -data['NOE'].x = torch.tensor([[1, 2, 3]], dtype=torch.float32) -data['NOE'].f = torch.zeros((2, 1), dtype=torch.float32) +data['NOE'].x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8 , 9]], dtype=torch.float32) +data['NOE'].f = torch.zeros((3, 1), dtype=torch.float32) data['RES'].x = torch.tensor([[7, 8, 9, 1, 2], [10, 11, 12, 7, 3]], dtype=torch.float32) data['RES'].f = torch.zeros((2, 1), dtype=torch.float32) data['SHIFT'].x = torch.tensor([[4, 5], [1, 2]], dtype=torch.float32) data['SHIFT'].f = torch.zeros((2, 1), dtype=torch.float32) # Triples (nothing is ever stored in these - they are placeholders to be used in edge types) -data['TRIPLE'].x = torch.zeros((1, 1), dtype=torch.float32) # easier just to use separate triple for self updates data['TRIPLE0'].x = torch.zeros((1, 1), dtype=torch.float32) data['TRIPLE1'].x = torch.zeros((1, 1), dtype=torch.float32) data['TRIPLE2'].x = torch.zeros((1, 1), dtype=torch.float32) data['TRIPLE3'].x = torch.zeros((1, 1), dtype=torch.float32) -# # Edges -# data['NOE', 'NOE_extract', 'TRIPLE'].edge_index = torch.tensor([[0, 0], [0, 0]]) -# data['RES', 'NH1_extract', 'TRIPLE'].edge_index = torch.tensor([[0, 1], [0, 0]]) -# data['RES', 'NH2_extract', 'TRIPLE'].edge_index = torch.tensor([[1, 0], [0, 0]]) #[[1, 0], [0, 1] -# data['SHIFT', 'NH1_extract', 'TRIPLE'].edge_index = torch.tensor([[0, 1], [0, 0]]) -# data['SHIFT', 'NH2_extract', 'TRIPLE'].edge_index = torch.tensor([[1, 0], [0, 0]]) +data['VALUE_NOE'].x = torch.zeros(1, 1, dtype=torch.float32) # has to be num graphs batched? +data['VALUE_SHIFT'].x = torch.zeros(1, 1, dtype=torch.float32) +data['VALUE_RES'].x = torch.zeros(1, 1, dtype=torch.float32) -# data['TRIPLE', 'update', 'TRIPLE'].edge_index = torch.tensor([[0], [0]]) # self loop -# data['TRIPLE', 'NOE_add', 'NOE'].edge_index = torch.tensor([[0, 0], [0, 0]]) -# data['TRIPLE', 'NH1_add', 'RES'].edge_index = torch.tensor([[0, 0], [0, 1]]) -# data['TRIPLE', 'NH2_add', 'RES'].edge_index = torch.tensor([[0, 0], [1, 0]]) -# data['TRIPLE', 'res1_add', 'RES'].edge_index = torch.tensor([[0, 0], [0, 1]]) -# data['TRIPLE', 'res2_add', 'RES'].edge_index = torch.tensor([[0, 0], [1, 0]]) -# data['TRIPLE', 'NH1_add', 'SHIFT'].edge_index = torch.tensor([[0, 0], [0, 1]]) -# data['TRIPLE', 'NH2_add', 'SHIFT'].edge_index = torch.tensor([[0, 0], [1, 0]]) class ConstructNodes(): - def __init__(self, data_history): - self.data_history = data_history + # def __init__(self, histories): + # self.histories = histories + + def construct_data_nodes(self, data, histories): + data['NOE'].x = torch.tensor(histories['noes'], dtype=torch.float32) + data['NOE'].f = torch.zeros((len(histories['noes']), 1), dtype=torch.float32) + data['SHIFT'].x = torch.tensor(histories['obs_chemical_shifts'], dtype=torch.float32) + data['SHIFT'].f = torch.zeros((len(histories['obs_chemical_shifts']), 1), dtype=torch.float32) + data['RES'].x = torch.tensor(histories['coordinates'], dtype=torch.float32) + data['RES'].f = torch.zeros((len(histories['coordinates']), 1), dtype=torch.float32) + return data - def construct_data_nodes(self, data): - # will need to extract each component of data from data_history - # data['NOE'].x - # data['NOE'].f - # data['SHIFT'].x - # data['SHIFT'].f - # data['RES'].x - # data['RES'].f + def construct_node_features(self, data, histories): + for i in range(len(data['SHIFT'].f)): + # is the shift being assigned right now? + if i == histories['shift_to_assign']: + data['SHIFT'].f[i, 0] = 1 + # has the shift already been assigned? + if i in histories['assignments'].keys(): + data['SHIFT'].f[i, 1] = 1 + + for i in range(len(data['RES'].f)): + # has the residue been assigned? + if i in histories['assignments'].values(): + data['RES'].f[i] = 1 return data def construct_triple_nodes(self, data): - # Triples (nothing is ever stored in these - they are placeholders to be used in edge types) - data['TRIPLE'].x = torch.zeros((1, 1), dtype=torch.float32) # easier just to use separate triple for self updates + """ + Generates triple node types. Nothing is ever stored in these - they are simply placeholders to construct/use in edge types. + """ data['TRIPLE0'].x = torch.zeros((1, 1), dtype=torch.float32) data['TRIPLE1'].x = torch.zeros((1, 1), dtype=torch.float32) data['TRIPLE2'].x = torch.zeros((1, 1), dtype=torch.float32) data['TRIPLE3'].x = torch.zeros((1, 1), dtype=torch.float32) return data - def construct_data(self): + def construct_value_nodes(self, data): + data['VALUE_NOE'].x = torch.zeros(1, 1, dtype=torch.float32) # has to be num graphs batched? + data['VALUE_SHIFT'].x = torch.zeros(1, 1, dtype=torch.float32) + data['VALUE_RES'].x = torch.zeros(1, 1, dtype=torch.float32) + return data + + def construct_data(self, histories): data = HeteroData() - self.construct_data_nodes(data) - self.construct_triple_nodes(data) + data = self.construct_data_nodes(data, histories) + data = self.construct_triple_nodes(data) + data = self.construct_value_nodes(data) + # data = self.construct_node_features(data, histories) # need to sort out all features first across types (consistent or not?) return data + + class ConstructEdges(): def __init__(self, data): self.data = data @@ -81,42 +94,81 @@ def __init__(self, data): self.num_shift = len(data['SHIFT'].x) self.num_res = len(data['RES'].x) - def get_triple_edges(self): + def get_triple_edges(self, source1, source2): """ - Grabs combinations from the three ranges. - Orders these into source and target indices. + Grabs all index combinations from the three ranges (residue, shift, noe) and orders these into source and target indices for the edges. + Combinations occur along columns into a triple node. """ - combo = list(product(range(self.num_noe), range(self.num_shift), range( self.num_res))) + combo = list(product(range(self.num_noe), range(source1), range(source2))) combo_tensor = torch.tensor(combo) - # switches tensor dimension (columns set up for each edge) + + # Switches tensor dimension for source (columns set up for each edge [0,0,0] --> [0],[0],[0]) source_nodes = torch.transpose(combo_tensor, 0, 1) - # repeats target edges for number of occurences in source (3 for the triple in this instance) + # Repeats target edges for number of occurences in source (3 for the triple in this instance, to be assigned to each incoming node type) target_nodes = torch.tensor(range(len(source_nodes[0]))).repeat(len(source_nodes), 1) + + # print(torch.stack([source_nodes, target_nodes], dim=0)) return torch.stack([source_nodes, target_nodes], dim=0) + def get_pairwise_edges(self): + """ + Grabs index combinations between shifts and residues and orders these into source and target indices for the edges. + """ + shift = torch.arange(0, self.num_shift) + resid = torch.arange(0, self.num_res) + + shift_repeats = shift.repeat_interleave(self.num_res) + resid_repeats = resid.repeat(self.num_shift) + return torch.stack((resid_repeats, shift_repeats), dim=0) + + def get_batch_edges(self): + shift = torch.arange(0, self.num_shift) + noe = torch.arange(0, self.num_noe) + resid = torch.arange(0, self.num_res) + + shift_repeats = torch.zeros(self.num_shift).long() + noe_repeats = torch.zeros(self.num_noe).long() + resid_repeats = torch.zeros(self.num_res).long() + + self.data['SHIFT', 'SHIFT_extract', f'VALUE_SHIFT'].edge_index = torch.stack((shift, shift_repeats), dim=0) + self.data['NOE', 'NOE_extract', f'VALUE_NOE'].edge_index = torch.stack((noe, noe_repeats), dim=0) + self.data['RES', 'RES_extract', f'VALUE_RES'].edge_index = torch.stack((resid, resid_repeats), dim=0) + + return self.data + def get_edges(self, triple_num, triple_type): + """ + Constructs all edges for the graph. + """ source1, source2, source3 = triple_type - edges_in = self.get_triple_edges() - self.data['NOE', 'NOE_extract', f'TRIPLE{triple_num}'].edge_index = edges_in[:, 0] + # Need to make sure I'm using the right node range (shift and residue number could vary) + num_node1 = self.num_res if source1 == 'RES' else self.num_shift + num_node2 = self.num_res if source2 == 'RES' else self.num_shift + + # Triple in + edges_in = self.get_triple_edges(num_node1, num_node2) + self.data[f'{source3}', 'NOE_extract', f'TRIPLE{triple_num}'].edge_index = edges_in[:, 0] self.data[f'{source1}', 'NH1_extract', f'TRIPLE{triple_num}'].edge_index = edges_in[:, 1] self.data[f'{source2}', 'NH2_extract', f'TRIPLE{triple_num}'].edge_index = edges_in[:, 2] - # reverse ordering for nodes going out + # Triple out (reverse in/out ordering) edges_out = torch.stack([edges_in[1], edges_in[0]], dim=0) - self.data[f'TRIPLE{triple_num}', 'NOE_add', 'NOE'].edge_index = edges_out[:, 0] + self.data[f'TRIPLE{triple_num}', 'NOE_add', f'{source3}'].edge_index = edges_out[:, 0] self.data[f'TRIPLE{triple_num}', 'NH1_add', f'{source1}'].edge_index = edges_out[:, 1] self.data[f'TRIPLE{triple_num}', 'NH2_add', f'{source2}'].edge_index = edges_out[:, 2] - # only residue nodes will have coordinate edges + # Only residue nodes will have coordinate edges if source1 == 'RES': self.data[f'TRIPLE{triple_num}', 'res1_add', 'RES'].edge_index = edges_out[:, 1] if source2 == 'RES': self.data[f'TRIPLE{triple_num}', 'res2_add', 'RES'].edge_index = edges_out[:, 2] - # self updates are the same across triple types - only need one + # Edges for self loop + self.data[f'TRIPLE{triple_num}', 'update', f'TRIPLE{triple_num}'].edge_index = torch.stack([edges_in[1, 0], edges_in[1, 0]], dim=0) + if triple_num == 0: - self.data['TRIPLE', 'update', 'TRIPLE'].edge_index = torch.stack([edges_in[1, 0], edges_in[1, 0]], dim=0) # self loop + self.data['RES', 'pair', 'SHIFT'].edge_index = self.get_pairwise_edges() # pairwise edges - no message passing return self.data def generate_edge_indices(self): @@ -124,6 +176,7 @@ def generate_edge_indices(self): self.data = self.get_edges(1, ('RES', 'SHIFT', 'NOE')) self.data = self.get_edges(2, ('SHIFT', 'RES', 'NOE')) self.data = self.get_edges(3, ('SHIFT', 'SHIFT', 'NOE')) + self.data = self.get_batch_edges() return self.data @@ -151,7 +204,7 @@ def grab_node(self, data, node_type, edge_type): return xi, xj, fi def construct_triple(self, data, edge_type1, edge_type2, edge_type3): - """ + """ Constructs triple based on type. """ # Measured shift or residue (shifts will return nonetype value) @@ -216,12 +269,14 @@ def __init__(self): self.mlp1 = nn.Sequential( nn.Linear(9, self.hidden), nn.ReLU(), - nn.Linear(self.hidden, 16)) + nn.Linear(self.hidden, 16), + nn.LayerNorm(16)) self.mlp2 = nn.Sequential( nn.Linear(8, self.hidden), nn.ReLU(), - nn.Linear(self.hidden, 10)) + nn.Linear(self.hidden, 10), + nn.LayerNorm(10)) def resresnoe_message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j, x22_j): """ @@ -307,13 +362,14 @@ def forward(self, x1, x12, x2, x22, x3, f1, f2, f3, edge_index): return out def message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j=None, x22_j=None): - # residue based triple type will have two sets of coordinates (shifts will have nonetype) + # residue based triple type will have two sets of coordinates (x12_j and x22_j) where shifts will have nonetype if x12_j != None and x22_j != None: return self.resresnoe_message(x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j, x22_j) else: return self.shiftshiftnoe_message(x1_j, x2_j, x3_j, f1_j, f2_j, f3_j) def update(self, aggr_out, x12, x22): + # residue based triple type will have two sets of coordinates (x12_j and x22_j) where shifts will have nonetype if x12 != None and x22 != None: # output shapes: noe[n, 3], shift1[n, 2], shift2[n, 2], dist1[n, 3], dist2[n, 3], f1[n, 1], f2[n, 1], f3[n, 1] delta_noe, delta_shift1, delta_shift2, delta_dist1, delta_dist2, delta_f1, delta_f2, delta_f3 = aggr_out[:, 0:3], aggr_out[:, 3:5], aggr_out[:, 5:7], aggr_out[:, 7:10], aggr_out[:, 10:13], aggr_out[:, 13:14], aggr_out[:, 14:15], aggr_out[:, 15:] @@ -329,8 +385,8 @@ class TripleMessagePass(MessagePassing): """ Standard message passing class for outgoing triple messages. """ - def __init__(self): - super().__init__(aggr='add') + def __init__(self, aggr): + super().__init__(aggr=aggr) def forward(self, x_source, x_target, edge_index): # SIZE (n, m) (source, target) @@ -348,11 +404,38 @@ def update(self, aggr_out, x): out = aggr_out + x[1] return out + + +class BatchMessagePass(MessagePassing): + """ + + """ + def __init__(self, aggr): + super().__init__(aggr=aggr) + self.noe_reduce = (nn.Linear(3, 1)) + self.shift_reduce = (nn.Linear(2, 1)) + self.res_reduce = (nn.Linear(5, 1)) + + def forward(self, x_source, x_target, edge_index): + # SIZE (n, m) (source, target) + return self.propagate(edge_index=edge_index, x=(x_source, x_target), size=(x_source.size(0), x_target.size(0))) + + def message(self, x_j): + if len(x_j[0]) == 3: + return self.noe_reduce(x_j) + if len(x_j[0]) == 2: + return self.shift_reduce(x_j) + if len(x_j[0]) == 5: + return self.res_reduce(x_j) + + def update(self, aggr_out, x): + return aggr_out + class TripleOut(): def __init__(self): - self.triple_messgage = TripleMessagePass() + self.triple_message = TripleMessagePass(aggr='add') def update_data(self, data, x_source, f_source, edge_type, x_index, update_f=True): """ @@ -363,20 +446,22 @@ def update_data(self, data, x_source, f_source, edge_type, x_index, update_f=Tru edge_index = data[edge_type].edge_index # Message passing/update for spatial features - data[target].x[:, x_index] = self.triple_messgage(x_source, data[target].x[:, x_index], edge_index) + data[target].x[:, x_index] = self.triple_message(x_source, data[target].x[:, x_index], edge_index) # Message passing/update for non-spatial features if needed (don't want a duplicate update for residue features) if update_f: - data[target].f = self.triple_messgage(f_source, data[target].f, edge_index) + data[target].f = self.triple_message(f_source, data[target].f, edge_index) return data -class ProteinGNN(): +class NMRLayer(nn.Module): def __init__(self): + super(NMRLayer, self).__init__() self.triple_in = TripleIn() self.triple_self = TripleUpdate() self.triple_out = TripleOut() + self.calc_manager = CalculationManager() self.NOE = slice(0,3) self.RES_XYZ = slice(0,3) @@ -384,6 +469,8 @@ def __init__(self): self.SHIFT_NH = slice(0,2) def update_noes(self, data, noe_delta, feature, i): + # 'i' is for the triple type in all of these cases (TRIPLE0 --> RES RES NOE, etc.) + # Should simplify this to be one or the other from the very start (string --> number identification) data = self.triple_out.update_data(data, noe_delta, feature, (f'TRIPLE{i}', 'NOE_add', 'NOE'), self.NOE) return data @@ -393,7 +480,7 @@ def update_coordinates(self, data, dist_delta1, dist_delta2, feature1, feature2, return data def update_shifts(self, data, shift_delta1, shift_delta2, feature1, feature2, target1, target2, i): - # selects target indexing for residue or shift node + # Selects target indexing for residue or shift node target_range1 = self.RES_NH if target1 == 'RES' else self.SHIFT_NH target_range2 = self.RES_NH if target2 == 'RES' else self.SHIFT_NH @@ -407,13 +494,13 @@ def do_updates(self, data, triple_type, i): # Self update and output if triple_type == ('RES', 'RES', 'NOE'): - delta_shift1, delta_shift2, delta_noe, delta_dist1, delta_dist2, deltaf1, deltaf2, deltaf3 = self.triple_self(shift_x1, coord_x1, shift_x2, coord_x2, noe_x, shift_f1, shift_f2, noe_f, data['TRIPLE', 'update', 'TRIPLE'].edge_index) + delta_shift1, delta_shift2, delta_noe, delta_dist1, delta_dist2, deltaf1, deltaf2, deltaf3 = self.triple_self(shift_x1, coord_x1, shift_x2, coord_x2, noe_x, shift_f1, shift_f2, noe_f, data[f'TRIPLE{i}', 'update', f'TRIPLE{i}'].edge_index) data = self.update_noes(data, delta_noe, deltaf3, i) data = self.update_coordinates(data, delta_dist1, delta_dist2, deltaf1, deltaf2, i) data = self.update_shifts(data, delta_shift1, delta_shift2, deltaf1, deltaf2, triple_type[0], triple_type[1], i) if triple_type in (('SHIFT', 'RES', 'NOE'), ('RES', 'SHIFT', 'NOE'), ('SHIFT', 'SHIFT', 'NOE')): - delta_shift1, delta_shift2, delta_noe, deltaf1, deltaf2, deltaf3 = self.triple_self(shift_x1, coord_x1, shift_x2, coord_x2, noe_x, shift_f1, shift_f2, noe_f, data['TRIPLE', 'update', 'TRIPLE'].edge_index) + delta_shift1, delta_shift2, delta_noe, deltaf1, deltaf2, deltaf3 = self.triple_self(shift_x1, coord_x1, shift_x2, coord_x2, noe_x, shift_f1, shift_f2, noe_f, data[f'TRIPLE{i}', 'update', f'TRIPLE{i}'].edge_index) data = self.update_noes(data, delta_noe, deltaf3, i) data = self.update_shifts(data, delta_shift1, delta_shift2, deltaf1, deltaf2, triple_type[0], triple_type[1], i) return data @@ -426,27 +513,109 @@ def forward(self, data): return data -if __name__ == "__main__": + +class ValueCalc(): + def __init__(self): + self.batch_message = BatchMessagePass(aggr='mean') + self.hidden = 64 + + self.testmlp = nn.Sequential( + nn.Linear(3, self.hidden), + nn.ReLU(), + nn.Linear(self.hidden, 1) + ) + + def get_aggr(self, x_source, x_target, edge_index): + aggr = self.batch_message(x_source, x_target, edge_index) + return aggr + + def calc_value(self, data): + noe = self.get_aggr(data['NOE'].x, data['VALUE_NOE'].x, data['NOE', 'NOE_extract', 'VALUE_NOE'].edge_index) + shift = self.get_aggr(data['SHIFT'].x, data['VALUE_SHIFT'].x, data['SHIFT', 'SHIFT_extract', 'VALUE_SHIFT'].edge_index) + resid = self.get_aggr(data['RES'].x, data['VALUE_RES'].x, data['RES', 'RES_extract', 'VALUE_RES'].edge_index) + + concat_aggr = torch.cat((noe, shift, resid), dim=-1) + values = self.testmlp(concat_aggr) + return values + + + +class PolicyCalc(): + def __init__(self): + self.RES_NH = slice(3,5) + + def pairwise_distance(self, data, edge_type): + node_type1, edge, node_type2 = edge_type + + # residue + resid_nodes = data[node_type1].x[data[edge_type].edge_index[0]][:, self.RES_NH] + # shift + shift_nodes = data[node_type2].x[data[edge_type].edge_index[1]] + + pair_dist = (shift_nodes - resid_nodes) + return -pair_dist**2 + + def calc_policy(self, data): + pair_dist2 = self.pairwise_distance(data, ('RES', 'pair', 'SHIFT')) + return pair_dist2 + + + +class TestLayer(): + def __init__(self): + self.value = ValueCalc() + self.policy = PolicyCalc() + + self.nmr = nn.Sequential( + NMRLayer(), + NMRLayer(), + ) + + def forward(self, data): + test_data = self.nmr(data) + value = self.value.calc_value(test_data) + policy = self.policy.calc_policy(test_data) + return value, policy + + + +def extract_data(pickle_file="fake_histories_r10_1.pkl"): + with open(pickle_file, "rb") as f: + histories = pickle.load(f) + return histories + +def construct_graph(history): + nodes = ConstructNodes() + data = nodes.construct_data(history) edges = ConstructEdges(data) data = edges.generate_edge_indices() - # testgnn = ProteinGNN() + return data - # data = testgnn.forward(data) - # print(f"updated RES: {data['RES']}") - # print(f"updated NOE: {data['NOE']}") - # print(f"updated SHIFT: {data['SHIFT']}") +def preprocess_data(histories, batch_size): + graphs = [] + for history in histories: + graphs.append(construct_graph(history)) + + data_loader = DataLoader(graphs, batch_size=batch_size, shuffle=False) + return data_loader - graph_list = [data, data] - data_loader = DataLoader(graph_list, batch_size=1, shuffle=False) - - unbatched_data_list = [] - testgnn = ProteinGNN() - for batch in data_loader: - test = testgnn.forward(batch) - unbatched_data_list += test.to_data_list() - # results = [] - # for subgraph in unbatched_data_list: - # results += torch.matmul(subgraph['RES'].x[:, 3:5], subgraph['SHIFT'].x.T) - # print(results) \ No newline at end of file +if __name__ == "__main__": + # nodes = ConstructNodes() + # data = nodes.construct_data(extract_data()) + edges = ConstructEdges(data) + data = edges.generate_edge_indices() + # data = preprocess_data(extract_data(), batch_size=2) + + value = ValueCalc() + + # create fake graph list (not from histories) to check + graph_list = [data, data, data, data] + data_loader = DataLoader(graph_list, batch_size=2, shuffle=False) + + testgnn = TestLayer() + for batch in data_loader: + value, policy = testgnn.forward(batch) + print(value.shape) # 2 values (1 per graph in batch) + print(policy.shape) # 8 policies (4 per graph in batch --> 4 comparisons being made 2shift/2resid) **THESE ARE NOT SPLIT BY GRAPH** From 8821099cd4503235569edcd6e0d682cb903c0469 Mon Sep 17 00:00:00 2001 From: abbie-p Date: Sun, 9 Nov 2025 21:38:17 -0700 Subject: [PATCH 38/43] Pkl example with single step --- .gitignore | 3 +++ fake_histories_r4_0.pkl | Bin 0 -> 6702 bytes 2 files changed, 3 insertions(+) create mode 100644 fake_histories_r4_0.pkl diff --git a/.gitignore b/.gitignore index 2f45271..beed6a1 100644 --- a/.gitignore +++ b/.gitignore @@ -5,7 +5,10 @@ __pycache__/ # Pickled data and visualization *.pkl +!fake_histories_r4_0.pkl *.png +runs/ +*.pdb # C extensions *.so diff --git a/fake_histories_r4_0.pkl b/fake_histories_r4_0.pkl new file mode 100644 index 0000000000000000000000000000000000000000..f1ccb4dffb372b746620f1152b1fade6e2a02cb8 GIT binary patch literal 6702 zcmZo*nW`+s00yyBG-{`4^l&HV=NF}9<|US-7Eb|+^>C&oW~audB$gyj>0u8j$}dUH z%$qV9C?T3xnp;q*mz-aes+U`uQ<9ljRFqgbrH8FJIWZ@(2&krq6|4kEr<7C{q=H14 z(k!NQcC=3knxf&&+{0)y#m~>r>pu{H32%myDM`*86XLyPV%$&KPiX_H%izr5cjoXC zmz;g!%u8gUb9IO7ciwq|EabbbhjE98oCA`^Wfl1a@=`Gnp^_<$Q!>~zAVwb9C}w}k z@)ELsht=0?BTQZ+3wa-0t7sB;8(D~Xua5ffJIctm)Ol!WoQ;dZVoOIa?}16Y_mTBq zIU-ZT@aP`0kW1ocgV6ko$U=-$ukARWgknqYp;vcaDkfmDrAU5n%&%}1zw}(BjWgeKY(#OmLs^i(*&t}}w1FZa5~6mKqvn*zGqT9m z=a$!>y(|HZFl4nWR6V)aqNR}4ew?8-v%C!IRAjZToPs;IR-m}yK>WF@JC&h<2U45b z!;+T|N#sC(u>1MDqIkn^5l4#gQ52^!U*g=}FAjAY!~^Dh6@5E5LoI3t`K%Tqaep7* zlSS)LTo4(2^SWm+G!!7J4E~$PCc!L0)?QRL{q`!Dxe&F7g#HIVhXo=;$SCJdWl1d- z+e@zhV|mhuV*AmudX{A{|3kE2Q)are7-k8w?FYlket9*aXuq0Lb5mRe8hjA#p2uXD zH-$qZ1X+9E`>LHQqEKS0@X5SK;joy3XsnCP0a)5nA#~lJP>|7u#8PD0{f6NgQtfn z6RA|p5WrR@LJDaH23XPd!kL2^DuyJ)h$+N`ER@>A56->uCHe6XmrTic)5BAeUjoe6 zsd=eI>6KI5VOI696{VIZ7NsCWN`W;*1`{Y=V1?*tQ4Oq|22oL+5eAAKXz4m+O5%(v zkgK63=V-x8_kuU$04V$q52>Ox<5BIDj3*#>b7lYwLs+4OR2Y&|K4yH&ItDC6Grq%1 z!Hgep=4kQf%rUrIs-wl Date: Sun, 9 Nov 2025 21:42:18 -0700 Subject: [PATCH 39/43] Training loop and visualization --- pyg_triple_update.py | 285 +++++++++++++++++++++++-------------------- training_loop_gnn.py | 119 ++++++++++++++++++ 2 files changed, 273 insertions(+), 131 deletions(-) create mode 100644 training_loop_gnn.py diff --git a/pyg_triple_update.py b/pyg_triple_update.py index 8315044..6f339e8 100644 --- a/pyg_triple_update.py +++ b/pyg_triple_update.py @@ -6,44 +6,51 @@ from torch_geometric.nn import MessagePassing from torch_geometric.data import HeteroData import torch.nn as nn +import numpy as np + +import matplotlib.pyplot as plt +import matplotlib.animation from itertools import product import pickle -# Set up data -data = HeteroData() +# # Set up data +# data = HeteroData() -# Spatial and non-spatial -data['NOE'].x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8 , 9]], dtype=torch.float32) -data['NOE'].f = torch.zeros((3, 1), dtype=torch.float32) -data['RES'].x = torch.tensor([[7, 8, 9, 1, 2], [10, 11, 12, 7, 3]], dtype=torch.float32) -data['RES'].f = torch.zeros((2, 1), dtype=torch.float32) -data['SHIFT'].x = torch.tensor([[4, 5], [1, 2]], dtype=torch.float32) -data['SHIFT'].f = torch.zeros((2, 1), dtype=torch.float32) +# # Spatial and non-spatial +# data['NOE'].x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8 , 9]], dtype=torch.float32) +# data['NOE'].f = torch.zeros((3, 2), dtype=torch.float32) +# data['RES'].x = torch.tensor([[7, 8, 9, 1, 2], [10, 11, 12, 7, 3]], dtype=torch.float32) +# data['RES'].f = torch.zeros((2, 2), dtype=torch.float32) +# data['SHIFT'].x = torch.tensor([[4, 5], [1, 2]], dtype=torch.float32) +# data['SHIFT'].f = torch.zeros((2, 2), dtype=torch.float32) -# Triples (nothing is ever stored in these - they are placeholders to be used in edge types) -data['TRIPLE0'].x = torch.zeros((1, 1), dtype=torch.float32) -data['TRIPLE1'].x = torch.zeros((1, 1), dtype=torch.float32) -data['TRIPLE2'].x = torch.zeros((1, 1), dtype=torch.float32) -data['TRIPLE3'].x = torch.zeros((1, 1), dtype=torch.float32) +# # Triples (nothing is ever stored in these - they are placeholders to be used in edge types) +# data['TRIPLE0'].x = torch.zeros((1, 1), dtype=torch.float32) +# data['TRIPLE1'].x = torch.zeros((1, 1), dtype=torch.float32) +# data['TRIPLE2'].x = torch.zeros((1, 1), dtype=torch.float32) +# data['TRIPLE3'].x = torch.zeros((1, 1), dtype=torch.float32) -data['VALUE_NOE'].x = torch.zeros(1, 1, dtype=torch.float32) # has to be num graphs batched? -data['VALUE_SHIFT'].x = torch.zeros(1, 1, dtype=torch.float32) -data['VALUE_RES'].x = torch.zeros(1, 1, dtype=torch.float32) +# data['VALUE_NOE'].x = torch.zeros(1, 1, dtype=torch.float32) # has to be num graphs batched? +# data['VALUE_SHIFT'].x = torch.zeros(1, 1, dtype=torch.float32) +# data['VALUE_RES'].x = torch.zeros(1, 1, dtype=torch.float32) class ConstructNodes(): - # def __init__(self, histories): - # self.histories = histories + """ + Builds graph nodes from input + """ + def __init__(self, device): + self.device = device def construct_data_nodes(self, data, histories): - data['NOE'].x = torch.tensor(histories['noes'], dtype=torch.float32) - data['NOE'].f = torch.zeros((len(histories['noes']), 1), dtype=torch.float32) - data['SHIFT'].x = torch.tensor(histories['obs_chemical_shifts'], dtype=torch.float32) - data['SHIFT'].f = torch.zeros((len(histories['obs_chemical_shifts']), 1), dtype=torch.float32) - data['RES'].x = torch.tensor(histories['coordinates'], dtype=torch.float32) - data['RES'].f = torch.zeros((len(histories['coordinates']), 1), dtype=torch.float32) + data['NOE'].x = torch.tensor(histories['noes'], dtype=torch.float32, device=self.device) + data['NOE'].f = torch.zeros((len(histories['noes']), 2), dtype=torch.float32, device=self.device) + data['SHIFT'].x = torch.tensor(histories['obs_chemical_shifts'], dtype=torch.float32, device=self.device) + data['SHIFT'].f = torch.zeros((len(histories['obs_chemical_shifts']), 2), dtype=torch.float32, device=self.device) + data['RES'].x = torch.tensor(histories['coordinates'], dtype=torch.float32, device=self.device) + data['RES'].f = torch.zeros((len(histories['coordinates']), 2), dtype=torch.float32, device=self.device) return data def construct_node_features(self, data, histories): @@ -65,16 +72,16 @@ def construct_triple_nodes(self, data): """ Generates triple node types. Nothing is ever stored in these - they are simply placeholders to construct/use in edge types. """ - data['TRIPLE0'].x = torch.zeros((1, 1), dtype=torch.float32) - data['TRIPLE1'].x = torch.zeros((1, 1), dtype=torch.float32) - data['TRIPLE2'].x = torch.zeros((1, 1), dtype=torch.float32) - data['TRIPLE3'].x = torch.zeros((1, 1), dtype=torch.float32) + data['TRIPLE0'].x = torch.zeros((1, 1), dtype=torch.float32, device=self.device) + data['TRIPLE1'].x = torch.zeros((1, 1), dtype=torch.float32, device=self.device) + data['TRIPLE2'].x = torch.zeros((1, 1), dtype=torch.float32, device=self.device) + data['TRIPLE3'].x = torch.zeros((1, 1), dtype=torch.float32, device=self.device) return data def construct_value_nodes(self, data): - data['VALUE_NOE'].x = torch.zeros(1, 1, dtype=torch.float32) # has to be num graphs batched? - data['VALUE_SHIFT'].x = torch.zeros(1, 1, dtype=torch.float32) - data['VALUE_RES'].x = torch.zeros(1, 1, dtype=torch.float32) + data['VALUE_NOE'].x = torch.zeros(1, 1, dtype=torch.float32, device=self.device) # has to be num graphs batched? + data['VALUE_SHIFT'].x = torch.zeros(1, 1, dtype=torch.float32, device=self.device) + data['VALUE_RES'].x = torch.zeros(1, 1, dtype=torch.float32, device=self.device) return data def construct_data(self, histories): @@ -82,17 +89,18 @@ def construct_data(self, histories): data = self.construct_data_nodes(data, histories) data = self.construct_triple_nodes(data) data = self.construct_value_nodes(data) - # data = self.construct_node_features(data, histories) # need to sort out all features first across types (consistent or not?) + data = self.construct_node_features(data, histories) # need to sort out all features first across types (consistent or not?) return data class ConstructEdges(): - def __init__(self, data): + def __init__(self, data, device): self.data = data self.num_noe = len(data['NOE'].x) self.num_shift = len(data['SHIFT'].x) self.num_res = len(data['RES'].x) + self.device = device def get_triple_edges(self, source1, source2): """ @@ -100,12 +108,12 @@ def get_triple_edges(self, source1, source2): Combinations occur along columns into a triple node. """ combo = list(product(range(self.num_noe), range(source1), range(source2))) - combo_tensor = torch.tensor(combo) + combo_tensor = torch.tensor(combo, device=self.device) # Switches tensor dimension for source (columns set up for each edge [0,0,0] --> [0],[0],[0]) source_nodes = torch.transpose(combo_tensor, 0, 1) # Repeats target edges for number of occurences in source (3 for the triple in this instance, to be assigned to each incoming node type) - target_nodes = torch.tensor(range(len(source_nodes[0]))).repeat(len(source_nodes), 1) + target_nodes = torch.tensor(range(len(source_nodes[0])), device=self.device).repeat(len(source_nodes), 1) # print(torch.stack([source_nodes, target_nodes], dim=0)) return torch.stack([source_nodes, target_nodes], dim=0) @@ -114,21 +122,23 @@ def get_pairwise_edges(self): """ Grabs index combinations between shifts and residues and orders these into source and target indices for the edges. """ - shift = torch.arange(0, self.num_shift) - resid = torch.arange(0, self.num_res) - + # shift = torch.arange(0, self.num_shift) + shift = np.nonzero(self.data['SHIFT'].f[:, 0])[0] # need to specify if it's the shift being assigned + resid = torch.arange(0, self.num_res, device=self.device) + shift_repeats = shift.repeat_interleave(self.num_res) - resid_repeats = resid.repeat(self.num_shift) + # resid_repeats = resid.repeat(self.num_shift) + resid_repeats = resid return torch.stack((resid_repeats, shift_repeats), dim=0) def get_batch_edges(self): - shift = torch.arange(0, self.num_shift) - noe = torch.arange(0, self.num_noe) - resid = torch.arange(0, self.num_res) + shift = torch.arange(0, self.num_shift, device=self.device) + noe = torch.arange(0, self.num_noe, device=self.device) + resid = torch.arange(0, self.num_res, device=self.device) - shift_repeats = torch.zeros(self.num_shift).long() - noe_repeats = torch.zeros(self.num_noe).long() - resid_repeats = torch.zeros(self.num_res).long() + shift_repeats = torch.zeros(self.num_shift, device=self.device).long() + noe_repeats = torch.zeros(self.num_noe, device=self.device).long() + resid_repeats = torch.zeros(self.num_res, device=self.device).long() self.data['SHIFT', 'SHIFT_extract', f'VALUE_SHIFT'].edge_index = torch.stack((shift, shift_repeats), dim=0) self.data['NOE', 'NOE_extract', f'VALUE_NOE'].edge_index = torch.stack((noe, noe_repeats), dim=0) @@ -166,7 +176,7 @@ def get_edges(self, triple_num, triple_type): # Edges for self loop self.data[f'TRIPLE{triple_num}', 'update', f'TRIPLE{triple_num}'].edge_index = torch.stack([edges_in[1, 0], edges_in[1, 0]], dim=0) - + if triple_num == 0: self.data['RES', 'pair', 'SHIFT'].edge_index = self.get_pairwise_edges() # pairwise edges - no message passing return self.data @@ -177,6 +187,7 @@ def generate_edge_indices(self): self.data = self.get_edges(2, ('SHIFT', 'RES', 'NOE')) self.data = self.get_edges(3, ('SHIFT', 'SHIFT', 'NOE')) self.data = self.get_batch_edges() + # print(self.data['RES', 'pair', 'SHIFT'].edge_index) return self.data @@ -260,23 +271,24 @@ class TripleUpdate(MessagePassing): """ Message passing class for triple self updates. """ - def __init__(self): + def __init__(self, device): super().__init__(aggr='add') self.calc_manager = CalculationManager() + self.device = device self.hidden = 64 self.mlp1 = nn.Sequential( - nn.Linear(9, self.hidden), + nn.Linear(12, self.hidden, device=self.device), nn.ReLU(), - nn.Linear(self.hidden, 16), - nn.LayerNorm(16)) + nn.Linear(self.hidden, 19, device=self.device), + nn.LayerNorm(19, device=self.device)) self.mlp2 = nn.Sequential( - nn.Linear(8, self.hidden), + nn.Linear(11, self.hidden, device=self.device), nn.ReLU(), - nn.Linear(self.hidden, 10), - nn.LayerNorm(10)) + nn.Linear(self.hidden, 13, device=self.device), + nn.LayerNorm(13, device=self.device)) def resresnoe_message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j, x22_j): """ @@ -316,10 +328,10 @@ def resresnoe_message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j, x22_j): delta13x = (rel_dist * (mlp_out[:, 10:13])) # FEATURE deltas [n, 1] - delta1f = mlp_out[:, 13:14] - delta2f = mlp_out[:, 14:15] - delta3f = mlp_out[:, 15:] - return torch.cat((delta1x, delta2x, delta3x, delta4x, delta5x, delta6x, delta7x, delta12x, delta13x, delta1f, delta2f, delta3f), dim=-1) # [n, 16] + delta1f = mlp_out[:, 13:15] + delta2f = mlp_out[:, 15:17] + delta3f = mlp_out[:, 17:] + return torch.cat((delta1x, delta2x, delta3x, delta4x, delta5x, delta6x, delta7x, delta12x, delta13x, delta1f, delta2f, delta3f), dim=-1) # [n, 16] [n, 19] def shiftshiftnoe_message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j): """ @@ -352,10 +364,10 @@ def shiftshiftnoe_message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j): delta7x = (diff5 * mlp_out[:, 6:7]) # H2 # FEATURE deltas [n, 1] - delta1f = mlp_out[:, 7:8] - delta2f = mlp_out[:, 8:9] - delta3f = mlp_out[:, 9:] - return torch.cat((delta1x, delta2x, delta3x, delta4x, delta5x, delta6x, delta7x, delta1f, delta2f, delta3f), dim=-1) # [n, 10] + delta1f = mlp_out[:, 7:9] + delta2f = mlp_out[:, 9:11] + delta3f = mlp_out[:, 11:] + return torch.cat((delta1x, delta2x, delta3x, delta4x, delta5x, delta6x, delta7x, delta1f, delta2f, delta3f), dim=-1) # [n, 10] [n, 13] def forward(self, x1, x12, x2, x22, x3, f1, f2, f3, edge_index): out = self.propagate(edge_index, x1=x1, x2=x2, x3=x3, f1=f1, f2=f2, f3=f3, x12=x12, x22=x22) # not sure how to deal with size here @@ -371,12 +383,12 @@ def message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j=None, x22_j=None): def update(self, aggr_out, x12, x22): # residue based triple type will have two sets of coordinates (x12_j and x22_j) where shifts will have nonetype if x12 != None and x22 != None: - # output shapes: noe[n, 3], shift1[n, 2], shift2[n, 2], dist1[n, 3], dist2[n, 3], f1[n, 1], f2[n, 1], f3[n, 1] - delta_noe, delta_shift1, delta_shift2, delta_dist1, delta_dist2, delta_f1, delta_f2, delta_f3 = aggr_out[:, 0:3], aggr_out[:, 3:5], aggr_out[:, 5:7], aggr_out[:, 7:10], aggr_out[:, 10:13], aggr_out[:, 13:14], aggr_out[:, 14:15], aggr_out[:, 15:] + # output shapes: noe[n, 3], shift1[n, 2], shift2[n, 2], dist1[n, 3], dist2[n, 3], f1[n, 2], f2[n, 2], f3[n, 2] + delta_noe, delta_shift1, delta_shift2, delta_dist1, delta_dist2, delta_f1, delta_f2, delta_f3 = aggr_out[:, 0:3], aggr_out[:, 3:5], aggr_out[:, 5:7], aggr_out[:, 7:10], aggr_out[:, 10:13], aggr_out[:, 13:15], aggr_out[:, 15:17], aggr_out[:, 17:] return delta_shift1, delta_shift2, delta_noe, delta_dist1, delta_dist2, delta_f1, delta_f2, delta_f3 else: - # output shapes: noe[n, 3], shift1[n, 2], shift2[n, 2], f1[n, 1], f2[n, 1], f3[n, 1] - delta_noe, delta_shift1, delta_shift2, delta_f1, delta_f2, delta_f3 = aggr_out[:, 0:3], aggr_out[:, 3:5], aggr_out[:, 5:7], aggr_out[:, 7:8], aggr_out[:, 8:9], aggr_out[:, 9:] + # output shapes: noe[n, 3], shift1[n, 2], shift2[n, 2], f1[n, 2], f2[n, 2], f3[n, 2] + delta_noe, delta_shift1, delta_shift2, delta_f1, delta_f2, delta_f3 = aggr_out[:, 0:3], aggr_out[:, 3:5], aggr_out[:, 5:7], aggr_out[:, 7:9], aggr_out[:, 9:11], aggr_out[:, 11:] return delta_shift1, delta_shift2, delta_noe, delta_f1, delta_f2, delta_f3 @@ -407,14 +419,13 @@ def update(self, aggr_out, x): class BatchMessagePass(MessagePassing): - """ - - """ - def __init__(self, aggr): + def __init__(self, aggr, device): super().__init__(aggr=aggr) - self.noe_reduce = (nn.Linear(3, 1)) - self.shift_reduce = (nn.Linear(2, 1)) - self.res_reduce = (nn.Linear(5, 1)) + self.device = device + + self.noe_reduce = (nn.Linear(3, 1, device=self.device)) + self.shift_reduce = (nn.Linear(2, 1, device=self.device)) + self.res_reduce = (nn.Linear(5, 1, device=self.device)) def forward(self, x_source, x_target, edge_index): # SIZE (n, m) (source, target) @@ -453,13 +464,13 @@ def update_data(self, data, x_source, f_source, edge_type, x_index, update_f=Tru data[target].f = self.triple_message(f_source, data[target].f, edge_index) return data - + class NMRLayer(nn.Module): - def __init__(self): + def __init__(self, device): super(NMRLayer, self).__init__() self.triple_in = TripleIn() - self.triple_self = TripleUpdate() + self.triple_self = TripleUpdate(device) self.triple_out = TripleOut() self.calc_manager = CalculationManager() @@ -470,7 +481,7 @@ def __init__(self): def update_noes(self, data, noe_delta, feature, i): # 'i' is for the triple type in all of these cases (TRIPLE0 --> RES RES NOE, etc.) - # Should simplify this to be one or the other from the very start (string --> number identification) + # Should really simplify this to be one or the other from the very start (string --> number identification) data = self.triple_out.update_data(data, noe_delta, feature, (f'TRIPLE{i}', 'NOE_add', 'NOE'), self.NOE) return data @@ -515,14 +526,16 @@ def forward(self, data): class ValueCalc(): - def __init__(self): - self.batch_message = BatchMessagePass(aggr='mean') + def __init__(self, device): + self.device = device + + self.batch_message = BatchMessagePass(aggr='mean', device=self.device) self.hidden = 64 self.testmlp = nn.Sequential( - nn.Linear(3, self.hidden), + nn.Linear(3, self.hidden, device=self.device), nn.ReLU(), - nn.Linear(self.hidden, 1) + nn.Linear(self.hidden, 1, device=self.device) ) def get_aggr(self, x_source, x_target, edge_index): @@ -535,14 +548,15 @@ def calc_value(self, data): resid = self.get_aggr(data['RES'].x, data['VALUE_RES'].x, data['RES', 'RES_extract', 'VALUE_RES'].edge_index) concat_aggr = torch.cat((noe, shift, resid), dim=-1) - values = self.testmlp(concat_aggr) - return values + value = self.testmlp(concat_aggr) + return value class PolicyCalc(): - def __init__(self): + def __init__(self, device): self.RES_NH = slice(3,5) + self.device = device def pairwise_distance(self, data, edge_type): node_type1, edge, node_type2 = edge_type @@ -552,70 +566,79 @@ def pairwise_distance(self, data, edge_type): # shift shift_nodes = data[node_type2].x[data[edge_type].edge_index[1]] - pair_dist = (shift_nodes - resid_nodes) - return -pair_dist**2 + rel_dist = (shift_nodes - resid_nodes) + pair_dist2 = -(torch.norm(rel_dist, dim=-1, keepdim=True)**2) + return pair_dist2.squeeze() # reduce dimension [num_res, 1] --> [num_res] def calc_policy(self, data): - pair_dist2 = self.pairwise_distance(data, ('RES', 'pair', 'SHIFT')) - return pair_dist2 + data_unbatched = data.to_data_list() + policy = [self.pairwise_distance(batch, ('RES', 'pair', 'SHIFT')).tolist() for batch in data_unbatched] + return torch.tensor(policy, dtype=torch.float32, requires_grad=True, device=self.device).unsqueeze(1) # add dimension back, now transposed [1, num_res] -class TestLayer(): - def __init__(self): - self.value = ValueCalc() - self.policy = PolicyCalc() +class TestLayer(nn.Module): + def __init__(self, device): + super(TestLayer, self).__init__() + self.value = ValueCalc(device) + self.policy = PolicyCalc(device) self.nmr = nn.Sequential( - NMRLayer(), - NMRLayer(), + NMRLayer(device), ) def forward(self, data): - test_data = self.nmr(data) - value = self.value.calc_value(test_data) - policy = self.policy.calc_policy(test_data) + out_data = self.nmr(data) + value = self.value.calc_value(out_data) + policy = self.policy.calc_policy(out_data) return value, policy + +if __name__ == "__main__": -def extract_data(pickle_file="fake_histories_r10_1.pkl"): - with open(pickle_file, "rb") as f: - histories = pickle.load(f) - return histories + def extract_data(pickle_file="fake_histories_r4_0.pkl"): + with open(pickle_file, "rb") as f: + histories = pickle.load(f) + return histories -def construct_graph(history): - nodes = ConstructNodes() - data = nodes.construct_data(history) - edges = ConstructEdges(data) - data = edges.generate_edge_indices() - return data + def construct_graph(history, device): + nodes = ConstructNodes(device) + data = nodes.construct_data(history) + edges = ConstructEdges(data, device) + data = edges.generate_edge_indices() + return data -def preprocess_data(histories, batch_size): - graphs = [] - for history in histories: - graphs.append(construct_graph(history)) + def preprocess_data(histories, batch_size, device): + graphs = [] + for history in histories: + graphs.append(construct_graph(history, device)) - data_loader = DataLoader(graphs, batch_size=batch_size, shuffle=False) - return data_loader + data_loader = DataLoader(graphs, batch_size=batch_size, shuffle=False) + return data_loader + + data_loader = preprocess_data(extract_data(), batch_size=1, device='cuda') + + plt.ion() + fig, ax = plt.subplots() + epochs = 5 + + testgnn = TestLayer(device='cuda') + for epoch in range(epochs): + batch_num = 0 + for batch in data_loader: + batch_num += 1 -if __name__ == "__main__": - # nodes = ConstructNodes() - # data = nodes.construct_data(extract_data()) - edges = ConstructEdges(data) - data = edges.generate_edge_indices() - # data = preprocess_data(extract_data(), batch_size=2) - - value = ValueCalc() - - # create fake graph list (not from histories) to check - graph_list = [data, data, data, data] - data_loader = DataLoader(graph_list, batch_size=2, shuffle=False) - - testgnn = TestLayer() - for batch in data_loader: - value, policy = testgnn.forward(batch) - print(value.shape) # 2 values (1 per graph in batch) - print(policy.shape) # 8 policies (4 per graph in batch --> 4 comparisons being made 2shift/2resid) **THESE ARE NOT SPLIT BY GRAPH** + value, policy = testgnn.forward(batch) + + ax.cla() + ax.scatter(batch['SHIFT'].x[:, 0].tolist(), batch['SHIFT'].x[:, 1].tolist(), color='red') + ax.scatter(batch['RES'].x[:, 3].tolist(), batch['RES'].x[:, 4].tolist(), color='blue') + ax.set_title(f"Epoch {epoch} Batch {batch_num}") + + plt.pause(0.01) + + plt.ioff() + plt.show() \ No newline at end of file diff --git a/training_loop_gnn.py b/training_loop_gnn.py new file mode 100644 index 0000000..8917f6d --- /dev/null +++ b/training_loop_gnn.py @@ -0,0 +1,119 @@ +import torch +import torch_geometric +import pickle +import numpy as np + +from pyg_triple_update import * +from torch.utils.tensorboard import SummaryWriter +import random + +def extract_data(pickle_file="fake_histories_r4_0.pkl"): + with open(pickle_file, "rb") as f: + histories = pickle.load(f) + return histories + +def construct_graph(history, device): + nodes = ConstructNodes(device) + data = nodes.construct_data(history) + edges = ConstructEdges(data, device) + data = edges.generate_edge_indices() + return data + +def preprocess_data(histories, device): + graphs = [] + for history in histories: + graphs.append(construct_graph(history, device)) + # print(history['assign_order']) + return graphs + +if __name__ == "__main__": + + # Set device + device = 'cuda' + + # Set up TensorBoard + # writer = SummaryWriter() + + # Load data + histories = extract_data() + nmr_graphs = preprocess_data(extract_data(), device) + + # Split dataset manually into training and testing + train_size = int(0.8 * len(nmr_graphs)) + train_data = nmr_graphs[:train_size] + test_data = nmr_graphs[train_size:] + train_target = [inp['shift_to_assign'] for inp in histories][:train_size] + test_target = [inp['shift_to_assign'] for inp in histories][train_size:] + + batch_size = 1 + epochs = 50_000 # Define the number of epochs + eval_interval = 100 # Evaluate the model every 100 iterations + + train_data_loader = DataLoader(train_data, batch_size=batch_size, shuffle=False) + test_data_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False) + + # Create our network and optimizer + net = TestLayer(device) + opt = torch.optim.AdamW(net.parameters(), lr=1e-4, weight_decay=0.01) + + # Interactive plotting + plt.ion() + fig, ax = plt.subplots() + + iteration = 0 + for epoch in range(epochs): + + for i, xs in enumerate(train_data_loader): + net.train() + + # Break targets into batches manually (if in same order - no shuffling) + ys = train_target[i*batch_size:(i+1)*batch_size] + + iteration += 1 + + _, y_hats = net(xs) + + # Plots current iteration in batch ########################################## + ax.cla() + x_val1, y_val1 = xs['SHIFT'].x[:, 0].tolist(), xs['SHIFT'].x[:, 1].tolist() + x_val2, y_val2 = xs['RES'].x[:, 3].tolist(), xs['RES'].x[:, 4].tolist() + + ax.scatter(x_val1, y_val1, color='red', label='shift') + ax.scatter(x_val2, y_val2, color='blue', label='resid') + + for j in range(len(x_val1)): + ax.annotate(j, (x_val1[j], y_val1[j]+0.5), color='red') + ax.annotate(j, (x_val2[j], y_val2[j]+0.5), color='blue') + + ax.legend() + ax.set_title(f"Iteration {iteration} - Batch {i}") + plt.pause(0.01) + ############################################################################# + + loss = 0 + correct = 0 + count = 0 + + for y_hat, y in zip(y_hats, ys): + y = torch.tensor([y], dtype=torch.long, device=device) + loss += torch.nn.functional.cross_entropy(y_hat, y) + y_pred = torch.argmax(y_hat) + + if y_pred == y: + correct += 1 + count += 1 + + loss = loss / count + accuracy = correct / count + + print(f"Iteration {iteration} - Loss: {loss.item()} - Accuracy: {accuracy}") + + # writer.add_scalar("Loss/train", loss.item(), iteration) + # writer.add_scalar("Accuracy/train", accuracy, iteration) + + opt.zero_grad() + loss.backward() + opt.step() + + plt.ioff() + plt.show() From d62b09f4c03d5de7d4d45e95e3f63cf86acb9008 Mon Sep 17 00:00:00 2001 From: Justin MacCallum Date: Tue, 11 Nov 2025 20:27:01 -0700 Subject: [PATCH 40/43] Reorganize and cleanup --- README.md | 46 +- assignment_unittest.py | 68 -- fake_data.py | 244 ------- fake_histories.py | 89 --- fake_histories_agent.py | 97 --- fractional_activations_unittest.py | 68 -- nmr/__init__.py | 32 + nmr/construct.py | 291 ++++++++ nmr/models/__init__.py | 37 + nmr/models/heads.py | 145 ++++ nmr/models/network.py | 205 ++++++ nmr/models/triple.py | 342 ++++++++++ nmr/nmr_gym/__init__.py | 24 + .../nmr_gym/assignment_order.py | 0 nmr/nmr_gym/data_structures.py | 45 ++ energy.py => nmr/nmr_gym/energy.py | 0 nmr/nmr_gym/fake_data.py | 409 +++++++++++ nmr/nmr_gym/fake_histories.py | 243 +++++++ nmr_gym_env.py => nmr/nmr_gym/gym_env.py | 15 +- nmr/nmr_gym/io.py | 279 ++++++++ nmr/nmr_gym/state.py | 132 ++++ nmr_text_adventure.py | 140 ---- nmr_transformer.py | 256 ------- pyg_triple_update.py | 644 ------------------ scripts/generate_dataset.py | 151 ++++ scripts/generate_histories.py | 196 ++++++ scripts/text_adventure.py | 131 ++++ scripts/train_gnn.py | 164 +++++ .../visualize_data.py | 0 write_pdb.py => scripts/write_pdb.py | 8 +- energy_unittest.py => tests/test_energy.py | 7 +- tests/test_fake_data.py | 171 +++++ tests/test_fake_histories.py | 191 ++++++ .../test_fractional_activation.py | 7 +- tests/test_generate_dataset.py | 289 ++++++++ tests/test_generate_histories_algorithm.py | 265 +++++++ tests/test_generate_histories_cli.py | 339 +++++++++ gym_env_unittest.py => tests/test_gym_env.py | 54 +- tests/test_integration_end_to_end.py | 303 ++++++++ training_loop.py | 205 ------ training_loop_gnn.py | 119 ---- 41 files changed, 4494 insertions(+), 1957 deletions(-) delete mode 100644 assignment_unittest.py delete mode 100644 fake_data.py delete mode 100644 fake_histories.py delete mode 100644 fake_histories_agent.py delete mode 100644 fractional_activations_unittest.py create mode 100644 nmr/__init__.py create mode 100644 nmr/construct.py create mode 100644 nmr/models/__init__.py create mode 100644 nmr/models/heads.py create mode 100644 nmr/models/network.py create mode 100644 nmr/models/triple.py create mode 100644 nmr/nmr_gym/__init__.py rename assignment_order.py => nmr/nmr_gym/assignment_order.py (100%) create mode 100644 nmr/nmr_gym/data_structures.py rename energy.py => nmr/nmr_gym/energy.py (100%) create mode 100644 nmr/nmr_gym/fake_data.py create mode 100644 nmr/nmr_gym/fake_histories.py rename nmr_gym_env.py => nmr/nmr_gym/gym_env.py (93%) create mode 100644 nmr/nmr_gym/io.py create mode 100644 nmr/nmr_gym/state.py delete mode 100644 nmr_text_adventure.py delete mode 100644 nmr_transformer.py delete mode 100644 pyg_triple_update.py create mode 100644 scripts/generate_dataset.py create mode 100644 scripts/generate_histories.py create mode 100644 scripts/text_adventure.py create mode 100644 scripts/train_gnn.py rename visualize_data.py => scripts/visualize_data.py (100%) rename write_pdb.py => scripts/write_pdb.py (89%) rename energy_unittest.py => tests/test_energy.py (93%) create mode 100644 tests/test_fake_data.py create mode 100644 tests/test_fake_histories.py rename fractional_activation_unittest.py => tests/test_fractional_activation.py (93%) create mode 100644 tests/test_generate_dataset.py create mode 100644 tests/test_generate_histories_algorithm.py create mode 100644 tests/test_generate_histories_cli.py rename gym_env_unittest.py => tests/test_gym_env.py (59%) create mode 100644 tests/test_integration_end_to_end.py delete mode 100644 training_loop.py delete mode 100644 training_loop_gnn.py diff --git a/README.md b/README.md index 6dfa5fb..2703e48 100644 --- a/README.md +++ b/README.md @@ -1 +1,45 @@ -# nmr \ No newline at end of file +# nmr + +NMR Chemical Shift Assignment using Graph Neural Networks and Reinforcement Learning + +## Quick Start + +### Generating Training Data + +Create synthetic datasets for supervised pre-training: + +```bash +# Generate a dataset +python scripts/generate_dataset.py --num-resid 10 --output dataset_10.pkl + +# Generate training histories +python scripts/generate_histories.py \ + --dataset dataset_10.pkl \ + --num-histories 100 +``` + +For detailed documentation, see [docs/fake-data-guide.md](docs/fake-data-guide.md) + +## Documentation + +- **[Fake Data Generation Guide](docs/fake-data-guide.md)** - Complete guide to generating synthetic training data +- **[Development Guide](docs/development-guide.md)** - Setup, development patterns, and testing +- **[Architecture Documentation](docs/architecture.md)** - System architecture and design +- **[Project Overview](docs/project-overview.md)** - Executive summary + +## Project Structure + +``` +nmr/ +├── nmr/ # Main package +│ ├── construct.py # Graph construction +│ ├── models/ # GNN architecture +│ └── env/ # RL environment +├── scripts/ # Executable scripts +│ ├── generate_dataset.py # Dataset generation +│ ├── generate_histories.py # History generation +│ ├── train_gnn.py # Training loop +│ └── visualize_data.py # Data visualization +├── tests/ # Unit tests +└── docs/ # Documentation +``` diff --git a/assignment_unittest.py b/assignment_unittest.py deleted file mode 100644 index ce63fc7..0000000 --- a/assignment_unittest.py +++ /dev/null @@ -1,68 +0,0 @@ -#unittest for fractional activations -#has to be made into a class for it to work - -import unittest - -class TestFracAct(unittest.TestCase): - - def setUp(self): - self.frac_act = FracAct() - - #test for activation loop - #expected to return the noe peak divided by the total number of different/unique noe peaks - - def test_activation_loop(self): - assign_order = [] - noe = [(0, 1), (0, 2)] - target = 0 - - result = self.frac_act.activation_loop(target, noe, assign_order) - expected = 1/3 - - self.assertEqual(result, expected) - - #test if each noe peak will be evaluated to give the fractional activation - - def test_fractional_activation_of_all_noe_peaks(self): - num_resid = 3 - noes = [[(0, 1), (0, 2)]] - - result = self.frac_act.fractional_activation_summation(num_resid, noes) - expected = [1/3, 1/3, 1/3] - - self.assertEqual(result, expected) - - #for more then one noe, the probabilities are summed - - def test_fractional_activation_summation_of_noes(self): - num_resid = 4 - noes = [(0, 1), (0, 2)], [(2,1),(1,1)], [(1,0),(3,2)] - - result = self.frac_act.fractional_activation_summation(num_resid, noes) - expected = [7/12, 13/12, 13/12, 1/4] - - self.assertEqual(result, expected) - - #according to the fractional activations calculated, should give the correct order - #if it is a tie, then lowest value goes first - - def test_fractional_activation_for_assignment_order(self): - num_resid = 4 - noes = [(0, 1), (0, 2)], [(2,1),(1,1)], [(1,0),(3,2)] - - result = self.frac_act.fractional_activation(num_resid, noes) - expected = [1, 2, 0, 3] - - self.assertEqual(result, expected) - - #if number of residues is higher than the number of peaks, then it will loop through to give first/lowest value peak ex. will give 0 for every additional residue - #if number of residues is lower than the number of peaks, then it will return the first value again (0, in this case) - - def test_fractional_activation_for_num_resid(self): - num_resid = 5 - noes = [[(0, 1), (0, 2)]], [(2,1),(1,1)], [(1,0),(3,2)] - - result = self.frac_act.fractional_activation(num_resid, noes) - expected = [1, 2, 0, 3] - - self.assertEqual(result, expected) \ No newline at end of file diff --git a/fake_data.py b/fake_data.py deleted file mode 100644 index 38ae6f8..0000000 --- a/fake_data.py +++ /dev/null @@ -1,244 +0,0 @@ -import numpy as np -from itertools import product -from typing import NamedTuple -import random -import argparse -import pickle -import glob as glob - - -class HSQCPeak(NamedTuple): - H1: float - N15: float - - -class NOEPeak(NamedTuple): - H1: float - N15: float - H2: float - - -class Protein(NamedTuple): - x: float - y: float - z: float - H1: float - N15: float - - -class Connectivity(NamedTuple): - atom1: float - atom2: float - distance: float - - -class FakeDataGenerator(): - - def __init__(self, num_resid): - self.num_resid: int = num_resid - self.nshift_min = 100 - self.nshift_max = 135 - self.hshift_min = 6 - self.hshift_max = 10 - self.cutoff = 0.5 - - def sample_unit(self, num_points, num_sides, min_len=0, max_len=1): - """ - Samples points from unit square or cube. - """ - return np.random.uniform(low=min_len, high=max_len, size=(num_points, num_sides)) - - def scale_unit(self, coordinates): - """ - Scaling box size to reflect a globular protein. - PDB:1CRC (~100 resid, globular, ~3.2 nm diameter) - - Rg = RN**v - Scaling factor (v) of 0.4, instead of 0.3 (cube-root) - R of 0.2 nm used in paper for fit - DOI: 10.1142/S021972002050050X - """ - radius_gyration = 0.2*(self.num_resid**0.4) - - com = np.mean(coordinates, axis=0) # should be ~0.5 - distances_sampled = np.linalg.norm(coordinates - com, axis=1) # distances around origin - radius_sampled = np.sqrt(np.mean((distances_sampled**2))) # radius of gyration based on these distances - - # Expected is radius_gyration, current is radius_sampled --> scale to fit - # Will overlap occur? - scale = radius_gyration / radius_sampled - coordinates_scaled = com + ((coordinates - com) * scale) - - return coordinates_scaled - - def create_hsqc(self): - """ - Sample points within NMR window for H1 and N15. Stack values to create HSQC peaks. - """ - # H shifts - h_shifts = self.sample_unit(self.num_resid, num_sides=1, min_len=self.hshift_min, max_len=self.hshift_max) - # N shifts - n_shifts = self.sample_unit(self.num_resid, num_sides=1, min_len=self.nshift_min, max_len=self.nshift_max) - - return np.hstack((h_shifts, n_shifts)) - - def add_noise(self, points, scale=0.1): - """ - Random selection from normal (gaussian) distribution of 'scale' width from 0 (center). - 'size' makes sure that it's the same shape as the point we are adding noises to. - ex. point = [0,1,2], noise = [1,1,1], noisy_point = [1,2,3] - """ - noise = np.random.normal(loc=0, scale=scale, size=points.shape) - noisy_point = points + noise - - # fold over points outside of boundaries - # for point in noisy_point: - # for i, axis in enumerate(point): - # if max_len < axis: - # point[i] = (max_len - (axis - max_len)) - # if axis < min_len: - # point[i] = (min_len - axis) - # else: - # pass - - return noisy_point - - def calculate_dist(self, p1, p2): - """ - Euclidian distance between two points. - """ - return np.linalg.norm((p2 - p1)) - - def create_noes(self, coordinates, shifts, random_key=False): - """ - Grabs close coordinates and 'associated' HSQC shift peaks (randomized index or 1:1 correlation) to create noes. - Predicted shifts assigned as shift order - this order is associated with the coordinate order if randomized (ie. the answer key). - Adds gaussian noise to all points at the end. - - **Can end up with no noes depending on the cutoff** - """ - if random_key: - shifts = np.random.permutation(shifts) - - noe_list = [] - for i, atom1 in enumerate(coordinates): - for j, atom2 in enumerate(coordinates): - if i != j: - dist = self.calculate_dist(atom1, atom2) - # dist = dist_grid[i][j] - if dist < self.cutoff: - #print(dist, shifts[i], shifts[j]) - noe = list(shifts[i][:]) # H1, N15 - noe.append(shifts[j][0]) # H2 - noe_list.append(noe) - - noisy_noes = self.add_noise(np.array(noe_list), scale=0.01) - pred_chemical_shifts = self.add_noise(np.array(shifts), scale=0.1) - - return noisy_noes, pred_chemical_shifts - - def calculate_connectivity(self, coordinates): - """ - Calculates close contacts based on protein coordinates. - """ - contacts = [] - - for i, atom1 in enumerate(coordinates): - for j, atom2 in enumerate(coordinates): - if i != j: - dist = self.calculate_dist(atom1, atom2) - # dist = dist_grid[i][j] - if dist < self.cutoff: - contacts.append((i, j, dist)) - - return contacts - - # def pdist_test(self, coordinates): - # pairwise_dists = pdist(coordinates) - # square_dists = squareform(pairwise_dists) - # return square_dists - - def generate_data_arrays(self, random_key): - """ - Generates all qualities of the system as individual arrays. - """ - # 3D structure [x,y,z] - pred_coordinates = self.sample_unit(self.num_resid, num_sides=3) - pred_coordinates = self.scale_unit(pred_coordinates) - - # "Actual" shifts [H1,N15] - obs_chemical_shifts = self.create_hsqc() - - # Distance grid - # dist_grid = self.pdist_test(pred_coordinates) - - # NOEs [H1,N15,H2] and predicted shifts [H1,N15] - noes, pred_chemical_shifts = self.create_noes(pred_coordinates, obs_chemical_shifts, random_key=random_key) - - # Connectivity [atom1,atom2,dist] - connectivity = self.calculate_connectivity(pred_coordinates) - - return pred_coordinates, obs_chemical_shifts, pred_chemical_shifts, noes, connectivity - - def order_data(self, pred_coordinates, obs_chemical_shifts, pred_chemical_shifts, noes, connectivity): - """ - Compiles data in a list of named tuples with one object per residue. - """ - pred_coordinates = [Protein(x=resid[0], y=resid[1], z=resid[2], H1=shift[0], N15=shift[1]) for resid, shift in zip(pred_coordinates, pred_chemical_shifts)] - # pred_chemical_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in pred_chemical_shifts] - obs_chemical_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in obs_chemical_shifts] - noes = [NOEPeak(H1=shift[0], N15=shift[1], H2=shift[2]) for shift in noes] - connectivity = [Connectivity(atom1=contact[0], atom2=contact[1], distance=contact[2]) for contact in connectivity] - - return pred_coordinates, obs_chemical_shifts, noes, connectivity - - def dump_pickle(self, pred_coordinates, obs_chemical_shifts, noes, connectivity, example=True): - """ - Saves data to disk. Run is given a proper name if used as a saved example. - """ - name = f'fakedata_r{self.num_resid}.pkl' if example else f'current_run.pkl' - - with open(name, 'wb') as f: - pickle.dump(pred_coordinates, f) - pickle.dump(obs_chemical_shifts, f) - pickle.dump(noes, f) - pickle.dump(connectivity, f) - - def load_pickle(self, example): - - name = f"./*r{self.num_resid}.pkl" if example else "./current_run.pkl" - - pickle_file = glob.glob(name) - with open(pickle_file[0], 'rb') as f: - coordinates = pickle.load(f) - obs_chemical_shifts = pickle.load(f) - noes = pickle.load(f) - connectivity = pickle.load(f) - - return coordinates, obs_chemical_shifts, noes, connectivity - - def generate_data(self, pickle_data=True, example=True, random_key=False): - """ - Generates all fake data, orders it in lists of namedtuples, and pickles if required. - """ - pred_coordinates, obs_chemical_shifts, pred_chemical_shifts, noes, connectivity = self.generate_data_arrays(random_key=random_key) - pred_coordinates, obs_chemical_shifts, noes, connectivity = self.order_data(pred_coordinates, obs_chemical_shifts, pred_chemical_shifts, noes, connectivity) - - if pickle_data: - self.dump_pickle(pred_coordinates, obs_chemical_shifts, noes, connectivity, example=example) - - else: - return pred_coordinates, obs_chemical_shifts, noes, connectivity - - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('num_resid', type=int, help='Number of residues') - args = parser.parse_args() - - num_resid = args.num_resid - - fakedata = FakeDataGenerator(num_resid) - pred_coordinates, obs_chemical_shifts, noes, connectivity = fakedata.generate_data(pickle_data=False, example=False) - \ No newline at end of file diff --git a/fake_histories.py b/fake_histories.py deleted file mode 100644 index 0e19108..0000000 --- a/fake_histories.py +++ /dev/null @@ -1,89 +0,0 @@ -import numpy as np - -from fake_data import FakeDataGenerator, HSQCPeak, Protein -from energy import Energy -from assignment_order import FractionalActivation - - - -class FakeHistoryGenerator(): - - def __init__(self, num_resid): - self.num_resid = num_resid - self.fakedata = FakeDataGenerator(num_resid) - self.cutoff = 0.5 - - def perturb_data(self, coordinates): - """ - Adds noise to original fake data and generates new 'actual' shifts based on predicted shifts. - """ - coords = self.fakedata.add_noise(np.array(coordinates)[:, :3], scale=0.1) - synthetic_shifts = self.fakedata.add_noise(np.array(coordinates)[:, 3:], scale=0.1) - - return coords, synthetic_shifts - - def weighted_error(self, scale=1): - """ - Generates weights for number of errors added based on exponential distribution. - Chooses error number to add based on weights. - """ - # Weight each point - weights = np.exp(-np.arange(self.num_resid)/scale) - # Normalized - weights = weights/np.sum(weights) - # Number of errors based on weight - error_num = np.random.choice(self.num_resid, 1, p=weights) - - # Can't make a singular error - needs to be replaced by another choice - return error_num if error_num > 1 else error_num*2 - - def add_errors(self, answer): - """ - If errors are present chooses that number of random shifts and move answer one over from true. - """ - error_num = self.weighted_error(scale=1) - - if error_num == 0: - return answer - else: - selections = np.random.choice(list(answer.values()), error_num, replace=False) - # print(selections) - mutations = np.roll(selections, 1) - # print(mutations) - for i, j in enumerate(selections): - # print(answer, answer[j], mutations[i], error_num) - answer[j] = mutations[i] - - return answer - - def generate_history(self, original): - """ - Adds noise to original data, reorganizes, and edits answer key if errors are introduced. - """ - - coordinates, obs_chemical_shifts = self.perturb_data(original['coordinates']) - - noes, pred_chemical_shifts = self.fakedata.create_noes(coordinates, obs_chemical_shifts) - - connectivity = self.fakedata.calculate_connectivity(coordinates) - - coordinates, obs_chemical_shifts, noes, connectivity = self.fakedata.order_data(coordinates, obs_chemical_shifts, pred_chemical_shifts, noes, connectivity) - - # Errors not included currently - # answer = original['assignments'] - # answer = add_errors(answer) - - - state = { - "coordinates": coordinates, - "obs_chemical_shifts": obs_chemical_shifts, - "noes": noes, - "connectivity": connectivity, - "assignments": {}, - "assign_order": [], - "shift_to_assign": 0, - "total_energy": 0.0, - "reward": 0.0 - } - - return state diff --git a/fake_histories_agent.py b/fake_histories_agent.py deleted file mode 100644 index 0289c0f..0000000 --- a/fake_histories_agent.py +++ /dev/null @@ -1,97 +0,0 @@ -import gymnasium as gym -import argparse -import pickle -import copy -import numpy as np - -from nmr_gym_env import GymEnv -from fake_data import FakeDataGenerator -from fake_histories import FakeHistoryGenerator -from visualize_data import Visualization - -from nmr_text_adventure import get_data - -import time - - - -if __name__ == '__main__': - start = time.process_time() - """ - Runs through protein examples of one size following the 1:1 shift to atom assignment. - Once example is assigned, generates history and repeats assignment process. - All histories are then saved to disk. - """ - parser = argparse.ArgumentParser() - parser.add_argument('num_resid', type=int, help='Max protein size') - parser.add_argument('history_length', type=int, help='Number of synthetic examples to generate') - args = parser.parse_args() - - num_resid = args.num_resid - history_length = args.history_length - - gym_env = GymEnv(num_resid) - - fakedata = FakeDataGenerator(num_resid) - fakehistory = FakeHistoryGenerator(num_resid) - - ###################################################### - - print("\nAvailable options to work with:\n1. Generate and save an example\n2. Grab a previous saved example\n3. Run a new example (not saved)") - - selection = int(input("Option Selection: ")) - - if selection == 1: - # 1. Need to generate and save an example? Saved in all instances as fakedata_r{num_resid}.pkl - coordinates, obs_chemical_shifts, noes, connectivity = get_data(num_resid, pickled=False, pickle_data=True, example=True, random_key=True) - if selection == 2: - # 2. Grab one of the previous examples? Make sure example exists with desired resid number - coordinates, obs_chemical_shifts, noes, connectivity = get_data(num_resid, pickled=True, pickle_data=False, random_key=True) - if selection == 3: - # 3. Run a new example? Saved in all instances as current_run.pkl - coordinates, obs_chemical_shifts, noes, connectivity = get_data(num_resid, pickled=False, pickle_data=True, example=False, random_key=True) - - ###################################################### - - print(f'\nCumulative Time:') - print(f'{time.process_time() - start}s --> Data Generated') - - observation = gym_env.reset(coordinates, obs_chemical_shifts, noes, connectivity) - - print(f'{time.process_time() - start}s --> Observation Reset') - - histories = [] - - for i in range(history_length+1): - - # If solving for fake histories off the original - if i > 0: - observation = gym_env.custom_state(fake_observation) - print(f'{time.process_time() - start}s --> Custom Reset ') - - if i <= 0: - original_observation = copy.deepcopy(observation) - - # Initial observation (nothing assigned) - histories.append(copy.deepcopy(observation)) - - terminated = False - while not terminated: - for action in observation['assign_order']: - observation, reward, terminated, total_energy = gym_env.step(action) - if not terminated: - histories.append(copy.deepcopy(observation)) - - # Generate new history - if i != history_length + 1: - fake_observation = fakehistory.generate_history(original_observation) - - # Final observation (everything assigned) - histories.append(copy.deepcopy(observation)) - - print(f"{time.process_time() - start} --> {i} Histories Complete") - - with open(f'fake_histories_r{num_resid}_{history_length}.pkl', 'wb') as f: - pickle.dump(histories, f) - - print(f'{time.process_time() - start}s --> {history_length} Histories Saved\n') \ No newline at end of file diff --git a/fractional_activations_unittest.py b/fractional_activations_unittest.py deleted file mode 100644 index b5e264e..0000000 --- a/fractional_activations_unittest.py +++ /dev/null @@ -1,68 +0,0 @@ -#unittest for fractional activations -#has to be made into a class for it to work - -import unittest - -class TestFracAct(unittest.TestCase): - - def setUp(self): - self.frac_act = FracAct() - -#test for activation loop -#expected to return the noe peak divided by the total number of different/unique noe peaks - - def test_activation_loop(self): - assign_order = [] - noe = [(0, 1), (0, 2)] - target = 0 - - result = self.frac_act.activation_loop(target, noe, assign_order) - expected = 1/3 - - self.assertEqual(result, expected) - -#test if each noe peak will be evaluated to give the fractional activation - - def test_fractional_activation_of_all_noe_peaks(self): - num_resid = 3 - noes = [[(0, 1), (0, 2)]] - - result = self.frac_act.fractional_activation_summation(num_resid, noes) - expected = [1/3, 1/3, 1/3] - - self.assertEqual(result, expected) - -#for more then one noe, the probabilities are summed - - def test_fractional_activation_summation_of_noes(self): - num_resid = 4 - noes = [(0, 1), (0, 2)], [(2,1),(1,1)], [(1,0),(3,2)] - - result = self.frac_act.fractional_activation_summation(num_resid, noes) - expected = [7/12, 13/12, 13/12, 1/4] - - self.assertEqual(result, expected) - -#according to the fractional activations calculated, should give the correct order -#if it is a tie, then lowest value goes first - - def test_fractional_activation_for_assignment_order(self): - num_resid = 4 - noes = [(0, 1), (0, 2)], [(2,1),(1,1)], [(1,0),(3,2)] - - result = self.frac_act.fractional_activation(num_resid, noes) - expected = [1, 2, 0, 3] - - self.assertEqual(result, expected) - -#if number of residues is higher than the number of peaks, then it will loop through to give first/lowest value peak ex. will give 0 for every additional residue -#if number of residues is lower than the number of peaks, then it will return the first value again (0, in this case) - - def test_fractional_activation_for_num_resid(self): - num_resid = 5 - noes = [[(0, 1), (0, 2)]], [(2,1),(1,1)], [(1,0),(3,2)] - - result = self.frac_act.fractional_activation(num_resid, noes) - expected = [1, 2, 0, 3] - - self.assertEqual(result, expected) \ No newline at end of file diff --git a/nmr/__init__.py b/nmr/__init__.py new file mode 100644 index 0000000..ee25033 --- /dev/null +++ b/nmr/__init__.py @@ -0,0 +1,32 @@ +"""NMR chemical shift assignment package.""" + +from nmr.construct import construct_graph +from nmr.models import NMRNet +from nmr.nmr_gym.fake_data import FakeDataGenerator +from nmr.nmr_gym.energy import Energy +from nmr.nmr_gym.assignment_order import FractionalActivation +from nmr.nmr_gym.gym_env import GymEnv +from nmr.nmr_gym.fake_histories import FakeHistoryGenerator +from nmr.nmr_gym.data_structures import Connectivity, HSQCPeak, NOEPeak, Protein +from nmr.nmr_gym.io import load_dataset, load_histories, save_dataset, save_histories +from nmr.nmr_gym.state import create_state_dict, validate_state_dict + +__all__ = [ + "construct_graph", + "NMRNet", + "FakeDataGenerator", + "Energy", + "FractionalActivation", + "GymEnv", + "FakeHistoryGenerator", + "Connectivity", + "HSQCPeak", + "NOEPeak", + "Protein", + "load_dataset", + "load_histories", + "save_dataset", + "save_histories", + "create_state_dict", + "validate_state_dict", +] diff --git a/nmr/construct.py b/nmr/construct.py new file mode 100644 index 0000000..0dbaabe --- /dev/null +++ b/nmr/construct.py @@ -0,0 +1,291 @@ +"""Graph construction utilities for NMR assignment.""" + +from itertools import product +from typing import Any, Dict, Tuple + +import numpy as np +import torch +from torch_geometric.data import HeteroData + + +def construct_graph(history: Dict[str, Any], device: torch.device | str) -> HeteroData: + """ + Constructs a complete heterogeneous graph from a history state. + + Args: + history: Dictionary containing coordinates, shifts, NOEs, and assignment state + device: Device to place tensors on ('cpu' or 'cuda') + + Returns: + HeteroData graph with all nodes and edges constructed + """ + data = construct_node_data(history, device) + data = construct_edges(data, device) + return data + + +def construct_node_data( + histories: Dict[str, Any], device: torch.device | str +) -> HeteroData: + """ + Builds graph nodes from input histories. + + Args: + histories: Dictionary containing coordinates, shifts, NOEs, and assignment state + device: Device to place tensors on ('cpu' or 'cuda') + + Returns: + HeteroData graph with all nodes constructed + """ + data = HeteroData() + data = _construct_data_nodes(data, histories, device) + data = _construct_triple_nodes(data, device) + data = _construct_value_nodes(data, device) + data = _construct_node_features(data, histories) + data.shift_to_assign = torch.tensor( + [int(histories["shift_to_assign"])], dtype=torch.long, device=device + ) + return data + + +def construct_edges(data: HeteroData, device: torch.device | str) -> HeteroData: + """ + Constructs all edge indices for the heterogeneous graph. + + Args: + data: HeteroData graph with nodes already constructed + device: Device to place tensors on ('cpu' or 'cuda') + + Returns: + HeteroData graph with all edges constructed + """ + num_noe = len(data["NOE"].x) + num_shift = len(data["SHIFT"].x) + num_res = len(data["RES"].x) + + # Add all edge types + _add_triple_edges( + data, 0, ("RES", "RES", "NOE"), num_noe, num_shift, num_res, device + ) + _add_triple_edges( + data, 1, ("RES", "SHIFT", "NOE"), num_noe, num_shift, num_res, device + ) + _add_triple_edges( + data, 2, ("SHIFT", "RES", "NOE"), num_noe, num_shift, num_res, device + ) + _add_triple_edges( + data, 3, ("SHIFT", "SHIFT", "NOE"), num_noe, num_shift, num_res, device + ) + _add_value_aggregation_edges(data, num_noe, num_shift, num_res, device) + + return data + + +def _construct_data_nodes( + data: HeteroData, histories: Dict[str, Any], device: torch.device | str +) -> HeteroData: + """Constructs NOE, SHIFT, and RES nodes with initial features.""" + data["NOE"].x = torch.tensor(histories["noes"], dtype=torch.float32, device=device) + data["NOE"].f = torch.zeros( + (len(histories["noes"]), 2), dtype=torch.float32, device=device + ) + data["SHIFT"].x = torch.tensor( + histories["obs_chemical_shifts"], dtype=torch.float32, device=device + ) + data["SHIFT"].f = torch.zeros( + (len(histories["obs_chemical_shifts"]), 2), + dtype=torch.float32, + device=device, + ) + data["RES"].x = torch.tensor( + histories["coordinates"], dtype=torch.float32, device=device + ) + data["RES"].f = torch.zeros( + (len(histories["coordinates"]), 2), dtype=torch.float32, device=device + ) + return data + + +def _construct_node_features(data: HeteroData, histories: Dict[str, Any]) -> HeteroData: + """Sets node features based on assignment status.""" + for i in range(len(data["SHIFT"].f)): + # is the shift being assigned right now? + if i == histories["shift_to_assign"]: + data["SHIFT"].f[i, 0] = 1 + # has the shift already been assigned? + if i in histories["assignments"].keys(): + data["SHIFT"].f[i, 1] = 1 + + for i in range(len(data["RES"].f)): + # has the residue been assigned? + if i in histories["assignments"].values(): + data["RES"].f[i] = 1 + return data + + +def _construct_triple_nodes(data: HeteroData, device: torch.device | str) -> HeteroData: + """ + Generates triple node types. Nothing is ever stored in these - they are simply placeholders to construct/use in edge types. + """ + data["TRIPLE0"].x = torch.zeros((1, 1), dtype=torch.float32, device=device) + data["TRIPLE1"].x = torch.zeros((1, 1), dtype=torch.float32, device=device) + data["TRIPLE2"].x = torch.zeros((1, 1), dtype=torch.float32, device=device) + data["TRIPLE3"].x = torch.zeros((1, 1), dtype=torch.float32, device=device) + return data + + +def _construct_value_nodes(data: HeteroData, device: torch.device | str) -> HeteroData: + """Constructs value aggregation nodes for NOE, SHIFT, and RES.""" + data["VALUE_NOE"].x = torch.zeros( + 1, 1, dtype=torch.float32, device=device + ) # has to be num graphs batched? + data["VALUE_SHIFT"].x = torch.zeros(1, 1, dtype=torch.float32, device=device) + data["VALUE_RES"].x = torch.zeros(1, 1, dtype=torch.float32, device=device) + return data + + +def _get_triple_edges( + num_noe: int, source1: int, source2: int, device: torch.device | str +) -> torch.Tensor: + """ + Grabs all index combinations from the three ranges (residue, shift, noe) and orders these into source and target indices for the edges. + Combinations occur along columns into a triple node. + + Args: + num_noe: Number of NOE nodes + source1: Number of nodes for first source type + source2: Number of nodes for second source type + device: Device to place tensors on + + Returns: + Stacked tensor of source and target node indices + """ + combo = list(product(range(num_noe), range(source1), range(source2))) + combo_tensor = torch.tensor(combo, device=device) + + # Switches tensor dimension for source (columns set up for each edge [0,0,0] --> [0],[0],[0]) + source_nodes = torch.transpose(combo_tensor, 0, 1) + # Repeats target edges for number of occurences in source (3 for the triple in this instance, to be assigned to each incoming node type) + target_nodes = torch.tensor(range(len(source_nodes[0])), device=device).repeat( + len(source_nodes), 1 + ) + + return torch.stack([source_nodes, target_nodes], dim=0) + + +def _get_pairwise_edges( + data: HeteroData, num_res: int, device: torch.device | str +) -> torch.Tensor: + """ + Grabs index combinations between shifts and residues and orders these into source and target indices for the edges. + + Args: + data: HeteroData graph with shift features + num_res: Number of residue nodes + device: Device to place tensors on + + Returns: + Stacked tensor of residue and shift indices for pairwise edges + """ + # shift = torch.arange(0, num_shift) + shift = np.nonzero(data["SHIFT"].f[:, 0])[ + 0 + ] # need to specify if it's the shift being assigned + resid = torch.arange(0, num_res, device=device) + + shift_repeats = shift.repeat_interleave(num_res) + # resid_repeats = resid.repeat(num_shift) + resid_repeats = resid + return torch.stack((resid_repeats, shift_repeats), dim=0) + + +def _add_value_aggregation_edges( + data: HeteroData, + num_noe: int, + num_shift: int, + num_res: int, + device: torch.device | str, +) -> None: + """ + Adds value aggregation edges from data nodes to value nodes. + + Args: + data: HeteroData graph to add edges to + num_noe: Number of NOE nodes + num_shift: Number of shift nodes + num_res: Number of residue nodes + device: Device to place tensors on + """ + shift = torch.arange(0, num_shift, device=device) + noe = torch.arange(0, num_noe, device=device) + resid = torch.arange(0, num_res, device=device) + + shift_repeats = torch.zeros(num_shift, device=device).long() + noe_repeats = torch.zeros(num_noe, device=device).long() + resid_repeats = torch.zeros(num_res, device=device).long() + + data["SHIFT", "SHIFT_extract", "VALUE_SHIFT"].edge_index = torch.stack( + (shift, shift_repeats), dim=0 + ) + data["NOE", "NOE_extract", "VALUE_NOE"].edge_index = torch.stack( + (noe, noe_repeats), dim=0 + ) + data["RES", "RES_extract", "VALUE_RES"].edge_index = torch.stack( + (resid, resid_repeats), dim=0 + ) + + +def _add_triple_edges( + data: HeteroData, + triple_num: int, + triple_type: Tuple[str, str, str], + num_noe: int, + num_shift: int, + num_res: int, + device: torch.device | str, +) -> None: + """ + Constructs all edges for a single triple type. + + Args: + data: HeteroData graph to add edges to + triple_num: Triple type number (0-3) + triple_type: Tuple of (source1, source2, source3) node type names + num_noe: Number of NOE nodes + num_shift: Number of shift nodes + num_res: Number of residue nodes + device: Device to place tensors on + """ + source1, source2, source3 = triple_type + + # Need to make sure I'm using the right node range (shift and residue number could vary) + num_node1 = num_res if source1 == "RES" else num_shift + num_node2 = num_res if source2 == "RES" else num_shift + + # Triple in + edges_in = _get_triple_edges(num_noe, num_node1, num_node2, device) + data[f"{source3}", "NOE_extract", f"TRIPLE{triple_num}"].edge_index = edges_in[:, 0] + data[f"{source1}", "NH1_extract", f"TRIPLE{triple_num}"].edge_index = edges_in[:, 1] + data[f"{source2}", "NH2_extract", f"TRIPLE{triple_num}"].edge_index = edges_in[:, 2] + + # Triple out (reverse in/out ordering) + edges_out = torch.stack([edges_in[1], edges_in[0]], dim=0) + data[f"TRIPLE{triple_num}", "NOE_add", f"{source3}"].edge_index = edges_out[:, 0] + data[f"TRIPLE{triple_num}", "NH1_add", f"{source1}"].edge_index = edges_out[:, 1] + data[f"TRIPLE{triple_num}", "NH2_add", f"{source2}"].edge_index = edges_out[:, 2] + + # Only residue nodes will have coordinate edges + if source1 == "RES": + data[f"TRIPLE{triple_num}", "res1_add", "RES"].edge_index = edges_out[:, 1] + if source2 == "RES": + data[f"TRIPLE{triple_num}", "res2_add", "RES"].edge_index = edges_out[:, 2] + + # Edges for self loop + data[f"TRIPLE{triple_num}", "update", f"TRIPLE{triple_num}"].edge_index = ( + torch.stack([edges_in[1, 0], edges_in[1, 0]], dim=0) + ) + + if triple_num == 0: + data["RES", "pair", "SHIFT"].edge_index = _get_pairwise_edges( + data, num_res, device + ) # pairwise edges - no message passing diff --git a/nmr/models/__init__.py b/nmr/models/__init__.py new file mode 100644 index 0000000..0808e10 --- /dev/null +++ b/nmr/models/__init__.py @@ -0,0 +1,37 @@ +""" +NMR GNN Models Package + +Provides neural network components for NMR chemical shift assignment +using heterogeneous graph neural networks with triple-based message passing. +""" + +# Core network components +from .network import NMRLayer, NMRNet + +# Triple message passing +from .triple import ( + CalculationManager, + TripleIn, + TripleMessagePass, + TripleOut, + TripleUpdate, +) + +# Prediction heads +from .heads import BatchMessagePass, PolicyCalc, ValueCalc + +__all__ = [ + # Network + "NMRLayer", + "NMRNet", + # Triple components + "CalculationManager", + "TripleIn", + "TripleUpdate", + "TripleMessagePass", + "TripleOut", + # Heads + "BatchMessagePass", + "ValueCalc", + "PolicyCalc", +] diff --git a/nmr/models/heads.py b/nmr/models/heads.py new file mode 100644 index 0000000..555197f --- /dev/null +++ b/nmr/models/heads.py @@ -0,0 +1,145 @@ +""" +Value and policy prediction heads for the NMR GNN. + +Contains components for aggregating graph information and predicting +value estimates and policy distributions. +""" + +import torch +import torch.nn as nn +from torch_geometric.nn import MessagePassing + + +class BatchMessagePass(MessagePassing): + """ + Message passing for aggregating node features to batch-level representations. + Used by value and policy heads to collect information across the graph. + """ + + def __init__(self, aggr, device): + super().__init__(aggr=aggr) + self.device = device + + self.noe_reduce = nn.Linear(3, 1, device=self.device) + self.shift_reduce = nn.Linear(2, 1, device=self.device) + self.res_reduce = nn.Linear(5, 1, device=self.device) + + def forward(self, x_source, x_target, edge_index): + # SIZE (n, m) (source, target) + return self.propagate( + edge_index=edge_index, + x=(x_source, x_target), + size=(x_source.size(0), x_target.size(0)), + ) + + def message(self, x_j): + if len(x_j[0]) == 3: + return self.noe_reduce(x_j) + if len(x_j[0]) == 2: + return self.shift_reduce(x_j) + if len(x_j[0]) == 5: + return self.res_reduce(x_j) + + def update(self, aggr_out, x): + return aggr_out + + +class ValueCalc(nn.Module): + """ + Value function estimation head. + + Aggregates information from NOE, SHIFT, and RES nodes to predict + a scalar value representing the quality of the current state. + """ + + def __init__(self, device): + super().__init__() + self.device = device + + self.batch_message = BatchMessagePass(aggr="mean", device=self.device) + self.hidden = 64 + + self.testmlp = nn.Sequential( + nn.Linear(3, self.hidden, device=self.device), + nn.ReLU(), + nn.Linear(self.hidden, 1, device=self.device), + ) + + def get_aggr(self, x_source, x_target, edge_index): + aggr = self.batch_message(x_source, x_target, edge_index) + return aggr + + def calc_value(self, data): + noe = self.get_aggr( + data["NOE"].x, + data["VALUE_NOE"].x, + data["NOE", "NOE_extract", "VALUE_NOE"].edge_index, + ) + shift = self.get_aggr( + data["SHIFT"].x, + data["VALUE_SHIFT"].x, + data["SHIFT", "SHIFT_extract", "VALUE_SHIFT"].edge_index, + ) + resid = self.get_aggr( + data["RES"].x, + data["VALUE_RES"].x, + data["RES", "RES_extract", "VALUE_RES"].edge_index, + ) + + concat_aggr = torch.cat((noe, shift, resid), dim=-1) + value = self.testmlp(concat_aggr) + return value + + +class PolicyCalc(nn.Module): + """ + Policy prediction head. + + Computes action probabilities (which residue to assign to the current shift) + based on pairwise distances between residue and shift features. + """ + + def __init__(self, device): + super().__init__() + self.RES_NH = slice(3, 5) + self.device = device + + def calc_policy(self, data): + """ + Calculate policy logits for each graph in the batch. + + For each graph, computes squared distances between the shift being assigned + and all residues. Returns logits (not probabilities). + + Args: + data: Batched HeteroData with RES nodes, SHIFT nodes, and ("RES", "pair", "SHIFT") edges + + Returns: + List of logit tensors (negative squared distances), one per graph in batch + """ + # Unbatch the data + data_unbatched = data.to_data_list() + + # Compute the policy logits for each item + policies = [] + for item in data_unbatched: + # Extract the shift being assigned and all residues + # The pairwise edges already connect only the shift being assigned to all residues + edge_index = item["RES", "pair", "SHIFT"].edge_index + + # Get residue NH features (H1, N15) + resid_features = item["RES"].x[edge_index[0]][:, self.RES_NH] + + # Get shift NH features (H1, N15) + shift_features = item["SHIFT"].x[edge_index[1]] + + # Compute squared pairwise distances (one per residue) + # squared_distance = (H1_res - H1_shift)^2 + (N15_res - N15_shift)^2 + squared_distances = torch.sum((resid_features - shift_features) ** 2, dim=-1) + + # Use negative squared distances as logits (closer = higher logit) + logits = -squared_distances + + policies.append(logits.unsqueeze(0)) # Add batch dimension for consistency + + return policies diff --git a/nmr/models/network.py b/nmr/models/network.py new file mode 100644 index 0000000..56d3701 --- /dev/null +++ b/nmr/models/network.py @@ -0,0 +1,205 @@ +""" +Top-level NMR GNN network architecture. + +Combines triple-based message passing with value and policy heads +to create the complete neural network for NMR assignment. +""" + +import pickle + +import matplotlib.pyplot as plt +import torch +import torch.nn as nn +from torch_geometric.loader import DataLoader + +from .triple import CalculationManager, TripleIn, TripleOut, TripleUpdate +from .heads import ValueCalc, PolicyCalc + + +class NMRLayer(nn.Module): + """ + Single layer of NMR-specific message passing. + + Orchestrates the full triple update pipeline: + 1. Extract features from nodes into triples (TripleIn) + 2. Update triple representations via message passing (TripleUpdate) + 3. Propagate updated information back to nodes (TripleOut) + + Handles all four triple types: + - TRIPLE0: (RES, RES, NOE) + - TRIPLE1: (RES, SHIFT, NOE) + - TRIPLE2: (SHIFT, RES, NOE) + - TRIPLE3: (SHIFT, SHIFT, NOE) + """ + + def __init__(self, device): + super(NMRLayer, self).__init__() + self.triple_in = TripleIn() + self.triple_self = TripleUpdate(device) + self.triple_out = TripleOut() + self.calc_manager = CalculationManager() + + self.NOE = slice(0, 3) + self.RES_XYZ = slice(0, 3) + self.RES_NH = slice(3, 5) + self.SHIFT_NH = slice(0, 2) + + def update_noes(self, data, noe_delta, feature, i): + # 'i' is for the triple type in all of these cases (TRIPLE0 --> RES RES NOE, etc.) + # Should really simplify this to be one or the other from the very start (string --> number identification) + data = self.triple_out.update_data( + data, noe_delta, feature, (f"TRIPLE{i}", "NOE_add", "NOE"), self.NOE + ) + return data + + def update_coordinates(self, data, dist_delta1, dist_delta2, feature1, feature2, i): + data = self.triple_out.update_data( + data, + dist_delta1, + feature1, + (f"TRIPLE{i}", "res1_add", "RES"), + self.RES_XYZ, + update_f=False, + ) + data = self.triple_out.update_data( + data, + dist_delta2, + feature2, + (f"TRIPLE{i}", "res2_add", "RES"), + self.RES_XYZ, + update_f=False, + ) + return data + + def update_shifts( + self, data, shift_delta1, shift_delta2, feature1, feature2, target1, target2, i + ): + # Selects target indexing for residue or shift node + target_range1 = self.RES_NH if target1 == "RES" else self.SHIFT_NH + target_range2 = self.RES_NH if target2 == "RES" else self.SHIFT_NH + + data = self.triple_out.update_data( + data, + shift_delta1, + feature1, + (f"TRIPLE{i}", "NH1_add", target1), + target_range1, + ) + data = self.triple_out.update_data( + data, + shift_delta2, + feature2, + (f"TRIPLE{i}", "NH2_add", target2), + target_range2, + ) + return data + + def do_updates(self, data, triple_type, i): + # Triple in + shift_x1, coord_x1, shift_x2, coord_x2, noe_x, shift_f1, shift_f2, noe_f = ( + self.triple_in.construct_triple( + data, + (triple_type[0], "NH1_extract", f"TRIPLE{i}"), + (triple_type[1], "NH2_extract", f"TRIPLE{i}"), + ("NOE", "NOE_extract", f"TRIPLE{i}"), + ) + ) + + # Self update and output + if triple_type == ("RES", "RES", "NOE"): + ( + delta_shift1, + delta_shift2, + delta_noe, + delta_dist1, + delta_dist2, + deltaf1, + deltaf2, + deltaf3, + ) = self.triple_self( + shift_x1, + coord_x1, + shift_x2, + coord_x2, + noe_x, + shift_f1, + shift_f2, + noe_f, + data[f"TRIPLE{i}", "update", f"TRIPLE{i}"].edge_index, + ) + data = self.update_noes(data, delta_noe, deltaf3, i) + data = self.update_coordinates( + data, delta_dist1, delta_dist2, deltaf1, deltaf2, i + ) + data = self.update_shifts( + data, + delta_shift1, + delta_shift2, + deltaf1, + deltaf2, + triple_type[0], + triple_type[1], + i, + ) + + if triple_type in ( + ("SHIFT", "RES", "NOE"), + ("RES", "SHIFT", "NOE"), + ("SHIFT", "SHIFT", "NOE"), + ): + delta_shift1, delta_shift2, delta_noe, deltaf1, deltaf2, deltaf3 = ( + self.triple_self( + shift_x1, + coord_x1, + shift_x2, + coord_x2, + noe_x, + shift_f1, + shift_f2, + noe_f, + data[f"TRIPLE{i}", "update", f"TRIPLE{i}"].edge_index, + ) + ) + data = self.update_noes(data, delta_noe, deltaf3, i) + data = self.update_shifts( + data, + delta_shift1, + delta_shift2, + deltaf1, + deltaf2, + triple_type[0], + triple_type[1], + i, + ) + return data + + def forward(self, data): + data = self.do_updates(data, ("RES", "RES", "NOE"), 0) + data = self.do_updates(data, ("RES", "SHIFT", "NOE"), 1) + data = self.do_updates(data, ("SHIFT", "RES", "NOE"), 2) + data = self.do_updates(data, ("SHIFT", "SHIFT", "NOE"), 3) + return data + + +class NMRNet(nn.Module): + """ + Complete NMR GNN model combining message passing with prediction heads. + + Stacks NMRLayer(s) for graph message passing, then uses ValueCalc and + PolicyCalc heads to predict state value and action probabilities. + """ + + def __init__(self, device): + super().__init__() + self.value = ValueCalc(device) + self.policy = PolicyCalc(device) + + self.nmr = nn.Sequential( + NMRLayer(device), + ) + + def forward(self, data): + out_data = self.nmr(data) + value = self.value.calc_value(out_data) + policy = self.policy.calc_policy(out_data) + return out_data, value, policy diff --git a/nmr/models/triple.py b/nmr/models/triple.py new file mode 100644 index 0000000..4e2115f --- /dev/null +++ b/nmr/models/triple.py @@ -0,0 +1,342 @@ +""" +Triple-based graph message passing components. + +Contains classes for constructing, updating, and propagating messages +through triple nodes in the heterogeneous NMR graph. +""" + +import torch +import torch.nn as nn +from torch_geometric.nn import MessagePassing + + +class CalculationManager: + """ + Common calculations used across triple and pairwise comparisons. + """ + + def __init__(self): + # Input index organization - slices to keep proper tensor dimension + self.NOE_N1 = slice(0, 1) + self.NOE_H1 = slice(1, 2) + self.NOE_H2 = slice(2, 3) + self.RES_XYZ = slice(0, 3) + self.SHIFT_N = slice(0, 1) + self.SHIFT_H = slice(1, 2) + + def calc_noe_difference(self, x1, x2, noe): + """ + Calculates shift difference between NOE and residue/measured shifts (direct/indirect only reverse option is available by index). + """ + diff_N = noe[:, self.NOE_N1] - x1[:, self.SHIFT_N] # N [n, 1] + diff_H1 = noe[:, self.NOE_H1] - x1[:, self.SHIFT_H] # H' [n, 1] + diff_H2 = noe[:, self.NOE_H2] - x2[:, self.SHIFT_H] # H" [n, 1] + return diff_N, diff_H1, diff_H2 + + def calc_shift_difference(self, x1, x2): + """ + Calculates shift difference between residue/measured shifts. + """ + diff_N = x1[:, self.SHIFT_N] - x2[:, self.SHIFT_N] # N [n, 1] + diff_H = x1[:, self.SHIFT_H] - x2[:, self.SHIFT_H] # H [n, 1] + return diff_N, diff_H + + def calc_res_distance(self, x1, x2): + """ + Calculates relative distance between residues and this value squared for equivariant calculations. + """ + rel_dist = x1 - x2 # [n, 3] + dist2 = torch.norm(rel_dist, dim=-1, keepdim=True) ** 2 # [n, 1] + return rel_dist, dist2 + + +class TripleIn(nn.Module): + def __init__(self): + super().__init__() + # Input index organization - slices to keep proper tensor dimension + self.RES_XYZ = slice(0, 3) + self.RES_NH = slice(3, 5) + + def grab_node(self, data, node_type, edge_type): + """ + Source node to triple via indexing. + For residue data, extracts shift and coordinate information to process separately. + """ + # Spatial features + xi = data[node_type].x[data[edge_type].edge_index[0]] + # Non-spatial features + fi = data[node_type].f[data[edge_type].edge_index[0]] + + # Splits coordinate and shift features if combined + xj = None + if edge_type[0] == "RES": + xi, xj = xi[:, self.RES_NH], xi[:, self.RES_XYZ] + return xi, xj, fi + + def construct_triple(self, data, edge_type1, edge_type2, edge_type3): + """ + Constructs triple based on type. + """ + # Measured shift or residue (shifts will return nonetype value) + x1, x12, f1 = self.grab_node( + data, node_type=edge_type1[0], edge_type=edge_type1 + ) + x2, x22, f2 = self.grab_node( + data, node_type=edge_type2[0], edge_type=edge_type2 + ) + # NOE (does not require any value split) + x3, _, f3 = self.grab_node(data, node_type=edge_type3[0], edge_type=edge_type3) + return x1, x12, x2, x22, x3, f1, f2, f3 + + +class TripleUpdate(MessagePassing): + """ + Message passing class for triple self updates. + """ + + def __init__(self, device): + super().__init__(aggr="add") + self.calc_manager = CalculationManager() + self.device = device + + self.hidden = 64 + + self.mlp1 = nn.Sequential( + nn.Linear(12, self.hidden, device=self.device), + nn.ReLU(), + nn.Linear(self.hidden, 19, device=self.device), + nn.LayerNorm(19, device=self.device), + ) + + self.mlp2 = nn.Sequential( + nn.Linear(11, self.hidden, device=self.device), + nn.ReLU(), + nn.Linear(self.hidden, 13, device=self.device), + nn.LayerNorm(13, device=self.device), + ) + + def resresnoe_message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j, x22_j): + """ + Constructs message for (res, res, NOE) triple type. + """ + # Residue distances + rel_dist, dist2 = self.calc_manager.calc_res_distance(x12_j, x22_j) + + # Differences relative to NOE shifts (N, H', H") + diff1, diff2, diff3 = self.calc_manager.calc_noe_difference(x1_j, x2_j, x3_j) + + # Shift differences + diff4, diff5 = self.calc_manager.calc_shift_difference(x1_j, x2_j) + + # Input for MLP + # (N, H', H", N, H, dist2, features) + mlp_in = torch.cat( + (diff1, diff2, diff3, diff4, diff5, dist2, f1_j, f2_j, f3_j), dim=-1 + ) # [n, 9] + + # Output from MLP + # Out should include values for each 'change' wanting to make + # (N, H', H", N1, H1, N2, H2, dist1, dist2, features) + mlp_out = self.mlp1(mlp_in) # [n, 16] + + # NOE deltas [n, 1] + delta1x = diff1 * mlp_out[:, 0:1] # N + delta2x = diff2 * mlp_out[:, 1:2] # H' + delta3x = diff3 * mlp_out[:, 2:3] # H" + + # SHIFT deltas residue [n, 1] + delta4x = diff4 * mlp_out[:, 3:4] # N1 + delta5x = diff5 * mlp_out[:, 4:5] # H1 + delta6x = diff4 * mlp_out[:, 5:6] # N2 + delta7x = diff5 * mlp_out[:, 6:7] # H2 + + # DISTANCE deltas [n, 3] + delta12x = rel_dist * (mlp_out[:, 7:10]) + delta13x = rel_dist * (mlp_out[:, 10:13]) + + # FEATURE deltas [n, 1] + delta1f = mlp_out[:, 13:15] + delta2f = mlp_out[:, 15:17] + delta3f = mlp_out[:, 17:] + return torch.cat( + ( + delta1x, + delta2x, + delta3x, + delta4x, + delta5x, + delta6x, + delta7x, + delta12x, + delta13x, + delta1f, + delta2f, + delta3f, + ), + dim=-1, + ) # [n, 16] [n, 19] + + def shiftshiftnoe_message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j): + """ + Constructs message for (shift, shift, NOE), (shift, res, NOE), (res, shift, NOE) triple types. + """ + # Differences relative to NOE shifts (N, H', H") + diff1, diff2, diff3 = self.calc_manager.calc_noe_difference(x1_j, x2_j, x3_j) + + # Shift differences + diff4, diff5 = self.calc_manager.calc_shift_difference(x1_j, x2_j) + + # Input for MLP + # (N, H', H", N, H, dist2, features) + mlp_in = torch.cat( + (diff1, diff2, diff3, diff4, diff5, f1_j, f2_j, f3_j), dim=-1 + ) # [n, 8] + + # Output from MLP + # Out should include values for each 'change' wanting to make + # (N, H', H", N1, H1, N2, H2, features) + mlp_out = self.mlp2(mlp_in) # [n, 10] + + # NOE deltas [n, 1] + delta1x = diff1 * mlp_out[:, 0:1] # N + delta2x = diff2 * mlp_out[:, 1:2] # H' + delta3x = diff3 * mlp_out[:, 2:3] # H" + + # SHIFT deltas residue [n, 1] + delta4x = diff4 * mlp_out[:, 3:4] # N1 + delta5x = diff5 * mlp_out[:, 4:5] # H1 + delta6x = diff4 * mlp_out[:, 5:6] # N2 + delta7x = diff5 * mlp_out[:, 6:7] # H2 + + # FEATURE deltas [n, 1] + delta1f = mlp_out[:, 7:9] + delta2f = mlp_out[:, 9:11] + delta3f = mlp_out[:, 11:] + return torch.cat( + ( + delta1x, + delta2x, + delta3x, + delta4x, + delta5x, + delta6x, + delta7x, + delta1f, + delta2f, + delta3f, + ), + dim=-1, + ) # [n, 10] [n, 13] + + def forward(self, x1, x12, x2, x22, x3, f1, f2, f3, edge_index): + out = self.propagate( + edge_index, x1=x1, x2=x2, x3=x3, f1=f1, f2=f2, f3=f3, x12=x12, x22=x22 + ) # not sure how to deal with size here + return out + + def message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j=None, x22_j=None): + # residue based triple type will have two sets of coordinates (x12_j and x22_j) where shifts will have nonetype + if x12_j != None and x22_j != None: + return self.resresnoe_message( + x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j, x22_j + ) + else: + return self.shiftshiftnoe_message(x1_j, x2_j, x3_j, f1_j, f2_j, f3_j) + + def update(self, aggr_out, x12, x22): + # residue based triple type will have two sets of coordinates (x12_j and x22_j) where shifts will have nonetype + if x12 != None and x22 != None: + # output shapes: noe[n, 3], shift1[n, 2], shift2[n, 2], dist1[n, 3], dist2[n, 3], f1[n, 2], f2[n, 2], f3[n, 2] + ( + delta_noe, + delta_shift1, + delta_shift2, + delta_dist1, + delta_dist2, + delta_f1, + delta_f2, + delta_f3, + ) = ( + aggr_out[:, 0:3], + aggr_out[:, 3:5], + aggr_out[:, 5:7], + aggr_out[:, 7:10], + aggr_out[:, 10:13], + aggr_out[:, 13:15], + aggr_out[:, 15:17], + aggr_out[:, 17:], + ) + return ( + delta_shift1, + delta_shift2, + delta_noe, + delta_dist1, + delta_dist2, + delta_f1, + delta_f2, + delta_f3, + ) + else: + # output shapes: noe[n, 3], shift1[n, 2], shift2[n, 2], f1[n, 2], f2[n, 2], f3[n, 2] + delta_noe, delta_shift1, delta_shift2, delta_f1, delta_f2, delta_f3 = ( + aggr_out[:, 0:3], + aggr_out[:, 3:5], + aggr_out[:, 5:7], + aggr_out[:, 7:9], + aggr_out[:, 9:11], + aggr_out[:, 11:], + ) + return delta_shift1, delta_shift2, delta_noe, delta_f1, delta_f2, delta_f3 + + +class TripleMessagePass(MessagePassing): + """ + Standard message passing class for outgoing triple messages. + """ + + def __init__(self, aggr): + super().__init__(aggr=aggr) + + def forward(self, x_source, x_target, edge_index): + # SIZE (n, m) (source, target) + return self.propagate( + edge_index=edge_index, + x=(x_source, x_target), + size=(x_source.size(0), x_target.size(0)), + ) + + def message(self, x_j): + return x_j + + def update(self, aggr_out, x): + # non-spatial features can also be included in same call but requires more work if they're not being updated + # (don't want to do residue feature updates 2x - only call once on shifts or coordinates) + # x_val, f_val = aggr_out[:, :len(x[1][1])], aggr_out[:, len(x[1][1]):] + # outf = f_val + f[1] + # outx = x_val + x[1] + out = aggr_out + x[1] + return out + + +class TripleOut(nn.Module): + def __init__(self): + super().__init__() + self.triple_message = TripleMessagePass(aggr="add") + + def update_data(self, data, x_source, f_source, edge_type, x_index, update_f=True): + """ + Calls message passing and directly updates the heterodata object based on target. + Set to update non-spatial features in all cases, but can be turned off case by case to prevent duplicate updates. + """ + source, edge, target = edge_type + edge_index = data[edge_type].edge_index + + # Message passing/update for spatial features + data[target].x[:, x_index] = self.triple_message( + x_source, data[target].x[:, x_index], edge_index + ) + + # Message passing/update for non-spatial features if needed (don't want a duplicate update for residue features) + if update_f: + data[target].f = self.triple_message(f_source, data[target].f, edge_index) + return data diff --git a/nmr/nmr_gym/__init__.py b/nmr/nmr_gym/__init__.py new file mode 100644 index 0000000..d3e0468 --- /dev/null +++ b/nmr/nmr_gym/__init__.py @@ -0,0 +1,24 @@ +"""Environment components for NMR chemical shift assignment. + +This module contains all the components related to the reinforcement learning +environment, including synthetic data generation, energy calculations, and +the Gymnasium interface. +""" + +from .fake_data import Protein, HSQCPeak, NOEPeak, Connectivity, FakeDataGenerator +from .energy import Energy +from .assignment_order import FractionalActivation +from .gym_env import GymEnv +from .fake_histories import FakeHistoryGenerator + +__all__ = [ + 'Protein', + 'HSQCPeak', + 'NOEPeak', + 'Connectivity', + 'FakeDataGenerator', + 'Energy', + 'FractionalActivation', + 'GymEnv', + 'FakeHistoryGenerator', +] diff --git a/assignment_order.py b/nmr/nmr_gym/assignment_order.py similarity index 100% rename from assignment_order.py rename to nmr/nmr_gym/assignment_order.py diff --git a/nmr/nmr_gym/data_structures.py b/nmr/nmr_gym/data_structures.py new file mode 100644 index 0000000..1f04e4b --- /dev/null +++ b/nmr/nmr_gym/data_structures.py @@ -0,0 +1,45 @@ +""" +NMR data structures used across the package. + +This module defines the core NamedTuple data structures for representing +NMR experimental data, protein structures, and connectivity information. +These structures provide type-safe, immutable containers for data exchange +between modules. + +Data Structures: +- HSQCPeak: HSQC peak with H1 and N15 chemical shifts (ppm) +- NOEPeak: NOE crosspeak indicating spatial proximity +- Protein: Residue with 3D coordinates and predicted chemical shifts +- Connectivity: Spatial connectivity between two atoms +""" + +from typing import NamedTuple + + +class HSQCPeak(NamedTuple): + """HSQC peak with H1 and N15 chemical shifts in ppm.""" + H1: float + N15: float + + +class NOEPeak(NamedTuple): + """NOE crosspeak with two H1-N15 pairs indicating spatial proximity.""" + H1: float + N15: float + H2: float + + +class Protein(NamedTuple): + """Protein residue with 3D coordinates and predicted chemical shifts.""" + x: float + y: float + z: float + H1: float + N15: float + + +class Connectivity(NamedTuple): + """Spatial connectivity between two atoms within cutoff distance.""" + atom1: float + atom2: float + distance: float diff --git a/energy.py b/nmr/nmr_gym/energy.py similarity index 100% rename from energy.py rename to nmr/nmr_gym/energy.py diff --git a/nmr/nmr_gym/fake_data.py b/nmr/nmr_gym/fake_data.py new file mode 100644 index 0000000..f8fa048 --- /dev/null +++ b/nmr/nmr_gym/fake_data.py @@ -0,0 +1,409 @@ +""" +Consolidated fake NMR data generation module. + +This module provides utilities for generating synthetic NMR data including: +- 3D protein coordinates generation +- HSQC chemical shift peak generation +- NOE distance constraint generation +- Dataset generation with explicit file paths (see io.py for save/load) + +The generation process simulates a simplified NMR assignment scenario: +1. Generate 3D protein structure (scaled to realistic dimensions using Flory scaling) +2. Generate HSQC peaks (observed chemical shifts in ppm) +3. Generate NOE crosspeaks (distance constraints from spatial proximity) +4. Add Gaussian noise to simulate experimental uncertainty + +Physical Parameters: +- Cutoff distance: 0.5 nm (~5 Angstroms) for NOE generation +- H1 shift range: 6-10 ppm (typical for amide protons) +- N15 shift range: 100-135 ppm (typical for backbone nitrogens) +- Radius of gyration: Rg = 0.2 * N^0.4 (Flory scaling for globular proteins) + +Typical usage: + >>> # Generate dataset + >>> generator = FakeDataGenerator(num_resid=10) + >>> coords, shifts, noes, conn = generator.generate_data(pickle_data=False) + >>> + >>> # Save dataset with metadata (see io.py) + >>> from nmr.nmr_gym.io import save_dataset + >>> metadata = {"version": "1.0", "num_resid": 10} + >>> save_dataset("dataset.pkl", coords, shifts, noes, conn, metadata) + >>> + >>> # Load dataset (see io.py) + >>> from nmr.nmr_gym.io import load_dataset + >>> coords, shifts, noes, conn, metadata = load_dataset("dataset.pkl") + >>> + >>> # Create state dictionary for RL environment (see state.py) + >>> from nmr.nmr_gym.state import create_state_dict + >>> state = create_state_dict(coords, shifts, noes, conn) +""" + +import pickle +from typing import List, Optional, Tuple + +import numpy as np + +from .data_structures import Connectivity, HSQCPeak, NOEPeak, Protein + + +class FakeDataGenerator: + """ + Generates synthetic NMR data for protein chemical shift assignment training. + + Creates realistic synthetic datasets including 3D protein coordinates, HSQC chemical + shift peaks, NOE distance crosspeaks, and connectivity information. Used for supervised + pre-training of graph neural networks before reinforcement learning fine-tuning. + + The generation process simulates a simplified NMR assignment scenario where: + - pred_coordinates: 3D structure from a predicted model + - obs_chemical_shifts: Experimental/observed shifts that need to be assigned + - pred_chemical_shifts: Shifts calculated from the predicted structure + - noes: NOE crosspeaks from spatial proximity in predicted structure + - connectivity: Residue pairs within distance cutoff + + Typical usage: + >>> generator = FakeDataGenerator(num_resid=10) + >>> coords, shifts, noes, conn = generator.generate_data(pickle_data=False) + >>> print(f"Generated {len(coords)} residues with {len(noes)} NOEs") + + For detailed documentation, see docs/fake-data-guide.md + """ + + def __init__(self, num_resid: int): + """ + Initialize fake data generator. + + Args: + num_resid: Number of residues in the synthetic protein (must be > 0) + + Raises: + ValueError: If num_resid <= 0 + + Attributes: + num_resid: Number of residues + nshift_min: Minimum N15 chemical shift (ppm) - default 100 + nshift_max: Maximum N15 chemical shift (ppm) - default 135 + hshift_min: Minimum H1 chemical shift (ppm) - default 6 + hshift_max: Maximum H1 chemical shift (ppm) - default 10 + cutoff: Distance cutoff for NOE generation (nm) - default 0.5 + """ + if num_resid <= 0: + raise ValueError(f"num_resid must be > 0, got {num_resid}") + + self.num_resid: int = num_resid + self.nshift_min = 100 + self.nshift_max = 135 + self.hshift_min = 6 + self.hshift_max = 10 + self.cutoff = 0.5 + + def sample_unit(self, num_points: int, num_sides: int, min_len: float = 0, max_len: float = 1) -> np.ndarray: + """ + Samples points from unit square or cube. + + Args: + num_points: Number of points to sample + num_sides: Dimensionality (2 for square, 3 for cube) + min_len: Minimum coordinate value + max_len: Maximum coordinate value + + Returns: + Array of shape (num_points, num_sides) with uniformly sampled coordinates + """ + return np.random.uniform(low=min_len, high=max_len, size=(num_points, num_sides)) + + def scale_unit(self, coordinates: np.ndarray) -> np.ndarray: + """ + Scale coordinates to realistic protein dimensions using radius of gyration. + + Transforms uniformly sampled coordinates to match expected globular protein + compactness based on the Flory scaling relationship: Rg = R * N^v + + Args: + coordinates: Uniformly sampled coordinates in [0,1]³ + + Returns: + Scaled coordinates representing realistic protein structure + + Reference: + PDB:1CRC (~100 resid, globular, ~3.2 nm diameter) + DOI: 10.1142/S021972002050050X + Scaling factor (v) = 0.4 for globular proteins + R = 0.2 nm (empirical constant) + """ + # Calculate target radius of gyration based on residue count + # Rg = 0.2 * N^0.4 gives realistic protein dimensions + radius_gyration = 0.2 * (self.num_resid ** 0.4) + + # Calculate current center of mass (should be ~0.5 for uniform [0,1] sampling) + com = np.mean(coordinates, axis=0) + + # Calculate distances from COM for each residue + distances_sampled = np.linalg.norm(coordinates - com, axis=1) + + # Calculate current radius of gyration: sqrt(mean(r²)) + radius_sampled = np.sqrt(np.mean((distances_sampled ** 2))) + + # Scale factor to match target Rg + scale = radius_gyration / radius_sampled + + # Apply scaling: translate to origin, scale, translate back + coordinates_scaled = com + ((coordinates - com) * scale) + + return coordinates_scaled + + def create_hsqc(self) -> np.ndarray: + """ + Sample points within NMR window for H1 and N15 to create HSQC peaks. + + Returns: + Array of shape (num_resid, 2) containing [H1, N15] shifts in ppm + """ + # H shifts + h_shifts = self.sample_unit(self.num_resid, num_sides=1, min_len=self.hshift_min, max_len=self.hshift_max) + # N shifts + n_shifts = self.sample_unit(self.num_resid, num_sides=1, min_len=self.nshift_min, max_len=self.nshift_max) + + return np.hstack((h_shifts, n_shifts)) + + def add_noise(self, points: np.ndarray, scale: float = 0.1) -> np.ndarray: + """ + Add Gaussian noise to simulate experimental uncertainty. + + Random selection from normal (Gaussian) distribution of 'scale' width centered at 0. + The noise has the same shape as the input points. + + Args: + points: Array of values to perturb + scale: Standard deviation of Gaussian noise (must be > 0) + + Returns: + Perturbed points with added noise + + Raises: + ValueError: If scale <= 0 + + Example: + point = [0, 1, 2], noise = [0.1, -0.05, 0.2], noisy_point = [0.1, 0.95, 2.2] + """ + if scale <= 0: + raise ValueError(f"scale must be > 0, got {scale}") + + noise = np.random.normal(loc=0, scale=scale, size=points.shape) + noisy_point = points + noise + + return noisy_point + + def calculate_dist(self, p1: np.ndarray, p2: np.ndarray) -> float: + """ + Calculate Euclidean distance between two points. + + Args: + p1: First point coordinates + p2: Second point coordinates + + Returns: + Euclidean distance between p1 and p2 + """ + return np.linalg.norm((p2 - p1)) + + def create_noes(self, coordinates: np.ndarray, shifts: np.ndarray, random_key: bool = False) -> Tuple[np.ndarray, np.ndarray]: + """ + Generate NOE crosspeaks from spatially close residues. + + NOE (Nuclear Overhauser Effect) crosspeaks indicate spatial proximity + between residues. This generates synthetic NOEs for all residue pairs + within the distance cutoff. + + Args: + coordinates: 3D coordinates for each residue (from predicted structure) + shifts: Chemical shift values [H1, N15] for each residue + random_key: If True, shuffle shifts to create non-identity mapping + (creates assignment problem). If False, use 1:1 mapping + where shift index equals residue index. + + Returns: + Tuple of (noisy_noes, pred_chemical_shifts): + - noisy_noes: NOE crosspeaks [H1_i, N15_i, H1_j] with noise + - pred_chemical_shifts: Predicted shifts with noise + + Note: + May return empty NOE list if protein is small or linear. + Cutoff is self.cutoff (default 0.5 nm = ~5 Angstroms). + """ + # If random_key=True, shuffle shifts to create assignment problem + if random_key: + shifts = np.random.permutation(shifts) + + # Generate NOEs for all residue pairs within cutoff distance + noe_list = [] + for i, atom1 in enumerate(coordinates): + for j, atom2 in enumerate(coordinates): + if i != j: + dist = self.calculate_dist(atom1, atom2) + # If residues are close, create NOE crosspeak + if dist < self.cutoff: + # NOE format: [H1 of res_i, N15 of res_i, H1 of res_j] + noe = list(shifts[i][:]) # H1, N15 from residue i + noe.append(shifts[j][0]) # H1 from residue j + noe_list.append(noe) + + # Add experimental noise to NOEs and predicted shifts + noisy_noes = self.add_noise(np.array(noe_list), scale=0.01) + pred_chemical_shifts = self.add_noise(np.array(shifts), scale=0.1) + + return noisy_noes, pred_chemical_shifts + + def calculate_connectivity(self, coordinates: np.ndarray) -> List[Tuple[int, int, float]]: + """ + Calculate close contacts based on protein coordinates. + + Finds all residue pairs within cutoff distance and returns their + connectivity information. + + Args: + coordinates: 3D coordinates for each residue + + Returns: + List of tuples (atom1_index, atom2_index, distance) for pairs + within cutoff distance + """ + contacts = [] + + for i, atom1 in enumerate(coordinates): + for j, atom2 in enumerate(coordinates): + if i != j: + dist = self.calculate_dist(atom1, atom2) + if dist < self.cutoff: + contacts.append((i, j, dist)) + + return contacts + + def generate_data_arrays(self, random_key: bool) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, List[Tuple]]: + """ + Generate all components of the synthetic NMR system as individual arrays. + + Args: + random_key: If True, shuffle shifts to create non-identity mapping + + Returns: + Tuple of (pred_coordinates, obs_chemical_shifts, pred_chemical_shifts, noes, connectivity): + - pred_coordinates: 3D structure [x, y, z] from predicted model + - obs_chemical_shifts: Observed/experimental shifts [H1, N15] to be assigned + - pred_chemical_shifts: Predicted shifts [H1, N15] from structure + - noes: NOE crosspeaks [H1, N15, H2] + - connectivity: List of (atom1, atom2, distance) tuples + """ + # Generate 3D structure from predicted model + pred_coordinates = self.sample_unit(self.num_resid, num_sides=3) + pred_coordinates = self.scale_unit(pred_coordinates) + + # Generate observed/experimental chemical shifts + obs_chemical_shifts = self.create_hsqc() + + # Generate NOEs and predicted shifts from structure + noes, pred_chemical_shifts = self.create_noes(pred_coordinates, obs_chemical_shifts, random_key=random_key) + + # Calculate connectivity from structure + connectivity = self.calculate_connectivity(pred_coordinates) + + return pred_coordinates, obs_chemical_shifts, pred_chemical_shifts, noes, connectivity + + def order_data(self, pred_coordinates: np.ndarray, obs_chemical_shifts: np.ndarray, + pred_chemical_shifts: np.ndarray, noes: np.ndarray, + connectivity: List[Tuple]) -> Tuple[List[Protein], List[HSQCPeak], List[NOEPeak], List[Connectivity]]: + """ + Compile data into lists of named tuples with one object per residue. + + Args: + pred_coordinates: 3D structure coordinates + obs_chemical_shifts: Observed shift values + pred_chemical_shifts: Predicted shift values + noes: NOE crosspeak data + connectivity: Connectivity tuples + + Returns: + Tuple of (pred_coordinates, obs_chemical_shifts, noes, connectivity) as named tuples + """ + pred_coordinates = [Protein(x=resid[0], y=resid[1], z=resid[2], H1=shift[0], N15=shift[1]) + for resid, shift in zip(pred_coordinates, pred_chemical_shifts)] + obs_chemical_shifts = [HSQCPeak(H1=shift[0], N15=shift[1]) for shift in obs_chemical_shifts] + noes = [NOEPeak(H1=shift[0], N15=shift[1], H2=shift[2]) for shift in noes] + connectivity = [Connectivity(atom1=contact[0], atom2=contact[1], distance=contact[2]) + for contact in connectivity] + + return pred_coordinates, obs_chemical_shifts, noes, connectivity + + def generate_data(self, pickle_data: bool = True, example: bool = True, + random_key: bool = False) -> Optional[Tuple[List, List, List, List]]: + """ + Generate all fake data, order it in lists of namedtuples, and optionally pickle. + + NOTE: This method uses deprecated dump_pickle/load_pickle with default filenames. + For new code, use generate_data_arrays() followed by save_dataset() with explicit path. + + Args: + pickle_data: If True, save to disk using deprecated dump_pickle + example: Passed to dump_pickle (deprecated) + random_key: If True, shuffle shifts to create non-identity mapping + + Returns: + If pickle_data=False, returns tuple of (pred_coordinates, obs_chemical_shifts, noes, connectivity) + If pickle_data=True, returns None (data is saved to disk) + """ + pred_coordinates, obs_chemical_shifts, pred_chemical_shifts, noes, connectivity = \ + self.generate_data_arrays(random_key=random_key) + pred_coordinates, obs_chemical_shifts, noes, connectivity = \ + self.order_data(pred_coordinates, obs_chemical_shifts, pred_chemical_shifts, noes, connectivity) + + if pickle_data: + self.dump_pickle(pred_coordinates, obs_chemical_shifts, noes, connectivity, example=example) + else: + return pred_coordinates, obs_chemical_shifts, noes, connectivity + + def dump_pickle(self, pred_coordinates: List[Protein], obs_chemical_shifts: List[HSQCPeak], + noes: List[NOEPeak], connectivity: List[Connectivity], example: bool = True): + """ + Save data to disk using default filename. + + DEPRECATED: Use save_dataset() from nmr.nmr_gym.io with explicit filepath instead. + + Args: + pred_coordinates: Protein structures with coordinates and shifts + obs_chemical_shifts: Observed HSQC peaks + noes: NOE crosspeaks + connectivity: Connectivity information + example: If True, uses 'fakedata_r{num_resid}.pkl', else 'current_run.pkl' + """ + name = f'fakedata_r{self.num_resid}.pkl' if example else f'current_run.pkl' + + with open(name, 'wb') as f: + pickle.dump(pred_coordinates, f) + pickle.dump(obs_chemical_shifts, f) + pickle.dump(noes, f) + pickle.dump(connectivity, f) + + def load_pickle(self, example: bool) -> Tuple[List[Protein], List[HSQCPeak], List[NOEPeak], List[Connectivity]]: + """ + Load data from disk using glob pattern. + + DEPRECATED: Use load_dataset() from nmr.nmr_gym.io with explicit filepath instead. + + Args: + example: If True, searches for '*r{num_resid}.pkl', else 'current_run.pkl' + + Returns: + Tuple of (coordinates, obs_chemical_shifts, noes, connectivity) + """ + import glob + + name = f"./*r{self.num_resid}.pkl" if example else "./current_run.pkl" + + pickle_file = glob.glob(name) + with open(pickle_file[0], 'rb') as f: + coordinates = pickle.load(f) + obs_chemical_shifts = pickle.load(f) + noes = pickle.load(f) + connectivity = pickle.load(f) + + return coordinates, obs_chemical_shifts, noes, connectivity diff --git a/nmr/nmr_gym/fake_histories.py b/nmr/nmr_gym/fake_histories.py new file mode 100644 index 0000000..c1e8eb0 --- /dev/null +++ b/nmr/nmr_gym/fake_histories.py @@ -0,0 +1,243 @@ +""" +Generates training histories by perturbing synthetic NMR datasets. + +This module implements the correct history generation algorithm for creating +training data from synthetic NMR datasets. The key principle is that spatial +relationships (NOEs) must reflect the true structure, while introducing +variability only in the observed chemical shifts. + +Correct Algorithm: +1. Load base dataset with original coordinates and shifts +2. For each trajectory: + a. Keep original coordinates unchanged (defines true structure) + b. Perturb ONLY observed chemical shifts (simulates experimental variation) + c. Regenerate NOEs from original coordinates (spatial relationships) + d. Play through with identity mapping for perfect play + e. Record trajectory as list of (state_dict, action, reward) tuples + +Why This Matters: +- Coordinates define the protein structure and spatial relationships +- NOEs are distance constraints derived from 3D structure +- Perturbing coordinates would change the structure and invalidate NOEs +- Only shifts should vary to simulate experimental uncertainty + +Identity Mapping for Perfect Play: +- Due to how synthetic data is generated (random_key=False), identity mapping holds +- When environment asks to assign shift index j, correct residue is also j +- This is shift_index → residue_index as 1:1 mapping +- Perfect play uses: action = state["shift_to_assign"] + +Trajectory Format: +Each trajectory is a list of tuples: +(state_dict, action_taken, value) + +Where: +- state_dict: Complete environment state at this step +- action_taken: Residue index chosen for assignment (integer) +- value: Sum of rewards from this point onwards (return) +""" + +import copy +import numpy as np + +from .fake_data import FakeDataGenerator + + +class FakeHistoryGenerator(): + """ + Generates training histories by perturbing synthetic NMR datasets. + + Creates multiple training examples from a base dataset by adding noise to + observed shifts and generating "perfect play" assignment sequences. Each + history represents one complete assignment episode where the model makes + correct choices at each step. + + The correct perturbation process: + 1. Keep original 3D coordinates unchanged (defines structure) + 2. Perturb only observed chemical shifts (σ=0.1 ppm) + 3. Regenerate NOEs from original (unperturbed) coordinates + 4. Create assignment history playing through with identity mapping + + This provides training diversity while maintaining structural consistency, + enabling supervised pre-training before RL fine-tuning. + + Typical usage: + >>> generator = FakeHistoryGenerator(num_resid=10) + >>> perturbed_state = generator.generate_history(original_state) + >>> trajectory = generate_trajectory(env, perturbed_state) + """ + + def __init__(self, num_resid): + """ + Initialize history generator. + + Args: + num_resid: Number of residues (must match base dataset) + """ + self.num_resid = num_resid + self.fakedata = FakeDataGenerator(num_resid) + self.cutoff = 0.5 + + def perturb_data(self, obs_chemical_shifts): + """ + Perturb only observed chemical shifts with Gaussian noise. + + IMPORTANT: Coordinates are NOT perturbed because they define the + protein structure. Perturbing coordinates would change spatial + relationships and invalidate the NOE distance constraints. Only + chemical shifts vary to simulate experimental uncertainty. + + Args: + obs_chemical_shifts: List of HSQCPeak NamedTuples with H1, N15 + + Returns: + Perturbed chemical shifts [H1, N15] with σ=0.1 ppm noise + """ + # Extract shifts and convert to numpy array + shifts_array = np.array([[s.H1, s.N15] for s in obs_chemical_shifts]) + + # Add chemical shift noise (σ=0.1 ppm) + perturbed_shifts = self.fakedata.add_noise(shifts_array, scale=0.1) + + return perturbed_shifts + + def generate_history(self, original): + """ + Generate a perturbed training history from a base dataset. + + Creates a new training example by: + 1. Extracting original coordinates (UNCHANGED) + 2. Perturbing ONLY observed shifts with Gaussian noise + 3. Regenerating NOEs from original coordinates + perturbed shifts + 4. Recalculating connectivity from original coordinates + 5. Packaging into a clean initial state dictionary + + The history is initialized with empty assignments, ready to be played + through by the RL environment making "perfect" assignment choices + using the identity mapping (action = shift_to_assign). + + Args: + original: Base state dictionary containing: + - 'coordinates': List of Protein NamedTuples (x,y,z,H1,N15) + - 'obs_chemical_shifts': List of HSQCPeak NamedTuples + + Returns: + State dictionary ready for RL environment with structure: + { + "coordinates": List[Protein], # UNCHANGED from original + "obs_chemical_shifts": List[HSQCPeak], # Perturbed + "noes": List[NOEPeak], # Regenerated from original coords + "connectivity": List[Connectivity], # From original coords + "assignments": {}, # Empty - to be filled during play + "assign_order": [], # To be set by environment + "shift_to_assign": 0, + "total_energy": 0.0, + "reward": 0.0 + } + """ + # Step 1: Extract ORIGINAL coordinates (do NOT perturb) + # Coordinates define the structure and must remain unchanged + original_coords = np.array([ + [p.x, p.y, p.z] for p in original['coordinates'] + ]) + + # Step 2: Perturb ONLY observed shifts + perturbed_shifts = self.perturb_data(original['obs_chemical_shifts']) + + # Step 3: Regenerate NOEs using ORIGINAL coordinates + perturbed shifts + # This ensures NOEs reflect true spatial relationships + noes, pred_chemical_shifts = self.fakedata.create_noes( + original_coords, perturbed_shifts, random_key=False + ) + + # Step 4: Recalculate connectivity using ORIGINAL coordinates + connectivity = self.fakedata.calculate_connectivity(original_coords) + + # Step 5: Convert arrays to NamedTuples + coordinates, obs_chemical_shifts, noes, connectivity = ( + self.fakedata.order_data( + original_coords, perturbed_shifts, pred_chemical_shifts, + noes, connectivity + ) + ) + + # Create clean initial state + state = { + "coordinates": coordinates, + "obs_chemical_shifts": obs_chemical_shifts, + "noes": noes, + "connectivity": connectivity, + "assignments": {}, + "assign_order": [], + "shift_to_assign": 0, + "total_energy": 0.0, + "reward": 0.0 + } + + return state + + +def generate_trajectory(env, state_dict): + """ + Generate a perfect play trajectory using identity mapping. + + Plays through a complete assignment episode recording each step as a + tuple of (state_dict, action_taken, value). Uses identity mapping for + perfect play where action equals shift_to_assign. + + Values are computed as the sum of rewards from each point onwards. + For example, if rewards are [0, 1, 1, 0], values are [2, 1, 0, 0]. + + Args: + env: GymEnv instance configured with num_resid + state_dict: Initial state dictionary from generate_history or + env.reset() + + Returns: + List of trajectory tuples: [(state_dict, action, value), ...] + Each tuple contains: + - state_dict: Complete environment state at this step (deep copy) + - action: Residue index chosen for assignment (int) + - value: Sum of rewards from this point onwards (float) + + Example: + >>> env = GymEnv(num_resid=10) + >>> state = env.reset(coords, shifts, noes, connectivity) + >>> trajectory = generate_trajectory(env, state) + >>> print(f"Trajectory has {len(trajectory)} steps") + """ + # Initialize environment with the provided state + # Use custom_state to set up the environment with this specific state + current_state = env.custom_state(state_dict) + + trajectory_with_rewards = [] + terminated = False + + # First pass: collect states, actions, and rewards + while not terminated: + # Perfect play: use identity mapping + # For synthetic data with random_key=False, shift_index = residue_index + action = current_state["shift_to_assign"] + + # IMPORTANT: Deep copy the state BEFORE taking the action to preserve it + state_before_action = copy.deepcopy(current_state) + + # Take the action + current_state, reward, terminated, _ = env.step(action) + + # Store the state before action, the action taken, and the reward + trajectory_with_rewards.append((state_before_action, action, reward)) + + # Second pass: convert rewards to values (sum of future rewards) + trajectory = [] + cumulative_value = 0.0 + + # Iterate backwards to compute cumulative values + for state, action, reward in reversed(trajectory_with_rewards): + cumulative_value += reward + trajectory.append((state, action, cumulative_value)) + + # Reverse to restore original order + trajectory.reverse() + + return trajectory diff --git a/nmr_gym_env.py b/nmr/nmr_gym/gym_env.py similarity index 93% rename from nmr_gym_env.py rename to nmr/nmr_gym/gym_env.py index 049b259..16da01d 100644 --- a/nmr_gym_env.py +++ b/nmr/nmr_gym/gym_env.py @@ -4,9 +4,9 @@ import pickle import math -from energy import Energy -from assignment_order import FractionalActivation -from visualize_data import Visualization +from .energy import Energy +from .assignment_order import FractionalActivation +from scripts.visualize_data import Visualization import cProfile import pstats @@ -58,12 +58,9 @@ def step(self, action): # temp_intermediate_energy = 0 ########################################## - # Calculate reward based on previous energy and temporary energy - self.state['reward'] = (self.intermediate_energy - temp_intermediate_energy) - - # Before correction, +ve means energy went down - we DO NOT want this, -ve means energy went up - assert self.state['reward'] <= 0 - self.state['reward'] = abs(self.state['reward']) + # Calculate reward based on previous energy and temporary energy + # Positive reward when energy goes down (good), negative when energy goes up (bad) + self.state['reward'] = self.intermediate_energy - temp_intermediate_energy # Store energy as intermediate self.intermediate_energy = temp_intermediate_energy diff --git a/nmr/nmr_gym/io.py b/nmr/nmr_gym/io.py new file mode 100644 index 0000000..8565ed2 --- /dev/null +++ b/nmr/nmr_gym/io.py @@ -0,0 +1,279 @@ +""" +I/O utilities for saving and loading NMR datasets and training histories. + +This module provides functions for persisting and loading NMR datasets and +training histories with explicit file paths and metadata. All functions use +Python's pickle format with well-defined structures and validation. + +Dataset Format: +- Single Dataset namedtuple containing: + - pred_coordinates (list of Protein named tuples) + - obs_chemical_shifts (list of HSQCPeak named tuples) + - noes (list of NOEPeak named tuples) + - connectivity (list of Connectivity named tuples) + - metadata (dictionary) + +Histories Format: +- Single Histories namedtuple containing: + - trajectories (list of trajectory lists) + - metadata (dictionary) + +Typical usage: + >>> # Save dataset + >>> metadata = {"version": "1.0", "num_resid": 10} + >>> save_dataset("dataset.pkl", coords, shifts, noes, conn, metadata) + >>> + >>> # Load dataset + >>> dataset = load_dataset("dataset.pkl") + >>> coords = dataset.pred_coordinates + >>> metadata = dataset.metadata + >>> + >>> # Save histories + >>> metadata = {"base_dataset_path": "dataset.pkl", "num_trajectories": 100} + >>> save_histories("histories.pkl", trajectories, metadata) + >>> + >>> # Load histories + >>> histories = load_histories("histories.pkl") + >>> trajectories = histories.trajectories + >>> metadata = histories.metadata +""" + +import pickle +from collections import namedtuple +from pathlib import Path +from typing import Dict, List, Tuple, Union + +from .data_structures import Connectivity, HSQCPeak, NOEPeak, Protein + +# Named tuples for packaging dataset and histories data +Dataset = namedtuple('Dataset', ['pred_coordinates', 'obs_chemical_shifts', 'noes', 'connectivity', 'metadata']) +Histories = namedtuple('Histories', ['trajectories', 'metadata']) + + +def save_dataset(filepath: Union[str, Path], pred_coordinates: List[Protein], + obs_chemical_shifts: List[HSQCPeak], noes: List[NOEPeak], + connectivity: List[Connectivity], metadata: Dict) -> None: + """ + Save dataset to disk with explicit file path and metadata. + + Pickle Format: + - Single Dataset namedtuple containing: + - pred_coordinates (list of Protein named tuples) + - obs_chemical_shifts (list of HSQCPeak named tuples) + - noes (list of NOEPeak named tuples) + - connectivity (list of Connectivity named tuples) + - metadata (dictionary) + + Args: + filepath: Explicit path where dataset will be saved + pred_coordinates: Protein structures with coordinates and predicted shifts + obs_chemical_shifts: Observed/experimental HSQC peaks to be assigned + noes: NOE crosspeaks indicating spatial proximity + connectivity: Residue pairs within distance cutoff + metadata: Dictionary with version, timestamp, generation parameters, etc. + + Raises: + OSError: If file cannot be written + TypeError: If data structures are invalid + + Example: + >>> metadata = { + ... "version": "1.0", + ... "timestamp": datetime.now().isoformat(), + ... "num_resid": 10, + ... "random_key": False, + ... "cutoff": 0.5 + ... } + >>> save_dataset("dataset.pkl", coords, shifts, noes, conn, metadata) + """ + filepath = Path(filepath) + + # Create parent directory if it doesn't exist + filepath.parent.mkdir(parents=True, exist_ok=True) + + # Package into Dataset namedtuple + dataset = Dataset( + pred_coordinates=pred_coordinates, + obs_chemical_shifts=obs_chemical_shifts, + noes=noes, + connectivity=connectivity, + metadata=metadata + ) + + try: + with open(filepath, 'wb') as f: + pickle.dump(dataset, f) + except Exception as e: + raise OSError(f"Failed to save dataset to {filepath}: {e}") from e + + +def load_dataset(filepath: Union[str, Path]) -> Dataset: + """ + Load dataset from disk with explicit file path and validate structure. + + Expected Pickle Format: + - Single Dataset namedtuple containing: + - pred_coordinates (list of Protein named tuples) + - obs_chemical_shifts (list of HSQCPeak named tuples) + - noes (list of NOEPeak named tuples) + - connectivity (list of Connectivity named tuples) + - metadata (dictionary) + + Args: + filepath: Explicit path to dataset pickle file + + Returns: + Dataset namedtuple with fields: pred_coordinates, obs_chemical_shifts, noes, connectivity, metadata + + Raises: + FileNotFoundError: If filepath does not exist + ValueError: If pickle structure is invalid + OSError: If file cannot be read + + Example: + >>> dataset = load_dataset("dataset.pkl") + >>> print(f"Loaded {dataset.metadata['num_resid']} residues") + >>> coords = dataset.pred_coordinates + """ + filepath = Path(filepath) + + if not filepath.exists(): + raise FileNotFoundError(f"Dataset file not found: {filepath}") + + try: + with open(filepath, 'rb') as f: + dataset = pickle.load(f) + except Exception as e: + raise OSError(f"Failed to load dataset from {filepath}: {e}") from e + + # Validate loaded data structure + if not isinstance(dataset, Dataset): + raise ValueError(f"Expected Dataset namedtuple, got {type(dataset)}") + + if not isinstance(dataset.pred_coordinates, list): + raise ValueError(f"pred_coordinates must be a list, got {type(dataset.pred_coordinates)}") + if not isinstance(dataset.obs_chemical_shifts, list): + raise ValueError(f"obs_chemical_shifts must be a list, got {type(dataset.obs_chemical_shifts)}") + if not isinstance(dataset.noes, list): + raise ValueError(f"noes must be a list, got {type(dataset.noes)}") + if not isinstance(dataset.connectivity, list): + raise ValueError(f"connectivity must be a list, got {type(dataset.connectivity)}") + if not isinstance(dataset.metadata, dict): + raise ValueError(f"metadata must be a dict, got {type(dataset.metadata)}") + + # Validate lengths match + num_coords = len(dataset.pred_coordinates) + num_shifts = len(dataset.obs_chemical_shifts) + if num_coords != num_shifts: + raise ValueError(f"Mismatch: {num_coords} coordinates but {num_shifts} shifts") + + return dataset + + +def save_histories(filepath: Union[str, Path], trajectories: List[List[Tuple]], metadata: Dict) -> None: + """ + Save training histories to disk with explicit file path and metadata. + + Each trajectory is a list of tuples with format: + (state_dict, action_taken, reward) + + Pickle Format: + - Single Histories namedtuple containing: + - trajectories (list of trajectory lists) + - metadata (dictionary) + + Args: + filepath: Explicit path where histories will be saved + trajectories: List of trajectories, where each trajectory is a list of + (state_dict, action, reward) tuples + metadata: Dictionary with base_dataset_path, num_trajectories, timestamp, etc. + + Raises: + OSError: If file cannot be written + ValueError: If trajectory structure is invalid + + Example: + >>> metadata = { + ... "base_dataset_path": "/path/to/dataset.pkl", + ... "num_trajectories": 100, + ... "timestamp": datetime.now().isoformat(), + ... "num_resid": 10 + ... } + >>> save_histories("histories.pkl", trajectories, metadata) + """ + filepath = Path(filepath) + + # Validate trajectory structure + if not isinstance(trajectories, list): + raise ValueError(f"trajectories must be a list, got {type(trajectories)}") + + for i, traj in enumerate(trajectories): + if not isinstance(traj, list): + raise ValueError(f"Trajectory {i} must be a list, got {type(traj)}") + + for j, step in enumerate(traj): + if not isinstance(step, tuple) or len(step) != 3: + raise ValueError(f"Trajectory {i}, step {j} must be a 3-tuple (state_dict, action, reward)") + + # Create parent directory if it doesn't exist + filepath.parent.mkdir(parents=True, exist_ok=True) + + # Package into Histories namedtuple + histories = Histories( + trajectories=trajectories, + metadata=metadata + ) + + try: + with open(filepath, 'wb') as f: + pickle.dump(histories, f) + except Exception as e: + raise OSError(f"Failed to save histories to {filepath}: {e}") from e + + +def load_histories(filepath: Union[str, Path]) -> Histories: + """ + Load training histories from disk with explicit file path and validate structure. + + Expected Pickle Format: + - Single Histories namedtuple containing: + - trajectories (list of trajectory lists) + - metadata (dictionary) + + Args: + filepath: Explicit path to histories pickle file + + Returns: + Histories namedtuple with fields: trajectories, metadata + + Raises: + FileNotFoundError: If filepath does not exist + ValueError: If pickle structure is invalid + OSError: If file cannot be read + + Example: + >>> histories = load_histories("histories.pkl") + >>> print(f"Loaded {histories.metadata['num_trajectories']} trajectories") + >>> trajectories = histories.trajectories + """ + filepath = Path(filepath) + + if not filepath.exists(): + raise FileNotFoundError(f"Histories file not found: {filepath}") + + try: + with open(filepath, 'rb') as f: + histories = pickle.load(f) + except Exception as e: + raise OSError(f"Failed to load histories from {filepath}: {e}") from e + + # Validate loaded data structure + if not isinstance(histories, Histories): + raise ValueError(f"Expected Histories namedtuple, got {type(histories)}") + + if not isinstance(histories.trajectories, list): + raise ValueError(f"trajectories must be a list, got {type(histories.trajectories)}") + if not isinstance(histories.metadata, dict): + raise ValueError(f"metadata must be a dict, got {type(histories.metadata)}") + + return histories diff --git a/nmr/nmr_gym/state.py b/nmr/nmr_gym/state.py new file mode 100644 index 0000000..f48cdbe --- /dev/null +++ b/nmr/nmr_gym/state.py @@ -0,0 +1,132 @@ +""" +State management utilities for RL environment. + +This module provides functions for creating and validating state dictionaries +used by the NMR chemical shift assignment RL environment. State dictionaries +contain all information needed to represent a point in the assignment process. + +State Dictionary Structure: +- coordinates: List of Protein named tuples (x, y, z, H1, N15) +- obs_chemical_shifts: List of HSQCPeak named tuples (H1, N15) to be assigned +- noes: List of NOEPeak named tuples (H1, N15, H2) representing spatial constraints +- connectivity: List of Connectivity named tuples (atom1, atom2, distance) +- assignments: Dict mapping shift indices to residue indices +- assign_order: List of shift indices in assignment order +- shift_to_assign: Int or None - current shift index to assign +- total_energy: Float - running total energy +- reward: Float - reward for current step + +Typical usage: + >>> # Create initial state from dataset + >>> state = create_state_dict(coords, shifts, noes, conn) + >>> validate_state_dict(state) # Raises ValueError if invalid + >>> env.custom_state(state) +""" + +from typing import Dict, List + +from .data_structures import Connectivity, HSQCPeak, NOEPeak, Protein + + +def create_state_dict(coordinates: List[Protein], obs_chemical_shifts: List[HSQCPeak], + noes: List[NOEPeak], connectivity: List[Connectivity]) -> Dict: + """ + Create state dictionary for RL environment from dataset components. + + State Dictionary Structure: + - coordinates: List of Protein named tuples (x, y, z, H1, N15) + - obs_chemical_shifts: List of HSQCPeak named tuples (H1, N15) to be assigned + - noes: List of NOEPeak named tuples (H1, N15, H2) representing spatial constraints + - connectivity: List of Connectivity named tuples (atom1, atom2, distance) + - assignments: Empty dict {} (will be filled during episode) + - assign_order: Empty list [] (will be filled by FractionalActivation) + - shift_to_assign: 0 (initial index before assignment order is determined) + - total_energy: 0.0 (running total energy) + - reward: 0.0 (reward for current step) + + Args: + coordinates: Protein structures with coordinates and predicted shifts + obs_chemical_shifts: Observed HSQC peaks to be assigned + noes: NOE crosspeaks indicating spatial proximity + connectivity: Residue pairs within distance cutoff + + Returns: + State dictionary matching GymEnv.reset() format + + Example: + >>> state = create_state_dict(coords, shifts, noes, conn) + >>> env.custom_state(state) + """ + return { + "coordinates": coordinates, + "obs_chemical_shifts": obs_chemical_shifts, + "noes": noes, + "connectivity": connectivity, + "assignments": {}, + "assign_order": [], + "shift_to_assign": 0, + "total_energy": 0.0, + "reward": 0.0 + } + + +def validate_state_dict(state: Dict) -> None: + """ + Validate state dictionary structure and contents. + + Checks for required keys, correct data types, and shape consistency. + Raises specific exceptions with clear messages on validation failure. + + Required Keys: + - coordinates: list of Protein named tuples + - obs_chemical_shifts: list of HSQCPeak named tuples + - noes: list of NOEPeak named tuples + - connectivity: list of Connectivity named tuples + - assignments: dict mapping shift indices to residue indices + - assign_order: list of shift indices in assignment order + - shift_to_assign: int or None (current shift index to assign) + - total_energy: float (running total energy) + - reward: float (reward for current step) + + Args: + state: State dictionary to validate + + Raises: + ValueError: If required key is missing, wrong type, or shape mismatch + + Example: + >>> validate_state_dict(state) # Raises ValueError if invalid + """ + required_keys = [ + "coordinates", "obs_chemical_shifts", "noes", "connectivity", + "assignments", "assign_order", "shift_to_assign", "total_energy", "reward" + ] + + # Check for missing keys + for key in required_keys: + if key not in state: + raise ValueError(f"Missing required key in state dictionary: '{key}'") + + # Check data types + if not isinstance(state["coordinates"], list): + raise ValueError(f"coordinates must be a list, got {type(state['coordinates'])}") + if not isinstance(state["obs_chemical_shifts"], list): + raise ValueError(f"obs_chemical_shifts must be a list, got {type(state['obs_chemical_shifts'])}") + if not isinstance(state["noes"], list): + raise ValueError(f"noes must be a list, got {type(state['noes'])}") + if not isinstance(state["connectivity"], list): + raise ValueError(f"connectivity must be a list, got {type(state['connectivity'])}") + if not isinstance(state["assignments"], dict): + raise ValueError(f"assignments must be a dict, got {type(state['assignments'])}") + if not isinstance(state["assign_order"], list): + raise ValueError(f"assign_order must be a list, got {type(state['assign_order'])}") + if not isinstance(state["total_energy"], (int, float)): + raise ValueError(f"total_energy must be a number, got {type(state['total_energy'])}") + if not isinstance(state["reward"], (int, float)): + raise ValueError(f"reward must be a number, got {type(state['reward'])}") + + # Check shape consistency + num_coords = len(state["coordinates"]) + num_shifts = len(state["obs_chemical_shifts"]) + if num_coords != num_shifts: + raise ValueError(f"Shape mismatch: {num_coords} coordinates but {num_shifts} obs_chemical_shifts") diff --git a/nmr_text_adventure.py b/nmr_text_adventure.py deleted file mode 100644 index cab1f37..0000000 --- a/nmr_text_adventure.py +++ /dev/null @@ -1,140 +0,0 @@ -import gymnasium as gym -import math -import argparse - -from nmr_gym_env import GymEnv -from fake_data import FakeDataGenerator -from visualize_data import Visualization - - - -def get_data(num_resid, pickled=False, pickle_data=True, example=True, random_key=False): - assert not pickled or not pickle_data - - fakedata = FakeDataGenerator(num_resid) - name = f"./*r{num_resid}.pkl" if example else "./current_run.pkl" - - # Grab a previous example (already pickled) - if pickled: - coordinates, obs_chemical_shifts, noes, connectivity = fakedata.load_pickle(example=example) - - # Create a new example and pickle - elif pickle_data: - fakedata.generate_data(pickle_data=True, example=example, random_key=random_key) - coordinates, obs_chemical_shifts, noes, connectivity = fakedata.load_pickle(example=example) - - # Run a new example and don't pickle - else: - coordinates, obs_chemical_shifts, noes, connectivity = fakedata.generate_data(pickle_data=False, example=example, random_key=random_key) - - return coordinates, obs_chemical_shifts, noes, connectivity - - -def text_adventure(): - - parser = argparse.ArgumentParser() - parser.add_argument('num_resid', type=int, help='Number of residues') - args = parser.parse_args() - - num_resid = args.num_resid - - gym_env = GymEnv(num_resid) - total_energy = math.inf - - ###################################################### - - print("\nAvailable options to work with:\n1. Generate and save an example\n2. Grab a previous saved example\n3. Run a new example (not saved)") - - selection = int(input("Option Selection: ")) - - if selection == 1: - # 1. Need to generate and save an example? Saved in all instances as fakedata_r{num_resid}.pkl - coordinates, obs_chemical_shifts, noes, connectivity = get_data(num_resid, pickled=False, pickle_data=True, example=True, random_key=True) - if selection == 2: - # 2. Grab one of the previous examples? Make sure example exists with desired resid number - coordinates, obs_chemical_shifts, noes, connectivity = get_data(num_resid, pickled=True, pickle_data=False, random_key=True) - if selection == 3: - # 3. Run a new example? Saved in all instances as current_run.pkl - coordinates, obs_chemical_shifts, noes, connectivity = get_data(num_resid, pickled=False, pickle_data=True, example=False, random_key=True) - - ###################################################### - - # Game repeats until energy of zero is reached - while total_energy > 0: - - observation = gym_env.reset(coordinates, obs_chemical_shifts, noes, connectivity) - terminated = False - - while not terminated: - gym_env.render() - action = int(input()) - if action < num_resid: - if action not in observation['assignments'].values(): - observation, reward, terminated, total_energy = gym_env.step(action) - - - -if __name__ == '__main__': - text_adventure() - -# parser = argparse.ArgumentParser() -# parser.add_argument('num_resid', type=int, help='Number of residues') -# args = parser.parse_args() - -# num_resid = args.num_resid - -# gym_env = GymEnv(num_resid) -# total_energy = math.inf - -# while total_energy > 0: - -# ###### Available options to work with right now ###### - -# # 1. Need to generate and save an example? Saved in all instances as fakedata_r{num_resid}.pkl -# # observation = gym_env.reset(pickled=False, pickle_data=True, random_key=False) - -# # 2. Grab one of these previous examples? Make sure example exists with desired resid number. -# # observation = gym_env.reset(pickled=True, pickle_data=False, random_key=False) - -# # 3. Run a new example? Saved in all instances as current_run.pkl -# # with cProfile.Profile() as pr: -# observation = gym_env.reset(pickled=False, pickle_data=True, example=False, random_key=False) -# # stats = pstats.Stats(pr) -# # stats.sort_stats('time').print_stats() - -# # 4. Use a custom hardcoded state? May be some issues with data types and visualization. -# # observation = gym_env.custom_state(state) - -# ###################################################### - -# terminated = False - -# while not terminated: -# action = int(input()) -# if action < num_resid: -# if action not in observation['assignments'].values(): -# observation, reward, terminated, total_energy = gym_env.step(action) - -# if __name__ == '__main__': - # cProfile.run('noe_combinations()') - # test_run() -# random generated set of 4 - # state = { - # "coords": [[0.35, 0.83, 0.89, 0.16, 0.034], [0.69, 0.12, 0.92, 0.41, 0.44], [0.05, 0.94, 0.51, 0.189, 0.62], [0.45, 0.06, 0.70, 0.63, 0.15]], - # "actual_shifts": [[0.17, 0.044], [0.41, 0.34], [0.18, 0.61], [0.64, 0.23]], - # "noes": [[0.16, 0.040, 0.17], [0.40, 0.34, 0.63], [0.18, 0.61, 0.16], [0.64, 0.24, 0.39]], - # "restraints": [[(0, 2)], [(1, 3)], [(0, 2)], [(1, 3)]], - # "assignments": {}, - # "total_energy": 0 - # } - - # random generated set of 10 with many restraints - # state = { - # "coords": [[0.8136200069440616, 0.7789831346675045, 0.6474707035745629, 0.8235029227185441, 0.6167276537744274], [0.9458251909740947, 0.8744248260731231, 0.16939643090298961, 0.54118846730192773, 0.7390592014969155], [0.07941814153270477, 0.8569710876203085, 0.6679287503567067, 0.3807330934400178, 0.713863583391123], [0.9295839450890829, 0.2591652267820922, 0.5767903211938266, 0.43720641978511525, 0.3031243342244849], [0.5686652108970609, 0.6600261693178074, 0.8353850759130969, 0.01585880729962143, 0.6545433843661965], [0.3049822956354953, 0.01597728824712763, 0.8007859163237383, 0.15284379187386453, 0.7329892824868469], [0.2707959067671988, 0.16592702782732538, 0.8530675522602625, 0.8257925033719939, 0.6225148209456146], [0.5906467171408669, 0.2275483180754293, 0.1179098748663463, 0.4240407057810011, 0.9558723658454763], [0.8089112838766771, 0.014346015197136075, 0.2651392664874339, 0.45659118663170697, 0.6579588797141084], [0.021085864510464902, 0.3582433470216929, 0.9002893595219581, 0.5229950482718816, 0.9970152102924102]], - # "actual_shifts": [[0.7235029227185441, 0.8167276537744274], [0.44118846730192773, 0.8390592014969155], [0.3607330934400178, 0.813863583391123], [0.35720641978511525, 0.2931243342244849], [0.03585880729962143, 0.6045433843661965], [0.13284379187386453, 0.7429892824868469], [0.8157925033719939, 0.5225148209456146], [0.3240407057810011, 0.9258723658454763], [0.35659118663170697, 0.4579588797141084], [0.6229950482718816, 0.9870152102924102]], - # "noes": [[0.7353272235348708, 0.7985692976248703, 0.024966821752980907], [0.3602327190429789, 0.27520736325235085, 0.3444804557069924], [0.030177646349364318, 0.6038110746526516, 0.7374671133868814], [0.1147855723335682, 0.7472335679962178, 0.8334688788656137], [0.14395802509472194, 0.7602008741457149, 0.6270020652720573], [0.8084393678158522, 0.49689451509103233, 0.1359092345145138], [0.8359026620374025, 0.5227806613419343, 0.6230942517135893], [0.3158886510580014, 0.9127784810952192, 0.36108155258485086], [0.35718811981030113, 0.46209771728576327, 0.36091347433509957], [0.3473018809331518, 0.45763566566297215, 0.3106537228148144], [0.63688033551314, 0.9916090993120613, 0.12731578248735478], [0.6196043419116839, 0.9903515010993935, 0.808008737938382]], - # # "restraints": [[(4, 9), (5, 9), (0, 4), (0, 5)], [(3, 8), (3, 7), (1, 8), (2, 3), (1, 3), (7, 8), (2, 8)], [(0, 4), (4, 9), (4, 6), (5, 6), (0, 5), (5, 9)], [(4, 6), (0, 4), (5, 6), (0, 5)], [(0, 7), (0, 4), (1, 5), (4, 9), (4, 6), (1, 4), (7, 9), (6, 7), (1, 7), (5, 6), (0, 5), (5, 9)], [(6, 7), (4, 6), (5, 6)], [(1, 6), (6, 9), (0, 6)], [(1, 2), (2, 7), (1, 5), (5, 8), (3, 7), (1, 8), (5, 7), (2, 3), (1, 7), (7, 8), (2, 5), (1, 3), (3, 5), (2, 8)], [(3, 8), (3, 7), (1, 8), (2, 3), (1, 3), (7, 8), (2, 8)], [(3, 8), (5, 8), (3, 7), (1, 8), (2, 3), (7, 8), (1, 3), (3, 5), (2, 8)], [(0, 7), (0, 4), (1, 5), (4, 9), (1, 4), (7, 9), (1, 7), (0, 5), (5, 9)], [(0, 1), (0, 9), (0, 6), (1, 6), (6, 9), (1, 9)]], - # "restraints": [[(3, 8), (5, 9)], [(5, 6), (0, 5), (5, 9)], [(0, 4), (4, 9), (4, 6), (5, 6)], [(3, 8), (3, 7), (1, 8)], [(1, 6), (6, 9)]], # there's something wrong with too many restraints, - # "assignments": {}, - # "total_energy": 0 - # } diff --git a/nmr_transformer.py b/nmr_transformer.py deleted file mode 100644 index 016f368..0000000 --- a/nmr_transformer.py +++ /dev/null @@ -1,256 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -from typing import NamedTuple, List, Optional, Tuple - - -class NMRInput(NamedTuple): - # input array of observed chemical shifts - # 2-dimensions: N_shift, H_shift - # (n_peaks, 2) - obs_chemical_shifts: torch.Tensor - - # input array of observed chemical shifts - # 2-dimensions: N_shift, H_shift - # (n_peaks, 2) - pred_chemical_shifts: torch.Tensor - - # input array of observed noes - # 3-dimensions: N_shift1, H_shift1, H_shift2 - # (n_noe, 3) - obs_noes: Optional[torch.Tensor] - - # input array of close distances - # 4-dimensions: N_shift1, H_shift1, H_shift2, H_shift2 - # note: each close distance should be listed twice, with the residue order swapped - # (n_close, 4) - close_distances: Optional[torch.Tensor] - - # input array of peaks assigned so far - # 4-dimensions: N_obs_shift, H_obs_shift, N_pred_shift, H_pred_shift - # (n_assigned, 4) - # note: can set to None if no peaks have been assigned - assigned_peaks: Optional[torch.Tensor] - - # which peak should be assigned? - peak_to_assign: int - - -class NMRTransformer(nn.Module): - """ - A transformer model for predicting the next peak to assign in an NMR spectrum. - """ - - def __init__( - self, - n_hidden: int = 128, - n_heads: int = 4, - n_layers: int = 4, - dropout: float = 0.1, - device=None, - ): - super(NMRTransformer, self).__init__() - self.embedder = NMRInitialEmbedding(n_hidden, device=device) - self.linearmapping = NMRInputMapping() - encoder_layers = nn.TransformerEncoderLayer( - d_model=n_hidden, - dim_feedforward=n_hidden * n_heads, - nhead=n_heads, - dropout=dropout, - device=device, - ) - self.encoder = nn.TransformerEncoder(encoder_layers, num_layers=n_layers) - self.layer_norm = nn.LayerNorm(n_hidden, device=device)############### - - self.policy_linear1 = nn.Linear(2 * n_hidden, 2 * n_hidden, device=device) - self.policy_linear2 = nn.Linear(2 * n_hidden, 1, device=device) - self.value_linear1 = nn.Linear(n_hidden, n_hidden, device=device) - self.value_linear2 = nn.Linear(n_hidden, 1, device=device) - - def forward(self, nmr_inputs: List[NMRInput]) -> Tuple[torch.Tensor, torch.Tensor]: - - linear_values = self.linearmapping(nmr_inputs) - embedded, mask = self.embedder(linear_values) - - # embedded, mask = self.embedder(nmr_inputs) # (N, S, D), (N, S) - embedded = embedded.transpose( - 0, 1 - ) # (S, N, D) for nn.Transformer compatibility - - # embedded_norm = self.layer_norm(embedded) - - x = self.encoder( - embedded, src_key_padding_mask=mask, is_causal=False - ) # (S, N, D) - x = x.transpose(0, 1) # (N, S, D) - - policies = [] - for i, nmr_input in enumerate(linear_values): - n_res = len(nmr_input.pred_chemical_shifts) - residue_embeddings = x[i, 1 : (n_res + 1), :] - peak_to_assign = nmr_input.peak_to_assign - peak_embedding = x[i, n_res + peak_to_assign + 1, :] - policy = (torch.matmul(residue_embeddings, peak_embedding)) - policies.append(policy) - - embeddings = x[:, 0, :] # Global tokens for each input in the batch - values = self.value_linear2(F.relu(self.value_linear1(embeddings))).reshape(-1) - - return policies, values - - -class NMRInputMapping(nn.Module): - def __init__( - self, - xmin_h: int = 6, - xmax_h: int = 10, - xmin_n: int = 100, - xmax_n: int = 135 - ): - super(NMRInputMapping, self).__init__() - self.xmin_h = xmin_h - self.xmax_h = xmax_h - self.xmin_n = xmin_n - self.xmax_n = xmax_n - - def linear_transform(self, x, xmin, xmax): - value = 2*((x - xmin)/(xmax - xmin)) - 1 - return value - - def forward(self, nmr_inputs: List[NMRInput]): #-> Tuple[torch.Tensor, torch.Tensor]: - - linear_values = [] - for i, nmr_input in enumerate(nmr_inputs): - pred_shift1 = self.linear_transform(nmr_input.pred_chemical_shifts[:,0], self.xmin_h, self.xmax_h) - pred_shift2 = self.linear_transform(nmr_input.pred_chemical_shifts[:,1], self.xmin_n, self.xmax_n) - pred_chemical_shifts = torch.stack((pred_shift1, pred_shift2), dim=-1) - # print(pred_chemical_shifts.shape) - - obs_shift1 = self.linear_transform(nmr_input.obs_chemical_shifts[:,0], self.xmin_h, self.xmax_h) - obs_shift2 = self.linear_transform(nmr_input.obs_chemical_shifts[:,1], self.xmin_n, self.xmax_n) - obs_chemical_shifts = torch.stack((obs_shift1, obs_shift2), dim=-1) - # print(obs_chemical_shifts.shape) - - if nmr_input.obs_noes is not None: - noe1 = self.linear_transform(nmr_input.obs_noes[:,0], self.xmin_h, self.xmax_h) - noe2 = self.linear_transform(nmr_input.obs_noes[:,1], self.xmin_n, self.xmax_n) - noe3 = self.linear_transform(nmr_input.obs_noes[:,2], self.xmin_h, self.xmax_h) - obs_noes = torch.stack((noe1, noe2, noe3), dim=-1) - else: - obs_noes = nmr_input.obs_noes - # print(obs_noes.shape) - - if nmr_input.close_distances is not None: - dist1 = self.linear_transform(nmr_input.close_distances[:,0], self.xmin_h, self.xmax_h) - dist2 = self.linear_transform(nmr_input.close_distances[:,1], self.xmin_n, self.xmax_n) - dist3 = self.linear_transform(nmr_input.close_distances[:,2], self.xmin_h, self.xmax_h) - dist4 = self.linear_transform(nmr_input.close_distances[:,3], self.xmin_n, self.xmax_n) - close_distances = torch.stack((dist1, dist2, dist3, dist4), dim=-1) - else: - close_distances = nmr_input.close_distances - # print(close_distances.shape) - - if nmr_input.assigned_peaks is not None: - peak1 = self.linear_transform(nmr_input.assigned_peaks[:,0], self.xmin_h, self.xmax_h) - peak2 = self.linear_transform(nmr_input.assigned_peaks[:,1], self.xmin_n, self.xmax_n) - peak3 = self.linear_transform(nmr_input.assigned_peaks[:,2], self.xmin_h, self.xmax_h) - peak4 = self.linear_transform(nmr_input.assigned_peaks[:,3], self.xmin_n, self.xmax_n) - assigned_peaks = torch.stack((peak1, peak2, peak3, peak4), dim=-1) - else: - assigned_peaks = nmr_input.assigned_peaks - - peak_to_assign = nmr_input.peak_to_assign - - linear_values.append(NMRInput(obs_chemical_shifts=obs_chemical_shifts, - pred_chemical_shifts=pred_chemical_shifts, - obs_noes=obs_noes, - close_distances=close_distances, - assigned_peaks=assigned_peaks, - peak_to_assign=peak_to_assign)) - - return linear_values - - - -class NMRInitialEmbedding(nn.Module): - def __init__(self, n_hidden: int = 128, device=None): - super(NMRInitialEmbedding, self).__init__() - self.device = device - self.n_hidden = n_hidden - self.linear_obs_shifts = nn.Linear(2, n_hidden, device=device) - self.linear_pred_shifts = nn.Linear(2, n_hidden, device=device) - self.linear_obs_noes = nn.Linear(3, n_hidden, device=device) - self.linear_close_distances = nn.Linear(4, n_hidden, device=device) - self.linear_assigned_peaks = nn.Linear(4, n_hidden, device=device) - self.embeddings = nn.Embedding(6, n_hidden, device=device) - - def forward(self, nmr_inputs: List[NMRInput]): - batch_embeddings = [] - lengths = [] - - for nmr in nmr_inputs: - parts = [] - - global_token = self.embeddings( - torch.tensor([5], device=nmr.pred_chemical_shifts.device) - ).expand(1, -1) - parts.append(global_token) - - x = self.linear_pred_shifts(nmr.pred_chemical_shifts) - x += self.embeddings(torch.tensor([0], device=x.device)) - parts.append(x) - - x = self.linear_obs_shifts(nmr.obs_chemical_shifts) - x += self.embeddings(torch.tensor([1], device=x.device)) - parts.append(x) - - if nmr.obs_noes is not None: - x = self.linear_obs_noes(nmr.obs_noes) - x += self.embeddings(torch.tensor([2], device=x.device)) - parts.append(x) - - if nmr.close_distances is not None: - x = self.linear_close_distances(nmr.close_distances) - x += self.embeddings(torch.tensor([3], device=x.device)) - parts.append(x) - - if nmr.assigned_peaks is not None: - x = self.linear_assigned_peaks(nmr.assigned_peaks) - x += self.embeddings(torch.tensor([4], device=x.device)) - parts.append(x) - - full_seq = torch.cat(parts, dim=0) - batch_embeddings.append(full_seq) - lengths.append(full_seq.size(0)) - - max_len = max(lengths) - padded = [] - padding_mask = [] - - for emb in batch_embeddings: - emb_len = emb.size(0) - - if emb_len < max_len: - # Create a padding tensor to append - pad_size = max_len - emb_len - pad_tensor = torch.zeros(pad_size, self.n_hidden, device=emb.device) - padded_emb = torch.cat([emb, pad_tensor], dim=0) - - # Create padding mask - mask = torch.cat( - [ - torch.zeros(emb_len, dtype=torch.bool, device=emb.device), - torch.ones(pad_size, dtype=torch.bool, device=emb.device), - ] - ) - else: - padded_emb = emb - mask = torch.zeros(emb_len, dtype=torch.bool, device=emb.device) - - padded.append(padded_emb) - padding_mask.append(mask) - - padded = torch.stack(padded) - padding_mask = torch.stack(padding_mask) - - return padded, padding_mask diff --git a/pyg_triple_update.py b/pyg_triple_update.py deleted file mode 100644 index 6f339e8..0000000 --- a/pyg_triple_update.py +++ /dev/null @@ -1,644 +0,0 @@ -import torch -import torch_geometric -from torch_geometric.data import Data - -from torch_geometric.loader import DataLoader -from torch_geometric.nn import MessagePassing -from torch_geometric.data import HeteroData -import torch.nn as nn -import numpy as np - -import matplotlib.pyplot as plt -import matplotlib.animation - -from itertools import product -import pickle - -# # Set up data -# data = HeteroData() - -# # Spatial and non-spatial -# data['NOE'].x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8 , 9]], dtype=torch.float32) -# data['NOE'].f = torch.zeros((3, 2), dtype=torch.float32) -# data['RES'].x = torch.tensor([[7, 8, 9, 1, 2], [10, 11, 12, 7, 3]], dtype=torch.float32) -# data['RES'].f = torch.zeros((2, 2), dtype=torch.float32) -# data['SHIFT'].x = torch.tensor([[4, 5], [1, 2]], dtype=torch.float32) -# data['SHIFT'].f = torch.zeros((2, 2), dtype=torch.float32) - -# # Triples (nothing is ever stored in these - they are placeholders to be used in edge types) -# data['TRIPLE0'].x = torch.zeros((1, 1), dtype=torch.float32) -# data['TRIPLE1'].x = torch.zeros((1, 1), dtype=torch.float32) -# data['TRIPLE2'].x = torch.zeros((1, 1), dtype=torch.float32) -# data['TRIPLE3'].x = torch.zeros((1, 1), dtype=torch.float32) - -# data['VALUE_NOE'].x = torch.zeros(1, 1, dtype=torch.float32) # has to be num graphs batched? -# data['VALUE_SHIFT'].x = torch.zeros(1, 1, dtype=torch.float32) -# data['VALUE_RES'].x = torch.zeros(1, 1, dtype=torch.float32) - - - -class ConstructNodes(): - """ - Builds graph nodes from input - """ - def __init__(self, device): - self.device = device - - def construct_data_nodes(self, data, histories): - data['NOE'].x = torch.tensor(histories['noes'], dtype=torch.float32, device=self.device) - data['NOE'].f = torch.zeros((len(histories['noes']), 2), dtype=torch.float32, device=self.device) - data['SHIFT'].x = torch.tensor(histories['obs_chemical_shifts'], dtype=torch.float32, device=self.device) - data['SHIFT'].f = torch.zeros((len(histories['obs_chemical_shifts']), 2), dtype=torch.float32, device=self.device) - data['RES'].x = torch.tensor(histories['coordinates'], dtype=torch.float32, device=self.device) - data['RES'].f = torch.zeros((len(histories['coordinates']), 2), dtype=torch.float32, device=self.device) - return data - - def construct_node_features(self, data, histories): - for i in range(len(data['SHIFT'].f)): - # is the shift being assigned right now? - if i == histories['shift_to_assign']: - data['SHIFT'].f[i, 0] = 1 - # has the shift already been assigned? - if i in histories['assignments'].keys(): - data['SHIFT'].f[i, 1] = 1 - - for i in range(len(data['RES'].f)): - # has the residue been assigned? - if i in histories['assignments'].values(): - data['RES'].f[i] = 1 - return data - - def construct_triple_nodes(self, data): - """ - Generates triple node types. Nothing is ever stored in these - they are simply placeholders to construct/use in edge types. - """ - data['TRIPLE0'].x = torch.zeros((1, 1), dtype=torch.float32, device=self.device) - data['TRIPLE1'].x = torch.zeros((1, 1), dtype=torch.float32, device=self.device) - data['TRIPLE2'].x = torch.zeros((1, 1), dtype=torch.float32, device=self.device) - data['TRIPLE3'].x = torch.zeros((1, 1), dtype=torch.float32, device=self.device) - return data - - def construct_value_nodes(self, data): - data['VALUE_NOE'].x = torch.zeros(1, 1, dtype=torch.float32, device=self.device) # has to be num graphs batched? - data['VALUE_SHIFT'].x = torch.zeros(1, 1, dtype=torch.float32, device=self.device) - data['VALUE_RES'].x = torch.zeros(1, 1, dtype=torch.float32, device=self.device) - return data - - def construct_data(self, histories): - data = HeteroData() - data = self.construct_data_nodes(data, histories) - data = self.construct_triple_nodes(data) - data = self.construct_value_nodes(data) - data = self.construct_node_features(data, histories) # need to sort out all features first across types (consistent or not?) - return data - - - -class ConstructEdges(): - def __init__(self, data, device): - self.data = data - self.num_noe = len(data['NOE'].x) - self.num_shift = len(data['SHIFT'].x) - self.num_res = len(data['RES'].x) - self.device = device - - def get_triple_edges(self, source1, source2): - """ - Grabs all index combinations from the three ranges (residue, shift, noe) and orders these into source and target indices for the edges. - Combinations occur along columns into a triple node. - """ - combo = list(product(range(self.num_noe), range(source1), range(source2))) - combo_tensor = torch.tensor(combo, device=self.device) - - # Switches tensor dimension for source (columns set up for each edge [0,0,0] --> [0],[0],[0]) - source_nodes = torch.transpose(combo_tensor, 0, 1) - # Repeats target edges for number of occurences in source (3 for the triple in this instance, to be assigned to each incoming node type) - target_nodes = torch.tensor(range(len(source_nodes[0])), device=self.device).repeat(len(source_nodes), 1) - - # print(torch.stack([source_nodes, target_nodes], dim=0)) - return torch.stack([source_nodes, target_nodes], dim=0) - - def get_pairwise_edges(self): - """ - Grabs index combinations between shifts and residues and orders these into source and target indices for the edges. - """ - # shift = torch.arange(0, self.num_shift) - shift = np.nonzero(self.data['SHIFT'].f[:, 0])[0] # need to specify if it's the shift being assigned - resid = torch.arange(0, self.num_res, device=self.device) - - shift_repeats = shift.repeat_interleave(self.num_res) - # resid_repeats = resid.repeat(self.num_shift) - resid_repeats = resid - return torch.stack((resid_repeats, shift_repeats), dim=0) - - def get_batch_edges(self): - shift = torch.arange(0, self.num_shift, device=self.device) - noe = torch.arange(0, self.num_noe, device=self.device) - resid = torch.arange(0, self.num_res, device=self.device) - - shift_repeats = torch.zeros(self.num_shift, device=self.device).long() - noe_repeats = torch.zeros(self.num_noe, device=self.device).long() - resid_repeats = torch.zeros(self.num_res, device=self.device).long() - - self.data['SHIFT', 'SHIFT_extract', f'VALUE_SHIFT'].edge_index = torch.stack((shift, shift_repeats), dim=0) - self.data['NOE', 'NOE_extract', f'VALUE_NOE'].edge_index = torch.stack((noe, noe_repeats), dim=0) - self.data['RES', 'RES_extract', f'VALUE_RES'].edge_index = torch.stack((resid, resid_repeats), dim=0) - - return self.data - - def get_edges(self, triple_num, triple_type): - """ - Constructs all edges for the graph. - """ - source1, source2, source3 = triple_type - - # Need to make sure I'm using the right node range (shift and residue number could vary) - num_node1 = self.num_res if source1 == 'RES' else self.num_shift - num_node2 = self.num_res if source2 == 'RES' else self.num_shift - - # Triple in - edges_in = self.get_triple_edges(num_node1, num_node2) - self.data[f'{source3}', 'NOE_extract', f'TRIPLE{triple_num}'].edge_index = edges_in[:, 0] - self.data[f'{source1}', 'NH1_extract', f'TRIPLE{triple_num}'].edge_index = edges_in[:, 1] - self.data[f'{source2}', 'NH2_extract', f'TRIPLE{triple_num}'].edge_index = edges_in[:, 2] - - # Triple out (reverse in/out ordering) - edges_out = torch.stack([edges_in[1], edges_in[0]], dim=0) - self.data[f'TRIPLE{triple_num}', 'NOE_add', f'{source3}'].edge_index = edges_out[:, 0] - self.data[f'TRIPLE{triple_num}', 'NH1_add', f'{source1}'].edge_index = edges_out[:, 1] - self.data[f'TRIPLE{triple_num}', 'NH2_add', f'{source2}'].edge_index = edges_out[:, 2] - - # Only residue nodes will have coordinate edges - if source1 == 'RES': - self.data[f'TRIPLE{triple_num}', 'res1_add', 'RES'].edge_index = edges_out[:, 1] - if source2 == 'RES': - self.data[f'TRIPLE{triple_num}', 'res2_add', 'RES'].edge_index = edges_out[:, 2] - - # Edges for self loop - self.data[f'TRIPLE{triple_num}', 'update', f'TRIPLE{triple_num}'].edge_index = torch.stack([edges_in[1, 0], edges_in[1, 0]], dim=0) - - if triple_num == 0: - self.data['RES', 'pair', 'SHIFT'].edge_index = self.get_pairwise_edges() # pairwise edges - no message passing - return self.data - - def generate_edge_indices(self): - self.data = self.get_edges(0, ('RES', 'RES', 'NOE')) - self.data = self.get_edges(1, ('RES', 'SHIFT', 'NOE')) - self.data = self.get_edges(2, ('SHIFT', 'RES', 'NOE')) - self.data = self.get_edges(3, ('SHIFT', 'SHIFT', 'NOE')) - self.data = self.get_batch_edges() - # print(self.data['RES', 'pair', 'SHIFT'].edge_index) - return self.data - - - -class TripleIn(): - def __init__(self): - # Input index organization - slices to keep proper tensor dimension - self.RES_XYZ = slice(0,3) - self.RES_NH = slice(3,5) - - def grab_node(self, data, node_type, edge_type): - """ - Source node to triple via indexing. - For residue data, extracts shift and coordinate information to process separately. - """ - # Spatial features - xi = data[node_type].x[data[edge_type].edge_index[0]] - # Non-spatial features - fi = data[node_type].f[data[edge_type].edge_index[0]] - - # Splits coordinate and shift features if combined - xj = None - if edge_type[0] == 'RES': - xi, xj = xi[:, self.RES_NH], xi[:, self.RES_XYZ] - return xi, xj, fi - - def construct_triple(self, data, edge_type1, edge_type2, edge_type3): - """ - Constructs triple based on type. - """ - # Measured shift or residue (shifts will return nonetype value) - x1, x12, f1 = self.grab_node(data, node_type=edge_type1[0], edge_type=edge_type1) - x2, x22, f2 = self.grab_node(data, node_type=edge_type2[0], edge_type=edge_type2) - # NOE (does not require any value split) - x3, _, f3 = self.grab_node(data, node_type=edge_type3[0], edge_type=edge_type3) - return x1, x12, x2, x22, x3, f1, f2, f3 - - - -class CalculationManager(): - """ - Common calculations used across triple and pairwise comparisons. - """ - def __init__(self): - # Input index organization - slices to keep proper tensor dimension - self.NOE_N1 = slice(0,1) - self.NOE_H1 = slice(1,2) - self.NOE_H2 = slice(2,3) - self.RES_XYZ = slice(0,3) - self.SHIFT_N = slice(0,1) - self.SHIFT_H = slice(1,2) - - def calc_noe_difference(self, x1, x2, noe): - """ - Calculates shift difference between NOE and residue/measured shifts (direct/indirect only reverse option is available by index). - """ - diff_N = (noe[:, self.NOE_N1] - x1[:, self.SHIFT_N]) # N [n, 1] - diff_H1 = (noe[:, self.NOE_H1] - x1[:, self.SHIFT_H]) # H' [n, 1] - diff_H2 = (noe[:, self.NOE_H2] - x2[:, self.SHIFT_H]) # H" [n, 1] - return diff_N, diff_H1, diff_H2 - - def calc_shift_difference(self, x1, x2): - """ - Calculates shift difference between residue/measured shifts. - """ - diff_N = (x1[:, self.SHIFT_N] - x2[:, self.SHIFT_N]) # N [n, 1] - diff_H = (x1[:, self.SHIFT_H] - x2[:, self.SHIFT_H]) # H [n, 1] - return diff_N, diff_H - - def calc_res_distance(self, x1, x2): - """ - Calculates relative distance between residues and this value squared for equivariant calculations. - """ - rel_dist = (x1 - x2) # [n, 3] - dist2 = torch.norm(rel_dist, dim=-1, keepdim=True)**2 # [n, 1] - return rel_dist, dist2 - - - -class TripleUpdate(MessagePassing): - """ - Message passing class for triple self updates. - """ - def __init__(self, device): - super().__init__(aggr='add') - self.calc_manager = CalculationManager() - self.device = device - - self.hidden = 64 - - self.mlp1 = nn.Sequential( - nn.Linear(12, self.hidden, device=self.device), - nn.ReLU(), - nn.Linear(self.hidden, 19, device=self.device), - nn.LayerNorm(19, device=self.device)) - - self.mlp2 = nn.Sequential( - nn.Linear(11, self.hidden, device=self.device), - nn.ReLU(), - nn.Linear(self.hidden, 13, device=self.device), - nn.LayerNorm(13, device=self.device)) - - def resresnoe_message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j, x22_j): - """ - Constructs message for (res, res, NOE) triple type. - """ - # Residue distances - rel_dist, dist2 = self.calc_manager.calc_res_distance(x12_j, x22_j) - - # Differences relative to NOE shifts (N, H', H") - diff1, diff2, diff3 = self.calc_manager.calc_noe_difference(x1_j, x2_j, x3_j) - - # Shift differences - diff4, diff5 = self.calc_manager.calc_shift_difference(x1_j, x2_j) - - # Input for MLP - # (N, H', H", N, H, dist2, features) - mlp_in = (torch.cat((diff1, diff2, diff3, diff4, diff5, dist2, f1_j, f2_j, f3_j), dim=-1)) # [n, 9] - - # Output from MLP - # Out should include values for each 'change' wanting to make - # (N, H', H", N1, H1, N2, H2, dist1, dist2, features) - mlp_out = self.mlp1(mlp_in) # [n, 16] - - # NOE deltas [n, 1] - delta1x = (diff1 * mlp_out[:, 0:1]) # N - delta2x = (diff2 * mlp_out[:, 1:2]) # H' - delta3x = (diff3 * mlp_out[:, 2:3]) # H" - - # SHIFT deltas residue [n, 1] - delta4x = (diff4 * mlp_out[:, 3:4]) # N1 - delta5x = (diff5 * mlp_out[:, 4:5]) # H1 - delta6x = (diff4 * mlp_out[:, 5:6]) # N2 - delta7x = (diff5 * mlp_out[:, 6:7]) # H2 - - # DISTANCE deltas [n, 3] - delta12x = (rel_dist * (mlp_out[:, 7:10])) - delta13x = (rel_dist * (mlp_out[:, 10:13])) - - # FEATURE deltas [n, 1] - delta1f = mlp_out[:, 13:15] - delta2f = mlp_out[:, 15:17] - delta3f = mlp_out[:, 17:] - return torch.cat((delta1x, delta2x, delta3x, delta4x, delta5x, delta6x, delta7x, delta12x, delta13x, delta1f, delta2f, delta3f), dim=-1) # [n, 16] [n, 19] - - def shiftshiftnoe_message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j): - """ - Constructs message for (shift, shift, NOE), (shift, res, NOE), (res, shift, NOE) triple types. - """ - # Differences relative to NOE shifts (N, H', H") - diff1, diff2, diff3 = self.calc_manager.calc_noe_difference(x1_j, x2_j, x3_j) - - # Shift differences - diff4, diff5 = self.calc_manager.calc_shift_difference(x1_j, x2_j) - - # Input for MLP - # (N, H', H", N, H, dist2, features) - mlp_in = (torch.cat((diff1, diff2, diff3, diff4, diff5, f1_j, f2_j, f3_j), dim=-1)) # [n, 8] - - # Output from MLP - # Out should include values for each 'change' wanting to make - # (N, H', H", N1, H1, N2, H2, features) - mlp_out = self.mlp2(mlp_in) # [n, 10] - - # NOE deltas [n, 1] - delta1x = (diff1 * mlp_out[:, 0:1]) # N - delta2x = (diff2 * mlp_out[:, 1:2]) # H' - delta3x = (diff3 * mlp_out[:, 2:3]) # H" - - # SHIFT deltas residue [n, 1] - delta4x = (diff4 * mlp_out[:, 3:4]) # N1 - delta5x = (diff5 * mlp_out[:, 4:5]) # H1 - delta6x = (diff4 * mlp_out[:, 5:6]) # N2 - delta7x = (diff5 * mlp_out[:, 6:7]) # H2 - - # FEATURE deltas [n, 1] - delta1f = mlp_out[:, 7:9] - delta2f = mlp_out[:, 9:11] - delta3f = mlp_out[:, 11:] - return torch.cat((delta1x, delta2x, delta3x, delta4x, delta5x, delta6x, delta7x, delta1f, delta2f, delta3f), dim=-1) # [n, 10] [n, 13] - - def forward(self, x1, x12, x2, x22, x3, f1, f2, f3, edge_index): - out = self.propagate(edge_index, x1=x1, x2=x2, x3=x3, f1=f1, f2=f2, f3=f3, x12=x12, x22=x22) # not sure how to deal with size here - return out - - def message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j=None, x22_j=None): - # residue based triple type will have two sets of coordinates (x12_j and x22_j) where shifts will have nonetype - if x12_j != None and x22_j != None: - return self.resresnoe_message(x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j, x22_j) - else: - return self.shiftshiftnoe_message(x1_j, x2_j, x3_j, f1_j, f2_j, f3_j) - - def update(self, aggr_out, x12, x22): - # residue based triple type will have two sets of coordinates (x12_j and x22_j) where shifts will have nonetype - if x12 != None and x22 != None: - # output shapes: noe[n, 3], shift1[n, 2], shift2[n, 2], dist1[n, 3], dist2[n, 3], f1[n, 2], f2[n, 2], f3[n, 2] - delta_noe, delta_shift1, delta_shift2, delta_dist1, delta_dist2, delta_f1, delta_f2, delta_f3 = aggr_out[:, 0:3], aggr_out[:, 3:5], aggr_out[:, 5:7], aggr_out[:, 7:10], aggr_out[:, 10:13], aggr_out[:, 13:15], aggr_out[:, 15:17], aggr_out[:, 17:] - return delta_shift1, delta_shift2, delta_noe, delta_dist1, delta_dist2, delta_f1, delta_f2, delta_f3 - else: - # output shapes: noe[n, 3], shift1[n, 2], shift2[n, 2], f1[n, 2], f2[n, 2], f3[n, 2] - delta_noe, delta_shift1, delta_shift2, delta_f1, delta_f2, delta_f3 = aggr_out[:, 0:3], aggr_out[:, 3:5], aggr_out[:, 5:7], aggr_out[:, 7:9], aggr_out[:, 9:11], aggr_out[:, 11:] - return delta_shift1, delta_shift2, delta_noe, delta_f1, delta_f2, delta_f3 - - - -class TripleMessagePass(MessagePassing): - """ - Standard message passing class for outgoing triple messages. - """ - def __init__(self, aggr): - super().__init__(aggr=aggr) - - def forward(self, x_source, x_target, edge_index): - # SIZE (n, m) (source, target) - return self.propagate(edge_index=edge_index, x=(x_source, x_target), size=(x_source.size(0), x_target.size(0))) - - def message(self, x_j): - return x_j - - def update(self, aggr_out, x): - # non-spatial features can also be included in same call but requires more work if they're not being updated - # (don't want to do residue feature updates 2x - only call once on shifts or coordinates) - # x_val, f_val = aggr_out[:, :len(x[1][1])], aggr_out[:, len(x[1][1]):] - # outf = f_val + f[1] - # outx = x_val + x[1] - out = aggr_out + x[1] - return out - - - -class BatchMessagePass(MessagePassing): - def __init__(self, aggr, device): - super().__init__(aggr=aggr) - self.device = device - - self.noe_reduce = (nn.Linear(3, 1, device=self.device)) - self.shift_reduce = (nn.Linear(2, 1, device=self.device)) - self.res_reduce = (nn.Linear(5, 1, device=self.device)) - - def forward(self, x_source, x_target, edge_index): - # SIZE (n, m) (source, target) - return self.propagate(edge_index=edge_index, x=(x_source, x_target), size=(x_source.size(0), x_target.size(0))) - - def message(self, x_j): - if len(x_j[0]) == 3: - return self.noe_reduce(x_j) - if len(x_j[0]) == 2: - return self.shift_reduce(x_j) - if len(x_j[0]) == 5: - return self.res_reduce(x_j) - - def update(self, aggr_out, x): - return aggr_out - - - -class TripleOut(): - def __init__(self): - self.triple_message = TripleMessagePass(aggr='add') - - def update_data(self, data, x_source, f_source, edge_type, x_index, update_f=True): - """ - Calls message passing and directly updates the heterodata object based on target. - Set to update non-spatial features in all cases, but can be turned off case by case to prevent duplicate updates. - """ - source, edge, target = edge_type - edge_index = data[edge_type].edge_index - - # Message passing/update for spatial features - data[target].x[:, x_index] = self.triple_message(x_source, data[target].x[:, x_index], edge_index) - - # Message passing/update for non-spatial features if needed (don't want a duplicate update for residue features) - if update_f: - data[target].f = self.triple_message(f_source, data[target].f, edge_index) - return data - - - -class NMRLayer(nn.Module): - def __init__(self, device): - super(NMRLayer, self).__init__() - self.triple_in = TripleIn() - self.triple_self = TripleUpdate(device) - self.triple_out = TripleOut() - self.calc_manager = CalculationManager() - - self.NOE = slice(0,3) - self.RES_XYZ = slice(0,3) - self.RES_NH = slice(3,5) - self.SHIFT_NH = slice(0,2) - - def update_noes(self, data, noe_delta, feature, i): - # 'i' is for the triple type in all of these cases (TRIPLE0 --> RES RES NOE, etc.) - # Should really simplify this to be one or the other from the very start (string --> number identification) - data = self.triple_out.update_data(data, noe_delta, feature, (f'TRIPLE{i}', 'NOE_add', 'NOE'), self.NOE) - return data - - def update_coordinates(self, data, dist_delta1, dist_delta2, feature1, feature2, i): - data = self.triple_out.update_data(data, dist_delta1, feature1, (f'TRIPLE{i}', 'res1_add', 'RES'), self.RES_XYZ, update_f=False) - data = self.triple_out.update_data(data, dist_delta2, feature2, (f'TRIPLE{i}', 'res2_add', 'RES'), self.RES_XYZ, update_f=False) - return data - - def update_shifts(self, data, shift_delta1, shift_delta2, feature1, feature2, target1, target2, i): - # Selects target indexing for residue or shift node - target_range1 = self.RES_NH if target1 == 'RES' else self.SHIFT_NH - target_range2 = self.RES_NH if target2 == 'RES' else self.SHIFT_NH - - data = self.triple_out.update_data(data, shift_delta1, feature1, (f'TRIPLE{i}', 'NH1_add', target1), target_range1) - data = self.triple_out.update_data(data, shift_delta2, feature2, (f'TRIPLE{i}', 'NH2_add', target2), target_range2) - return data - - def do_updates(self, data, triple_type, i): - # Triple in - shift_x1, coord_x1, shift_x2, coord_x2, noe_x, shift_f1, shift_f2, noe_f = self.triple_in.construct_triple(data, (triple_type[0], 'NH1_extract', f'TRIPLE{i}'), (triple_type[1], 'NH2_extract', f'TRIPLE{i}'), ('NOE', 'NOE_extract', f'TRIPLE{i}')) - - # Self update and output - if triple_type == ('RES', 'RES', 'NOE'): - delta_shift1, delta_shift2, delta_noe, delta_dist1, delta_dist2, deltaf1, deltaf2, deltaf3 = self.triple_self(shift_x1, coord_x1, shift_x2, coord_x2, noe_x, shift_f1, shift_f2, noe_f, data[f'TRIPLE{i}', 'update', f'TRIPLE{i}'].edge_index) - data = self.update_noes(data, delta_noe, deltaf3, i) - data = self.update_coordinates(data, delta_dist1, delta_dist2, deltaf1, deltaf2, i) - data = self.update_shifts(data, delta_shift1, delta_shift2, deltaf1, deltaf2, triple_type[0], triple_type[1], i) - - if triple_type in (('SHIFT', 'RES', 'NOE'), ('RES', 'SHIFT', 'NOE'), ('SHIFT', 'SHIFT', 'NOE')): - delta_shift1, delta_shift2, delta_noe, deltaf1, deltaf2, deltaf3 = self.triple_self(shift_x1, coord_x1, shift_x2, coord_x2, noe_x, shift_f1, shift_f2, noe_f, data[f'TRIPLE{i}', 'update', f'TRIPLE{i}'].edge_index) - data = self.update_noes(data, delta_noe, deltaf3, i) - data = self.update_shifts(data, delta_shift1, delta_shift2, deltaf1, deltaf2, triple_type[0], triple_type[1], i) - return data - - def forward(self, data): - data = self.do_updates(data, ('RES', 'RES', 'NOE'), 0) - data = self.do_updates(data, ('RES', 'SHIFT', 'NOE'), 1) - data = self.do_updates(data, ('SHIFT', 'RES', 'NOE'), 2) - data = self.do_updates(data, ('SHIFT', 'SHIFT', 'NOE'), 3) - return data - - - -class ValueCalc(): - def __init__(self, device): - self.device = device - - self.batch_message = BatchMessagePass(aggr='mean', device=self.device) - self.hidden = 64 - - self.testmlp = nn.Sequential( - nn.Linear(3, self.hidden, device=self.device), - nn.ReLU(), - nn.Linear(self.hidden, 1, device=self.device) - ) - - def get_aggr(self, x_source, x_target, edge_index): - aggr = self.batch_message(x_source, x_target, edge_index) - return aggr - - def calc_value(self, data): - noe = self.get_aggr(data['NOE'].x, data['VALUE_NOE'].x, data['NOE', 'NOE_extract', 'VALUE_NOE'].edge_index) - shift = self.get_aggr(data['SHIFT'].x, data['VALUE_SHIFT'].x, data['SHIFT', 'SHIFT_extract', 'VALUE_SHIFT'].edge_index) - resid = self.get_aggr(data['RES'].x, data['VALUE_RES'].x, data['RES', 'RES_extract', 'VALUE_RES'].edge_index) - - concat_aggr = torch.cat((noe, shift, resid), dim=-1) - value = self.testmlp(concat_aggr) - return value - - - -class PolicyCalc(): - def __init__(self, device): - self.RES_NH = slice(3,5) - self.device = device - - def pairwise_distance(self, data, edge_type): - node_type1, edge, node_type2 = edge_type - - # residue - resid_nodes = data[node_type1].x[data[edge_type].edge_index[0]][:, self.RES_NH] - # shift - shift_nodes = data[node_type2].x[data[edge_type].edge_index[1]] - - rel_dist = (shift_nodes - resid_nodes) - pair_dist2 = -(torch.norm(rel_dist, dim=-1, keepdim=True)**2) - return pair_dist2.squeeze() # reduce dimension [num_res, 1] --> [num_res] - - def calc_policy(self, data): - data_unbatched = data.to_data_list() - policy = [self.pairwise_distance(batch, ('RES', 'pair', 'SHIFT')).tolist() for batch in data_unbatched] - return torch.tensor(policy, dtype=torch.float32, requires_grad=True, device=self.device).unsqueeze(1) # add dimension back, now transposed [1, num_res] - - - -class TestLayer(nn.Module): - def __init__(self, device): - super(TestLayer, self).__init__() - self.value = ValueCalc(device) - self.policy = PolicyCalc(device) - - self.nmr = nn.Sequential( - NMRLayer(device), - ) - - def forward(self, data): - out_data = self.nmr(data) - value = self.value.calc_value(out_data) - policy = self.policy.calc_policy(out_data) - return value, policy - - - -if __name__ == "__main__": - - def extract_data(pickle_file="fake_histories_r4_0.pkl"): - with open(pickle_file, "rb") as f: - histories = pickle.load(f) - return histories - - def construct_graph(history, device): - nodes = ConstructNodes(device) - data = nodes.construct_data(history) - edges = ConstructEdges(data, device) - data = edges.generate_edge_indices() - return data - - def preprocess_data(histories, batch_size, device): - graphs = [] - for history in histories: - graphs.append(construct_graph(history, device)) - - data_loader = DataLoader(graphs, batch_size=batch_size, shuffle=False) - return data_loader - - data_loader = preprocess_data(extract_data(), batch_size=1, device='cuda') - - plt.ion() - fig, ax = plt.subplots() - - epochs = 5 - - testgnn = TestLayer(device='cuda') - for epoch in range(epochs): - batch_num = 0 - for batch in data_loader: - - batch_num += 1 - - value, policy = testgnn.forward(batch) - - ax.cla() - ax.scatter(batch['SHIFT'].x[:, 0].tolist(), batch['SHIFT'].x[:, 1].tolist(), color='red') - ax.scatter(batch['RES'].x[:, 3].tolist(), batch['RES'].x[:, 4].tolist(), color='blue') - ax.set_title(f"Epoch {epoch} Batch {batch_num}") - - plt.pause(0.01) - - plt.ioff() - plt.show() \ No newline at end of file diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py new file mode 100644 index 0000000..9b7f514 --- /dev/null +++ b/scripts/generate_dataset.py @@ -0,0 +1,151 @@ +""" +Generate a single fake NMR dataset with explicit output path. + +This script creates synthetic NMR data including protein coordinates, HSQC chemical +shifts, NOE distance crosspeaks, and connectivity information. Used for supervised +pre-training of graph neural networks before reinforcement learning fine-tuning. + +The script uses the consolidated nmr.nmr_gym.fake_data module and saves datasets with +metadata for versioning and reproducibility. + +Output Format: + The pickle file contains 5 objects in sequence: + 1. pred_coordinates: List of Protein named tuples (x, y, z, H1, N15) + 2. obs_chemical_shifts: List of HSQCPeak named tuples (H1, N15) + 3. noes: List of NOEPeak named tuples (H1, N15, H2) + 4. connectivity: List of Connectivity named tuples (atom1, atom2, distance) + 5. metadata: Dictionary with version, timestamp, num_resid, seed (if provided) + +Usage: + python scripts/generate_dataset.py --num-resid 10 --output my_dataset.pkl + python scripts/generate_dataset.py --num-resid 10 --output test.pkl --seed 42 + +For more information, see docs/fake-data-guide.md +""" + +import argparse +import random +import sys +from datetime import datetime +from pathlib import Path + +import numpy as np + +# Add parent directory to path to allow imports from nmr package +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.nmr_gym.fake_data import FakeDataGenerator +from nmr.nmr_gym.io import save_dataset + + +def set_seed(seed: int) -> None: + """Set random seed for reproducible generation.""" + np.random.seed(seed) + random.seed(seed) + + +def main(): + """Main entry point for dataset generation script.""" + parser = argparse.ArgumentParser( + description="Generate a single fake NMR dataset with explicit output path." + ) + parser.add_argument( + "--num-resid", + type=int, + required=True, + help="Number of residues in the synthetic protein (must be > 0)" + ) + parser.add_argument( + "--output", + type=str, + required=True, + help="Output filename for the dataset (e.g., my_dataset.pkl)" + ) + parser.add_argument( + "--seed", + type=int, + help="Random seed for reproducible generation (optional)" + ) + + args = parser.parse_args() + + # Validate num_resid + if args.num_resid < 1: + print( + f"Error: num_resid must be >= 1 (got {args.num_resid})", + file=sys.stderr + ) + sys.exit(1) + + # Set seed if provided + if args.seed is not None: + set_seed(args.seed) + + # Validate output path + try: + output_path = Path(args.output) + except (TypeError, ValueError) as e: + print( + f"Error: Invalid output path '{args.output}': {e}", + file=sys.stderr + ) + sys.exit(1) + + # Generate dataset + print(f"Generating dataset with {args.num_resid} residues...") + + try: + generator = FakeDataGenerator(args.num_resid) + + # Generate data arrays (coordinates, shifts, noes, connectivity) + # Note: generate_data_arrays returns 5 values including pred_shifts + pred_coordinates, obs_chemical_shifts, pred_chemical_shifts, noes, \ + connectivity = generator.generate_data_arrays(random_key=False) + + # Convert arrays to named tuples using order_data + pred_coordinates, obs_chemical_shifts, noes, connectivity = \ + generator.order_data( + pred_coordinates, obs_chemical_shifts, + pred_chemical_shifts, noes, connectivity + ) + + # Prepare metadata + metadata = { + "version": "1.0", + "timestamp": datetime.now().isoformat(), + "num_resid": args.num_resid, + "random_key": False, + "cutoff": generator.cutoff + } + + # Add seed to metadata if provided + if args.seed is not None: + metadata["seed"] = args.seed + + # Save dataset using consolidated save_dataset utility + save_dataset( + output_path, pred_coordinates, obs_chemical_shifts, + noes, connectivity, metadata + ) + + print(f"Saved to: {output_path}") + + except OSError as e: + # Handle file I/O errors specifically + print( + f"Error: Failed to save dataset to '{args.output}'. {e}", + file=sys.stderr + ) + print( + "Ensure the directory exists and you have write permissions.", + file=sys.stderr + ) + sys.exit(1) + except (ValueError, TypeError) as e: + # Handle validation and data generation errors + print(f"Error during dataset generation: {e}", file=sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/scripts/generate_histories.py b/scripts/generate_histories.py new file mode 100644 index 0000000..f12dd8a --- /dev/null +++ b/scripts/generate_histories.py @@ -0,0 +1,196 @@ +""" +Generate synthetic NMR assignment histories from a dataset. + +This script implements the correct history generation algorithm: +1. Load base dataset ONCE (coordinates and shifts) +2. For N iterations: + - Perturb ONLY observed shifts (keep original coordinates unchanged) + - Regenerate NOEs from original coordinates + - Create environment with perturbed state + - Generate trajectory using identity mapping (action = shift_to_assign) + - Record trajectory as list of (state_dict, action, reward) tuples +3. Save all trajectories to output file + +The key principle is that coordinates define the protein structure and must remain +constant across all trajectories. Only observed shifts vary to simulate experimental +uncertainty, while NOEs are regenerated from the true (original) structure. + +Trajectory Format: +Each trajectory is a list of tuples: (state_dict, action_taken, reward) +- state_dict: Complete environment state at this step +- action_taken: Residue index chosen for assignment (integer) +- reward: Negative energy increase from env.step() (float) + +Identity Mapping: +Due to how synthetic data is generated (random_key=False), the identity mapping +holds: when environment asks to assign shift index j, the correct residue is also j. +This is shift_index → residue_index as 1:1 mapping for perfect play. + +Usage Examples: + # Generate 100 histories from dataset + python generate_histories.py --dataset data.pkl --num-histories 100 --output histories.pkl + + # Generate with reproducible seed + python generate_histories.py --dataset data.pkl --num-histories 50 --output hist.pkl --seed 42 +""" + +import argparse +import sys +from datetime import datetime +from pathlib import Path + +import numpy as np + +# Add parent directory to path to allow imports from nmr package +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.nmr_gym.io import load_dataset, save_histories +from nmr.nmr_gym.fake_histories import FakeHistoryGenerator, generate_trajectory +from nmr.nmr_gym.gym_env import GymEnv + + +def set_seed(seed: int) -> None: + """Set random seed for reproducibility. + + Args: + seed: Random seed value + """ + import random + random.seed(seed) + np.random.seed(seed) + + +def main(): + """Main execution function.""" + parser = argparse.ArgumentParser( + description="Generate synthetic NMR assignment histories from a dataset", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + # Generate 100 histories + python generate_histories.py --dataset data.pkl --num-histories 100 --output histories.pkl + + # Generate with reproducible seed + python generate_histories.py --dataset data.pkl --num-histories 50 --output hist.pkl --seed 42 +""", + ) + + parser.add_argument( + "--dataset", + type=str, + required=True, + help="Path to input dataset pickle file (created by generate_dataset.py)", + ) + parser.add_argument( + "--num-histories", + type=int, + required=True, + help="Number of histories to generate (must be > 0)", + ) + parser.add_argument( + "--output", + type=str, + required=True, + help="Output filepath for histories pickle file", + ) + parser.add_argument( + "--seed", + type=int, + default=None, + help="Random seed for reproducibility (optional)", + ) + + args = parser.parse_args() + + # Validate num_histories + if args.num_histories < 1: + print("Error: num-histories must be at least 1") + sys.exit(1) + + # Set seed if provided + if args.seed is not None: + set_seed(args.seed) + print(f"Random seed set to: {args.seed}") + + # Load base dataset ONCE + print(f"Loading base dataset from: {args.dataset}") + try: + dataset = load_dataset(args.dataset) + coords = dataset.pred_coordinates + obs_shifts = dataset.obs_chemical_shifts + noes = dataset.noes + conn = dataset.connectivity + except FileNotFoundError as e: + print(f"Error: {e}") + sys.exit(1) + except ValueError as e: + print(f"Error: {e}") + sys.exit(1) + except Exception as e: + print(f"Error loading dataset: {e}") + sys.exit(1) + + # Extract num_resid from the data + num_resid = len(coords) + print(f"Dataset loaded successfully ({num_resid} residues)") + + # Initialize environment and history generator + env = GymEnv(num_resid) + history_gen = FakeHistoryGenerator(num_resid) + + # Generate trajectories + print(f"Generating {args.num_histories} trajectories...") + + # Initialize environment with original data to get base state + original_state = env.reset(coords, obs_shifts, noes, conn) + + trajectories = [] + trajectory_lengths = [] + + # Calculate progress reporting interval (every 25%) + progress_interval = max(1, args.num_histories // 4) + + for i in range(args.num_histories): + # Perturb ONLY observed shifts, regenerate NOEs from original coordinates + perturbed_state = history_gen.generate_history(original_state) + + # Generate trajectory using identity mapping (action = shift_to_assign) + trajectory = generate_trajectory(env, perturbed_state) + + # Record trajectory and its length + trajectories.append(trajectory) + trajectory_lengths.append(len(trajectory)) + + # Report progress every 25% + if (i + 1) % progress_interval == 0 or (i + 1) == args.num_histories: + print(f"Progress: Generated trajectory {i + 1}/{args.num_histories}") + + # Calculate statistics + avg_length = np.mean(trajectory_lengths) + print(f"\nGeneration complete:") + print(f" Total trajectories: {args.num_histories}") + print(f" Average trajectory length: {avg_length:.1f} steps") + + # Save all trajectories with metadata + history_metadata = { + "base_dataset_path": str(Path(args.dataset).resolve()), + "num_trajectories": args.num_histories, + "num_resid": num_resid, + "timestamp": datetime.now().isoformat(), + "seed": args.seed, + "avg_trajectory_length": float(avg_length) + } + + try: + save_histories(args.output, trajectories, history_metadata) + print(f"Histories saved to: {args.output}") + except OSError as e: + print(f"Error saving histories: {e}") + sys.exit(1) + except ValueError as e: + print(f"Error validating trajectories: {e}") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/scripts/text_adventure.py b/scripts/text_adventure.py new file mode 100644 index 0000000..95dc274 --- /dev/null +++ b/scripts/text_adventure.py @@ -0,0 +1,131 @@ +""" +Interactive text adventure for NMR chemical shift assignment. + +This script provides an interactive command-line interface for manually testing +the NMR chemical shift assignment environment. Users load a previously saved dataset, +then interactively assign chemical shifts to residues while observing energy changes. + +Usage: + python scripts/text_adventure.py + +Example: + python scripts/text_adventure.py dataset_10.pkl + +Interactive Gameplay: +- The environment will display the current state and available actions +- Enter a residue index (0-based) to assign the current shift to that residue +- The game continues until all shifts are assigned to residues +- Invalid assignments (already used residues) are rejected +- Final energy and assignments are displayed at the end +""" + +import argparse +import math +import sys +from pathlib import Path + +# Add parent directory to path to allow imports from nmr package +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.nmr_gym.gym_env import GymEnv +from nmr.nmr_gym.io import load_dataset + + +def load_data(filepath: str) -> tuple: + """ + Load NMR dataset from file. + + Args: + filepath: Path to the dataset file + + Returns: + Tuple of (coordinates, obs_chemical_shifts, noes, connectivity) + + Raises: + FileNotFoundError: If file doesn't exist + OSError: If file I/O fails + """ + print(f"Loading dataset from {filepath}...") + dataset = load_dataset(filepath) + print(f"Loaded dataset with {len(dataset.pred_coordinates)} residues") + return dataset.pred_coordinates, dataset.obs_chemical_shifts, dataset.noes, dataset.connectivity + + +def text_adventure(): + """ + Main entry point for interactive text adventure. + + Parses command-line arguments, loads dataset from file, + and runs interactive assignment game. + """ + parser = argparse.ArgumentParser( + description="Interactive NMR chemical shift assignment game", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + python scripts/text_adventure.py dataset_10.pkl + python scripts/text_adventure.py /path/to/my_dataset.pkl + """ + ) + parser.add_argument("dataset", type=str, help="Path to dataset file (.pkl)") + args = parser.parse_args() + + # Load data from file + try: + coordinates, obs_chemical_shifts, noes, connectivity = load_data(args.dataset) + except (FileNotFoundError, OSError, ValueError) as e: + print(f"Error loading dataset: {e}") + sys.exit(1) + + # Determine num_resid from loaded data + num_resid = len(coordinates) + gym_env = GymEnv(num_resid) + + # Run one episode + observation = gym_env.reset( + coordinates, obs_chemical_shifts, noes, connectivity + ) + terminated = False + + while not terminated: + # Display current state to the user + shift_idx = observation["shift_to_assign"] + shift_val = observation["obs_chemical_shifts"][shift_idx] + assigned = sorted(observation["assignments"].values()) + unassigned = [i for i in range(num_resid) if i not in assigned] + + print("\n" + "="*60) + print(f"Current shift to assign: #{shift_idx}") + print(f" H1: {shift_val.H1:.2f} ppm, N15: {shift_val.N15:.2f} ppm") + print(f"Assigned residues: {assigned}") + print(f"Unassigned residues: {unassigned}") + print(f"Total energy: {observation['total_energy']:.2f}") + print(f"Shift to assign: {shift_idx}") + print("="*60) + + try: + action = int(input("Enter residue index to assign: ")) + if action < 0 or action >= num_resid: + print(f"Invalid action {action}. Must be between 0 and {num_resid - 1}") + continue + if action not in observation["assignments"].values(): + observation, reward, terminated, total_energy = gym_env.step(action) + print(f"\nAssigned shift {shift_idx} to residue {action}. Reward: {reward:.2f}") + else: + print(f"\nResidue {action} has already been assigned. Choose a different residue.") + except ValueError: + print("Please enter a valid integer.") + except KeyboardInterrupt: + print("\nExiting game...") + sys.exit(0) + + # Display final results + print("\n" + "="*60) + print("Episode complete!") + print(f"Final energy: {total_energy:.2f}") + print(f"Final assignments (shift -> residue): {observation['assignments']}") + print("="*60) + + +if __name__ == "__main__": + text_adventure() diff --git a/scripts/train_gnn.py b/scripts/train_gnn.py new file mode 100644 index 0000000..5bc0bda --- /dev/null +++ b/scripts/train_gnn.py @@ -0,0 +1,164 @@ +import torch +from pprint import pp +import sys +from pathlib import Path +import argparse + +# Add parent directory to path to allow imports from nmr package +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.construct import construct_graph +from nmr.models import NMRNet +from nmr.nmr_gym.io import load_histories +from torch_geometric.loader import DataLoader +from matplotlib import pyplot as plt + + +def extract_data(pickle_file): + """Load histories and flatten into list of (state_dict, action, reward) tuples.""" + histories = load_histories(pickle_file) + + # Flatten trajectories into individual training examples + # Each trajectory contains steps: (state_dict, action, reward) + examples = [] + for trajectory in histories.trajectories: + for state_dict, action, value in trajectory: + examples.append((state_dict, action, value)) + + return examples + + +def preprocess_data(examples, device): + """Convert state dictionaries to graphs and return actions and values separately.""" + graphs = [] + actions = [] + values = [] + + for state_dict, action, value in examples: + graph = construct_graph(state_dict, device) + graphs.append(graph) + actions.append(action) + values.append(value) + + return graphs, actions, values + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Train GNN model on NMR assignment histories") + parser.add_argument( + "--histories", + type=str, + required=True, + help="Path to the pickle file containing training histories" + ) + parser.add_argument( + "--device", + type=str, + default="cpu", + choices=["cpu", "cuda"], + help="Device to use for training (default: cpu)" + ) + parser.add_argument( + "--epochs", + type=int, + default=50000, + help="Number of training epochs (default: 50000)" + ) + parser.add_argument( + "--batch-size", + type=int, + default=1, + help="Batch size for training (default: 1)" + ) + parser.add_argument( + "--learning-rate", + type=float, + default=1e-4, + help="Learning rate for optimizer (default: 1e-4)" + ) + parser.add_argument( + "--eval-interval", + type=int, + default=100, + help="Evaluate model every N iterations (default: 100)" + ) + + args = parser.parse_args() + + # Set device + device = args.device + + # Load data + examples = extract_data(args.histories) + nmr_graphs, actions, rewards = preprocess_data(examples, device) + + batch_size = args.batch_size + epochs = args.epochs + eval_interval = args.eval_interval + + data_loader = DataLoader(nmr_graphs, batch_size=batch_size, shuffle=False) + + # Create our network and optimizer + net = NMRNet(device) + opt = torch.optim.AdamW(net.parameters(), lr=args.learning_rate, weight_decay=0.01) + + # Interactive plotting + plt.ion() + fig, ax = plt.subplots() + + iteration = 0 + for epoch in range(epochs): + + for i, xs in enumerate(data_loader): + net.train() + iteration += 1 + + # Break targets into batches manually (if in same order - no shuffling) + ys = actions[i * batch_size : (i + 1) * batch_size] + + _, _, policies = net(xs) + + # Plots current iteration in batch ########################################## + ax.cla() + x_val1, y_val1 = xs["SHIFT"].x[:, 0].tolist(), xs["SHIFT"].x[:, 1].tolist() + x_val2, y_val2 = xs["RES"].x[:, 3].tolist(), xs["RES"].x[:, 4].tolist() + + ax.scatter(x_val1, y_val1, color="red", label="shift") + ax.scatter(x_val2, y_val2, color="blue", label="resid") + + for j in range(len(x_val1)): + ax.annotate(j, (x_val1[j], y_val1[j] + 0.5), color="red") + ax.annotate(j, (x_val2[j], y_val2[j] + 0.5), color="blue") + + ax.legend() + ax.set_title(f"Iteration {iteration} - Batch {i}") + plt.pause(0.01) + ############################################################################# + + loss = 0 + correct = 0 + count = 0 + + for policy, y in zip(policies, ys): + y = torch.tensor([y], dtype=torch.long, device=device) + loss += torch.nn.functional.cross_entropy(policy, y) + y_pred = torch.argmax(policy) + + if y_pred == y: + correct += 1 + count += 1 + + loss = loss / count + accuracy = correct / count + + print(f"Iteration {iteration} - Loss: {loss.item()} - Accuracy: {accuracy}") + + # writer.add_scalar("Loss/train", loss.item(), iteration) + # writer.add_scalar("Accuracy/train", accuracy, iteration) + + opt.zero_grad() + loss.backward() + opt.step() + + plt.ioff() + plt.show() diff --git a/visualize_data.py b/scripts/visualize_data.py similarity index 100% rename from visualize_data.py rename to scripts/visualize_data.py diff --git a/write_pdb.py b/scripts/write_pdb.py similarity index 89% rename from write_pdb.py rename to scripts/write_pdb.py index 0b291e2..22cff6d 100644 --- a/write_pdb.py +++ b/scripts/write_pdb.py @@ -1,7 +1,13 @@ -from fake_data import FakeDataGenerator from typing import NamedTuple import numpy as np import argparse +import sys +from pathlib import Path + +# Add parent directory to path to allow imports from nmr package +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.nmr_gym.fake_data import FakeDataGenerator diff --git a/energy_unittest.py b/tests/test_energy.py similarity index 93% rename from energy_unittest.py rename to tests/test_energy.py index 6edd4bd..74409ed 100644 --- a/energy_unittest.py +++ b/tests/test_energy.py @@ -1,7 +1,12 @@ import unittest import numpy as np +import sys +from pathlib import Path -from energy import Energy +# Add parent directory to path to allow imports from nmr package +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.nmr_gym.energy import Energy diff --git a/tests/test_fake_data.py b/tests/test_fake_data.py new file mode 100644 index 0000000..9836999 --- /dev/null +++ b/tests/test_fake_data.py @@ -0,0 +1,171 @@ +""" +Unit tests for consolidated fake data module functions. + +Tests the new utility functions added to nmr/env/fake_data.py: +- Dataset saving/loading with explicit paths +- State dictionary creation and validation +- Shift-only perturbation (not coordinates) +""" + +import pickle +import sys +import tempfile +import unittest +from datetime import datetime +from pathlib import Path + +import numpy as np + +# Add parent directory to path to allow imports from nmr package +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.nmr_gym.fake_data import FakeDataGenerator +from nmr.nmr_gym.io import load_dataset, save_dataset +from nmr.nmr_gym.state import create_state_dict, validate_state_dict + + +class TestDatasetSaveLoad(unittest.TestCase): + """Test save_dataset and load_dataset utilities.""" + + def setUp(self): + """Set up test fixtures.""" + self.test_dir = tempfile.mkdtemp() + self.test_path = Path(self.test_dir) + + def tearDown(self): + """Clean up temporary files.""" + import shutil + if self.test_path.exists(): + shutil.rmtree(self.test_dir) + + def test_save_and_load_dataset_roundtrip(self): + """Test that dataset can be saved and loaded with explicit path.""" + # Generate test data + generator = FakeDataGenerator(num_resid=5) + pred_coords, obs_shifts, pred_shifts, noes, connectivity = generator.generate_data_arrays(random_key=False) + pred_coords, obs_shifts, noes, connectivity = generator.order_data( + pred_coords, obs_shifts, pred_shifts, noes, connectivity + ) + + # Save dataset + filepath = self.test_path / "test_dataset.pkl" + metadata = { + "version": "1.0", + "timestamp": datetime.now().isoformat(), + "num_resid": 5, + "random_key": False, + } + save_dataset(filepath, pred_coords, obs_shifts, noes, connectivity, metadata) + + # Load dataset + dataset = load_dataset(filepath) + + # Verify data matches + self.assertEqual(len(dataset.pred_coordinates), 5) + self.assertEqual(len(dataset.obs_chemical_shifts), 5) + self.assertEqual(dataset.metadata["num_resid"], 5) + self.assertEqual(dataset.metadata["version"], "1.0") + + def test_load_dataset_validates_structure(self): + """Test that load_dataset validates pickle structure.""" + # Create invalid pickle with a list instead of Dataset namedtuple + filepath = self.test_path / "invalid_dataset.pkl" + with open(filepath, 'wb') as f: + pickle.dump([1, 2, 3], f) + + # Should raise exception due to invalid structure + with self.assertRaises(ValueError) as context: + load_dataset(filepath) + + self.assertIn("Expected Dataset namedtuple", str(context.exception)) + + +class TestStateDict(unittest.TestCase): + """Test state dictionary creation and validation utilities.""" + + def test_create_state_dict_structure(self): + """Test that create_state_dict returns correct structure.""" + # Create minimal test data + generator = FakeDataGenerator(num_resid=3) + pred_coords, obs_shifts, pred_shifts, noes, connectivity = generator.generate_data_arrays(random_key=False) + pred_coords, obs_shifts, noes, connectivity = generator.order_data( + pred_coords, obs_shifts, pred_shifts, noes, connectivity + ) + + # Create state dict + state = create_state_dict(pred_coords, obs_shifts, noes, connectivity) + + # Verify structure + self.assertIn("coordinates", state) + self.assertIn("obs_chemical_shifts", state) + self.assertIn("noes", state) + self.assertIn("connectivity", state) + self.assertIn("assignments", state) + self.assertIn("assign_order", state) + self.assertIn("shift_to_assign", state) + self.assertIn("total_energy", state) + self.assertIn("reward", state) + + # Verify initial values + self.assertEqual(state["assignments"], {}) + self.assertEqual(state["assign_order"], []) + self.assertEqual(state["shift_to_assign"], 0) + self.assertEqual(state["total_energy"], 0.0) + self.assertEqual(state["reward"], 0.0) + + def test_validate_state_dict_accepts_valid(self): + """Test that validate_state_dict accepts valid state.""" + # Create valid state + generator = FakeDataGenerator(num_resid=3) + pred_coords, obs_shifts, pred_shifts, noes, connectivity = generator.generate_data_arrays(random_key=False) + pred_coords, obs_shifts, noes, connectivity = generator.order_data( + pred_coords, obs_shifts, pred_shifts, noes, connectivity + ) + state = create_state_dict(pred_coords, obs_shifts, noes, connectivity) + + # Should not raise exception + validate_state_dict(state) + + def test_validate_state_dict_rejects_missing_keys(self): + """Test that validate_state_dict catches missing keys.""" + # Create incomplete state + invalid_state = { + "coordinates": [], + "obs_chemical_shifts": [], + # Missing other required keys + } + + # Should raise exception + with self.assertRaises(ValueError) as context: + validate_state_dict(invalid_state) + + self.assertIn("Missing required key", str(context.exception)) + + +class TestShiftPerturbation(unittest.TestCase): + """Test that perturbation affects only shifts, not coordinates.""" + + def test_add_noise_perturbs_shifts_not_coordinates(self): + """Test that add_noise only modifies shift values.""" + generator = FakeDataGenerator(num_resid=5) + + # Generate coordinates and shifts + pred_coords = generator.sample_unit(5, num_sides=3) + pred_coords = generator.scale_unit(pred_coords) + obs_shifts = generator.create_hsqc() + + # Store original coordinates + original_coords = pred_coords.copy() + + # Perturb only shifts + perturbed_shifts = generator.add_noise(obs_shifts, scale=0.1) + + # Verify coordinates unchanged + np.testing.assert_array_equal(pred_coords, original_coords) + + # Verify shifts changed (with high probability) + self.assertFalse(np.array_equal(obs_shifts, perturbed_shifts)) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_fake_histories.py b/tests/test_fake_histories.py new file mode 100644 index 0000000..35f77a5 --- /dev/null +++ b/tests/test_fake_histories.py @@ -0,0 +1,191 @@ +""" +Unit tests for corrected fake history generation algorithm. + +Tests verify that the history generation algorithm correctly: +- Keeps coordinates unchanged across perturbations +- Perturbs only observed shifts, not coordinates +- Regenerates NOEs from original (unperturbed) coordinates +- Uses identity mapping for perfect play (action = shift_to_assign) +- Creates proper trajectory format: (state_dict, action, reward) tuples +""" + +import unittest +import sys +import copy +from pathlib import Path + +import numpy as np + +# Add parent directory to path +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.nmr_gym.fake_data import FakeDataGenerator +from nmr.nmr_gym.fake_histories import FakeHistoryGenerator +from nmr.nmr_gym.gym_env import GymEnv + + +class TestFakeHistoryFix(unittest.TestCase): + """Test suite for corrected fake history generation.""" + + def setUp(self): + """Create test dataset for use in tests.""" + self.num_resid = 5 + self.fake_data_gen = FakeDataGenerator(self.num_resid) + + # Generate base dataset + self.coordinates, self.obs_chemical_shifts, self.noes, self.connectivity = ( + self.fake_data_gen.generate_data(pickle_data=False, random_key=False) + ) + + # Create state dict + self.state = { + "coordinates": self.coordinates, + "obs_chemical_shifts": self.obs_chemical_shifts, + "noes": self.noes, + "connectivity": self.connectivity, + "assignments": {}, + "assign_order": [], + "shift_to_assign": 0, + "total_energy": 0.0, + "reward": 0.0 + } + + def test_coordinates_unchanged_across_perturbations(self): + """Test that coordinates remain unchanged when perturbing data.""" + history_gen = FakeHistoryGenerator(self.num_resid) + + # Extract original coordinates as numpy array for comparison + original_coords = np.array([[p.x, p.y, p.z] for p in self.coordinates]) + + # Generate multiple perturbed versions + for _ in range(3): + perturbed_state = history_gen.generate_history(self.state) + perturbed_coords = np.array([[p.x, p.y, p.z] for p in perturbed_state['coordinates']]) + + # Coordinates should be identical to original + np.testing.assert_array_equal( + original_coords, perturbed_coords, + err_msg="Coordinates should remain unchanged across perturbations" + ) + + def test_only_shifts_perturbed(self): + """Test that only observed shifts are perturbed, not coordinates.""" + history_gen = FakeHistoryGenerator(self.num_resid) + + # Extract original values + original_coords = np.array([[p.x, p.y, p.z] for p in self.coordinates]) + original_shifts = np.array([[s.H1, s.N15] for s in self.obs_chemical_shifts]) + + # Generate perturbed version + perturbed_state = history_gen.generate_history(self.state) + + # Extract perturbed values + perturbed_coords = np.array([[p.x, p.y, p.z] for p in perturbed_state['coordinates']]) + perturbed_shifts = np.array([[s.H1, s.N15] for s in perturbed_state['obs_chemical_shifts']]) + + # Coordinates should be unchanged + np.testing.assert_array_equal( + original_coords, perturbed_coords, + err_msg="Coordinates should NOT be perturbed" + ) + + # Shifts should be different (perturbed with noise) + self.assertFalse( + np.array_equal(original_shifts, perturbed_shifts), + "Shifts should be perturbed with noise" + ) + + def test_noes_regenerated_from_original_coordinates(self): + """Test that NOEs are regenerated from original coordinates, not perturbed ones.""" + history_gen = FakeHistoryGenerator(self.num_resid) + + # Extract original coordinates + original_coords = np.array([[p.x, p.y, p.z] for p in self.coordinates]) + + # Generate perturbed version + perturbed_state = history_gen.generate_history(self.state) + + # NOEs should be based on spatial distances from original coordinates + # We verify this by checking that coordinates used for NOE generation are unchanged + perturbed_coords = np.array([[p.x, p.y, p.z] for p in perturbed_state['coordinates']]) + + np.testing.assert_array_equal( + original_coords, perturbed_coords, + err_msg="NOEs must be generated from original (unperturbed) coordinates" + ) + + # Also verify that NOEs were actually regenerated (not empty) + self.assertGreater( + len(perturbed_state['noes']), 0, + "NOEs should be regenerated and not empty" + ) + + def test_identity_mapping_perfect_play(self): + """Test that identity mapping works for perfect play: action = shift_to_assign.""" + # Initialize environment + env = GymEnv(self.num_resid) + state = env.reset(self.coordinates, self.obs_chemical_shifts, self.noes, self.connectivity) + + # Perfect play: use identity mapping + terminated = False + while not terminated: + # For synthetic data with random_key=False, identity mapping holds + # The correct action is the shift index itself + action = state["shift_to_assign"] + + # Verify action is valid (not already assigned) + self.assertNotIn( + action, state["assignments"].values(), + f"Action {action} should not be already assigned" + ) + + state, reward, terminated, _ = env.step(action) + + # Reward should be non-negative (energy should not increase) + self.assertGreaterEqual( + reward, 0.0, + "Reward should be non-negative for perfect play" + ) + + def test_trajectory_format(self): + """Test that trajectory has correct format: list of (state_dict, action, reward) tuples.""" + from nmr.nmr_gym.fake_histories import generate_trajectory + + # Initialize environment + env = GymEnv(self.num_resid) + initial_state = env.reset(self.coordinates, self.obs_chemical_shifts, self.noes, self.connectivity) + + # Generate trajectory + trajectory = generate_trajectory(env, initial_state) + + # Should be a list + self.assertIsInstance(trajectory, list, "Trajectory should be a list") + + # Should have num_resid steps (one per assignment) + self.assertEqual( + len(trajectory), self.num_resid, + f"Trajectory should have {self.num_resid} steps" + ) + + # Each element should be a tuple of (state_dict, action, reward) + for i, step in enumerate(trajectory): + self.assertIsInstance(step, tuple, f"Step {i} should be a tuple") + self.assertEqual(len(step), 3, f"Step {i} should have 3 elements") + + state_dict, action, reward = step + + # Verify types + self.assertIsInstance(state_dict, dict, f"Step {i}: state should be dict") + self.assertIsInstance(action, (int, np.integer), f"Step {i}: action should be int") + self.assertIsInstance(reward, (int, float, np.number), f"Step {i}: reward should be numeric") + + # Verify state dict has required keys + required_keys = ['coordinates', 'obs_chemical_shifts', 'noes', + 'connectivity', 'assignments', 'assign_order', + 'shift_to_assign', 'total_energy', 'reward'] + for key in required_keys: + self.assertIn(key, state_dict, f"Step {i}: missing key '{key}'") + + +if __name__ == '__main__': + unittest.main() diff --git a/fractional_activation_unittest.py b/tests/test_fractional_activation.py similarity index 93% rename from fractional_activation_unittest.py rename to tests/test_fractional_activation.py index e1ceb04..bcb1ac2 100644 --- a/fractional_activation_unittest.py +++ b/tests/test_fractional_activation.py @@ -1,7 +1,12 @@ import unittest from itertools import chain +import sys +from pathlib import Path -from assignment_order import FractionalActivation +# Add parent directory to path to allow imports from nmr package +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.nmr_gym.assignment_order import FractionalActivation diff --git a/tests/test_generate_dataset.py b/tests/test_generate_dataset.py new file mode 100644 index 0000000..4957b8d --- /dev/null +++ b/tests/test_generate_dataset.py @@ -0,0 +1,289 @@ +""" +Unit tests for scripts/generate_dataset.py + +Tests dataset generation script functionality including: +- Required output argument validation +- Specified filename correctness +- Seed reproducibility +- Error handling for invalid inputs +- Pickle file round-trip (save and load) +- Metadata validation (version, timestamp, num_resid, seed) +""" + +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +import numpy as np + +# Add parent directory to path to allow imports from nmr package +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.nmr_gym.io import load_dataset + + +class TestGenerateDataset(unittest.TestCase): + """Test suite for generate_dataset.py script with comprehensive coverage.""" + + def setUp(self): + """Set up test fixtures.""" + self.test_dir = tempfile.mkdtemp() + self.test_path = Path(self.test_dir) + self.script_path = Path(__file__).parent.parent / "scripts" / "generate_dataset.py" + + def tearDown(self): + """Clean up temporary files.""" + # Clean up test directory + import shutil + if self.test_path.exists(): + shutil.rmtree(self.test_dir) + + def test_missing_output_argument_errors(self): + """Test that missing --output argument causes error.""" + result = subprocess.run( + [sys.executable, str(self.script_path), "--num-resid", "10"], + capture_output=True, + text=True + ) + # argparse exits with code 2 for missing required arguments + self.assertEqual(result.returncode, 2) + self.assertIn("required", result.stderr.lower()) + + def test_specified_filename_works(self): + """Test that specified output filename is created correctly.""" + output_file = self.test_path / "my_test_dataset.pkl" + + result = subprocess.run( + [ + sys.executable, + str(self.script_path), + "--num-resid", "5", + "--output", str(output_file) + ], + capture_output=True, + text=True + ) + + self.assertEqual(result.returncode, 0) + self.assertTrue(output_file.exists()) + self.assertIn("Generating dataset with 5 residues", result.stdout) + self.assertIn(f"Saved to: {output_file}", result.stdout) + + def test_seed_reproducibility(self): + """Test that same seed produces identical datasets.""" + output_file1 = self.test_path / "dataset1.pkl" + output_file2 = self.test_path / "dataset2.pkl" + + # Generate first dataset with seed 42 + result1 = subprocess.run( + [ + sys.executable, + str(self.script_path), + "--num-resid", "5", + "--output", str(output_file1), + "--seed", "42" + ], + capture_output=True, + text=True + ) + self.assertEqual(result1.returncode, 0) + + # Generate second dataset with same seed + result2 = subprocess.run( + [ + sys.executable, + str(self.script_path), + "--num-resid", "5", + "--output", str(output_file2), + "--seed", "42" + ], + capture_output=True, + text=True + ) + self.assertEqual(result2.returncode, 0) + + # Load both datasets using load_dataset utility + dataset1 = load_dataset(output_file1) + dataset2 = load_dataset(output_file2) + + # Verify data are identical + np.testing.assert_array_equal(dataset1.pred_coordinates, dataset2.pred_coordinates) + np.testing.assert_array_equal(dataset1.obs_chemical_shifts, dataset2.obs_chemical_shifts) + np.testing.assert_array_equal(dataset1.noes, dataset2.noes) + self.assertEqual(len(dataset1.connectivity), len(dataset2.connectivity)) + + # Verify metadata contains seed + self.assertEqual(dataset1.metadata["seed"], 42) + self.assertEqual(dataset2.metadata["seed"], 42) + + def test_num_resid_less_than_one_errors(self): + """Test that num_resid < 1 causes error with clear message.""" + output_file = self.test_path / "invalid.pkl" + + result = subprocess.run( + [ + sys.executable, + str(self.script_path), + "--num-resid", "0", + "--output", str(output_file) + ], + capture_output=True, + text=True + ) + + self.assertEqual(result.returncode, 1) + self.assertIn("num_resid must be >= 1", result.stderr) + + def test_invalid_output_path_errors(self): + """Test that invalid output path causes graceful error.""" + # Use a path that cannot be created (parent doesn't exist and is a file) + temp_file = self.test_path / "blockfile" + temp_file.touch() + invalid_path = temp_file / "subdir" / "output.pkl" + + result = subprocess.run( + [ + sys.executable, + str(self.script_path), + "--num-resid", "5", + "--output", str(invalid_path) + ], + capture_output=True, + text=True + ) + + self.assertNotEqual(result.returncode, 0) + self.assertIn("Error", result.stderr) + + def test_pickle_round_trip(self): + """Test that generated pickle can be loaded and contains expected structure.""" + output_file = self.test_path / "roundtrip_test.pkl" + num_resid = 10 + + # Generate dataset + result = subprocess.run( + [ + sys.executable, + str(self.script_path), + "--num-resid", str(num_resid), + "--output", str(output_file) + ], + capture_output=True, + text=True + ) + self.assertEqual(result.returncode, 0) + + # Load and verify structure using new format + dataset = load_dataset(output_file) + + # Verify coordinates are list of Protein named tuples + self.assertEqual(len(dataset.pred_coordinates), num_resid) + self.assertTrue(hasattr(dataset.pred_coordinates[0], 'x')) + self.assertTrue(hasattr(dataset.pred_coordinates[0], 'y')) + self.assertTrue(hasattr(dataset.pred_coordinates[0], 'z')) + + # Verify obs_chemical_shifts are list of HSQCPeak named tuples + self.assertEqual(len(dataset.obs_chemical_shifts), num_resid) + self.assertTrue(hasattr(dataset.obs_chemical_shifts[0], 'H1')) + self.assertTrue(hasattr(dataset.obs_chemical_shifts[0], 'N15')) + + # Verify NOEs exist and have correct structure + self.assertGreater(len(dataset.noes), 0, "Should generate at least some NOEs") + self.assertTrue(hasattr(dataset.noes[0], 'H1')) + self.assertTrue(hasattr(dataset.noes[0], 'N15')) + self.assertTrue(hasattr(dataset.noes[0], 'H2')) + + # Verify connectivity is a list of tuples + self.assertIsInstance(dataset.connectivity, list) + if len(dataset.connectivity) > 0: + self.assertIsInstance(dataset.connectivity[0], tuple) + self.assertEqual(len(dataset.connectivity[0]), 3) # (atom1, atom2, distance) + + def test_output_directory_creation(self): + """Test that parent directories are created if they don't exist.""" + nested_path = self.test_path / "subdir1" / "subdir2" / "dataset.pkl" + + result = subprocess.run( + [ + sys.executable, + str(self.script_path), + "--num-resid", "5", + "--output", str(nested_path) + ], + capture_output=True, + text=True + ) + + self.assertEqual(result.returncode, 0) + self.assertTrue(nested_path.exists()) + self.assertTrue(nested_path.parent.exists()) + + def test_metadata_validation(self): + """Test that generated pickle includes valid metadata with all required fields.""" + output_file = self.test_path / "metadata_test.pkl" + num_resid = 8 + + # Generate dataset with seed + result = subprocess.run( + [ + sys.executable, + str(self.script_path), + "--num-resid", str(num_resid), + "--output", str(output_file), + "--seed", "123" + ], + capture_output=True, + text=True + ) + + # Verify script succeeded + self.assertEqual(result.returncode, 0, f"Script failed: {result.stderr}") + self.assertTrue(output_file.exists(), "Output file was not created") + + # Load using load_dataset utility + dataset = load_dataset(output_file) + + # Verify metadata exists and contains expected keys + self.assertIsInstance(dataset.metadata, dict) + self.assertIn("version", dataset.metadata) + self.assertIn("timestamp", dataset.metadata) + self.assertIn("num_resid", dataset.metadata) + self.assertEqual(dataset.metadata["num_resid"], num_resid) + + # Verify seed is in metadata when provided + self.assertIn("seed", dataset.metadata) + self.assertEqual(dataset.metadata["seed"], 123) + + def test_metadata_without_seed(self): + """Test that metadata is created correctly when no seed is provided.""" + output_file = self.test_path / "no_seed_test.pkl" + + # Generate dataset without seed + result = subprocess.run( + [ + sys.executable, + str(self.script_path), + "--num-resid", "5", + "--output", str(output_file) + ], + capture_output=True, + text=True + ) + + self.assertEqual(result.returncode, 0) + + # Load and verify metadata + dataset = load_dataset(output_file) + + # Verify metadata exists with required fields + self.assertIsInstance(dataset.metadata, dict) + self.assertIn("version", dataset.metadata) + self.assertIn("timestamp", dataset.metadata) + self.assertIn("num_resid", dataset.metadata) + # Seed may or may not be present when not explicitly provided + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_generate_histories_algorithm.py b/tests/test_generate_histories_algorithm.py new file mode 100644 index 0000000..b4933df --- /dev/null +++ b/tests/test_generate_histories_algorithm.py @@ -0,0 +1,265 @@ +""" +Unit tests for history generation algorithm. + +These tests verify the core algorithm logic for history generation: +- Load base dataset once +- Generate N trajectories by perturbing only shifts +- Keep coordinates constant across all trajectories +- Record trajectories in correct format: (state_dict, action, reward) +""" + +import pickle +import sys +import tempfile +import unittest +from pathlib import Path + +import numpy as np + +# Add parent directory to path to allow imports from nmr package +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.nmr_gym.fake_data import FakeDataGenerator +from nmr.nmr_gym.io import load_dataset, load_histories, save_dataset, save_histories +from nmr.nmr_gym.fake_histories import FakeHistoryGenerator, generate_trajectory +from nmr.nmr_gym.gym_env import GymEnv + + +class TestGenerateHistoriesAlgorithm(unittest.TestCase): + """Unit tests for the core history generation algorithm.""" + + def setUp(self): + """Create a temporary dataset for testing.""" + self.num_resid = 5 + self.temp_dir = tempfile.mkdtemp() + self.dataset_path = Path(self.temp_dir) / "test_dataset.pkl" + self.histories_path = Path(self.temp_dir) / "test_histories.pkl" + + # Generate and save a base dataset + generator = FakeDataGenerator(self.num_resid) + coords, obs_shifts, pred_shifts, noes, conn = generator.generate_data_arrays(random_key=False) + coords, obs_shifts, noes, conn = generator.order_data(coords, obs_shifts, pred_shifts, noes, conn) + + metadata = { + "version": "1.0", + "num_resid": self.num_resid, + "random_key": False + } + save_dataset(self.dataset_path, coords, obs_shifts, noes, conn, metadata) + + def test_script_loads_dataset_correctly(self): + """Test that script loads dataset with explicit path.""" + # Load the dataset + dataset = load_dataset(self.dataset_path) + + # Verify loaded data + self.assertEqual(len(dataset.pred_coordinates), self.num_resid) + self.assertEqual(len(dataset.obs_chemical_shifts), self.num_resid) + self.assertIsInstance(dataset.noes, list) + self.assertIsInstance(dataset.connectivity, list) + self.assertEqual(dataset.metadata["num_resid"], self.num_resid) + + def test_n_trajectories_from_one_base_dataset(self): + """Test that N trajectories are generated from one base dataset.""" + # Load base dataset once + dataset = load_dataset(self.dataset_path) + coords = dataset.pred_coordinates + obs_shifts = dataset.obs_chemical_shifts + noes = dataset.noes + conn = dataset.connectivity + + # Store original coordinates for comparison + original_coords = np.array([[p.x, p.y, p.z] for p in coords]) + + # Generate multiple trajectories + num_trajectories = 3 + env = GymEnv(self.num_resid) + history_gen = FakeHistoryGenerator(self.num_resid) + + trajectories = [] + + # Initialize environment once with original data + original_state = env.reset(coords, obs_shifts, noes, conn) + + for i in range(num_trajectories): + # Perturb only shifts, keep coordinates + perturbed_state = history_gen.generate_history(original_state) + + # Verify coordinates unchanged + perturbed_coords = np.array([[p.x, p.y, p.z] for p in perturbed_state['coordinates']]) + np.testing.assert_array_almost_equal(original_coords, perturbed_coords) + + # Generate trajectory + trajectory = generate_trajectory(env, perturbed_state) + trajectories.append(trajectory) + + # Verify we got correct number of trajectories + self.assertEqual(len(trajectories), num_trajectories) + + # Verify each trajectory is non-empty + for traj in trajectories: + self.assertGreater(len(traj), 0) + + def test_trajectory_format_correct(self): + """Test that trajectory format is (state_dict, action, reward) tuples.""" + # Load dataset + dataset = load_dataset(self.dataset_path) + coords = dataset.pred_coordinates + obs_shifts = dataset.obs_chemical_shifts + noes = dataset.noes + conn = dataset.connectivity + + # Generate one trajectory + env = GymEnv(self.num_resid) + history_gen = FakeHistoryGenerator(self.num_resid) + + original_state = env.reset(coords, obs_shifts, noes, conn) + perturbed_state = history_gen.generate_history(original_state) + trajectory = generate_trajectory(env, perturbed_state) + + # Verify trajectory is a list + self.assertIsInstance(trajectory, list) + self.assertGreater(len(trajectory), 0) + + # Verify each element is a tuple with 3 elements + for step in trajectory: + self.assertIsInstance(step, tuple) + self.assertEqual(len(step), 3) + + state_dict, action, reward = step + + # Verify state_dict is a dictionary with required keys + self.assertIsInstance(state_dict, dict) + required_keys = ['coordinates', 'obs_chemical_shifts', 'noes', + 'connectivity', 'assignments', 'assign_order', + 'shift_to_assign', 'total_energy', 'reward'] + for key in required_keys: + self.assertIn(key, state_dict) + + # Verify action is an integer + self.assertIsInstance(action, (int, np.integer)) + + # Verify reward is a number + self.assertIsInstance(reward, (int, float, np.number)) + + def test_coordinates_constant_across_trajectories(self): + """Test that coordinates remain constant across all trajectories.""" + # Load dataset + dataset = load_dataset(self.dataset_path) + coords = dataset.pred_coordinates + obs_shifts = dataset.obs_chemical_shifts + noes = dataset.noes + conn = dataset.connectivity + + # Extract original coordinates + original_coords = np.array([[p.x, p.y, p.z] for p in coords]) + + # Generate multiple trajectories and collect coordinates from each + num_trajectories = 3 + env = GymEnv(self.num_resid) + history_gen = FakeHistoryGenerator(self.num_resid) + + original_state = env.reset(coords, obs_shifts, noes, conn) + + all_coords_match = True + for i in range(num_trajectories): + # Generate perturbed state + perturbed_state = history_gen.generate_history(original_state) + + # Generate trajectory + trajectory = generate_trajectory(env, perturbed_state) + + # Check coordinates in every step of this trajectory + for state_dict, action, reward in trajectory: + step_coords = np.array([[p.x, p.y, p.z] for p in state_dict['coordinates']]) + + # Verify coordinates match original + if not np.allclose(step_coords, original_coords): + all_coords_match = False + break + + if not all_coords_match: + break + + self.assertTrue(all_coords_match, "Coordinates changed across trajectories") + + def test_save_and_load_histories(self): + """Test that histories can be saved and loaded correctly.""" + # Load dataset + dataset = load_dataset(self.dataset_path) + coords = dataset.pred_coordinates + obs_shifts = dataset.obs_chemical_shifts + noes = dataset.noes + conn = dataset.connectivity + + # Generate a few trajectories + num_trajectories = 2 + env = GymEnv(self.num_resid) + history_gen = FakeHistoryGenerator(self.num_resid) + + original_state = env.reset(coords, obs_shifts, noes, conn) + + trajectories = [] + for i in range(num_trajectories): + perturbed_state = history_gen.generate_history(original_state) + trajectory = generate_trajectory(env, perturbed_state) + trajectories.append(trajectory) + + # Save histories + history_metadata = { + "base_dataset_path": str(self.dataset_path), + "num_trajectories": num_trajectories, + "num_resid": self.num_resid + } + save_histories(self.histories_path, trajectories, history_metadata) + + # Load histories + histories = load_histories(self.histories_path) + + # Verify metadata + self.assertEqual(histories.metadata["num_trajectories"], num_trajectories) + self.assertEqual(histories.metadata["num_resid"], self.num_resid) + + # Verify structure + self.assertEqual(len(histories.trajectories), num_trajectories) + for traj in histories.trajectories: + self.assertIsInstance(traj, list) + self.assertGreater(len(traj), 0) + + def test_action_matches_shift_to_assign(self): + """Test that action always matches shift_to_assign for perfect play.""" + # Load dataset + dataset = load_dataset(self.dataset_path) + coords = dataset.pred_coordinates + obs_shifts = dataset.obs_chemical_shifts + noes = dataset.noes + conn = dataset.connectivity + + # Generate multiple trajectories to verify consistency + num_trajectories = 3 + env = GymEnv(self.num_resid) + history_gen = FakeHistoryGenerator(self.num_resid) + + original_state = env.reset(coords, obs_shifts, noes, conn) + + for i in range(num_trajectories): + # Generate perturbed state + perturbed_state = history_gen.generate_history(original_state) + + # Generate trajectory + trajectory = generate_trajectory(env, perturbed_state) + + # Verify each step: action should equal shift_to_assign + for step_num, (state_dict, action, reward) in enumerate(trajectory): + shift_to_assign = state_dict['shift_to_assign'] + + self.assertEqual( + action, shift_to_assign, + f"Trajectory {i}, step {step_num}: action ({action}) does not match " + f"shift_to_assign ({shift_to_assign}). Perfect play requires " + f"action = shift_to_assign for identity mapping." + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_generate_histories_cli.py b/tests/test_generate_histories_cli.py new file mode 100644 index 0000000..2e84bbe --- /dev/null +++ b/tests/test_generate_histories_cli.py @@ -0,0 +1,339 @@ +""" +CLI/Integration tests for the generate_histories.py script. + +Tests cover: +- Command-line interface and argument parsing +- History generation from dataset via CLI +- Default filename generation with timestamp +- Custom filename specification +- Seed reproducibility via CLI +- Error handling (missing files, invalid formats, invalid arguments) +- Progress messages and help text +- History structure validity +""" + +import unittest +import sys +import tempfile +import subprocess +from pathlib import Path +import numpy as np + +# Add parent directory to path +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.nmr_gym.fake_data import FakeDataGenerator +from nmr.nmr_gym.io import load_dataset +from scripts.generate_histories import set_seed + + +class TestGenerateHistoriesCLI(unittest.TestCase): + """CLI and integration tests for generate_histories.py script.""" + + def setUp(self): + """Create test dataset for use in tests.""" + self.temp_dir = tempfile.mkdtemp() + self.test_dataset_path = Path(self.temp_dir) / "test_dataset.pkl" + + # Generate a small test dataset + self.num_resid = 5 + fake_data_gen = FakeDataGenerator(self.num_resid) + coordinates, obs_chemical_shifts, noes, connectivity = ( + fake_data_gen.generate_data(pickle_data=False, random_key=False) + ) + + # Save test dataset with metadata in new format + from nmr.nmr_gym.io import save_dataset + metadata = {"num_resid": self.num_resid, "test": True} + save_dataset(self.test_dataset_path, coordinates, obs_chemical_shifts, noes, connectivity, metadata) + + def tearDown(self): + """Clean up temporary files.""" + import shutil + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def test_set_seed(self): + """Test that set_seed() sets the random seed correctly.""" + import random + + set_seed(42) + val1 = random.random() + + set_seed(42) + val2 = random.random() + + self.assertEqual(val1, val2, "Same seed should produce same random values") + + def test_load_dataset_success(self): + """Test successful dataset loading.""" + dataset = load_dataset(str(self.test_dataset_path)) + + self.assertIsNotNone(dataset.pred_coordinates) + self.assertIsNotNone(dataset.obs_chemical_shifts) + self.assertIsNotNone(dataset.noes) + self.assertIsNotNone(dataset.connectivity) + self.assertEqual(len(dataset.pred_coordinates), self.num_resid) + + def test_load_dataset_file_not_found(self): + """Test error handling for missing dataset file.""" + with self.assertRaises(FileNotFoundError) as context: + load_dataset("nonexistent_file.pkl") + + self.assertIn("Dataset file not found", str(context.exception)) + + def test_load_dataset_invalid_pickle(self): + """Test error handling for invalid pickle format.""" + invalid_file = Path(self.temp_dir) / "invalid.pkl" + + # Create a file with invalid pickle data + with open(invalid_file, 'w') as f: + f.write("This is not a pickle file") + + with self.assertRaises(OSError) as context: + load_dataset(str(invalid_file)) + + self.assertIn("Failed to load dataset", str(context.exception)) + + def test_history_generation_cli(self): + """Test history generation via CLI with small dataset.""" + output_file = Path(self.temp_dir) / "test_histories.pkl" + + # Run the script + result = subprocess.run([ + sys.executable, + "scripts/generate_histories.py", + "--dataset", str(self.test_dataset_path), + "--num-histories", "3", + "--output", str(output_file) + ], capture_output=True, text=True, cwd=Path(__file__).parent.parent) + + self.assertEqual(result.returncode, 0, + f"Script failed: {result.stderr}") + self.assertTrue(output_file.exists(), "Output file not created") + + # Load and verify structure using load_histories + from nmr.nmr_gym.io import load_histories + histories = load_histories(str(output_file)) + + self.assertIsInstance(histories.trajectories, list) + self.assertGreater(len(histories.trajectories), 0, "No trajectories generated") + + # Check first trajectory structure - should be list of tuples + first_trajectory = histories.trajectories[0] + self.assertIsInstance(first_trajectory, list) + self.assertGreater(len(first_trajectory), 0, "Empty trajectory") + + # Check first step structure: (state_dict, action, reward) + first_step = first_trajectory[0] + self.assertIsInstance(first_step, tuple) + self.assertEqual(len(first_step), 3, "Step should be 3-tuple") + + state_dict, action, reward = first_step + self.assertIsInstance(state_dict, dict) + self.assertIsInstance(action, (int, np.integer)) + self.assertIsInstance(reward, (int, float, np.number)) + + # Verify expected keys in state dictionary + expected_keys = ['coordinates', 'obs_chemical_shifts', 'noes', + 'connectivity', 'assignments'] + for key in expected_keys: + self.assertIn(key, state_dict, + f"Missing expected key: {key}") + + def test_custom_filename(self): + """Test that custom output filename is used correctly.""" + custom_name = "my_custom_histories.pkl" + output_file = Path(self.temp_dir) / custom_name + + result = subprocess.run([ + sys.executable, + "scripts/generate_histories.py", + "--dataset", str(self.test_dataset_path), + "--num-histories", "2", + "--output", str(output_file) + ], capture_output=True, text=True, cwd=Path(__file__).parent.parent) + + self.assertEqual(result.returncode, 0) + self.assertTrue(output_file.exists()) + self.assertEqual(output_file.name, custom_name) + + def test_seed_reproducibility(self): + """Test that same seed produces identical histories.""" + output1 = Path(self.temp_dir) / "histories1.pkl" + output2 = Path(self.temp_dir) / "histories2.pkl" + + # Generate with seed=42 + subprocess.run([ + sys.executable, + "scripts/generate_histories.py", + "--dataset", str(self.test_dataset_path), + "--num-histories", "3", + "--output", str(output1), + "--seed", "42" + ], capture_output=True, cwd=Path(__file__).parent.parent) + + # Generate again with same seed + subprocess.run([ + sys.executable, + "scripts/generate_histories.py", + "--dataset", str(self.test_dataset_path), + "--num-histories", "3", + "--output", str(output2), + "--seed", "42" + ], capture_output=True, cwd=Path(__file__).parent.parent) + + # Load both outputs using load_histories + from nmr.nmr_gym.io import load_histories + histories1 = load_histories(str(output1)) + histories2 = load_histories(str(output2)) + + self.assertEqual(len(histories1.trajectories), len(histories2.trajectories)) + + # Check structure - trajectories are lists of tuples + for i in range(min(3, len(histories1.trajectories))): + self.assertEqual(len(histories1.trajectories[i]), len(histories2.trajectories[i])) + # Check that each step has same structure + for step1, step2 in zip(histories1.trajectories[i], histories2.trajectories[i]): + self.assertEqual(len(step1), len(step2), "Steps should be 3-tuples") + + # Check actual content reproducibility - observed chemical shifts and NOEs have noise + # This tests that numpy random state is properly seeded + import numpy as np + if len(histories1.trajectories) > 0 and len(histories1.trajectories[0]) > 0: + # Get first state from first trajectory + state1, _, _ = histories1.trajectories[0][0] + state2, _, _ = histories2.trajectories[0][0] + + # Compare observed chemical shifts which have random noise added + obs_shifts1 = state1['obs_chemical_shifts'] + obs_shifts2 = state2['obs_chemical_shifts'] + + shifts_array1 = np.array([[s.H1, s.N15] for s in obs_shifts1]) + shifts_array2 = np.array([[s.H1, s.N15] for s in obs_shifts2]) + + np.testing.assert_array_equal( + shifts_array1, shifts_array2, + err_msg="Chemical shifts should be identical with same seed (numpy RNG not seeded?)" + ) + + # Also check NOEs which have random noise added + noes1 = state1['noes'] + noes2 = state2['noes'] + + noes_array1 = np.array([[n.H1, n.N15, n.H2] for n in noes1]) + noes_array2 = np.array([[n.H1, n.N15, n.H2] for n in noes2]) + + np.testing.assert_array_equal( + noes_array1, noes_array2, + err_msg="NOEs should be identical with same seed (numpy RNG not seeded?)" + ) + + def test_invalid_num_histories(self): + """Test error handling for invalid num_histories value.""" + output_file = Path(self.temp_dir) / "invalid_test.pkl" + result = subprocess.run([ + sys.executable, + "scripts/generate_histories.py", + "--dataset", str(self.test_dataset_path), + "--num-histories", "0", + "--output", str(output_file) + ], capture_output=True, text=True, cwd=Path(__file__).parent.parent) + + self.assertNotEqual(result.returncode, 0) + self.assertIn("must be at least 1", result.stderr + result.stdout) + + def test_help_output(self): + """Test that --help produces clear, descriptive help text.""" + result = subprocess.run([ + sys.executable, + "scripts/generate_histories.py", + "--help" + ], capture_output=True, text=True, cwd=Path(__file__).parent.parent) + + self.assertEqual(result.returncode, 0) + + # Check for key arguments in help text + self.assertIn("--dataset", result.stdout) + self.assertIn("--num-histories", result.stdout) + self.assertIn("--output", result.stdout) + self.assertIn("--seed", result.stdout) + self.assertIn("Examples:", result.stdout) + + def test_progress_messages(self): + """Test that progress messages are printed during generation.""" + result = subprocess.run([ + sys.executable, + "scripts/generate_histories.py", + "--dataset", str(self.test_dataset_path), + "--num-histories", "4", + "--output", str(Path(self.temp_dir) / "test.pkl") + ], capture_output=True, text=True, cwd=Path(__file__).parent.parent) + + self.assertEqual(result.returncode, 0) + + # Check for expected progress messages + output = result.stdout + result.stderr + self.assertIn("Loading base dataset from:", output) + self.assertIn("Dataset loaded successfully", output) + self.assertIn("Generating", output) + self.assertIn("Generation complete", output) + self.assertIn("Histories saved to:", output) + + def test_history_structure_validity(self): + """Test that generated histories have correct and valid structure.""" + output_file = Path(self.temp_dir) / "structure_test.pkl" + + subprocess.run([ + sys.executable, + "scripts/generate_histories.py", + "--dataset", str(self.test_dataset_path), + "--num-histories", "2", + "--output", str(output_file) + ], capture_output=True, cwd=Path(__file__).parent.parent) + + # Load using load_histories + from nmr.nmr_gym.io import load_histories + histories = load_histories(str(output_file)) + + # All trajectories should be lists of tuples + for i, trajectory in enumerate(histories.trajectories): + self.assertIsInstance(trajectory, list, + f"Trajectory {i} is not a list") + self.assertGreater(len(trajectory), 0, + f"Trajectory {i} is empty") + + # Each step should be a 3-tuple (state_dict, action, reward) + for j, step in enumerate(trajectory): + self.assertIsInstance(step, tuple, + f"Trajectory {i}, step {j} is not a tuple") + self.assertEqual(len(step), 3, + f"Trajectory {i}, step {j} is not a 3-tuple") + + state_dict, action, reward = step + + # State should be a dictionary + self.assertIsInstance(state_dict, dict, + f"Trajectory {i}, step {j} state is not a dict") + + # Check for required keys + required_keys = ['coordinates', 'obs_chemical_shifts', + 'noes', 'connectivity', 'assignments'] + for key in required_keys: + self.assertIn(key, state_dict, + f"Trajectory {i}, step {j} missing key: {key}") + + # Assignments should be a dictionary + self.assertIsInstance(state_dict['assignments'], dict, + f"Trajectory {i}, step {j} assignments not a dict") + + # Action should be an integer (Python int or numpy integer) + self.assertIsInstance(action, (int, np.integer), + f"Trajectory {i}, step {j} action not an int") + + # Reward should be numeric + self.assertIsInstance(reward, (int, float, np.number), + f"Trajectory {i}, step {j} reward not numeric") + + +if __name__ == '__main__': + unittest.main() diff --git a/gym_env_unittest.py b/tests/test_gym_env.py similarity index 59% rename from gym_env_unittest.py rename to tests/test_gym_env.py index 89f73ff..020c6d1 100644 --- a/gym_env_unittest.py +++ b/tests/test_gym_env.py @@ -1,36 +1,63 @@ import unittest import numpy as np +import sys +from pathlib import Path -from nmr_gym_env import GymEnv +# Add parent directory to path to allow imports from nmr package +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.nmr_gym.gym_env import GymEnv +from nmr.nmr_gym.data_structures import Connectivity, HSQCPeak, NOEPeak, Protein class TestGymEnv(unittest.TestCase): def setUp(self): self.num_resid = 4 + # Create proper named tuples instead of plain lists + coordinates = [ + Protein(0.35, 0.83, 0.89, 0.17, 0.044), + Protein(0.69, 0.12, 0.92, 0.41, 0.34), + Protein(0.05, 0.94, 0.51, 0.18, 0.61), + Protein(0.45, 0.06, 0.70, 0.64, 0.23) + ] + obs_chemical_shifts = [ + HSQCPeak(0.17, 0.044), + HSQCPeak(0.41, 0.34), + HSQCPeak(0.18, 0.61), + HSQCPeak(0.64, 0.23) + ] + noes = [ + NOEPeak(0.16, 0.040, 0.17), + NOEPeak(0.40, 0.34, 0.63), + NOEPeak(0.18, 0.61, 0.16), + NOEPeak(0.64, 0.24, 0.39) + ] + connectivity = [] # Add connectivity if needed + self.state = { - "coords": [[0.35, 0.83, 0.89], [0.69, 0.12, 0.92], [0.05, 0.94, 0.51], [0.45, 0.06, 0.70]], - "actual_shifts": [[0.17, 0.044], [0.41, 0.34], [0.18, 0.61], [0.64, 0.23]], - "noes": [[0.16, 0.040, 0.17], [0.40, 0.34, 0.63], [0.18, 0.61, 0.16], [0.64, 0.24, 0.39]], + "coordinates": coordinates, + "obs_chemical_shifts": obs_chemical_shifts, + "noes": noes, + "connectivity": connectivity, "restraints": [[(0, 2)], [(1, 3)], [(0, 2)], [(1, 3)]], "assignments": {}, "total_energy": 0 - } + } self.gym_env = GymEnv(self.num_resid) def test_step_energy_zeros(self): actions = [3, 1, 0, 2] - + expected_energy = [0, 0, 0, 0] expected_reward = [0, 0, 0, 0] observation = self.gym_env.custom_state(self.state) - terminated = False - + for i, action in enumerate(actions): - observation, reward, terminated = self.gym_env.step(action) + observation, reward, terminated, total_energy = self.gym_env.step(action) - self.assertEqual(observation['total_energy'], expected_energy[i]) + self.assertEqual(observation['total_energy'], expected_energy[i]) self.assertEqual(reward, expected_reward[i]) def test_step_energy_goes_up(self): @@ -55,12 +82,11 @@ def test_step_energy_goes_up(self): expected_reward = [0, -0.9558604159815726, 0, -0.4534183043507546] observation = self.gym_env.custom_state(self.state) - terminated = False - + for i, action in enumerate(actions): - observation, reward, terminated = self.gym_env.step(action) + observation, reward, terminated, total_energy = self.gym_env.step(action) - self.assertEqual(observation['total_energy'], expected_energy[i]) + self.assertEqual(observation['total_energy'], expected_energy[i]) self.assertEqual(reward, expected_reward[i]) unittest.main(argv=[''], verbosity=2, exit=False) \ No newline at end of file diff --git a/tests/test_integration_end_to_end.py b/tests/test_integration_end_to_end.py new file mode 100644 index 0000000..91e77b4 --- /dev/null +++ b/tests/test_integration_end_to_end.py @@ -0,0 +1,303 @@ +""" +Integration tests for end-to-end dataset and history generation workflow. + +These tests verify critical integration points: +- Dataset generation → history generation workflow +- Coordinates remain identical across all trajectories from same base dataset +- Perturbed shifts differ across trajectories +- Complete workflow with save/load operations +- Data format consistency throughout pipeline +""" + +import pickle +import sys +import tempfile +import unittest +from pathlib import Path + +import numpy as np + +# Add parent directory to path +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.nmr_gym.fake_data import FakeDataGenerator +from nmr.nmr_gym.fake_histories import FakeHistoryGenerator, generate_trajectory +from nmr.nmr_gym.gym_env import GymEnv +from nmr.nmr_gym.io import load_dataset, load_histories, save_dataset, save_histories + + +class TestEndToEndWorkflow(unittest.TestCase): + """Integration tests for complete dataset → histories workflow.""" + + def setUp(self): + """Set up temporary directory for test files.""" + self.temp_dir = tempfile.mkdtemp() + self.temp_path = Path(self.temp_dir) + + def tearDown(self): + """Clean up temporary files.""" + import shutil + if self.temp_path.exists(): + shutil.rmtree(self.temp_dir) + + def test_dataset_to_histories_workflow(self): + """Test complete workflow: generate dataset → generate histories → verify format.""" + num_resid = 5 + num_trajectories = 3 + + # Step 1: Generate dataset + generator = FakeDataGenerator(num_resid) + pred_coords, obs_shifts, pred_shifts, noes, connectivity = generator.generate_data_arrays( + random_key=False + ) + pred_coords, obs_shifts, noes, connectivity = generator.order_data( + pred_coords, obs_shifts, pred_shifts, noes, connectivity + ) + + # Step 2: Save dataset + dataset_path = self.temp_path / "test_dataset.pkl" + dataset_metadata = { + "version": "1.0", + "num_resid": num_resid, + "random_key": False, + } + save_dataset(dataset_path, pred_coords, obs_shifts, noes, connectivity, dataset_metadata) + + # Step 3: Load dataset + dataset = load_dataset(dataset_path) + + # Step 4: Generate histories from loaded dataset + env = GymEnv(num_resid) + history_gen = FakeHistoryGenerator(num_resid) + + # Initialize environment once with loaded data + original_state = env.reset( + dataset.pred_coordinates, + dataset.obs_chemical_shifts, + dataset.noes, + dataset.connectivity + ) + + trajectories = [] + for i in range(num_trajectories): + # Perturb state and generate trajectory + perturbed_state = history_gen.generate_history(original_state) + trajectory = generate_trajectory(env, perturbed_state) + trajectories.append(trajectory) + + # Step 5: Save histories + histories_path = self.temp_path / "test_histories.pkl" + histories_metadata = { + "base_dataset_path": str(dataset_path), + "num_trajectories": num_trajectories, + "num_resid": num_resid, + } + save_histories(histories_path, trajectories, histories_metadata) + + # Step 6: Load and verify histories + histories = load_histories(histories_path) + + # Verify structure + self.assertEqual(len(histories.trajectories), num_trajectories) + self.assertEqual(histories.metadata["num_trajectories"], num_trajectories) + + # Verify each trajectory has correct format + for traj in histories.trajectories: + self.assertIsInstance(traj, list) + self.assertEqual(len(traj), num_resid) # One step per residue + + for step in traj: + self.assertIsInstance(step, tuple) + self.assertEqual(len(step), 3) + + state_dict, action, reward = step + self.assertIsInstance(state_dict, dict) + self.assertIsInstance(action, (int, np.integer)) + self.assertIsInstance(reward, (int, float, np.number)) + + def test_coordinates_identical_across_trajectories(self): + """Test that histories from same base dataset have identical coordinates.""" + num_resid = 6 + num_trajectories = 4 + + # Generate base dataset + generator = FakeDataGenerator(num_resid) + pred_coords, obs_shifts, pred_shifts, noes, connectivity = generator.generate_data_arrays( + random_key=False + ) + pred_coords, obs_shifts, noes, connectivity = generator.order_data( + pred_coords, obs_shifts, pred_shifts, noes, connectivity + ) + + # Extract original coordinates as numpy array + original_coords = np.array([[p.x, p.y, p.z] for p in pred_coords]) + + # Generate multiple trajectories + env = GymEnv(num_resid) + history_gen = FakeHistoryGenerator(num_resid) + original_state = env.reset(pred_coords, obs_shifts, noes, connectivity) + + all_coords_from_trajectories = [] + + for i in range(num_trajectories): + perturbed_state = history_gen.generate_history(original_state) + trajectory = generate_trajectory(env, perturbed_state) + + # Extract coordinates from first step of this trajectory + first_state, _, _ = trajectory[0] + traj_coords = np.array([[p.x, p.y, p.z] for p in first_state["coordinates"]]) + all_coords_from_trajectories.append(traj_coords) + + # Verify all trajectories have identical coordinates to original + for i, traj_coords in enumerate(all_coords_from_trajectories): + np.testing.assert_array_almost_equal( + original_coords, + traj_coords, + decimal=10, + err_msg=f"Trajectory {i} has different coordinates from original", + ) + + # Verify all trajectories have identical coordinates to each other + for i in range(1, len(all_coords_from_trajectories)): + np.testing.assert_array_almost_equal( + all_coords_from_trajectories[0], + all_coords_from_trajectories[i], + decimal=10, + err_msg=f"Trajectory 0 and {i} have different coordinates", + ) + + def test_perturbed_shifts_differ_across_trajectories(self): + """Test that perturbed shifts differ across trajectories.""" + num_resid = 5 + num_trajectories = 3 + + # Generate base dataset + generator = FakeDataGenerator(num_resid) + pred_coords, obs_shifts, pred_shifts, noes, connectivity = generator.generate_data_arrays( + random_key=False + ) + pred_coords, obs_shifts, noes, connectivity = generator.order_data( + pred_coords, obs_shifts, pred_shifts, noes, connectivity + ) + + # Generate multiple trajectories + env = GymEnv(num_resid) + history_gen = FakeHistoryGenerator(num_resid) + original_state = env.reset(pred_coords, obs_shifts, noes, connectivity) + + all_shifts_from_trajectories = [] + + for i in range(num_trajectories): + perturbed_state = history_gen.generate_history(original_state) + trajectory = generate_trajectory(env, perturbed_state) + + # Extract shifts from first step of this trajectory + first_state, _, _ = trajectory[0] + traj_shifts = np.array([[s.H1, s.N15] for s in first_state["obs_chemical_shifts"]]) + all_shifts_from_trajectories.append(traj_shifts) + + # Verify that shifts differ between at least some trajectories + # (they should all be different due to random perturbation) + shifts_are_different = False + for i in range(len(all_shifts_from_trajectories) - 1): + for j in range(i + 1, len(all_shifts_from_trajectories)): + if not np.allclose( + all_shifts_from_trajectories[i], all_shifts_from_trajectories[j] + ): + shifts_are_different = True + break + if shifts_are_different: + break + + self.assertTrue( + shifts_are_different, + "Perturbed shifts should differ across trajectories due to random noise", + ) + + def test_noes_regenerated_consistently_from_coordinates(self): + """Test that NOEs are regenerated from original coordinates, not perturbed ones.""" + num_resid = 5 + num_trajectories = 3 + + # Generate base dataset + generator = FakeDataGenerator(num_resid) + pred_coords, obs_shifts, pred_shifts, noes, connectivity = generator.generate_data_arrays( + random_key=False + ) + pred_coords, obs_shifts, noes, connectivity = generator.order_data( + pred_coords, obs_shifts, pred_shifts, noes, connectivity + ) + + # Generate multiple perturbed states and check NOEs + env = GymEnv(num_resid) + history_gen = FakeHistoryGenerator(num_resid) + original_state = env.reset(pred_coords, obs_shifts, noes, connectivity) + + all_noes_from_trajectories = [] + + for i in range(num_trajectories): + perturbed_state = history_gen.generate_history(original_state) + trajectory = generate_trajectory(env, perturbed_state) + + # Extract NOEs from first step + first_state, _, _ = trajectory[0] + traj_noes = first_state["noes"] + all_noes_from_trajectories.append(traj_noes) + + # Since coordinates are the same across all trajectories, + # and NOEs are generated from coordinates + shifts with noise, + # the NOE list structure should be consistent (same number of NOEs) + noe_counts = [len(noes) for noes in all_noes_from_trajectories] + + # All trajectories should have the same number of NOEs + # (since they come from same coordinates with same cutoff) + self.assertTrue( + all(count == noe_counts[0] for count in noe_counts), + f"NOE counts differ across trajectories: {noe_counts}. " + "This suggests NOEs are not being regenerated consistently from coordinates.", + ) + + def test_trajectory_reproducibility_with_seed(self): + """Test that setting seed produces reproducible trajectories.""" + num_resid = 5 + + # Generate base dataset + generator = FakeDataGenerator(num_resid) + pred_coords, obs_shifts, pred_shifts, noes, connectivity = generator.generate_data_arrays( + random_key=False + ) + pred_coords, obs_shifts, noes, connectivity = generator.order_data( + pred_coords, obs_shifts, pred_shifts, noes, connectivity + ) + + # Generate trajectory 1 with seed + np.random.seed(42) + env1 = GymEnv(num_resid) + history_gen1 = FakeHistoryGenerator(num_resid) + original_state1 = env1.reset(pred_coords, obs_shifts, noes, connectivity) + perturbed_state1 = history_gen1.generate_history(original_state1) + trajectory1 = generate_trajectory(env1, perturbed_state1) + + # Generate trajectory 2 with same seed + np.random.seed(42) + env2 = GymEnv(num_resid) + history_gen2 = FakeHistoryGenerator(num_resid) + original_state2 = env2.reset(pred_coords, obs_shifts, noes, connectivity) + perturbed_state2 = history_gen2.generate_history(original_state2) + trajectory2 = generate_trajectory(env2, perturbed_state2) + + # Extract shifts from first step of each trajectory + shifts1 = np.array([[s.H1, s.N15] for s in trajectory1[0][0]["obs_chemical_shifts"]]) + shifts2 = np.array([[s.H1, s.N15] for s in trajectory2[0][0]["obs_chemical_shifts"]]) + + # Verify identical perturbations with same seed + np.testing.assert_array_almost_equal( + shifts1, + shifts2, + decimal=10, + err_msg="Same seed should produce identical perturbed shifts", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/training_loop.py b/training_loop.py deleted file mode 100644 index 0381600..0000000 --- a/training_loop.py +++ /dev/null @@ -1,205 +0,0 @@ -import torch -import pickle -import numpy as np -from nmr_transformer import * -from torch.utils.tensorboard import SummaryWriter -import random - - - -def extract_data(pickle_file="fake_histories_r20_1.pkl"): - with open(pickle_file, "rb") as f: - histories = pickle.load(f) - return histories - - -def preprocess_data(histories): - """ - Loops over chosen histories and answers, splits up components, and adds shifts if indices were used. - """ - predicted_shifts = [] - obs_chemical_shifts = [] - noes = [] - close_dist = [] - assignments = [] - peak_to_assign = [] - correct_answer = [] - - for history in histories: - # Predicted shifts - predicted_shifts_temp = [(i.H1, i.N15) for i in history["coordinates"]] - predicted_shifts.append(predicted_shifts_temp) - # Observed shifts - obs_chemical_shifts_temp = [(i.H1, i.N15) for i in history["obs_chemical_shifts"]] - obs_chemical_shifts.append(obs_chemical_shifts_temp) - # NOEs - noes.append(history["noes"]) - - # Close distances (duplicated for reverse pair order) - close = [] - for i in history["connectivity"]: - atom1 = predicted_shifts_temp[i.atom1] - atom2 = predicted_shifts_temp[i.atom2] - close.append((atom1[0], atom1[1], atom2[0], atom2[1])) - close.append((atom2[0], atom2[1], atom1[0], atom1[1])) - close_dist.append(close) - - # Assignments - assign = [] - for i, j in list(history["assignments"].items()): - shift1 = obs_chemical_shifts_temp[i] - shift2 = predicted_shifts_temp[j] - assign.append((shift1[0], shift1[1], shift2[0], shift2[1])) - if assign: - assignments.append(assign) - else: - assignments.append(None) - - # Peak to assign - peak_to_assign.append(history["shift_to_assign"]) - correct_answer.append(history["shift_to_assign"]) - - return ( - predicted_shifts, - obs_chemical_shifts, - noes, - close_dist, - assignments, - peak_to_assign, - correct_answer - ) - - -def organize_nmr_inputs(histories, device): - """ - Sets up NMR inputs as list of named tuples. - """ - ( - predicted_shifts, - actual_shifts, - noes, - close_dist, - assignments, - assign_peak, - correct_answer, - ) = preprocess_data(histories) - - nmr_inputs = [ - ( - NMRInput( - obs_chemical_shifts=torch.tensor(actual_shifts[i], dtype=torch.float, device=device), - pred_chemical_shifts=torch.tensor(predicted_shifts[i], dtype=torch.float, device=device), - obs_noes=torch.tensor(noes[i], dtype=torch.float, device=device), - close_distances=torch.tensor(close_dist[i], dtype=torch.float, device=device), - assigned_peaks=torch.tensor(assignments[i], dtype=torch.float, device=device) if assignments[i] else None, - peak_to_assign=assign_peak[i], - ), - correct_answer[i], - ) - for i in range(len(histories)) - ] - return nmr_inputs - - - -if __name__ == "__main__": - - # Set device - # device = "mps" - device ='cuda' - - # Set up TensorBoard - writer = SummaryWriter() - - # Load data - nmr_inputs = organize_nmr_inputs(extract_data(), device=device) - - # Shuffle the data - random.shuffle(nmr_inputs) - - # Split dataset manually into training and testing - train_size = int(0.8 * len(nmr_inputs)) - train_data = nmr_inputs[:train_size] - test_data = nmr_inputs[train_size:] - - # Create our network and optimizer - net = NMRTransformer(dropout=0.1, n_layers=4, device=device) - opt = torch.optim.AdamW(net.parameters(), lr=1e-4, weight_decay=0.01) - policy_loss = torch.nn.CrossEntropyLoss() - - # Training loop - batch_size = 10 - epochs = 500_000 # Define the number of epochs - eval_interval = 100 # Evaluate the model every 100 iterations - - iteration = 0 - for epoch in range(epochs): - random.shuffle(train_data) # Shuffle training data each epoch - - # Break training data into batches manually - for i in range(0, len(train_data), batch_size): - net.train() - batch = train_data[i:i + batch_size] - xs = [inp[0] for inp in batch] - ys = [inp[1] for inp in batch] - - iteration += 1 - - y_hats, _ = net(xs) - - loss = 0 - correct = 0 - count = 0 - - for y_hat, y in zip(y_hats, ys): - y = torch.tensor(y, device=device) - loss += policy_loss(y_hat, y) - - y_pred = torch.argmax(y_hat) - if y_pred == y: - correct += 1 - count += 1 - - loss = loss / count - accuracy = correct / count - - print(f"Iteration {iteration} - Loss: {loss.item()} - Accuracy: {accuracy}") - - writer.add_scalar("Loss/train", loss.item(), iteration) - writer.add_scalar("Accuracy/train", accuracy, iteration) - - opt.zero_grad() - loss.backward() - opt.step() - - # Evaluate on test set periodically - if iteration % eval_interval == 0: - net.eval() - test_loss = 0 - test_correct = 0 - test_trials = 0 - - with torch.no_grad(): - for j in range(0, len(test_data), batch_size): - test_batch = test_data[j:j + batch_size] - xs = [inp[0] for inp in test_batch] - ys = [inp[1] for inp in test_batch] - - test_y_hats, _ = net(xs) - - batch_loss = 0 - for y_hat, y in zip(test_y_hats, ys): - y = torch.tensor(y, device=device) - batch_loss += policy_loss(y_hat, y).item() - if torch.argmax(y_hat) == y: - test_correct += 1 - test_trials += 1 - - test_loss += batch_loss / len(test_batch) - - test_loss /= (len(test_data) / batch_size) - test_accuracy = test_correct / test_trials - - print(f"Test Loss: {test_loss} - Test Accuracy: {test_accuracy}") - writer.add_scalar("Loss/test", test_loss, iteration) - writer.add_scalar("Accuracy/test", test_accuracy, iteration) diff --git a/training_loop_gnn.py b/training_loop_gnn.py deleted file mode 100644 index 8917f6d..0000000 --- a/training_loop_gnn.py +++ /dev/null @@ -1,119 +0,0 @@ -import torch -import torch_geometric -import pickle -import numpy as np - -from pyg_triple_update import * -from torch.utils.tensorboard import SummaryWriter -import random - -def extract_data(pickle_file="fake_histories_r4_0.pkl"): - with open(pickle_file, "rb") as f: - histories = pickle.load(f) - return histories - -def construct_graph(history, device): - nodes = ConstructNodes(device) - data = nodes.construct_data(history) - edges = ConstructEdges(data, device) - data = edges.generate_edge_indices() - return data - -def preprocess_data(histories, device): - graphs = [] - for history in histories: - graphs.append(construct_graph(history, device)) - # print(history['assign_order']) - return graphs - -if __name__ == "__main__": - - # Set device - device = 'cuda' - - # Set up TensorBoard - # writer = SummaryWriter() - - # Load data - histories = extract_data() - nmr_graphs = preprocess_data(extract_data(), device) - - # Split dataset manually into training and testing - train_size = int(0.8 * len(nmr_graphs)) - train_data = nmr_graphs[:train_size] - test_data = nmr_graphs[train_size:] - train_target = [inp['shift_to_assign'] for inp in histories][:train_size] - test_target = [inp['shift_to_assign'] for inp in histories][train_size:] - - batch_size = 1 - epochs = 50_000 # Define the number of epochs - eval_interval = 100 # Evaluate the model every 100 iterations - - train_data_loader = DataLoader(train_data, batch_size=batch_size, shuffle=False) - test_data_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False) - - # Create our network and optimizer - net = TestLayer(device) - opt = torch.optim.AdamW(net.parameters(), lr=1e-4, weight_decay=0.01) - - # Interactive plotting - plt.ion() - fig, ax = plt.subplots() - - iteration = 0 - for epoch in range(epochs): - - for i, xs in enumerate(train_data_loader): - net.train() - - # Break targets into batches manually (if in same order - no shuffling) - ys = train_target[i*batch_size:(i+1)*batch_size] - - iteration += 1 - - _, y_hats = net(xs) - - # Plots current iteration in batch ########################################## - ax.cla() - x_val1, y_val1 = xs['SHIFT'].x[:, 0].tolist(), xs['SHIFT'].x[:, 1].tolist() - x_val2, y_val2 = xs['RES'].x[:, 3].tolist(), xs['RES'].x[:, 4].tolist() - - ax.scatter(x_val1, y_val1, color='red', label='shift') - ax.scatter(x_val2, y_val2, color='blue', label='resid') - - for j in range(len(x_val1)): - ax.annotate(j, (x_val1[j], y_val1[j]+0.5), color='red') - ax.annotate(j, (x_val2[j], y_val2[j]+0.5), color='blue') - - ax.legend() - ax.set_title(f"Iteration {iteration} - Batch {i}") - plt.pause(0.01) - ############################################################################# - - loss = 0 - correct = 0 - count = 0 - - for y_hat, y in zip(y_hats, ys): - y = torch.tensor([y], dtype=torch.long, device=device) - loss += torch.nn.functional.cross_entropy(y_hat, y) - y_pred = torch.argmax(y_hat) - - if y_pred == y: - correct += 1 - count += 1 - - loss = loss / count - accuracy = correct / count - - print(f"Iteration {iteration} - Loss: {loss.item()} - Accuracy: {accuracy}") - - # writer.add_scalar("Loss/train", loss.item(), iteration) - # writer.add_scalar("Accuracy/train", accuracy, iteration) - - opt.zero_grad() - loss.backward() - opt.step() - - plt.ioff() - plt.show() From ebdd27ccf5b8ebcf67f103a70a54afaa16ac7d71 Mon Sep 17 00:00:00 2001 From: Justin MacCallum Date: Tue, 11 Nov 2025 20:27:30 -0700 Subject: [PATCH 41/43] Remove pkl file --- fake_histories_r4_0.pkl | Bin 6702 -> 0 bytes 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 fake_histories_r4_0.pkl diff --git a/fake_histories_r4_0.pkl b/fake_histories_r4_0.pkl deleted file mode 100644 index f1ccb4dffb372b746620f1152b1fade6e2a02cb8..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 6702 zcmZo*nW`+s00yyBG-{`4^l&HV=NF}9<|US-7Eb|+^>C&oW~audB$gyj>0u8j$}dUH z%$qV9C?T3xnp;q*mz-aes+U`uQ<9ljRFqgbrH8FJIWZ@(2&krq6|4kEr<7C{q=H14 z(k!NQcC=3knxf&&+{0)y#m~>r>pu{H32%myDM`*86XLyPV%$&KPiX_H%izr5cjoXC zmz;g!%u8gUb9IO7ciwq|EabbbhjE98oCA`^Wfl1a@=`Gnp^_<$Q!>~zAVwb9C}w}k z@)ELsht=0?BTQZ+3wa-0t7sB;8(D~Xua5ffJIctm)Ol!WoQ;dZVoOIa?}16Y_mTBq zIU-ZT@aP`0kW1ocgV6ko$U=-$ukARWgknqYp;vcaDkfmDrAU5n%&%}1zw}(BjWgeKY(#OmLs^i(*&t}}w1FZa5~6mKqvn*zGqT9m z=a$!>y(|HZFl4nWR6V)aqNR}4ew?8-v%C!IRAjZToPs;IR-m}yK>WF@JC&h<2U45b z!;+T|N#sC(u>1MDqIkn^5l4#gQ52^!U*g=}FAjAY!~^Dh6@5E5LoI3t`K%Tqaep7* zlSS)LTo4(2^SWm+G!!7J4E~$PCc!L0)?QRL{q`!Dxe&F7g#HIVhXo=;$SCJdWl1d- z+e@zhV|mhuV*AmudX{A{|3kE2Q)are7-k8w?FYlket9*aXuq0Lb5mRe8hjA#p2uXD zH-$qZ1X+9E`>LHQqEKS0@X5SK;joy3XsnCP0a)5nA#~lJP>|7u#8PD0{f6NgQtfn z6RA|p5WrR@LJDaH23XPd!kL2^DuyJ)h$+N`ER@>A56->uCHe6XmrTic)5BAeUjoe6 zsd=eI>6KI5VOI696{VIZ7NsCWN`W;*1`{Y=V1?*tQ4Oq|22oL+5eAAKXz4m+O5%(v zkgK63=V-x8_kuU$04V$q52>Ox<5BIDj3*#>b7lYwLs+4OR2Y&|K4yH&ItDC6Grq%1 z!Hgep=4kQf%rUrIs-wl Date: Tue, 11 Nov 2025 22:48:21 -0700 Subject: [PATCH 42/43] Add normalization layer --- nmr/models/network.py | 59 ++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 55 insertions(+), 4 deletions(-) diff --git a/nmr/models/network.py b/nmr/models/network.py index 56d3701..7d940f5 100644 --- a/nmr/models/network.py +++ b/nmr/models/network.py @@ -5,12 +5,8 @@ to create the complete neural network for NMR assignment. """ -import pickle - -import matplotlib.pyplot as plt import torch import torch.nn as nn -from torch_geometric.loader import DataLoader from .triple import CalculationManager, TripleIn, TripleOut, TripleUpdate from .heads import ValueCalc, PolicyCalc @@ -181,6 +177,59 @@ def forward(self, data): return data +class NormalizeShifts(nn.Module): + """ + Normalize chemical shifts to have a typical range between -1 and 1 + """ + + def __init__(self, H_lower=6.0, H_upper=10.0, N_lower=100.0, N_upper=135.0): + super().__init__() + self.H_lower = H_lower + self.H_upper = H_upper + self.H_delta = H_upper - H_lower + self.N_lower = N_lower + self.N_upper = N_upper + self.N_delta = N_upper - N_lower + + def forward(self, data): + # split RES into coords and shifts + res_x = data["RES"].x + res_xyz = res_x[:, :3] + res_H = res_x[:, 3].unsqueeze(-1) + res_N = res_x[:, 4].unsqueeze(-1) + # split SHIFTS into 1H and 15N + shift_H = data["SHIFT"].x[:, 0].unsqueeze(-1) + shift_N = data["SHIFT"].x[:, 1].unsqueeze(-1) + # split NOEs into 1H and 15N + noe_H1 = data["NOE"].x[:, 0].unsqueeze(-1) + noe_N1 = data["NOE"].x[:, 1].unsqueeze(-1) + noe_H2 = data["NOE"].x[:, 2].unsqueeze(-1) + # transform all 1H + res_H = self._transform_H(res_H) + shift_H = self._transform_H(shift_H) + noe_H1 = self._transform_H(noe_H1) + noe_H2 = self._transform_H(noe_H2) + # transofrm all 15N + res_N = self._transform_N(res_N) + shift_N = self._transform_N(shift_N) + noe_N1 = self._transform_N(noe_N1) + # reassemble tensors + res_x = torch.cat([res_xyz, res_H, res_N], dim=-1) + shift_x = torch.cat([shift_H, shift_N], dim=-1) + noe_x = torch.cat([noe_H1, noe_N1, noe_H2], dim=-1) + + data["RES"].x = res_x + data["SHIFT"].x = shift_x + data["NOE"].x = noe_x + return data + + def _transform_H(self, value): + return 2 * (value - self.H_lower) / self.H_delta - 1 + + def _transform_N(self, value): + return 2 * (value - self.N_lower) / self.N_delta - 1 + + class NMRNet(nn.Module): """ Complete NMR GNN model combining message passing with prediction heads. @@ -194,11 +243,13 @@ def __init__(self, device): self.value = ValueCalc(device) self.policy = PolicyCalc(device) + self.normalize = NormalizeShifts() self.nmr = nn.Sequential( NMRLayer(device), ) def forward(self, data): + data = self.normalize(data) out_data = self.nmr(data) value = self.value.calc_value(out_data) policy = self.policy.calc_policy(out_data) From a35c957e1eb9de77f39024f68753f91f60deddae Mon Sep 17 00:00:00 2001 From: Justin MacCallum Date: Fri, 14 Nov 2025 13:36:53 -0700 Subject: [PATCH 43/43] Updated architecture --- nmr/construct.py | 390 ++++-- nmr/models/__init__.py | 75 +- nmr/models/heads.py | 113 +- nmr/models/network.py | 587 ++++++--- nmr/models/pair.py | 429 ++++++ nmr/models/triple.py | 1735 +++++++++++++++++++++---- scripts/train_gnn.py | 154 ++- tests/test_gather_components.py | 303 +++++ tests/test_gather_components.py.bak | 303 +++++ tests/test_network_orchestration.py | 146 +++ tests/test_scatter_operations.py | 218 ++++ tests/test_scatter_operations.py.bak | 212 +++ tests/test_terminology_updates.py | 98 ++ tests/test_triple_compositions.py | 351 +++++ tests/test_triple_compositions.py.bak | 345 +++++ tests/test_triple_helpers.py | 224 ++++ tests/test_triple_integration.py | 564 ++++++++ tests/test_update_components.py | 490 +++++++ 18 files changed, 6034 insertions(+), 703 deletions(-) create mode 100644 nmr/models/pair.py create mode 100644 tests/test_gather_components.py create mode 100644 tests/test_gather_components.py.bak create mode 100644 tests/test_network_orchestration.py create mode 100644 tests/test_scatter_operations.py create mode 100644 tests/test_scatter_operations.py.bak create mode 100644 tests/test_terminology_updates.py create mode 100644 tests/test_triple_compositions.py create mode 100644 tests/test_triple_compositions.py.bak create mode 100644 tests/test_triple_helpers.py create mode 100644 tests/test_triple_integration.py create mode 100644 tests/test_update_components.py diff --git a/nmr/construct.py b/nmr/construct.py index 0dbaabe..e3a35ce 100644 --- a/nmr/construct.py +++ b/nmr/construct.py @@ -1,11 +1,35 @@ -"""Graph construction utilities for NMR assignment.""" +""" +Graph construction utilities for NMR assignment. + +This module constructs heterogeneous graphs from NMR data for chemical shift assignment +using Graph Neural Networks. The graph structure includes: + +Node Types: +- Residue: Protein residues with coordinates [x,y,z] and predicted shifts [H,N] +- Peak: Observed chemical shifts [H,N] +- Noe: NOE distance constraints [N, H', H"] +- Triple nodes: Four types representing different relationship configurations: + * ResidueResidueNoeTriple: (Residue, Residue, Noe) - Updates coordinates and shifts + * ResiduePeakNoeTriple: (Residue, Peak, Noe) - Updates shifts only + * PeakResidueNoeTriple: (Peak, Residue, Noe) - Updates shifts only + * PeakPeakNoeTriple: (Peak, Peak, Noe) - Updates shifts only +- Value aggregation nodes: VALUE_NOE, VALUE_SHIFT, VALUE_RES + +Edge Types: +- Bidirectional propagation edges: Used for both gather (source → triple) and scatter (triple → source) + * (source_node, "prop_first", triple_type) - First position connections + * (source_node, "prop_second", triple_type) - Second position connections + * ("Noe", "prop_noe", triple_type) - NOE constraint connections + Direction is controlled by flow parameter in MessagePassing layers +- Value aggregation edges: Aggregate features for value prediction + * (node_type, "aggregate", value_node_type) +""" -from itertools import product from typing import Any, Dict, Tuple -import numpy as np import torch from torch_geometric.data import HeteroData +from time import time def construct_graph(history: Dict[str, Any], device: torch.device | str) -> HeteroData: @@ -30,6 +54,8 @@ def construct_node_data( """ Builds graph nodes from input histories. + Creates nodes for Residue, Peak, Noe, triple types, and value aggregation nodes. + Args: histories: Dictionary containing coordinates, shifts, NOEs, and assignment state device: Device to place tensors on ('cpu' or 'cuda') @@ -52,6 +78,9 @@ def construct_edges(data: HeteroData, device: torch.device | str) -> HeteroData: """ Constructs all edge indices for the heterogeneous graph. + Creates bidirectional propagation edges, value aggregation edges, and policy edges + for all triple configurations. + Args: data: HeteroData graph with nodes already constructed device: Device to place tensors on ('cpu' or 'cuda') @@ -59,233 +88,328 @@ def construct_edges(data: HeteroData, device: torch.device | str) -> HeteroData: Returns: HeteroData graph with all edges constructed """ - num_noe = len(data["NOE"].x) - num_shift = len(data["SHIFT"].x) - num_res = len(data["RES"].x) + num_noe = len(data["Noe"].x) + num_peak = len(data["Peak"].x) + num_residue = len(data["Residue"].x) - # Add all edge types + # Add all edge types for each triple configuration _add_triple_edges( - data, 0, ("RES", "RES", "NOE"), num_noe, num_shift, num_res, device + data, ("Residue", "Residue", "Noe"), num_noe, num_peak, num_residue, device ) _add_triple_edges( - data, 1, ("RES", "SHIFT", "NOE"), num_noe, num_shift, num_res, device + data, ("Residue", "Peak", "Noe"), num_noe, num_peak, num_residue, device ) _add_triple_edges( - data, 2, ("SHIFT", "RES", "NOE"), num_noe, num_shift, num_res, device + data, ("Peak", "Residue", "Noe"), num_noe, num_peak, num_residue, device ) _add_triple_edges( - data, 3, ("SHIFT", "SHIFT", "NOE"), num_noe, num_shift, num_res, device + data, ("Peak", "Peak", "Noe"), num_noe, num_peak, num_residue, device ) - _add_value_aggregation_edges(data, num_noe, num_shift, num_res, device) - + _add_value_aggregation_edges(data, num_noe, num_peak, num_residue, device) return data def _construct_data_nodes( data: HeteroData, histories: Dict[str, Any], device: torch.device | str ) -> HeteroData: - """Constructs NOE, SHIFT, and RES nodes with initial features.""" - data["NOE"].x = torch.tensor(histories["noes"], dtype=torch.float32, device=device) - data["NOE"].f = torch.zeros( + """ + Constructs Noe, Peak, and Residue nodes with initial features. + + Args: + data: HeteroData graph to add nodes to + histories: Dictionary containing coordinates, shifts, and NOEs + device: Device to place tensors on + + Returns: + HeteroData with data nodes added + """ + data["Noe"].x = torch.tensor(histories["noes"], dtype=torch.float32, device=device) + data["Noe"].f = torch.zeros( (len(histories["noes"]), 2), dtype=torch.float32, device=device ) - data["SHIFT"].x = torch.tensor( + data["Peak"].x = torch.tensor( histories["obs_chemical_shifts"], dtype=torch.float32, device=device ) - data["SHIFT"].f = torch.zeros( + data["Peak"].f = torch.zeros( (len(histories["obs_chemical_shifts"]), 2), dtype=torch.float32, device=device, ) - data["RES"].x = torch.tensor( + # Residue.x = [coordinates (3), predicted_shifts (2)] = [5 features total] + # Coordinates already contains [x, y, z, H, N] as 5 features + data["Residue"].x = torch.tensor( histories["coordinates"], dtype=torch.float32, device=device ) - data["RES"].f = torch.zeros( + data["Residue"].f = torch.zeros( (len(histories["coordinates"]), 2), dtype=torch.float32, device=device ) return data -def _construct_node_features(data: HeteroData, histories: Dict[str, Any]) -> HeteroData: - """Sets node features based on assignment status.""" - for i in range(len(data["SHIFT"].f)): - # is the shift being assigned right now? - if i == histories["shift_to_assign"]: - data["SHIFT"].f[i, 0] = 1 - # has the shift already been assigned? - if i in histories["assignments"].keys(): - data["SHIFT"].f[i, 1] = 1 - - for i in range(len(data["RES"].f)): - # has the residue been assigned? - if i in histories["assignments"].values(): - data["RES"].f[i] = 1 - return data - - def _construct_triple_nodes(data: HeteroData, device: torch.device | str) -> HeteroData: """ - Generates triple node types. Nothing is ever stored in these - they are simply placeholders to construct/use in edge types. + Constructs triple nodes for all four triple types. + + Triple nodes act as intermediaries for message passing between residues, peaks, and NOEs. + Each triple type represents a different configuration of source nodes. + + Args: + data: HeteroData graph to add triple nodes to + device: Device to place tensors on + + Returns: + HeteroData with triple nodes added """ - data["TRIPLE0"].x = torch.zeros((1, 1), dtype=torch.float32, device=device) - data["TRIPLE1"].x = torch.zeros((1, 1), dtype=torch.float32, device=device) - data["TRIPLE2"].x = torch.zeros((1, 1), dtype=torch.float32, device=device) - data["TRIPLE3"].x = torch.zeros((1, 1), dtype=torch.float32, device=device) + num_noe = len(data["Noe"].x) + num_peak = len(data["Peak"].x) + num_residue = len(data["Residue"].x) + + # ResidueResidueNoeTriple: All combinations of (residue_i, residue_j, noe_k) + num_res_res_noe = num_residue * num_residue * num_noe + data["ResidueResidueNoeTriple"].x = torch.zeros( + (num_res_res_noe, 1), dtype=torch.float32, device=device + ) + data["ResidueResidueNoeTriple"].f = torch.zeros( + (num_res_res_noe, 2), dtype=torch.float32, device=device + ) + + # ResiduePeakNoeTriple: All combinations of (residue_i, peak_j, noe_k) + num_res_peak_noe = num_residue * num_peak * num_noe + data["ResiduePeakNoeTriple"].x = torch.zeros( + (num_res_peak_noe, 1), dtype=torch.float32, device=device + ) + data["ResiduePeakNoeTriple"].f = torch.zeros( + (num_res_peak_noe, 2), dtype=torch.float32, device=device + ) + + # PeakResidueNoeTriple: All combinations of (peak_i, residue_j, noe_k) + num_peak_res_noe = num_peak * num_residue * num_noe + data["PeakResidueNoeTriple"].x = torch.zeros( + (num_peak_res_noe, 1), dtype=torch.float32, device=device + ) + data["PeakResidueNoeTriple"].f = torch.zeros( + (num_peak_res_noe, 2), dtype=torch.float32, device=device + ) + + # PeakPeakNoeTriple: All combinations of (peak_i, peak_j, noe_k) + num_peak_peak_noe = num_peak * num_peak * num_noe + data["PeakPeakNoeTriple"].x = torch.zeros( + (num_peak_peak_noe, 1), dtype=torch.float32, device=device + ) + data["PeakPeakNoeTriple"].f = torch.zeros( + (num_peak_peak_noe, 2), dtype=torch.float32, device=device + ) + return data def _construct_value_nodes(data: HeteroData, device: torch.device | str) -> HeteroData: - """Constructs value aggregation nodes for NOE, SHIFT, and RES.""" - data["VALUE_NOE"].x = torch.zeros( - 1, 1, dtype=torch.float32, device=device - ) # has to be num graphs batched? - data["VALUE_SHIFT"].x = torch.zeros(1, 1, dtype=torch.float32, device=device) - data["VALUE_RES"].x = torch.zeros(1, 1, dtype=torch.float32, device=device) + """ + Constructs value aggregation nodes for value prediction head. + + Value nodes aggregate information from Noe, Peak, and Residue nodes + to predict the quality of the current assignment state. + + Args: + data: HeteroData graph to add value nodes to + device: Device to place tensors on + + Returns: + HeteroData with value nodes added + """ + data["VALUE_NOE"].x = torch.zeros((1, 1), dtype=torch.float32, device=device) + data["VALUE_SHIFT"].x = torch.zeros((1, 1), dtype=torch.float32, device=device) + data["VALUE_RES"].x = torch.zeros((1, 1), dtype=torch.float32, device=device) return data -def _get_triple_edges( - num_noe: int, source1: int, source2: int, device: torch.device | str -) -> torch.Tensor: +def _construct_node_features(data: HeteroData, histories: Dict[str, Any]) -> HeteroData: """ - Grabs all index combinations from the three ranges (residue, shift, noe) and orders these into source and target indices for the edges. - Combinations occur along columns into a triple node. + Initializes node features based on assignment state. + + Sets feature flags on Peak and Residue nodes to indicate which nodes + are currently being assigned or have been assigned. Args: - num_noe: Number of NOE nodes - source1: Number of nodes for first source type - source2: Number of nodes for second source type - device: Device to place tensors on + data: HeteroData graph with nodes + histories: Dictionary containing assignment state Returns: - Stacked tensor of source and target node indices + HeteroData with node features initialized """ - combo = list(product(range(num_noe), range(source1), range(source2))) - combo_tensor = torch.tensor(combo, device=device) - - # Switches tensor dimension for source (columns set up for each edge [0,0,0] --> [0],[0],[0]) - source_nodes = torch.transpose(combo_tensor, 0, 1) - # Repeats target edges for number of occurences in source (3 for the triple in this instance, to be assigned to each incoming node type) - target_nodes = torch.tensor(range(len(source_nodes[0])), device=device).repeat( - len(source_nodes), 1 - ) + # Mark the peak being assigned + shift_to_assign = int(histories["shift_to_assign"]) + data["Peak"].f[shift_to_assign, 0] = 1.0 + + # Mark assigned peaks and residues + for shift_idx, residue_idx in histories["assignments"].items(): + data["Peak"].f[int(shift_idx), 1] = 1.0 + data["Residue"].f[int(residue_idx), 1] = 1.0 + + # Add edges for existing assignments (peak -> residue mappings) + # Always add the edge type, even if empty, for consistent graph structure + if histories["assignments"]: + peak_indices = list(histories["assignments"].keys()) + residue_indices = list(histories["assignments"].values()) + data["Peak", "assigned_to", "Residue"].edge_index = torch.tensor( + [peak_indices, residue_indices], + dtype=torch.long, + device=data["Peak"].x.device, + ) + else: + # Create empty edge_index with shape [2, 0] + data["Peak", "assigned_to", "Residue"].edge_index = torch.empty( + (2, 0), dtype=torch.long, device=data["Peak"].x.device + ) - return torch.stack([source_nodes, target_nodes], dim=0) + return data -def _get_pairwise_edges( - data: HeteroData, num_res: int, device: torch.device | str +def _get_triple_edges( + num_noe: int, source1: int, source2: int, device: torch.device | str ) -> torch.Tensor: """ - Grabs index combinations between shifts and residues and orders these into source and target indices for the edges. + Generates edge indices for connecting triple nodes to their source nodes. + + Each triple node represents a combination of (NOE, source1, source2). This function + creates the edge connectivity to link each triple back to its three constituent nodes. Args: - data: HeteroData graph with shift features - num_res: Number of residue nodes + num_noe: Number of NOE nodes + source1: Number of nodes in first source type (Residue or Peak) + source2: Number of nodes in second source type (Residue or Peak) device: Device to place tensors on Returns: - Stacked tensor of residue and shift indices for pairwise edges + Tensor of shape [2, 3, num_edges] where: + - dim 0: [source_indices, triple_indices] + - dim 1: [noe_position, first_position, second_position] + - dim 2: all num_noe * source1 * source2 edge combinations + + For example, edges[0, 1, k] gives the source1 node index for triple k, + and edges[1, 1, k] gives the triple node index k. """ - # shift = torch.arange(0, num_shift) - shift = np.nonzero(data["SHIFT"].f[:, 0])[ - 0 - ] # need to specify if it's the shift being assigned - resid = torch.arange(0, num_res, device=device) + total_edges = num_noe * source1 * source2 - shift_repeats = shift.repeat_interleave(num_res) - # resid_repeats = resid.repeat(num_shift) - resid_repeats = resid - return torch.stack((resid_repeats, shift_repeats), dim=0) + # Create source node indices for all triple combinations + # Each triple (i, j, k) corresponds to edge index: i*(source1*source2) + j*source2 + k + + # NOE indices: [0,0,...,0, 1,1,...,1, 2,2,...,2, ...] (each repeated source1*source2 times) + noe_indices = torch.arange(num_noe, device=device).repeat_interleave(source1 * source2) + + # Source1 indices: [0,0,...,0, 1,1,...,1, ...] (each repeated source2 times, pattern repeats num_noe times) + source1_indices = torch.arange(source1, device=device).repeat_interleave(source2).repeat(num_noe) + + # Source2 indices: [0,1,2,...,source2-1, 0,1,2,...,source2-1, ...] (cycles continuously) + source2_indices = torch.arange(source2, device=device).repeat(num_noe * source1) + + # Stack into source_nodes tensor [3, total_edges] + source_nodes = torch.stack([noe_indices, source1_indices, source2_indices], dim=0) + + # Create target indices [0, 1, 2, ..., total_edges-1] repeated 3 times + target_nodes = torch.arange(total_edges, device=device).unsqueeze(0).repeat(3, 1) + + return torch.stack([source_nodes, target_nodes], dim=0) def _add_value_aggregation_edges( data: HeteroData, num_noe: int, - num_shift: int, - num_res: int, + num_peak: int, + num_residue: int, device: torch.device | str, ) -> None: """ Adds value aggregation edges from data nodes to value nodes. + These edges allow the value prediction head to aggregate information + from all Noe, Peak, and Residue nodes. + Args: data: HeteroData graph to add edges to num_noe: Number of NOE nodes - num_shift: Number of shift nodes - num_res: Number of residue nodes + num_peak: Number of peak nodes + num_residue: Number of residue nodes device: Device to place tensors on """ - shift = torch.arange(0, num_shift, device=device) + peak = torch.arange(0, num_peak, device=device) noe = torch.arange(0, num_noe, device=device) - resid = torch.arange(0, num_res, device=device) + resid = torch.arange(0, num_residue, device=device) - shift_repeats = torch.zeros(num_shift, device=device).long() + peak_repeats = torch.zeros(num_peak, device=device).long() noe_repeats = torch.zeros(num_noe, device=device).long() - resid_repeats = torch.zeros(num_res, device=device).long() + resid_repeats = torch.zeros(num_residue, device=device).long() - data["SHIFT", "SHIFT_extract", "VALUE_SHIFT"].edge_index = torch.stack( - (shift, shift_repeats), dim=0 + data["Peak", "SHIFT_extract", "VALUE_SHIFT"].edge_index = torch.stack( + (peak, peak_repeats), dim=0 ) - data["NOE", "NOE_extract", "VALUE_NOE"].edge_index = torch.stack( + data["Noe", "aggregate", "VALUE_NOE"].edge_index = torch.stack( (noe, noe_repeats), dim=0 ) - data["RES", "RES_extract", "VALUE_RES"].edge_index = torch.stack( + data["Residue", "RES_extract", "VALUE_RES"].edge_index = torch.stack( (resid, resid_repeats), dim=0 ) def _add_triple_edges( data: HeteroData, - triple_num: int, triple_type: Tuple[str, str, str], num_noe: int, - num_shift: int, - num_res: int, + num_peak: int, + num_residue: int, device: torch.device | str, ) -> None: """ - Constructs all edges for a single triple type. + Constructs bidirectional propagation edges for a single triple type. + + Creates edges that can be traversed in both directions for message passing: + - Gather direction: source → triple (default flow='source_to_target') + - Scatter direction: triple → source (with flow='target_to_source') + + Triple types use descriptive names: + - ResidueResidueNoeTriple: Residue-Residue-Noe configuration + - ResiduePeakNoeTriple: Residue-Peak-Noe configuration + - PeakResidueNoeTriple: Peak-Residue-Noe configuration + - PeakPeakNoeTriple: Peak-Peak-Noe configuration + + Edge naming convention: + - (source_node, "prop_first", triple_name) - First position bidirectional edges + - (source_node, "prop_second", triple_name) - Second position bidirectional edges + - ("Noe", "prop_noe", triple_name) - NOE bidirectional edges Args: data: HeteroData graph to add edges to - triple_num: Triple type number (0-3) triple_type: Tuple of (source1, source2, source3) node type names num_noe: Number of NOE nodes - num_shift: Number of shift nodes - num_res: Number of residue nodes + num_peak: Number of peak nodes + num_residue: Number of residue nodes device: Device to place tensors on """ - source1, source2, source3 = triple_type - - # Need to make sure I'm using the right node range (shift and residue number could vary) - num_node1 = num_res if source1 == "RES" else num_shift - num_node2 = num_res if source2 == "RES" else num_shift - - # Triple in - edges_in = _get_triple_edges(num_noe, num_node1, num_node2, device) - data[f"{source3}", "NOE_extract", f"TRIPLE{triple_num}"].edge_index = edges_in[:, 0] - data[f"{source1}", "NH1_extract", f"TRIPLE{triple_num}"].edge_index = edges_in[:, 1] - data[f"{source2}", "NH2_extract", f"TRIPLE{triple_num}"].edge_index = edges_in[:, 2] - - # Triple out (reverse in/out ordering) - edges_out = torch.stack([edges_in[1], edges_in[0]], dim=0) - data[f"TRIPLE{triple_num}", "NOE_add", f"{source3}"].edge_index = edges_out[:, 0] - data[f"TRIPLE{triple_num}", "NH1_add", f"{source1}"].edge_index = edges_out[:, 1] - data[f"TRIPLE{triple_num}", "NH2_add", f"{source2}"].edge_index = edges_out[:, 2] - - # Only residue nodes will have coordinate edges - if source1 == "RES": - data[f"TRIPLE{triple_num}", "res1_add", "RES"].edge_index = edges_out[:, 1] - if source2 == "RES": - data[f"TRIPLE{triple_num}", "res2_add", "RES"].edge_index = edges_out[:, 2] - - # Edges for self loop - data[f"TRIPLE{triple_num}", "update", f"TRIPLE{triple_num}"].edge_index = ( - torch.stack([edges_in[1, 0], edges_in[1, 0]], dim=0) + source1, source2, _ = triple_type + + # Determine source counts based on node types + source1_count = num_residue if source1 == "Residue" else num_peak + source2_count = num_residue if source2 == "Residue" else num_peak + + # Generate triple name from types + triple_name = f"{source1}{source2}NoeTriple" + + # Get edge indices for all combinations + # edges shape: [2, 3, num_edges] where 2 = [source, target], 3 = [noe, first, second] + edges = _get_triple_edges(num_noe, source1_count, source2_count, device) + + # Bidirectional propagation edges: source → triple (can be reversed for scatter) + # First position (source1 ↔ triple) + data[source1, "prop_first", triple_name].edge_index = torch.stack( + [edges[0, 1], edges[1, 1]], dim=0 + ) + + # Second position (source2 ↔ triple) + data[source2, "prop_second", triple_name].edge_index = torch.stack( + [edges[0, 2], edges[1, 2]], dim=0 ) - if triple_num == 0: - data["RES", "pair", "SHIFT"].edge_index = _get_pairwise_edges( - data, num_res, device - ) # pairwise edges - no message passing + # NOE position (Noe ↔ triple) + data["Noe", "prop_noe", triple_name].edge_index = torch.stack( + [edges[0, 0], edges[1, 0]], dim=0 + ) diff --git a/nmr/models/__init__.py b/nmr/models/__init__.py index 0808e10..6fbdd53 100644 --- a/nmr/models/__init__.py +++ b/nmr/models/__init__.py @@ -6,15 +6,41 @@ """ # Core network components -from .network import NMRLayer, NMRNet +from .network import ( + NMRLayer, + NMRNet, + ModelConfig, + FeatureEmbedConfig, + ShiftEmbedConfig, + MLPConfig, +) -# Triple message passing +# Triple message passing - new modular architecture from .triple import ( - CalculationManager, - TripleIn, - TripleMessagePass, - TripleOut, - TripleUpdate, + # Gather components + FirstResidueGather, + FirstPeakGather, + SecondResidueGather, + SecondPeakGather, + NoeGather, + # Update components + ResidueUpdate, + PeakUpdate, + # Scatter components + FirstResidueScatter, + FirstPeakScatter, + SecondResidueScatter, + SecondPeakScatter, + NoeScatter, + # Triple composition + ResidueResidueNoeTriple, + ResiduePeakNoeTriple, + PeakResidueNoeTriple, + PeakPeakNoeTriple, + # Helper functions + calc_noe_difference, + calc_shift_difference, + calc_res_distance, ) # Prediction heads @@ -24,12 +50,35 @@ # Network "NMRLayer", "NMRNet", - # Triple components - "CalculationManager", - "TripleIn", - "TripleUpdate", - "TripleMessagePass", - "TripleOut", + # Configuration + "ModelConfig", + "FeatureEmbedConfig", + "ShiftEmbedConfig", + "MLPConfig", + # Gather components + "FirstResidueGather", + "FirstPeakGather", + "SecondResidueGather", + "SecondPeakGather", + "NoeGather", + # Update components + "ResidueUpdate", + "PeakUpdate", + # Scatter components + "FirstResidueScatter", + "FirstPeakScatter", + "SecondResidueScatter", + "SecondPeakScatter", + "NoeScatter", + # Triple composition + "ResidueResidueNoeTriple", + "ResiduePeakNoeTriple", + "PeakResidueNoeTriple", + "PeakPeakNoeTriple", + # Helper functions + "calc_noe_difference", + "calc_shift_difference", + "calc_res_distance", # Heads "BatchMessagePass", "ValueCalc", diff --git a/nmr/models/heads.py b/nmr/models/heads.py index 555197f..c986522 100644 --- a/nmr/models/heads.py +++ b/nmr/models/heads.py @@ -16,13 +16,19 @@ class BatchMessagePass(MessagePassing): Used by value and policy heads to collect information across the graph. """ - def __init__(self, aggr, device): + def __init__(self, aggr, device, config): super().__init__(aggr=aggr) self.device = device + self.config = config - self.noe_reduce = nn.Linear(3, 1, device=self.device) - self.shift_reduce = nn.Linear(2, 1, device=self.device) - self.res_reduce = nn.Linear(5, 1, device=self.device) + # After embeddings: + # - NOE: shift_dim (embedded from 3D) + # - Peak: shift_dim (embedded from 2D) + # - Residue: 3 (coords) + shift_dim + # Note: NOE and Peak both have shift_dim, so they share the same reduction layer + shift_dim = config.shift_embed.output_dim + self.shift_reduce = nn.Linear(shift_dim, 1, device=self.device) + self.res_reduce = nn.Linear(3 + shift_dim, 1, device=self.device) def forward(self, x_source, x_target, edge_index): # SIZE (n, m) (source, target) @@ -33,12 +39,18 @@ def forward(self, x_source, x_target, edge_index): ) def message(self, x_j): - if len(x_j[0]) == 3: - return self.noe_reduce(x_j) - if len(x_j[0]) == 2: + dim = x_j.size(1) + shift_dim = self.config.shift_embed.output_dim + + # Check dimension to determine node type + if dim == shift_dim: + # NOE and Peak both have shift_dim after embedding, use same reduction return self.shift_reduce(x_j) - if len(x_j[0]) == 5: + elif dim == 3 + shift_dim: + # Residue nodes: 3 coords + shift_dim return self.res_reduce(x_j) + else: + raise ValueError(f"Unexpected feature dimension: {dim}") def update(self, aggr_out, x): return aggr_out @@ -48,19 +60,21 @@ class ValueCalc(nn.Module): """ Value function estimation head. - Aggregates information from NOE, SHIFT, and RES nodes to predict + Aggregates information from Noe, Peak, and Residue nodes to predict a scalar value representing the quality of the current state. """ - def __init__(self, device): + def __init__(self, device, config): super().__init__() self.device = device + self.config = config - self.batch_message = BatchMessagePass(aggr="mean", device=self.device) + self.batch_message = BatchMessagePass(aggr="mean", device=self.device, config=config) self.hidden = 64 self.testmlp = nn.Sequential( nn.Linear(3, self.hidden, device=self.device), + nn.LayerNorm(self.hidden, device=self.device), nn.ReLU(), nn.Linear(self.hidden, 1, device=self.device), ) @@ -70,20 +84,31 @@ def get_aggr(self, x_source, x_target, edge_index): return aggr def calc_value(self, data): + """ + Calculate value estimate for the current state. + + Aggregates features from Noe, Peak, and Residue nodes using value aggregation edges. + + Args: + data: HeteroData graph with Noe, Peak, Residue nodes and value aggregation edges + + Returns: + Value tensor [batch_size, 1] + """ noe = self.get_aggr( - data["NOE"].x, + data["Noe"].x, data["VALUE_NOE"].x, - data["NOE", "NOE_extract", "VALUE_NOE"].edge_index, + data["Noe", "aggregate", "VALUE_NOE"].edge_index, ) shift = self.get_aggr( - data["SHIFT"].x, + data["Peak"].x, data["VALUE_SHIFT"].x, - data["SHIFT", "SHIFT_extract", "VALUE_SHIFT"].edge_index, + data["Peak", "SHIFT_extract", "VALUE_SHIFT"].edge_index, ) resid = self.get_aggr( - data["RES"].x, + data["Residue"].x, data["VALUE_RES"].x, - data["RES", "RES_extract", "VALUE_RES"].edge_index, + data["Residue", "RES_extract", "VALUE_RES"].edge_index, ) concat_aggr = torch.cat((noe, shift, resid), dim=-1) @@ -99,47 +124,65 @@ class PolicyCalc(nn.Module): based on pairwise distances between residue and shift features. """ - def __init__(self, device): + def __init__(self, device, config): super().__init__() - self.RES_NH = slice(3, 5) self.device = device + self.config = config + # After embeddings, residue shifts start at index 3 and go to the end + # No need for fixed slice since we'll use [:, 3:] dynamically def calc_policy(self, data): """ Calculate policy logits for each graph in the batch. - For each graph, computes squared distances between the shift being assigned + For each graph, computes the full pairwise distance matrix between all peaks and all residues. Returns logits (not probabilities). + The returned logits have shape [num_peaks, num_residues] where: + - logits[peak_i, residue_j] = logit for assigning peak i to residue j + - Higher logit = closer in chemical shift space (negative squared distance) + + For training: + - Use logits[shift_to_assign, :] with action as target for current assignment + - Use logits[assigned_peak, :] with assigned_residue as target for previous assignments + Args: - data: Batched HeteroData with RES nodes, SHIFT nodes, and ("RES", "pair", "SHIFT") edges + data: HeteroData (single or batched) with Residue nodes and Peak nodes Returns: List of logit tensors (negative squared distances), one per graph in batch + Shape per graph: [num_peaks, num_residues] """ - # Unbatch the data - data_unbatched = data.to_data_list() + # Check if data is batched or single + # Batched data has a 'batch' attribute on nodes + is_batched = hasattr(data["Residue"], 'batch') + + if is_batched: + # Use PyG's batching mechanism + data_unbatched = data.to_data_list() + else: + # Wrap single graph in a list + data_unbatched = [data] # Compute the policy logits for each item policies = [] for item in data_unbatched: - # Extract the shift being assigned and all residues - # The pairwise edges already connect only the shift being assigned to all residues - edge_index = item["RES", "pair", "SHIFT"].edge_index - - # Get residue NH features (H1, N15) - resid_features = item["RES"].x[edge_index[0]][:, self.RES_NH] + # Get all peak shift embeddings: shape [num_peaks, shift_dim] + peak_features = item["Peak"].x - # Get shift NH features (H1, N15) - shift_features = item["SHIFT"].x[edge_index[1]] + # Get all residue shift embeddings: shape [num_residues, shift_dim] + residue_features = item["Residue"].x[:, 3:] # Extract shift dimensions starting at index 3 - # Compute squared pairwise distances (one per residue) - # squared_distance = (H1_res - H1_shift)^2 + (N15_res - N15_shift)^2 - squared_distances = torch.sum((resid_features - shift_features) ** 2, dim=-1) + # Compute pairwise squared distances using broadcasting + # peak_features.unsqueeze(1): [num_peaks, 1, shift_dim] + # residue_features.unsqueeze(0): [1, num_residues, shift_dim] + # diff: [num_peaks, num_residues, shift_dim] + diff = peak_features.unsqueeze(1) - residue_features.unsqueeze(0) + squared_distances = torch.sum(diff ** 2, dim=-1) # [num_peaks, num_residues] # Use negative squared distances as logits (closer = higher logit) logits = -squared_distances - policies.append(logits.unsqueeze(0)) # Add batch dimension for consistency + policies.append(logits) return policies diff --git a/nmr/models/network.py b/nmr/models/network.py index 7d940f5..83601db 100644 --- a/nmr/models/network.py +++ b/nmr/models/network.py @@ -5,179 +5,114 @@ to create the complete neural network for NMR assignment. """ +from dataclasses import dataclass, field + import torch import torch.nn as nn +from torch_geometric.data import Batch + +from .heads import PolicyCalc, ValueCalc +from .pair import AssignedPair +from .triple import ( + PeakPeakNoeTriple, + PeakResidueNoeTriple, + ResiduePeakNoeTriple, + ResidueResidueNoeTriple, +) + + +@dataclass +class FeatureEmbedConfig: + """Configuration for embedding 2D assignment features to higher dimensions.""" + + output_dim: int = 128 + hidden_dim: int = 128 + num_layers: int = 1 + + +@dataclass +class ShiftEmbedConfig: + """Configuration for embedding 2D chemical shifts to higher dimensions.""" + + output_dim: int = 16 + hidden_dim: int = 128 + num_layers: int = 1 + -from .triple import CalculationManager, TripleIn, TripleOut, TripleUpdate -from .heads import ValueCalc, PolicyCalc +@dataclass +class MLPEmbedConfig: + """Internal configuration for generic MLP embedding.""" + + input_dim: int + output_dim: int + hidden_dim: int + num_layers: int + + +@dataclass +class MLPConfig: + """Configuration for MLP layers in message passing.""" + + hidden_size: int = 64 + num_layers: int = 1 + + +@dataclass +class ModelConfig: + """Top-level configuration for NMRNet model.""" + + num_nmr_layers: int = 1 + feature_embed: FeatureEmbedConfig = field(default_factory=FeatureEmbedConfig) + shift_embed: ShiftEmbedConfig = field(default_factory=ShiftEmbedConfig) + mlp: MLPConfig = field(default_factory=MLPConfig) class NMRLayer(nn.Module): """ Single layer of NMR-specific message passing. - Orchestrates the full triple update pipeline: - 1. Extract features from nodes into triples (TripleIn) - 2. Update triple representations via message passing (TripleUpdate) - 3. Propagate updated information back to nodes (TripleOut) + Orchestrates message passing through all four triple types and assigned pairs in sequence: + - ResidueResidueNoeTriple: (Residue, Residue, Noe) - Updates coordinates and shifts + - ResiduePeakNoeTriple: (Residue, Peak, Noe) - Updates shifts only + - PeakResidueNoeTriple: (Peak, Residue, Noe) - Updates shifts only + - PeakPeakNoeTriple: (Peak, Peak, Noe) - Updates shifts only + - AssignedPair: (Peak, Residue) assigned pairs - Updates shifts only - Handles all four triple types: - - TRIPLE0: (RES, RES, NOE) - - TRIPLE1: (RES, SHIFT, NOE) - - TRIPLE2: (SHIFT, RES, NOE) - - TRIPLE3: (SHIFT, SHIFT, NOE) + Each triple/pair type is instantiated as a separate attribute and called explicitly + in the forward pass, with no conditionals or branching logic. """ - def __init__(self, device): + def __init__(self, device, config: ModelConfig): super(NMRLayer, self).__init__() - self.triple_in = TripleIn() - self.triple_self = TripleUpdate(device) - self.triple_out = TripleOut() - self.calc_manager = CalculationManager() - - self.NOE = slice(0, 3) - self.RES_XYZ = slice(0, 3) - self.RES_NH = slice(3, 5) - self.SHIFT_NH = slice(0, 2) - - def update_noes(self, data, noe_delta, feature, i): - # 'i' is for the triple type in all of these cases (TRIPLE0 --> RES RES NOE, etc.) - # Should really simplify this to be one or the other from the very start (string --> number identification) - data = self.triple_out.update_data( - data, noe_delta, feature, (f"TRIPLE{i}", "NOE_add", "NOE"), self.NOE - ) - return data - - def update_coordinates(self, data, dist_delta1, dist_delta2, feature1, feature2, i): - data = self.triple_out.update_data( - data, - dist_delta1, - feature1, - (f"TRIPLE{i}", "res1_add", "RES"), - self.RES_XYZ, - update_f=False, - ) - data = self.triple_out.update_data( - data, - dist_delta2, - feature2, - (f"TRIPLE{i}", "res2_add", "RES"), - self.RES_XYZ, - update_f=False, - ) - return data - - def update_shifts( - self, data, shift_delta1, shift_delta2, feature1, feature2, target1, target2, i - ): - # Selects target indexing for residue or shift node - target_range1 = self.RES_NH if target1 == "RES" else self.SHIFT_NH - target_range2 = self.RES_NH if target2 == "RES" else self.SHIFT_NH - - data = self.triple_out.update_data( - data, - shift_delta1, - feature1, - (f"TRIPLE{i}", "NH1_add", target1), - target_range1, - ) - data = self.triple_out.update_data( - data, - shift_delta2, - feature2, - (f"TRIPLE{i}", "NH2_add", target2), - target_range2, - ) - return data - - def do_updates(self, data, triple_type, i): - # Triple in - shift_x1, coord_x1, shift_x2, coord_x2, noe_x, shift_f1, shift_f2, noe_f = ( - self.triple_in.construct_triple( - data, - (triple_type[0], "NH1_extract", f"TRIPLE{i}"), - (triple_type[1], "NH2_extract", f"TRIPLE{i}"), - ("NOE", "NOE_extract", f"TRIPLE{i}"), - ) - ) + # Instantiate all 4 triple classes as separate attributes, passing config + self.residue_residue_noe = ResidueResidueNoeTriple(device, config) + self.residue_peak_noe = ResiduePeakNoeTriple(device, config) + self.peak_residue_noe = PeakResidueNoeTriple(device, config) + self.peak_peak_noe = PeakPeakNoeTriple(device, config) + # Instantiate assigned pair processing, passing config + self.assigned_pair = AssignedPair(device, config) - # Self update and output - if triple_type == ("RES", "RES", "NOE"): - ( - delta_shift1, - delta_shift2, - delta_noe, - delta_dist1, - delta_dist2, - deltaf1, - deltaf2, - deltaf3, - ) = self.triple_self( - shift_x1, - coord_x1, - shift_x2, - coord_x2, - noe_x, - shift_f1, - shift_f2, - noe_f, - data[f"TRIPLE{i}", "update", f"TRIPLE{i}"].edge_index, - ) - data = self.update_noes(data, delta_noe, deltaf3, i) - data = self.update_coordinates( - data, delta_dist1, delta_dist2, deltaf1, deltaf2, i - ) - data = self.update_shifts( - data, - delta_shift1, - delta_shift2, - deltaf1, - deltaf2, - triple_type[0], - triple_type[1], - i, - ) + def forward(self, data): + """ + Process all four triple types and assigned pairs in sequence. - if triple_type in ( - ("SHIFT", "RES", "NOE"), - ("RES", "SHIFT", "NOE"), - ("SHIFT", "SHIFT", "NOE"), - ): - delta_shift1, delta_shift2, delta_noe, deltaf1, deltaf2, deltaf3 = ( - self.triple_self( - shift_x1, - coord_x1, - shift_x2, - coord_x2, - noe_x, - shift_f1, - shift_f2, - noe_f, - data[f"TRIPLE{i}", "update", f"TRIPLE{i}"].edge_index, - ) - ) - data = self.update_noes(data, delta_noe, deltaf3, i) - data = self.update_shifts( - data, - delta_shift1, - delta_shift2, - deltaf1, - deltaf2, - triple_type[0], - triple_type[1], - i, - ) - return data + Args: + data: HeteroData graph to process - def forward(self, data): - data = self.do_updates(data, ("RES", "RES", "NOE"), 0) - data = self.do_updates(data, ("RES", "SHIFT", "NOE"), 1) - data = self.do_updates(data, ("SHIFT", "RES", "NOE"), 2) - data = self.do_updates(data, ("SHIFT", "SHIFT", "NOE"), 3) + Returns: + Updated HeteroData graph with all node features updated + """ + # Call each triple class explicitly in sequence + data = self.residue_residue_noe(data) + data = self.residue_peak_noe(data) + data = self.peak_residue_noe(data) + data = self.peak_peak_noe(data) + # Process assigned pairs after triples + data = self.assigned_pair(data) return data -class NormalizeShifts(nn.Module): +class StandardizeShifts(nn.Module): """ Normalize chemical shifts to have a typical range between -1 and 1 """ @@ -192,18 +127,18 @@ def __init__(self, H_lower=6.0, H_upper=10.0, N_lower=100.0, N_upper=135.0): self.N_delta = N_upper - N_lower def forward(self, data): - # split RES into coords and shifts - res_x = data["RES"].x + # split Residue into coords and shifts + res_x = data["Residue"].x res_xyz = res_x[:, :3] res_H = res_x[:, 3].unsqueeze(-1) res_N = res_x[:, 4].unsqueeze(-1) - # split SHIFTS into 1H and 15N - shift_H = data["SHIFT"].x[:, 0].unsqueeze(-1) - shift_N = data["SHIFT"].x[:, 1].unsqueeze(-1) + # split Peaks into 1H and 15N + shift_H = data["Peak"].x[:, 0].unsqueeze(-1) + shift_N = data["Peak"].x[:, 1].unsqueeze(-1) # split NOEs into 1H and 15N - noe_H1 = data["NOE"].x[:, 0].unsqueeze(-1) - noe_N1 = data["NOE"].x[:, 1].unsqueeze(-1) - noe_H2 = data["NOE"].x[:, 2].unsqueeze(-1) + noe_H1 = data["Noe"].x[:, 0].unsqueeze(-1) + noe_N1 = data["Noe"].x[:, 1].unsqueeze(-1) + noe_H2 = data["Noe"].x[:, 2].unsqueeze(-1) # transform all 1H res_H = self._transform_H(res_H) shift_H = self._transform_H(shift_H) @@ -218,9 +153,9 @@ def forward(self, data): shift_x = torch.cat([shift_H, shift_N], dim=-1) noe_x = torch.cat([noe_H1, noe_N1, noe_H2], dim=-1) - data["RES"].x = res_x - data["SHIFT"].x = shift_x - data["NOE"].x = noe_x + data["Residue"].x = res_x + data["Peak"].x = shift_x + data["Noe"].x = noe_x return data def _transform_H(self, value): @@ -230,27 +165,325 @@ def _transform_N(self, value): return 2 * (value - self.N_lower) / self.N_delta - 1 +class MLPEmbedding(nn.Module): + """ + Generic MLP for embedding features of any dimensionality. + + Builds a configurable MLP: Linear → ReLU → ... → Linear based on + the provided MLPEmbedConfig. Used internally by semantic wrapper classes. + """ + + def __init__(self, config: MLPEmbedConfig, device): + super().__init__() + self.config = config + self.device = device + + layers = [] + + # Input layer: input_dim → hidden_dim + layers.append(nn.Linear(config.input_dim, config.hidden_dim, device=device)) + layers.append(nn.ReLU()) + + # Additional hidden layers + for _ in range(config.num_layers - 1): + layers.append( + nn.Linear(config.hidden_dim, config.hidden_dim, device=device) + ) + layers.append(nn.ReLU()) + + # Output layer: hidden_dim → output_dim + layers.append(nn.Linear(config.hidden_dim, config.output_dim, device=device)) + + self.mlp = nn.Sequential(*layers) + + def forward(self, x): + """ + Embed input features to higher dimensional space. + + Args: + x: Tensor of shape [n, input_dim] + + Returns: + Embedded features of shape [n, output_dim] + """ + return self.mlp(x) + + +class FeatureEmbedding(nn.Module): + """ + Embed 2D assignment features to higher dimensional space. + + Maps [assignment_indicator, already_assigned] → feature_dim using + a configurable MLP: Linear → ReLU → ... → Linear. + """ + + def __init__(self, config: FeatureEmbedConfig, device): + super().__init__() + mlp_config = MLPEmbedConfig( + input_dim=2, + output_dim=config.output_dim, + hidden_dim=config.hidden_dim, + num_layers=config.num_layers, + ) + self.mlp = MLPEmbedding(mlp_config, device) + + def forward(self, features): + """ + Embed 2D features to higher dimensional space. + + Args: + features: Tensor of shape [n, 2] with assignment indicators + + Returns: + Embedded features of shape [n, output_dim] + """ + return self.mlp(features) + + +class ShiftEmbedding(nn.Module): + """ + Embed 2D chemical shifts to higher dimensional space. + + Maps [H, N] → shift_dim using a configurable MLP: Linear → ReLU → ... → Linear. + This embedding is shared between Peak and Residue shift features. + """ + + def __init__(self, config: ShiftEmbedConfig, device): + super().__init__() + mlp_config = MLPEmbedConfig( + input_dim=2, + output_dim=config.output_dim, + hidden_dim=config.hidden_dim, + num_layers=config.num_layers, + ) + self.mlp = MLPEmbedding(mlp_config, device) + + def forward(self, shifts): + """ + Embed 2D shifts to higher dimensional space. + + Args: + shifts: Tensor of shape [n, 2] with [H, N] chemical shifts + + Returns: + Embedded shifts of shape [n, output_dim] + """ + return self.mlp(shifts) + + +class NoeEmbedding(nn.Module): + """ + Embed 3D NOE features to higher dimensional space. + + Maps [N, H', H"] → shift_dim using a configurable MLP: Linear → ReLU → ... → Linear. + Uses separate weights from ShiftEmbedding since NOE structure is different (3D vs 2D). + """ + + def __init__(self, config: ShiftEmbedConfig, device): + super().__init__() + mlp_config = MLPEmbedConfig( + input_dim=3, + output_dim=config.output_dim, + hidden_dim=config.hidden_dim, + num_layers=config.num_layers, + ) + self.mlp = MLPEmbedding(mlp_config, device) + + def forward(self, noe_features): + """ + Embed 3D NOE features to higher dimensional space. + + Args: + noe_features: Tensor of shape [n, 3] with [N, H', H"] NOE shifts + + Returns: + Embedded NOE features of shape [n, output_dim] + """ + return self.mlp(noe_features) + + +class EmbedFeatures(nn.Module): + """ + Apply all feature and shift embeddings to a heterogeneous graph. + + This module transforms all node features from their raw dimensions to learned + embeddings: + - Residue.x: [x, y, z, H, N] → [x, y, z, embedded_shifts...] + - Peak.x: [H, N] → [embedded_shifts...] + - Noe.x: [N, H', H"] → [embedded_noe...] + - All .f attributes: [2] → [feature_dim] + + Coordinates remain unchanged (still 3D). + """ + + def __init__(self, device, config: ModelConfig): + super().__init__() + self.device = device + self.config = config + + # Create embedding modules + self.res_feature_embed = FeatureEmbedding(config.feature_embed, device) + self.peak_feature_embed = FeatureEmbedding(config.feature_embed, device) + self.noe_feature_embed = FeatureEmbedding(config.feature_embed, device) + # use the same embedding for residue and peak shifts + self.shift_embed = ShiftEmbedding(config.shift_embed, device) + self.noe_embed = NoeEmbedding(config.shift_embed, device) + + def forward(self, data): + """ + Embed all features and shifts in the graph. + + Args: + data: HeteroData graph with standardized shifts + + Returns: + HeteroData graph with embedded features + """ + # Embed Residue shifts and concatenate with coordinates + res_coords = data["Residue"].x[:, 0:3] # [num_res, 3] + res_shifts = data["Residue"].x[:, 3:5] # [num_res, 2] + res_shifts_embedded = self.shift_embed(res_shifts) # [num_res, shift_dim] + data["Residue"].x = torch.cat([res_coords, res_shifts_embedded], dim=-1) + + # Embed Peak shifts + peak_shifts = data["Peak"].x # [num_peaks, 2] + data["Peak"].x = self.shift_embed(peak_shifts) # [num_peaks, shift_dim] + + # Embed NOE features + noe_features = data["Noe"].x # [num_noes, 3] + data["Noe"].x = self.noe_embed(noe_features) # [num_noes, shift_dim] + + # Embed all .f attributes (assignment features) + data["Residue"].f = self.res_feature_embed(data["Residue"].f) + data["Peak"].f = self.peak_feature_embed(data["Peak"].f) + data["Noe"].f = self.noe_feature_embed(data["Noe"].f) + + return data + + +class ChemicalShiftNorm(nn.Module): + """ + Per-graph normalization with separate normalization for each node type. + + Applied AFTER embeddings, so dimensions are: + - Residue.x: [coords (3D), embedded_shifts (shift_dim)] + - Peak.x: [embedded_shifts (shift_dim)] + - Noe.x: [embedded_noe (shift_dim)] + + For each graph independently: + - Standardizes Residue XYZ coordinates to zero mean, unit variance + - Standardizes Residue embedded shifts to zero mean, unit variance + - Standardizes Peak embedded shifts to zero mean, unit variance + - Standardizes NOE embedded features to zero mean, unit variance + + Each node type's shifts are normalized independently, allowing the network + to learn different representations for predicted vs observed shifts. + + Implementation unbatches graphs, normalizes each separately, then re-batches + to avoid in-place operations that break gradients. + """ + + def __init__(self, eps=1e-5): + super().__init__() + self.eps = eps + + def forward(self, data): + # Check if data is batched + is_batched = hasattr(data["Residue"], "batch") + + if not is_batched: + # Single graph - normalize directly + return self._normalize_single_graph(data) + else: + # Batched graphs - unbatch, normalize each, re-batch + graph_list = data.to_data_list() + normalized_list = [self._normalize_single_graph(g) for g in graph_list] + return Batch.from_data_list(normalized_list) + + def _normalize_single_graph(self, data): + """ + Normalize a single graph. + + Returns a new graph with normalized features (no in-place operations). + After embeddings, dimensions are: + - Residue.x: [coords (3D), embedded_shifts (shift_dim)] + - Peak.x: [embedded_shifts (shift_dim)] + - Noe.x: [embedded_noe (shift_dim)] + """ + # Normalize Residue coordinates + coords = data["Residue"].x[:, 0:3] + coords_mean = coords.mean(dim=0, keepdim=True) + coords_std = coords.std(dim=0, keepdim=True) + self.eps + normalized_coords = (coords - coords_mean) / coords_std + + # Normalize Residue shifts + residue_shifts = data["Residue"].x[:, 3:] # [shift_dim] + res_shift_mean = residue_shifts.mean(dim=0, keepdim=True) + res_shift_std = residue_shifts.std(dim=0, keepdim=True) + self.eps + normalized_residue_shifts = (residue_shifts - res_shift_mean) / res_shift_std + + # Normalize Peak shifts + peak_shifts = data["Peak"].x # [shift_dim] + peak_shift_mean = peak_shifts.mean(dim=0, keepdim=True) + peak_shift_std = peak_shifts.std(dim=0, keepdim=True) + self.eps + normalized_peak_shifts = (peak_shifts - peak_shift_mean) / peak_shift_std + + # Normalize NOE shifts + noe_shifts = data["Noe"].x # [shift_dim] + noe_shift_mean = noe_shifts.mean(dim=0, keepdim=True) + noe_shift_std = noe_shifts.std(dim=0, keepdim=True) + self.eps + normalized_noe_shifts = (noe_shifts - noe_shift_mean) / noe_shift_std + + # Reconstruct tensors + data["Residue"].x = torch.cat( + [normalized_coords, normalized_residue_shifts], dim=1 + ) + data["Peak"].x = normalized_peak_shifts + data["Noe"].x = normalized_noe_shifts + + return data + + class NMRNet(nn.Module): """ Complete NMR GNN model combining message passing with prediction heads. Stacks NMRLayer(s) for graph message passing, then uses ValueCalc and PolicyCalc heads to predict state value and action probabilities. + + Applies ChemicalShiftNorm after each layer to ensure consistent normalization + of chemical shifts across all node types. """ - def __init__(self, device): + def __init__(self, device, config: ModelConfig): super().__init__() - self.value = ValueCalc(device) - self.policy = PolicyCalc(device) + self.config = config + self.device = device - self.normalize = NormalizeShifts() - self.nmr = nn.Sequential( - NMRLayer(device), - ) + # Prediction heads - pass config for dimension calculations + self.value = ValueCalc(device, config) + self.policy = PolicyCalc(device, config) + + # Standardization and embedding layers + self.standardize = StandardizeShifts() + self.embed_features = EmbedFeatures(device, config) + + # Build sequential stack: [NMRLayer, ChemicalShiftNorm, NMRLayer, ChemicalShiftNorm, ...] + layers = [] + for i in range(config.num_nmr_layers): + layers.append(NMRLayer(device, config)) + layers.append(ChemicalShiftNorm()) + + self.nmr = nn.Sequential(*layers) def forward(self, data): - data = self.normalize(data) - out_data = self.nmr(data) - value = self.value.calc_value(out_data) - policy = self.policy.calc_policy(out_data) - return out_data, value, policy + # Standardize raw shifts to [-1, 1] range + data = self.standardize(data) + # Embed shifts and features to higher dimensions + data = self.embed_features(data) + # Message passing layers with normalization + data = self.nmr(data) + # Prediction heads + value = self.value.calc_value(data) + policy = self.policy.calc_policy(data) + return value, policy diff --git a/nmr/models/pair.py b/nmr/models/pair.py new file mode 100644 index 0000000..9584371 --- /dev/null +++ b/nmr/models/pair.py @@ -0,0 +1,429 @@ +""" +Pair-based graph message passing components for assigned peak-residue pairs. + +This module implements direct message passing between Peak and Residue nodes +that have been assigned to each other, without intermediate triple nodes. + +Architecture: +1. Message Operations (2 classes): Process edges directly with MLPs +2. Pair Composition (1 class): Wire both message passing directions together + +This is simpler than triple processing since pairs connect exactly 2 nodes +via edges, while triples simulate hypergraphs with 3 nodes. + +Pair Type: +- AssignedPair: (Peak, Residue) pairs connected by "assigned_to" edges + Updates shifts and features in both directions + +Node Type Naming: +- Residue: Protein residues with coordinates [x,y,z] and shifts [H,N] +- Peak: Observed chemical shifts [H,N] + +Edge Naming: +- Bidirectional edges: ("Peak", "assigned_to", "Residue") + * Used for Peak → Residue message passing with flow='source_to_target' + * Used for Residue → Peak message passing with flow='target_to_source' +""" + +import torch +import torch.nn as nn +from torch_geometric.nn import MessagePassing + + +# ============================================================================ +# MLP Architecture Constants +# ============================================================================ + +# AssignedPeakToResidueMessage MLP +PEAK_TO_RESIDUE_INPUT_SIZE = 6 # shift_diffs(2) + peak_features(2) + residue_features(2) +PEAK_TO_RESIDUE_OUTPUT_SIZE = 3 # shift_weight(1) + feature_deltas(2) for residue + +# AssignedResidueToPeakMessage MLP +RESIDUE_TO_PEAK_INPUT_SIZE = 6 # shift_diffs(2) + residue_features(2) + peak_features(2) +RESIDUE_TO_PEAK_OUTPUT_SIZE = 3 # shift_weight(1) + feature_deltas(2) for peak + + +# ============================================================================ +# SECTION 1: Message Passing Operations +# ============================================================================ + + +class AssignedPeakToResidueMessage(MessagePassing): + """ + Message passing from assigned Peak to Residue nodes. + + Uses edge type: ("Peak", "assigned_to", "Residue") with flow='source_to_target' + + Message computation: + 1. Extract peak shifts, residue shifts, and features + 2. Compute shift differences (peak - residue) + 3. Concatenate shift diffs and features + 4. Pass through MLP to get residue deltas + 5. Apply deltas to residue shifts and features (aggregated) + + Equivariance: + All calculations use shift differences (never absolute shifts) + """ + + def __init__(self, device, config, hidden_size: int = None, num_layers: int = None): + """ + Initialize AssignedPeakToResidueMessage. + + Args: + device: torch device (CPU or CUDA) + config: ModelConfig with dimension settings + hidden_size: Number of hidden units in MLP (default: from config) + num_layers: Number of hidden layers (default: from config) + """ + super().__init__(aggr="mean", flow="source_to_target") + self.device = device + self.config = config + self.edge_type = ("Peak", "assigned_to", "Residue") + + # Use config defaults if not specified + if hidden_size is None: + hidden_size = config.mlp.hidden_size + if num_layers is None: + num_layers = config.mlp.num_layers + + # Calculate input/output sizes from config + shift_dim = config.shift_embed.output_dim + feature_dim = config.feature_embed.output_dim + # Input: shift_diffs (shift_dim) + peak_features (feature_dim) + residue_features (feature_dim) + input_size = shift_dim + 2 * feature_dim + # Output: shift_weight (shift_dim) + feature_deltas (feature_dim) for residue + output_size = shift_dim + feature_dim + + # Store dimensions + self.shift_dim = shift_dim + self.feature_dim = feature_dim + + # Build MLP + layers = [] + + # First hidden layer + layers.append(nn.Linear(input_size, hidden_size)) + layers.append(nn.LayerNorm(hidden_size)) + layers.append(nn.ReLU()) + + # Additional hidden layers + for _ in range(num_layers - 1): + layers.append(nn.Linear(hidden_size, hidden_size)) + layers.append(nn.LayerNorm(hidden_size)) + layers.append(nn.ReLU()) + + # Output layer + layers.append(nn.Linear(hidden_size, output_size)) + + self.mlp = nn.Sequential(*layers).to(device) + + def forward(self, data): + """ + Execute message passing from Peak to Residue nodes. + + Args: + data: HeteroData graph with Peak and Residue nodes + + Returns: + Updated HeteroData with Residue.x and Residue.f modified + """ + # Get edge indices + edge_index = data[self.edge_type].edge_index + + # Extract peak shifts [n_peaks, shift_dim] and features [n_peaks, feature_dim] + peak_shifts = data["Peak"].x + peak_features = data["Peak"].f + + # Extract residue shifts [n_residues, shift_dim] and features [n_residues, feature_dim] + residue_shifts = data["Residue"].x[:, 3:] # Shifts start at index 3 + residue_features = data["Residue"].f + + # Determine sizes for message passing + num_residues = data["Residue"].x.size(0) + num_peaks = data["Peak"].x.size(0) + + # Propagate updates to residues + shift_updates, feature_updates = self.propagate( + edge_index, + peak_shifts=peak_shifts, + peak_features=peak_features, + residue_shifts=residue_shifts, + residue_features=residue_features, + size=(num_peaks, num_residues) + ) + + # Apply updates to residue nodes (assemble new tensors) + old_x = data["Residue"].x + new_coords = old_x[:, 0:3] # Keep coordinates unchanged + new_shifts = old_x[:, 3:] + shift_updates + data["Residue"].x = torch.cat([new_coords, new_shifts], dim=-1) + data["Residue"].f = data["Residue"].f + feature_updates + + return data + + def message(self, peak_shifts_j, peak_features_j, residue_shifts_i, residue_features_i): + """ + Compute messages from peak (source) to residue (target). + + Args: + peak_shifts_j: Peak shifts [n_edges, 2] (source) + peak_features_j: Peak features [n_edges, 2] (source) + residue_shifts_i: Residue shifts [n_edges, 2] (target) + residue_features_i: Residue features [n_edges, 2] (target) + + Returns: + Tuple of (shift_deltas, feature_deltas) for residues + """ + # Compute shift difference vector (peak - residue) for equivariance + # shift_diff is a 2D vector [H, N] + shift_diff = peak_shifts_j - residue_shifts_i # [n_edges, 2] + + # Concatenate features for MLP input (6 total) + mlp_input = torch.cat([ + shift_diff, # [n_edges, 2] + peak_features_j, # [n_edges, 2] + residue_features_i, # [n_edges, 2] + ], dim=-1) # [n_edges, 6] + + # Apply MLP to get scalar weight and feature deltas + mlp_output = self.mlp(mlp_input) # [n_edges, 3] + + # Parse output: shift_weight (shift_dim) + feature_deltas (feature_dim) + shift_weight = mlp_output[:, :self.shift_dim] # [n_edges, shift_dim] + feature_deltas = mlp_output[:, self.shift_dim:] # [n_edges, feature_dim] + + # Compute shift deltas by element-wise multiplication with shift difference + # This ensures movement is always in the direction of the difference + shift_deltas = shift_diff * shift_weight # [n_edges, shift_dim] + + # Concatenate shift and feature deltas for aggregation + # PyG will automatically apply mean aggregation (aggr='mean') + return torch.cat([shift_deltas, feature_deltas], dim=-1) # [n_edges, shift_dim + feature_dim] + + def update(self, aggr_out): + """ + Split aggregated updates back into shift and feature components. + + Args: + aggr_out: Aggregated tensor [n_nodes, shift_dim + feature_dim] + + Returns: + Tuple of (shift_updates, feature_updates) + """ + shift_updates = aggr_out[:, :self.shift_dim] # [n_nodes, shift_dim] + feature_updates = aggr_out[:, self.shift_dim:] # [n_nodes, feature_dim] + return shift_updates, feature_updates + + +class AssignedResidueToPeakMessage(MessagePassing): + """ + Message passing from Residue to assigned Peak nodes. + + Uses edge type: ("Peak", "assigned_to", "Residue") with flow='target_to_source' + + Message computation: + 1. Extract residue shifts, peak shifts, and features + 2. Compute shift differences (residue - peak) + 3. Concatenate shift diffs and features + 4. Pass through MLP to get peak deltas + 5. Apply deltas to peak shifts and features (aggregated) + + Equivariance: + All calculations use shift differences (never absolute shifts) + """ + + def __init__(self, device, config, hidden_size: int = None, num_layers: int = None): + """ + Initialize AssignedResidueToPeakMessage. + + Args: + device: torch device (CPU or CUDA) + config: ModelConfig with dimension settings + hidden_size: Number of hidden units in MLP (default: from config) + num_layers: Number of hidden layers (default: from config) + """ + super().__init__(aggr="mean", flow="target_to_source") + self.device = device + self.config = config + self.edge_type = ("Peak", "assigned_to", "Residue") + + # Use config defaults if not specified + if hidden_size is None: + hidden_size = config.mlp.hidden_size + if num_layers is None: + num_layers = config.mlp.num_layers + + # Calculate input/output sizes from config + shift_dim = config.shift_embed.output_dim + feature_dim = config.feature_embed.output_dim + # Input: shift_diffs (shift_dim) + residue_features (feature_dim) + peak_features (feature_dim) + input_size = shift_dim + 2 * feature_dim + # Output: shift_weight (shift_dim) + feature_deltas (feature_dim) for peak + output_size = shift_dim + feature_dim + + # Store dimensions + self.shift_dim = shift_dim + self.feature_dim = feature_dim + + # Build MLP + layers = [] + + # First hidden layer + layers.append(nn.Linear(input_size, hidden_size)) + layers.append(nn.LayerNorm(hidden_size)) + layers.append(nn.ReLU()) + + # Additional hidden layers + for _ in range(num_layers - 1): + layers.append(nn.Linear(hidden_size, hidden_size)) + layers.append(nn.LayerNorm(hidden_size)) + layers.append(nn.ReLU()) + + # Output layer + layers.append(nn.Linear(hidden_size, output_size)) + + self.mlp = nn.Sequential(*layers).to(device) + + def forward(self, data): + """ + Execute message passing from Residue to Peak nodes. + + Args: + data: HeteroData graph with Peak and Residue nodes + + Returns: + Updated HeteroData with Peak.x and Peak.f modified + """ + # Get edge indices + edge_index = data[self.edge_type].edge_index + + # Extract peak shifts [n_peaks, shift_dim] and features [n_peaks, feature_dim] + peak_shifts = data["Peak"].x + peak_features = data["Peak"].f + + # Extract residue shifts [n_residues, shift_dim] and features [n_residues, feature_dim] + residue_shifts = data["Residue"].x[:, 3:] # Shifts start at index 3 + residue_features = data["Residue"].f + + # Determine sizes for message passing + num_residues = data["Residue"].x.size(0) + num_peaks = data["Peak"].x.size(0) + + # Propagate updates to peaks (with reversed flow, size is (target, source)) + shift_updates, feature_updates = self.propagate( + edge_index, + residue_shifts=residue_shifts, + residue_features=residue_features, + peak_shifts=peak_shifts, + peak_features=peak_features, + size=(num_peaks, num_residues) + ) + + # Apply updates to peak nodes (assemble new tensors) + data["Peak"].x = data["Peak"].x + shift_updates + data["Peak"].f = data["Peak"].f + feature_updates + + return data + + def message(self, residue_shifts_j, residue_features_j, peak_shifts_i, peak_features_i): + """ + Compute messages from residue (source in reversed flow) to peak (target). + + Args: + residue_shifts_j: Residue shifts [n_edges, shift_dim] (source) + residue_features_j: Residue features [n_edges, feature_dim] (source) + peak_shifts_i: Peak shifts [n_edges, shift_dim] (target) + peak_features_i: Peak features [n_edges, feature_dim] (target) + + Returns: + Tuple of (shift_deltas, feature_deltas) for peaks + """ + # Compute shift difference vector (residue - peak) for equivariance + shift_diff = residue_shifts_j - peak_shifts_i # [n_edges, shift_dim] + + # Concatenate features for MLP input + mlp_input = torch.cat([ + shift_diff, # [n_edges, shift_dim] + residue_features_j, # [n_edges, feature_dim] + peak_features_i, # [n_edges, feature_dim] + ], dim=-1) + + # Apply MLP to get weight vector and feature deltas + mlp_output = self.mlp(mlp_input) + + # Parse output: shift_weight (shift_dim) + feature_deltas (feature_dim) + shift_weight = mlp_output[:, :self.shift_dim] # [n_edges, shift_dim] + feature_deltas = mlp_output[:, self.shift_dim:] # [n_edges, feature_dim] + + # Compute shift deltas by element-wise multiplication with shift difference + # This ensures movement is always in the direction of the difference + shift_deltas = shift_diff * shift_weight # [n_edges, shift_dim] + + # Concatenate shift and feature deltas for aggregation + # PyG will automatically apply mean aggregation (aggr='mean') + return torch.cat([shift_deltas, feature_deltas], dim=-1) # [n_edges, shift_dim + feature_dim] + + def update(self, aggr_out): + """ + Split aggregated updates back into shift and feature components. + + Args: + aggr_out: Aggregated tensor [n_nodes, shift_dim + feature_dim] + + Returns: + Tuple of (shift_updates, feature_updates) + """ + shift_updates = aggr_out[:, :self.shift_dim] # [n_nodes, shift_dim] + feature_updates = aggr_out[:, self.shift_dim:] # [n_nodes, feature_dim] + return shift_updates, feature_updates + + +# ============================================================================ +# SECTION 2: Pair Composition Class +# ============================================================================ + + +class AssignedPair(nn.Module): + """ + Wire together bidirectional message passing for assigned peak-residue pairs. + + This composition class processes assigned (Peak, Residue) pairs by running + message passing in both directions: Peak → Residue and Residue → Peak. + + Components: + - AssignedPeakToResidueMessage: Peak → Residue updates + - AssignedResidueToPeakMessage: Residue → Peak updates + + Forward pass sequence: + peak_to_residue → residue_to_peak + """ + + def __init__(self, device, config, hidden_size: int = None, num_layers: int = None): + """ + Initialize AssignedPair. + + Args: + device: torch device (CPU or CUDA) + config: ModelConfig with dimension settings + hidden_size: Hidden units in message MLPs (default: from config) + num_layers: Hidden layers in message MLPs (default: from config) + """ + super().__init__() + + # Instantiate both message passing operations - pass config + self.peak_to_residue = AssignedPeakToResidueMessage(device, config, hidden_size, num_layers) + self.residue_to_peak = AssignedResidueToPeakMessage(device, config, hidden_size, num_layers) + + def forward(self, data): + """ + Execute bidirectional message passing. + + Args: + data: HeteroData graph + + Returns: + Updated HeteroData graph + """ + data = self.peak_to_residue(data) + data = self.residue_to_peak(data) + return data diff --git a/nmr/models/triple.py b/nmr/models/triple.py index 4e2115f..2e69767 100644 --- a/nmr/models/triple.py +++ b/nmr/models/triple.py @@ -1,8 +1,34 @@ """ -Triple-based graph message passing components. +Modular triple-based graph message passing components. -Contains classes for constructing, updating, and propagating messages -through triple nodes in the heterogeneous NMR graph. +This module implements a three-stage message passing pipeline using explicit, +modular components instead of conditional logic: + +Architecture: +1. Gather Operations (5 classes): Extract features from source nodes to triple nodes +2. Update Operations (2 classes): Compute deltas via MLPs on triple nodes +3. Scatter Operations (5 classes): Propagate deltas back to source nodes +4. Triple Composition (4 classes): Wire gather/update/scatter for each triple type +5. Layer Orchestration (in network.py): Call all 4 triple types explicitly + +This architecture eliminates runtime conditionals by making all behavioral choices +explicit at construction time. + +Triple Types: +- ResidueResidueNoeTriple: (Residue, Residue, Noe) - Updates coordinates + shifts +- ResiduePeakNoeTriple: (Residue, Peak, Noe) - Updates shifts only +- PeakResidueNoeTriple: (Peak, Residue, Noe) - Updates shifts only +- PeakPeakNoeTriple: (Peak, Peak, Noe) - Updates shifts only + +Node Type Naming: +- Residue: Protein residues with coordinates [x,y,z] and shifts [H,N] +- Peak: Observed chemical shifts [H,N] +- Noe: NOE constraints [N, H', H"] + +Edge Naming Conventions: +- Bidirectional propagation edges: (source_node, "prop_first"/"prop_second"/"prop_noe", triple_type) + * Used for gather with default flow='source_to_target' + * Used for scatter with flow='target_to_source' """ import torch @@ -10,333 +36,1476 @@ from torch_geometric.nn import MessagePassing -class CalculationManager: +# ============================================================================ +# MLP Architecture Dimension Functions +# ============================================================================ +# These functions compute input/output sizes for MLP architectures based on +# configurable shift_dim and feature_dim parameters + + +def residue_update_input_size(shift_dim, feature_dim): """ - Common calculations used across triple and pairwise comparisons. + Calculate input size for ResidueUpdate MLP. + + Components: + - dist_squared: 1 + - shift differences: 3*shift_dim pairwise differences + (diff_first_to_second, diff_first_to_noe, diff_second_to_noe) + - features: 3*feature_dim (first_features, second_features, noe_features) + + Returns: 1 + 3*shift_dim + 3*feature_dim + """ + return 1 + 3 * shift_dim + 3 * feature_dim + + +def residue_update_output_size(shift_dim, feature_dim): + """ + Calculate output size for ResidueUpdate MLP. + + Components: + - coord_weights: 2 (w_coord_first, w_coord_second) + - shift_weights: 6 (w1-w6, scalar weights for each shift difference) + - feature_deltas: 3*feature_dim (delta_first_features, delta_second_features, delta_noe_features) + + Returns: coord_weights(2) + shift_weights(6) + feature_deltas(3*feature_dim) + """ + return 2 + 6 + 3 * feature_dim + + +def peak_update_input_size(shift_dim, feature_dim): + """ + Calculate input size for PeakUpdate MLP. + + Components: + - shift differences: 3*shift_dim pairwise differences + (diff_first_to_second, diff_first_to_noe, diff_second_to_noe) + - features: 3*feature_dim (first_features, second_features, noe_features) + + Returns: 3*shift_dim + 3*feature_dim + Note: NO dist_squared (peaks have no coordinates). + """ + return 3 * shift_dim + 3 * feature_dim + + +def peak_update_output_size(shift_dim, feature_dim): + """ + Calculate output size for PeakUpdate MLP. + + Components: + - shift_weights: 6 (w1-w6, scalar weights for each shift difference) + - feature_deltas: 3*feature_dim (delta_first_features, delta_second_features, delta_noe_features) + + Returns: shift_weights(6) + feature_deltas(3*feature_dim) + Note: NO coord_weights (peaks have no coordinates). + """ + return 6 + 3 * feature_dim + + +# ============================================================================ +# SECTION 1: Gather Operations +# ============================================================================ +# These classes extract features from source nodes to triple nodes using +# PyTorch Geometric MessagePassing with aggr="mean" + + +class FirstResidueGather(MessagePassing): """ + Extract coordinates, shifts, and features from residues in first position. - def __init__(self): - # Input index organization - slices to keep proper tensor dimension - self.NOE_N1 = slice(0, 1) - self.NOE_H1 = slice(1, 2) - self.NOE_H2 = slice(2, 3) - self.RES_XYZ = slice(0, 3) - self.SHIFT_N = slice(0, 1) - self.SHIFT_H = slice(1, 2) + Uses edge type: ("Residue", "prop_first", triple_type) + Sets attributes on triple nodes: + - first_coords: [n, 3] coordinates + - first_shifts: [n, shift_dim] chemical shift embeddings + - first_features: [n, feature_dim] assignment feature embeddings + """ - def calc_noe_difference(self, x1, x2, noe): + def __init__(self, triple_type: str): """ - Calculates shift difference between NOE and residue/measured shifts (direct/indirect only reverse option is available by index). + Initialize FirstResidueGather. + + Args: + triple_type: Name of target triple node type """ - diff_N = noe[:, self.NOE_N1] - x1[:, self.SHIFT_N] # N [n, 1] - diff_H1 = noe[:, self.NOE_H1] - x1[:, self.SHIFT_H] # H' [n, 1] - diff_H2 = noe[:, self.NOE_H2] - x2[:, self.SHIFT_H] # H" [n, 1] - return diff_N, diff_H1, diff_H2 + super().__init__(aggr="mean") + self.triple_type = triple_type + self.edge_type = ("Residue", "prop_first", triple_type) - def calc_shift_difference(self, x1, x2): + def forward(self, data): """ - Calculates shift difference between residue/measured shifts. + Extract features from Residue nodes to triple nodes. + + Args: + data: HeteroData graph with Residue nodes and gather edges + + Returns: + Updated HeteroData with first_coords, first_shifts, first_features set """ - diff_N = x1[:, self.SHIFT_N] - x2[:, self.SHIFT_N] # N [n, 1] - diff_H = x1[:, self.SHIFT_H] - x2[:, self.SHIFT_H] # H [n, 1] - return diff_N, diff_H + # Extract coordinates [n, 3] from Residue.x[:, 0:3] + coords = data["Residue"].x[:, 0:3] + + # Extract shifts [n, shift_dim] from Residue.x[:, 3:] + shifts = data["Residue"].x[:, 3:] + + # Extract features [n, feature_dim] from Residue.f + features = data["Residue"].f + + # Get edge indices for this gather operation + edge_index = data[self.edge_type].edge_index + + # Determine number of target nodes (triple nodes) + num_triples = data[self.triple_type].x.size(0) + + # Propagate coordinates with explicit size + data[self.triple_type].first_coords = self.propagate( + edge_index, x=coords, size=(coords.size(0), num_triples) + ) + + # Propagate shifts + data[self.triple_type].first_shifts = self.propagate( + edge_index, x=shifts, size=(shifts.size(0), num_triples) + ) - def calc_res_distance(self, x1, x2): + # Propagate features + data[self.triple_type].first_features = self.propagate( + edge_index, x=features, size=(features.size(0), num_triples) + ) + + return data + + def message(self, x_j): + """Pass through features from source nodes.""" + return x_j + + +class FirstPeakGather(MessagePassing): + """ + Extract shifts and features from peaks in first position. + + Uses edge type: ("Peak", "prop_first", triple_type) + Sets attributes on triple nodes: + - first_shifts: [n, shift_dim] chemical shift embeddings + - first_features: [n, feature_dim] assignment feature embeddings + Note: Peaks have NO coordinates + """ + + def __init__(self, triple_type: str): """ - Calculates relative distance between residues and this value squared for equivariant calculations. + Initialize FirstPeakGather. + + Args: + triple_type: Name of target triple node type """ - rel_dist = x1 - x2 # [n, 3] - dist2 = torch.norm(rel_dist, dim=-1, keepdim=True) ** 2 # [n, 1] - return rel_dist, dist2 + super().__init__(aggr="mean") + self.triple_type = triple_type + self.edge_type = ("Peak", "prop_first", triple_type) + def forward(self, data): + """ + Extract features from Peak nodes to triple nodes. -class TripleIn(nn.Module): - def __init__(self): - super().__init__() - # Input index organization - slices to keep proper tensor dimension - self.RES_XYZ = slice(0, 3) - self.RES_NH = slice(3, 5) + Args: + data: HeteroData graph with Peak nodes and gather edges - def grab_node(self, data, node_type, edge_type): + Returns: + Updated HeteroData with first_shifts, first_features set """ - Source node to triple via indexing. - For residue data, extracts shift and coordinate information to process separately. + # Extract shifts [n, shift_dim] from Peak.x + shifts = data["Peak"].x + + # Extract features [n, feature_dim] from Peak.f + features = data["Peak"].f + + # Get edge indices for this gather operation + edge_index = data[self.edge_type].edge_index + + # Determine number of target nodes (triple nodes) + num_triples = data[self.triple_type].x.size(0) + + # Propagate shifts with explicit size + data[self.triple_type].first_shifts = self.propagate( + edge_index, x=shifts, size=(shifts.size(0), num_triples) + ) + + # Propagate features + data[self.triple_type].first_features = self.propagate( + edge_index, x=features, size=(features.size(0), num_triples) + ) + + return data + + def message(self, x_j): + """Pass through features from source nodes.""" + return x_j + + +class SecondResidueGather(MessagePassing): + """ + Extract coordinates, shifts, and features from residues in second position. + + Uses edge type: ("Residue", "prop_second", triple_type) + Sets attributes on triple nodes: + - second_coords: [n, 3] coordinates + - second_shifts: [n, shift_dim] chemical shift embeddings + - second_features: [n, feature_dim] assignment feature embeddings + """ + + def __init__(self, triple_type: str): """ - # Spatial features - xi = data[node_type].x[data[edge_type].edge_index[0]] - # Non-spatial features - fi = data[node_type].f[data[edge_type].edge_index[0]] + Initialize SecondResidueGather. - # Splits coordinate and shift features if combined - xj = None - if edge_type[0] == "RES": - xi, xj = xi[:, self.RES_NH], xi[:, self.RES_XYZ] - return xi, xj, fi + Args: + triple_type: Name of target triple node type + """ + super().__init__(aggr="mean") + self.triple_type = triple_type + self.edge_type = ("Residue", "prop_second", triple_type) - def construct_triple(self, data, edge_type1, edge_type2, edge_type3): + def forward(self, data): """ - Constructs triple based on type. + Extract features from Residue nodes to triple nodes. + + Args: + data: HeteroData graph with Residue nodes and gather edges + + Returns: + Updated HeteroData with second_coords, second_shifts, second_features set """ - # Measured shift or residue (shifts will return nonetype value) - x1, x12, f1 = self.grab_node( - data, node_type=edge_type1[0], edge_type=edge_type1 + # Extract coordinates [n, 3] from Residue.x[:, 0:3] + coords = data["Residue"].x[:, 0:3] + + # Extract shifts [n, shift_dim] from Residue.x[:, 3:] + shifts = data["Residue"].x[:, 3:] + + # Extract features [n, feature_dim] from Residue.f + features = data["Residue"].f + + # Get edge indices for this gather operation + edge_index = data[self.edge_type].edge_index + + # Determine number of target nodes (triple nodes) + num_triples = data[self.triple_type].x.size(0) + + # Propagate coordinates with explicit size + data[self.triple_type].second_coords = self.propagate( + edge_index, x=coords, size=(coords.size(0), num_triples) + ) + + # Propagate shifts + data[self.triple_type].second_shifts = self.propagate( + edge_index, x=shifts, size=(shifts.size(0), num_triples) ) - x2, x22, f2 = self.grab_node( - data, node_type=edge_type2[0], edge_type=edge_type2 + + # Propagate features + data[self.triple_type].second_features = self.propagate( + edge_index, x=features, size=(features.size(0), num_triples) ) - # NOE (does not require any value split) - x3, _, f3 = self.grab_node(data, node_type=edge_type3[0], edge_type=edge_type3) - return x1, x12, x2, x22, x3, f1, f2, f3 + + return data + + def message(self, x_j): + """Pass through features from source nodes.""" + return x_j -class TripleUpdate(MessagePassing): +class SecondPeakGather(MessagePassing): """ - Message passing class for triple self updates. + Extract shifts and features from peaks in second position. + + Uses edge type: ("Peak", "prop_second", triple_type) + Sets attributes on triple nodes: + - second_shifts: [n, shift_dim] chemical shift embeddings + - second_features: [n, feature_dim] assignment feature embeddings + Note: Peaks have NO coordinates """ - def __init__(self, device): - super().__init__(aggr="add") - self.calc_manager = CalculationManager() - self.device = device + def __init__(self, triple_type: str): + """ + Initialize SecondPeakGather. + + Args: + triple_type: Name of target triple node type + """ + super().__init__(aggr="mean") + self.triple_type = triple_type + self.edge_type = ("Peak", "prop_second", triple_type) + + def forward(self, data): + """ + Extract features from Peak nodes to triple nodes. + + Args: + data: HeteroData graph with Peak nodes and gather edges + + Returns: + Updated HeteroData with second_shifts, second_features set + """ + # Extract shifts [n, shift_dim] from Peak.x + shifts = data["Peak"].x + + # Extract features [n, feature_dim] from Peak.f + features = data["Peak"].f + + # Get edge indices for this gather operation + edge_index = data[self.edge_type].edge_index + + # Determine number of target nodes (triple nodes) + num_triples = data[self.triple_type].x.size(0) + + # Propagate shifts with explicit size + data[self.triple_type].second_shifts = self.propagate( + edge_index, x=shifts, size=(shifts.size(0), num_triples) + ) + + # Propagate features + data[self.triple_type].second_features = self.propagate( + edge_index, x=features, size=(features.size(0), num_triples) + ) + + return data + + def message(self, x_j): + """Pass through features from source nodes.""" + return x_j + - self.hidden = 64 +class NoeGather(MessagePassing): + """ + Extract NOE shifts and features from NOE constraint nodes. + + Uses edge type: ("Noe", "prop_noe", triple_type) + Sets attributes on triple nodes: + - noe_shifts: [n, shift_dim] NOE shift embeddings + - noe_features: [n, feature_dim] NOE feature embeddings + """ + + def __init__(self, triple_type: str): + """ + Initialize NoeGather. + + Args: + triple_type: Name of target triple node type + """ + super().__init__(aggr="mean") + self.triple_type = triple_type + self.edge_type = ("Noe", "prop_noe", triple_type) + + def forward(self, data): + """ + Extract features from Noe nodes to triple nodes. + + Args: + data: HeteroData graph with Noe nodes and gather edges + + Returns: + Updated HeteroData with noe_shifts, noe_features set + """ + # Extract NOE shifts [n, shift_dim] from Noe.x (embedded from original 3D NOE features) + shifts = data["Noe"].x + + # Extract features [n, feature_dim] from Noe.f + features = data["Noe"].f + + # Get edge indices for this gather operation + edge_index = data[self.edge_type].edge_index + + # Determine number of target nodes (triple nodes) + num_triples = data[self.triple_type].x.size(0) - self.mlp1 = nn.Sequential( - nn.Linear(12, self.hidden, device=self.device), - nn.ReLU(), - nn.Linear(self.hidden, 19, device=self.device), - nn.LayerNorm(19, device=self.device), + # Propagate shifts with explicit size + data[self.triple_type].noe_shifts = self.propagate( + edge_index, x=shifts, size=(shifts.size(0), num_triples) ) - self.mlp2 = nn.Sequential( - nn.Linear(11, self.hidden, device=self.device), - nn.ReLU(), - nn.Linear(self.hidden, 13, device=self.device), - nn.LayerNorm(13, device=self.device), + # Propagate features + data[self.triple_type].noe_features = self.propagate( + edge_index, x=features, size=(features.size(0), num_triples) ) - def resresnoe_message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j, x22_j): - """ - Constructs message for (res, res, NOE) triple type. - """ - # Residue distances - rel_dist, dist2 = self.calc_manager.calc_res_distance(x12_j, x22_j) - - # Differences relative to NOE shifts (N, H', H") - diff1, diff2, diff3 = self.calc_manager.calc_noe_difference(x1_j, x2_j, x3_j) - - # Shift differences - diff4, diff5 = self.calc_manager.calc_shift_difference(x1_j, x2_j) - - # Input for MLP - # (N, H', H", N, H, dist2, features) - mlp_in = torch.cat( - (diff1, diff2, diff3, diff4, diff5, dist2, f1_j, f2_j, f3_j), dim=-1 - ) # [n, 9] - - # Output from MLP - # Out should include values for each 'change' wanting to make - # (N, H', H", N1, H1, N2, H2, dist1, dist2, features) - mlp_out = self.mlp1(mlp_in) # [n, 16] - - # NOE deltas [n, 1] - delta1x = diff1 * mlp_out[:, 0:1] # N - delta2x = diff2 * mlp_out[:, 1:2] # H' - delta3x = diff3 * mlp_out[:, 2:3] # H" - - # SHIFT deltas residue [n, 1] - delta4x = diff4 * mlp_out[:, 3:4] # N1 - delta5x = diff5 * mlp_out[:, 4:5] # H1 - delta6x = diff4 * mlp_out[:, 5:6] # N2 - delta7x = diff5 * mlp_out[:, 6:7] # H2 - - # DISTANCE deltas [n, 3] - delta12x = rel_dist * (mlp_out[:, 7:10]) - delta13x = rel_dist * (mlp_out[:, 10:13]) - - # FEATURE deltas [n, 1] - delta1f = mlp_out[:, 13:15] - delta2f = mlp_out[:, 15:17] - delta3f = mlp_out[:, 17:] - return torch.cat( - ( - delta1x, - delta2x, - delta3x, - delta4x, - delta5x, - delta6x, - delta7x, - delta12x, - delta13x, - delta1f, - delta2f, - delta3f, - ), + return data + + def message(self, x_j): + """Pass through features from source nodes.""" + return x_j + + +# ============================================================================ +# SECTION 2: Update Operations +# ============================================================================ +# These classes compute deltas on triple nodes using MLPs (standard nn.Module) + + +def _create_empty_deltas(device, num_triples=0, shift_dim=2, feature_dim=2): + """ + Create empty delta tensors for triple nodes with zero instances. + + This helper function is used when a triple set is empty to ensure all + delta attributes are properly initialized with correct shapes. + + Args: + device: torch device (CPU or CUDA) + num_triples: Number of triple nodes (typically 0 for empty sets) + shift_dim: Dimension of shift embeddings (default: 2 for backward compatibility) + feature_dim: Dimension of feature embeddings (default: 2 for backward compatibility) + + Returns: + Dictionary with 8 delta tensors: + - delta_first_coords: [num_triples, 3] + - delta_first_shifts: [num_triples, shift_dim] + - delta_first_features: [num_triples, feature_dim] + - delta_second_coords: [num_triples, 3] + - delta_second_shifts: [num_triples, shift_dim] + - delta_second_features: [num_triples, feature_dim] + - delta_noe_shifts: [num_triples, shift_dim] + - delta_noe_features: [num_triples, feature_dim] + """ + return { + 'delta_first_coords': torch.zeros((num_triples, 3), dtype=torch.float32, device=device), + 'delta_first_shifts': torch.zeros((num_triples, shift_dim), dtype=torch.float32, device=device), + 'delta_first_features': torch.zeros((num_triples, feature_dim), dtype=torch.float32, device=device), + 'delta_second_coords': torch.zeros((num_triples, 3), dtype=torch.float32, device=device), + 'delta_second_shifts': torch.zeros((num_triples, shift_dim), dtype=torch.float32, device=device), + 'delta_second_features': torch.zeros((num_triples, feature_dim), dtype=torch.float32, device=device), + 'delta_noe_shifts': torch.zeros((num_triples, shift_dim), dtype=torch.float32, device=device), + 'delta_noe_features': torch.zeros((num_triples, feature_dim), dtype=torch.float32, device=device), + } + + +class ResidueUpdate(nn.Module): + """ + Compute deltas for ResidueResidueNoeTriple using MLP. + + MLP Architecture: + Input: variable (dist_squared + 3*shift_dim + 3*feature_dim) + Hidden: configurable layers with ReLU activations + Output: variable (2 coord weights + 6*shift_dim + 3*feature_dim) + + Equivariance: + - Calculations are based on distances and differences between shifts + - Coordinate deltas are computed as difference * learned_weights to maintain + rotation/translation equivariance. + """ + + def __init__(self, triple_type: str, device, config): + """ + Initialize ResidueUpdate. + + Args: + triple_type: Name of triple node type (should be "ResidueResidueNoeTriple") + device: torch device (CPU or CUDA) + config: ModelConfig with shift_embed and feature_embed dimensions + """ + super().__init__() + self.triple_type = triple_type + self.device = device + self.config = config + + # Get MLP configuration from config + hidden_size = config.mlp.hidden_size + num_layers = config.mlp.num_layers + + # Calculate input/output sizes from config + shift_dim = config.shift_embed.output_dim + feature_dim = config.feature_embed.output_dim + input_size = residue_update_input_size(shift_dim, feature_dim) + output_size = residue_update_output_size(shift_dim, feature_dim) + + # Store dimensions for later use + self.shift_dim = shift_dim + self.feature_dim = feature_dim + + # Build MLP + layers = [] + + # First hidden layer + layers.append(nn.Linear(input_size, hidden_size)) + layers.append(nn.LayerNorm(hidden_size)) + layers.append(nn.ReLU()) + + # Additional hidden layers + for _ in range(num_layers - 1): + layers.append(nn.Linear(hidden_size, hidden_size)) + layers.append(nn.LayerNorm(hidden_size)) + layers.append(nn.ReLU()) + + # Output layer + layers.append(nn.Linear(hidden_size, output_size)) + layers.append(nn.LayerNorm(output_size)) + + self.mlp = nn.Sequential(*layers).to(device) + + def forward(self, data): + """ + Compute deltas for triple nodes. + + After embeddings, dimensions are: + - first_shifts: [n, shift_dim] + - second_shifts: [n, shift_dim] + - noe_shifts: [n, shift_dim] + - first_features: [n, feature_dim] + - second_features: [n, feature_dim] + - noe_features: [n, feature_dim] + + Args: + data: HeteroData with gathered attributes on triple nodes + + Returns: + Updated HeteroData with delta attributes set + """ + # Get gathered attributes (dimensions now variable) + first_coords = data[self.triple_type].first_coords # [n, 3] + first_shifts = data[self.triple_type].first_shifts # [n, shift_dim] + first_features = data[self.triple_type].first_features # [n, feature_dim] + second_coords = data[self.triple_type].second_coords # [n, 3] + second_shifts = data[self.triple_type].second_shifts # [n, shift_dim] + second_features = data[self.triple_type].second_features # [n, feature_dim] + noe_shifts = data[self.triple_type].noe_shifts # [n, shift_dim] + noe_features = data[self.triple_type].noe_features # [n, feature_dim] + + num_triples = first_shifts.size(0) + + # Handle empty triple sets using helper function + if num_triples == 0: + deltas = _create_empty_deltas(self.device, num_triples=0, shift_dim=self.shift_dim, feature_dim=self.feature_dim) + for key, value in deltas.items(): + setattr(data[self.triple_type], key, value) + return data + + # Calculate relative distance and dist_squared for equivariance + rel_dist, dist_squared = calc_res_distance(first_coords, second_coords) + + # Compute shift differences BEFORE MLP input (for translation equivariance) + diff_first_to_second = second_shifts - first_shifts # [n, shift_dim] + diff_first_to_noe = noe_shifts - first_shifts # [n, shift_dim] + diff_second_to_noe = noe_shifts - second_shifts # [n, shift_dim] + + # Concatenate all features for MLP input + # Input: dist_squared + 3*shift_dim (differences) + 3*feature_dim + mlp_input = torch.cat( + [ + dist_squared, # [n, 1] + diff_first_to_second, # [n, shift_dim] + diff_first_to_noe, # [n, shift_dim] + diff_second_to_noe, # [n, shift_dim] + first_features, # [n, feature_dim] + second_features, # [n, feature_dim] + noe_features, # [n, feature_dim] + ], dim=-1, - ) # [n, 16] [n, 19] - - def shiftshiftnoe_message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j): - """ - Constructs message for (shift, shift, NOE), (shift, res, NOE), (res, shift, NOE) triple types. - """ - # Differences relative to NOE shifts (N, H', H") - diff1, diff2, diff3 = self.calc_manager.calc_noe_difference(x1_j, x2_j, x3_j) - - # Shift differences - diff4, diff5 = self.calc_manager.calc_shift_difference(x1_j, x2_j) - - # Input for MLP - # (N, H', H", N, H, dist2, features) - mlp_in = torch.cat( - (diff1, diff2, diff3, diff4, diff5, f1_j, f2_j, f3_j), dim=-1 - ) # [n, 8] - - # Output from MLP - # Out should include values for each 'change' wanting to make - # (N, H', H", N1, H1, N2, H2, features) - mlp_out = self.mlp2(mlp_in) # [n, 10] - - # NOE deltas [n, 1] - delta1x = diff1 * mlp_out[:, 0:1] # N - delta2x = diff2 * mlp_out[:, 1:2] # H' - delta3x = diff3 * mlp_out[:, 2:3] # H" - - # SHIFT deltas residue [n, 1] - delta4x = diff4 * mlp_out[:, 3:4] # N1 - delta5x = diff5 * mlp_out[:, 4:5] # H1 - delta6x = diff4 * mlp_out[:, 5:6] # N2 - delta7x = diff5 * mlp_out[:, 6:7] # H2 - - # FEATURE deltas [n, 1] - delta1f = mlp_out[:, 7:9] - delta2f = mlp_out[:, 9:11] - delta3f = mlp_out[:, 11:] - return torch.cat( - ( - delta1x, - delta2x, - delta3x, - delta4x, - delta5x, - delta6x, - delta7x, - delta1f, - delta2f, - delta3f, - ), + ) + + # Apply MLP to get deltas + mlp_output = self.mlp(mlp_input) + + # Parse output: + # 2 coord weights + 6 scalar shift weights + 3*feature_dim feature deltas + coord_weight_first = 0 * mlp_output[:, 0:1] # [n, 1] - still zeroed out + coord_weight_second = 0 * mlp_output[:, 1:2] # [n, 1] - still zeroed out + + # Shift weights: 6 scalars + w1 = mlp_output[:, 2:3] # [n, 1] - weight for first -> second + w2 = mlp_output[:, 3:4] # [n, 1] - weight for second -> first + w3 = mlp_output[:, 4:5] # [n, 1] - weight for first -> noe + w4 = mlp_output[:, 5:6] # [n, 1] - weight for second -> noe + w5 = mlp_output[:, 6:7] # [n, 1] - weight for noe -> first + w6 = mlp_output[:, 7:8] # [n, 1] - weight for noe -> second + + # Feature deltas: 3 vectors of feature_dim each + feature_deltas = mlp_output[:, 8:] # [n, 3*feature_dim] + + # Compute equivariant coordinate deltas: rel_dist * learned_weights + delta_first_coords = rel_dist * coord_weight_first # [n, 3] + delta_second_coords = -rel_dist * coord_weight_second # [n, 3] + + # Compute shift deltas using scalar weights × difference vectors + # The scalar weights [n, 1] broadcast with difference vectors [n, shift_dim] + # + # Position 1 shifts - receives contributions from w1 and w3 + delta_first_shifts = w1 * diff_first_to_second + w3 * diff_first_to_noe # [n, shift_dim] + + # Position 2 shifts - receives contributions from w2 and w4 + # Note: diff_first_to_second = second - first, so -diff_first_to_second = first - second + delta_second_shifts = -w2 * diff_first_to_second + w4 * diff_second_to_noe # [n, shift_dim] + + # NOE shifts - receives contributions from w5 and w6 + delta_noe_shifts = -w5 * diff_first_to_noe + -w6 * diff_second_to_noe # [n, shift_dim] + + # Compute feature deltas (3 vectors of feature_dim each) + delta_first_features = feature_deltas[:, 0*self.feature_dim:1*self.feature_dim] # [n, feature_dim] + delta_second_features = feature_deltas[:, 1*self.feature_dim:2*self.feature_dim] # [n, feature_dim] + delta_noe_features = feature_deltas[:, 2*self.feature_dim:3*self.feature_dim] # [n, feature_dim] + + # Set delta attributes on triple nodes + data[self.triple_type].delta_first_coords = delta_first_coords + data[self.triple_type].delta_first_shifts = delta_first_shifts + data[self.triple_type].delta_first_features = delta_first_features + data[self.triple_type].delta_second_coords = delta_second_coords + data[self.triple_type].delta_second_shifts = delta_second_shifts + data[self.triple_type].delta_second_features = delta_second_features + data[self.triple_type].delta_noe_shifts = delta_noe_shifts + data[self.triple_type].delta_noe_features = delta_noe_features + + return data + + +class PeakUpdate(nn.Module): + """ + Compute deltas for Peak-based triples using MLP. + + MLP Architecture: + Input: variable (3*shift_dim + 3*feature_dim, NO dist_squared) + Hidden: configurable layers with ReLU activations + Output: variable (6*shift_dim + 3*feature_dim, NO coord weights) + + Coordinate Handling: + Peaks have no coordinates, so all coordinate deltas are zero. + + Equivariance: + Calculations are based on differences between shifts + """ + + def __init__(self, triple_type: str, device, config): + """ + Initialize PeakUpdate. + + Args: + triple_type: Name of triple node type + device: torch device (CPU or CUDA) + config: ModelConfig with shift_embed and feature_embed dimensions + """ + super().__init__() + self.triple_type = triple_type + self.device = device + self.config = config + + # Get MLP configuration from config + hidden_size = config.mlp.hidden_size + num_layers = config.mlp.num_layers + + # Calculate input/output sizes from config + shift_dim = config.shift_embed.output_dim + feature_dim = config.feature_embed.output_dim + input_size = peak_update_input_size(shift_dim, feature_dim) + output_size = peak_update_output_size(shift_dim, feature_dim) + + # Store dimensions for later use + self.shift_dim = shift_dim + self.feature_dim = feature_dim + + # Build MLP + layers = [] + + # First hidden layer + layers.append(nn.Linear(input_size, hidden_size)) + layers.append(nn.LayerNorm(hidden_size)) + layers.append(nn.ReLU()) + + # Additional hidden layers + for _ in range(num_layers - 1): + layers.append(nn.Linear(hidden_size, hidden_size)) + layers.append(nn.LayerNorm(hidden_size)) + layers.append(nn.ReLU()) + + # Output layer + layers.append(nn.Linear(hidden_size, output_size)) + layers.append(nn.LayerNorm(output_size)) + + self.mlp = nn.Sequential(*layers).to(device) + + def forward(self, data): + """ + Compute deltas for triple nodes. + + After embeddings, dimensions are: + - first_shifts: [n, shift_dim] + - second_shifts: [n, shift_dim] + - noe_shifts: [n, shift_dim] + - first_features: [n, feature_dim] + - second_features: [n, feature_dim] + - noe_features: [n, feature_dim] + + Args: + data: HeteroData with gathered attributes on triple nodes + + Returns: + Updated HeteroData with delta attributes set + """ + # Get gathered attributes (peaks have NO coordinates) + first_shifts = data[self.triple_type].first_shifts # [n, shift_dim] + first_features = data[self.triple_type].first_features # [n, feature_dim] + second_shifts = data[self.triple_type].second_shifts # [n, shift_dim] + second_features = data[self.triple_type].second_features # [n, feature_dim] + noe_shifts = data[self.triple_type].noe_shifts # [n, shift_dim] + noe_features = data[self.triple_type].noe_features # [n, feature_dim] + + num_triples = first_shifts.size(0) + + # Handle empty triple sets using helper function + if num_triples == 0: + deltas = _create_empty_deltas(self.device, num_triples=0, shift_dim=self.shift_dim, feature_dim=self.feature_dim) + for key, value in deltas.items(): + setattr(data[self.triple_type], key, value) + return data + + # Compute shift differences BEFORE MLP input (for translation equivariance) + diff_first_to_second = second_shifts - first_shifts # [n, shift_dim] + diff_first_to_noe = noe_shifts - first_shifts # [n, shift_dim] + diff_second_to_noe = noe_shifts - second_shifts # [n, shift_dim] + + # Concatenate all features for MLP input (NO distance calculations for peaks) + # Input: 3*shift_dim (differences) + 3*feature_dim + mlp_input = torch.cat( + [ + diff_first_to_second, # [n, shift_dim] + diff_first_to_noe, # [n, shift_dim] + diff_second_to_noe, # [n, shift_dim] + first_features, # [n, feature_dim] + second_features, # [n, feature_dim] + noe_features, # [n, feature_dim] + ], dim=-1, - ) # [n, 10] [n, 13] - - def forward(self, x1, x12, x2, x22, x3, f1, f2, f3, edge_index): - out = self.propagate( - edge_index, x1=x1, x2=x2, x3=x3, f1=f1, f2=f2, f3=f3, x12=x12, x22=x22 - ) # not sure how to deal with size here - return out - - def message(self, x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j=None, x22_j=None): - # residue based triple type will have two sets of coordinates (x12_j and x22_j) where shifts will have nonetype - if x12_j != None and x22_j != None: - return self.resresnoe_message( - x1_j, x2_j, x3_j, f1_j, f2_j, f3_j, x12_j, x22_j - ) - else: - return self.shiftshiftnoe_message(x1_j, x2_j, x3_j, f1_j, f2_j, f3_j) - - def update(self, aggr_out, x12, x22): - # residue based triple type will have two sets of coordinates (x12_j and x22_j) where shifts will have nonetype - if x12 != None and x22 != None: - # output shapes: noe[n, 3], shift1[n, 2], shift2[n, 2], dist1[n, 3], dist2[n, 3], f1[n, 2], f2[n, 2], f3[n, 2] - ( - delta_noe, - delta_shift1, - delta_shift2, - delta_dist1, - delta_dist2, - delta_f1, - delta_f2, - delta_f3, - ) = ( - aggr_out[:, 0:3], - aggr_out[:, 3:5], - aggr_out[:, 5:7], - aggr_out[:, 7:10], - aggr_out[:, 10:13], - aggr_out[:, 13:15], - aggr_out[:, 15:17], - aggr_out[:, 17:], - ) - return ( - delta_shift1, - delta_shift2, - delta_noe, - delta_dist1, - delta_dist2, - delta_f1, - delta_f2, - delta_f3, - ) - else: - # output shapes: noe[n, 3], shift1[n, 2], shift2[n, 2], f1[n, 2], f2[n, 2], f3[n, 2] - delta_noe, delta_shift1, delta_shift2, delta_f1, delta_f2, delta_f3 = ( - aggr_out[:, 0:3], - aggr_out[:, 3:5], - aggr_out[:, 5:7], - aggr_out[:, 7:9], - aggr_out[:, 9:11], - aggr_out[:, 11:], - ) - return delta_shift1, delta_shift2, delta_noe, delta_f1, delta_f2, delta_f3 - - -class TripleMessagePass(MessagePassing): - """ - Standard message passing class for outgoing triple messages. - """ - - def __init__(self, aggr): - super().__init__(aggr=aggr) - - def forward(self, x_source, x_target, edge_index): - # SIZE (n, m) (source, target) - return self.propagate( - edge_index=edge_index, - x=(x_source, x_target), - size=(x_source.size(0), x_target.size(0)), ) + # Apply MLP to get outputs + mlp_output = self.mlp(mlp_input) + + # Parse output: + # 6 scalar shift weights + 3*feature_dim feature deltas (NO coord weights) + w1 = mlp_output[:, 0:1] # [n, 1] - weight for first -> second + w2 = mlp_output[:, 1:2] # [n, 1] - weight for second -> first + w3 = mlp_output[:, 2:3] # [n, 1] - weight for first -> noe + w4 = mlp_output[:, 3:4] # [n, 1] - weight for second -> noe + w5 = mlp_output[:, 4:5] # [n, 1] - weight for noe -> first + w6 = mlp_output[:, 5:6] # [n, 1] - weight for noe -> second + + # Feature deltas: 3 vectors of feature_dim each + feature_deltas = mlp_output[:, 6:] # [n, 3*feature_dim] + + # Create ZERO coordinate deltas (peaks have no coordinates) + delta_first_coords = torch.zeros((num_triples, 3), dtype=torch.float32, device=self.device) + delta_second_coords = torch.zeros((num_triples, 3), dtype=torch.float32, device=self.device) + + # Compute shift deltas using scalar weights × difference vectors + # The scalar weights [n, 1] broadcast with difference vectors [n, shift_dim] + # + # Position 1 shifts - receives contributions from w1 and w3 + delta_first_shifts = w1 * diff_first_to_second + w3 * diff_first_to_noe # [n, shift_dim] + + # Position 2 shifts - receives contributions from w2 and w4 + # Note: diff_first_to_second = second - first, so -diff_first_to_second = first - second + delta_second_shifts = -w2 * diff_first_to_second + w4 * diff_second_to_noe # [n, shift_dim] + + # NOE shifts - receives contributions from w5 and w6 + delta_noe_shifts = -w5 * diff_first_to_noe + -w6 * diff_second_to_noe # [n, shift_dim] + + # Compute feature deltas (3 vectors of feature_dim each) + delta_first_features = feature_deltas[:, 0*self.feature_dim:1*self.feature_dim] # [n, feature_dim] + delta_second_features = feature_deltas[:, 1*self.feature_dim:2*self.feature_dim] # [n, feature_dim] + delta_noe_features = feature_deltas[:, 2*self.feature_dim:3*self.feature_dim] # [n, feature_dim] + + # Set delta attributes on triple nodes + data[self.triple_type].delta_first_coords = delta_first_coords + data[self.triple_type].delta_first_shifts = delta_first_shifts + data[self.triple_type].delta_first_features = delta_first_features + data[self.triple_type].delta_second_coords = delta_second_coords + data[self.triple_type].delta_second_shifts = delta_second_shifts + data[self.triple_type].delta_second_features = delta_second_features + data[self.triple_type].delta_noe_shifts = delta_noe_shifts + data[self.triple_type].delta_noe_features = delta_noe_features + + return data + + +# ============================================================================ +# SECTION 3: Scatter Operations +# ============================================================================ +# These classes propagate deltas from triple nodes back to source nodes using +# PyTorch Geometric MessagePassing with aggr="mean" + + +class FirstResidueScatter(MessagePassing): + """ + Propagate deltas from triple nodes to residues in first position. + + Uses edge type: ("Residue", "prop_first", triple_type) with reversed flow + Reads delta attributes from triple nodes: + - delta_first_coords: [n, 3] coordinate deltas + - delta_first_shifts: [n, shift_dim] shift deltas + - delta_first_features: [n, feature_dim] feature deltas + Updates Residue nodes: + - Residue.x[:, 0:3] with coordinate deltas + - Residue.x[:, 3:] with shift deltas + - Residue.f with feature deltas + """ + + def __init__(self, triple_type: str): + """ + Initialize FirstResidueScatter. + + Args: + triple_type: Name of source triple node type + """ + super().__init__(aggr="mean", flow="target_to_source") + self.triple_type = triple_type + self.edge_type = ("Residue", "prop_first", triple_type) + + def forward(self, data): + """ + Propagate deltas from triple nodes to Residue nodes. + + Args: + data: HeteroData with delta attributes on triple nodes + + Returns: + Updated HeteroData with Residue.x and Residue.f modified + """ + # Get delta attributes from triple nodes + delta_coords = data[self.triple_type].delta_first_coords # [n, 3] + delta_shifts = data[self.triple_type].delta_first_shifts # [n, shift_dim] + delta_features = data[self.triple_type].delta_first_features # [n, feature_dim] + + # Get edge indices for this scatter operation + edge_index = data[self.edge_type].edge_index + + # Determine number of target nodes (Residue nodes - targets when using reversed flow) + num_residues = data["Residue"].x.size(0) + + # Propagate coordinate deltas (with reversed flow, size is (target, source)) + coord_updates = self.propagate( + edge_index, x=delta_coords, size=(num_residues, delta_coords.size(0)) + ) + + # Propagate shift deltas + shift_updates = self.propagate( + edge_index, x=delta_shifts, size=(num_residues, delta_shifts.size(0)) + ) + + # Propagate feature deltas + feature_updates = self.propagate( + edge_index, x=delta_features, size=(num_residues, delta_features.size(0)) + ) + + # Assemble complete new tensors (avoid in-place updates for gradient preservation) + # Update Residue.x by concatenating updated coords and shifts + old_x = data["Residue"].x + new_coords = old_x[:, 0:3] + coord_updates + new_shifts = old_x[:, 3:] + shift_updates + data["Residue"].x = torch.cat([new_coords, new_shifts], dim=-1) + + # Update Residue.f + data["Residue"].f = data["Residue"].f + feature_updates + + return data + def message(self, x_j): + """Pass through deltas from triple nodes.""" return x_j - def update(self, aggr_out, x): - # non-spatial features can also be included in same call but requires more work if they're not being updated - # (don't want to do residue feature updates 2x - only call once on shifts or coordinates) - # x_val, f_val = aggr_out[:, :len(x[1][1])], aggr_out[:, len(x[1][1]):] - # outf = f_val + f[1] - # outx = x_val + x[1] - out = aggr_out + x[1] - return out +class FirstPeakScatter(MessagePassing): + """ + Propagate deltas from triple nodes to peaks in first position. -class TripleOut(nn.Module): - def __init__(self): - super().__init__() - self.triple_message = TripleMessagePass(aggr="add") + Uses edge type: ("Peak", "prop_first", triple_type) with reversed flow + Reads delta attributes from triple nodes: + - delta_first_shifts: [n, shift_dim] shift deltas + - delta_first_features: [n, feature_dim] feature deltas + Updates Peak nodes: + - Peak.x with shift deltas + - Peak.f with feature deltas + Note: Peaks have NO coordinates + """ + + def __init__(self, triple_type: str): + """ + Initialize FirstPeakScatter. + + Args: + triple_type: Name of source triple node type + """ + super().__init__(aggr="mean", flow="target_to_source") + self.triple_type = triple_type + self.edge_type = ("Peak", "prop_first", triple_type) + + def forward(self, data): + """ + Propagate deltas from triple nodes to Peak nodes. + + Args: + data: HeteroData with delta attributes on triple nodes + + Returns: + Updated HeteroData with Peak.x and Peak.f modified + """ + # Get delta attributes from triple nodes (no coords for peaks) + delta_shifts = data[self.triple_type].delta_first_shifts # [n, shift_dim] + delta_features = data[self.triple_type].delta_first_features # [n, feature_dim] + + # Get edge indices for this scatter operation + edge_index = data[self.edge_type].edge_index + + # Determine number of target nodes (Peak nodes - targets when using reversed flow) + num_peaks = data["Peak"].x.size(0) + + # Propagate shift deltas (with reversed flow, size is (target, source)) + shift_updates = self.propagate( + edge_index, x=delta_shifts, size=(num_peaks, delta_shifts.size(0)) + ) + + # Propagate feature deltas + feature_updates = self.propagate( + edge_index, x=delta_features, size=(num_peaks, delta_features.size(0)) + ) + + # Update Peak.x and Peak.f (assemble new tensors) + data["Peak"].x = data["Peak"].x + shift_updates + data["Peak"].f = data["Peak"].f + feature_updates + + return data + + def message(self, x_j): + """Pass through deltas from triple nodes.""" + return x_j + + +class SecondResidueScatter(MessagePassing): + """ + Propagate deltas from triple nodes to residues in second position. + + Uses edge type: ("Residue", "prop_second", triple_type) with reversed flow + Reads delta attributes from triple nodes: + - delta_second_coords: [n, 3] coordinate deltas + - delta_second_shifts: [n, shift_dim] shift deltas + - delta_second_features: [n, feature_dim] feature deltas + Updates Residue nodes: + - Residue.x[:, 0:3] with coordinate deltas + - Residue.x[:, 3:] with shift deltas + - Residue.f with feature deltas + """ + + def __init__(self, triple_type: str): + """ + Initialize SecondResidueScatter. + + Args: + triple_type: Name of source triple node type + """ + super().__init__(aggr="mean", flow="target_to_source") + self.triple_type = triple_type + self.edge_type = ("Residue", "prop_second", triple_type) + + def forward(self, data): + """ + Propagate deltas from triple nodes to Residue nodes. + + Args: + data: HeteroData with delta attributes on triple nodes + + Returns: + Updated HeteroData with Residue.x and Residue.f modified + """ + # Get delta attributes from triple nodes + delta_coords = data[self.triple_type].delta_second_coords # [n, 3] + delta_shifts = data[self.triple_type].delta_second_shifts # [n, shift_dim] + delta_features = data[self.triple_type].delta_second_features # [n, feature_dim] + + # Get edge indices for this scatter operation + edge_index = data[self.edge_type].edge_index + + # Determine number of target nodes (Residue nodes - targets when using reversed flow) + num_residues = data["Residue"].x.size(0) + + # Propagate coordinate deltas (with reversed flow, size is (target, source)) + coord_updates = self.propagate( + edge_index, x=delta_coords, size=(num_residues, delta_coords.size(0)) + ) + + # Propagate shift deltas + shift_updates = self.propagate( + edge_index, x=delta_shifts, size=(num_residues, delta_shifts.size(0)) + ) + + # Propagate feature deltas + feature_updates = self.propagate( + edge_index, x=delta_features, size=(num_residues, delta_features.size(0)) + ) + + # Assemble complete new tensors (avoid in-place updates for gradient preservation) + # Update Residue.x by concatenating updated coords and shifts + old_x = data["Residue"].x + new_coords = old_x[:, 0:3] + coord_updates + new_shifts = old_x[:, 3:] + shift_updates + data["Residue"].x = torch.cat([new_coords, new_shifts], dim=-1) + + # Update Residue.f + data["Residue"].f = data["Residue"].f + feature_updates + + return data + + def message(self, x_j): + """Pass through deltas from triple nodes.""" + return x_j + + +class SecondPeakScatter(MessagePassing): + """ + Propagate deltas from triple nodes to peaks in second position. + + Uses edge type: ("Peak", "prop_second", triple_type) with reversed flow + Reads delta attributes from triple nodes: + - delta_second_shifts: [n, shift_dim] shift deltas + - delta_second_features: [n, feature_dim] feature deltas + Updates Peak nodes: + - Peak.x with shift deltas + - Peak.f with feature deltas + Note: Peaks have NO coordinates + """ + + def __init__(self, triple_type: str): + """ + Initialize SecondPeakScatter. + + Args: + triple_type: Name of source triple node type + """ + super().__init__(aggr="mean", flow="target_to_source") + self.triple_type = triple_type + self.edge_type = ("Peak", "prop_second", triple_type) + + def forward(self, data): + """ + Propagate deltas from triple nodes to Peak nodes. + + Args: + data: HeteroData with delta attributes on triple nodes - def update_data(self, data, x_source, f_source, edge_type, x_index, update_f=True): + Returns: + Updated HeteroData with Peak.x and Peak.f modified """ - Calls message passing and directly updates the heterodata object based on target. - Set to update non-spatial features in all cases, but can be turned off case by case to prevent duplicate updates. + # Get delta attributes from triple nodes (no coords for peaks) + delta_shifts = data[self.triple_type].delta_second_shifts # [n, shift_dim] + delta_features = data[self.triple_type].delta_second_features # [n, feature_dim] + + # Get edge indices for this scatter operation + edge_index = data[self.edge_type].edge_index + + # Determine number of target nodes (Peak nodes - targets when using reversed flow) + num_peaks = data["Peak"].x.size(0) + + # Propagate shift deltas (with reversed flow, size is (target, source)) + shift_updates = self.propagate( + edge_index, x=delta_shifts, size=(num_peaks, delta_shifts.size(0)) + ) + + # Propagate feature deltas + feature_updates = self.propagate( + edge_index, x=delta_features, size=(num_peaks, delta_features.size(0)) + ) + + # Update Peak.x and Peak.f (assemble new tensors) + data["Peak"].x = data["Peak"].x + shift_updates + data["Peak"].f = data["Peak"].f + feature_updates + + return data + + def message(self, x_j): + """Pass through deltas from triple nodes.""" + return x_j + + +class NoeScatter(MessagePassing): + """ + Propagate deltas from triple nodes to NOE constraint nodes. + + Uses edge type: ("Noe", "prop_noe", triple_type) with reversed flow + Reads delta attributes from triple nodes: + - delta_noe_shifts: [n, shift_dim] NOE shift deltas + - delta_noe_features: [n, feature_dim] feature deltas + Updates Noe nodes: + - Noe.x with shift deltas + - Noe.f with feature deltas + """ + + def __init__(self, triple_type: str): + """ + Initialize NoeScatter. + + Args: + triple_type: Name of source triple node type + """ + super().__init__(aggr="mean", flow="target_to_source") + self.triple_type = triple_type + self.edge_type = ("Noe", "prop_noe", triple_type) + + def forward(self, data): + """ + Propagate deltas from triple nodes to Noe nodes. + + Args: + data: HeteroData with delta attributes on triple nodes + + Returns: + Updated HeteroData with Noe.x and Noe.f modified """ - source, edge, target = edge_type - edge_index = data[edge_type].edge_index + # Get delta attributes from triple nodes + delta_shifts = data[self.triple_type].delta_noe_shifts # [n, shift_dim] + delta_features = data[self.triple_type].delta_noe_features # [n, feature_dim] - # Message passing/update for spatial features - data[target].x[:, x_index] = self.triple_message( - x_source, data[target].x[:, x_index], edge_index + # Get edge indices for this scatter operation + edge_index = data[self.edge_type].edge_index + + # Determine number of target nodes (Noe nodes - targets when using reversed flow) + num_noes = data["Noe"].x.size(0) + + # Propagate shift deltas (with reversed flow, size is (target, source)) + shift_updates = self.propagate( + edge_index, x=delta_shifts, size=(num_noes, delta_shifts.size(0)) + ) + + # Propagate feature deltas + feature_updates = self.propagate( + edge_index, x=delta_features, size=(num_noes, delta_features.size(0)) ) - # Message passing/update for non-spatial features if needed (don't want a duplicate update for residue features) - if update_f: - data[target].f = self.triple_message(f_source, data[target].f, edge_index) + # Update Noe.x and Noe.f (assemble new tensors) + data["Noe"].x = data["Noe"].x + shift_updates + data["Noe"].f = data["Noe"].f + feature_updates + + return data + + def message(self, x_j): + """Pass through deltas from triple nodes.""" + return x_j + + +# ============================================================================ +# SECTION 4: Triple Composition Classes +# ============================================================================ +# These classes wire together gather, update, and scatter operations for each +# triple type + + +class ResidueResidueNoeTriple(nn.Module): + """ + Wire together components for ResidueResidueNoeTriple. + + This triple type handles (Residue, Residue, Noe) relationships and is the + only triple type that updates coordinates (equivariant). + + Components: + - FirstResidueGather: Extract from first residue + - SecondResidueGather: Extract from second residue + - NoeGather: Extract from NOE constraint + - ResidueUpdate: Compute coordinate and shift deltas + - FirstResidueScatter: Propagate to first residue + - SecondResidueScatter: Propagate to second residue + - NoeScatter: Propagate to NOE constraint + + Forward pass sequence: + gather_first → gather_second → gather_noe → update → + scatter_first → scatter_second → scatter_noe + """ + + def __init__(self, device, config): + """ + Initialize ResidueResidueNoeTriple. + + Args: + device: torch device (CPU or CUDA) + config: ModelConfig with dimension settings + """ + super().__init__() + triple_type = "ResidueResidueNoeTriple" + + # Instantiate gather operations + self.first_gather = FirstResidueGather(triple_type) + self.second_gather = SecondResidueGather(triple_type) + self.noe_gather = NoeGather(triple_type) + + # Instantiate update operation + self.update = ResidueUpdate(triple_type, device, config) + + # Instantiate scatter operations + self.first_scatter = FirstResidueScatter(triple_type) + self.second_scatter = SecondResidueScatter(triple_type) + self.noe_scatter = NoeScatter(triple_type) + + def forward(self, data): + """ + Execute gather → update → scatter pipeline. + + Args: + data: HeteroData graph + + Returns: + Updated HeteroData graph + """ + data = self.first_gather(data) + data = self.second_gather(data) + data = self.noe_gather(data) + data = self.update(data) + data = self.first_scatter(data) + data = self.second_scatter(data) + data = self.noe_scatter(data) + return data + + +class ResiduePeakNoeTriple(nn.Module): + """ + Wire together components for ResiduePeakNoeTriple. + + This triple type handles (Residue, Peak, Noe) relationships. Updates shifts + only, NO coordinate updates. + + Components: + - FirstResidueGather: Extract from first residue + - SecondPeakGather: Extract from second peak + - NoeGather: Extract from NOE constraint + - PeakUpdate: Compute shift deltas only + - FirstResidueScatter: Propagate to first residue + - SecondPeakScatter: Propagate to second peak + - NoeScatter: Propagate to NOE constraint + + Forward pass sequence: + gather_first → gather_second → gather_noe → update → + scatter_first → scatter_second → scatter_noe + """ + + def __init__(self, device, config): + """ + Initialize ResiduePeakNoeTriple. + + Args: + device: torch device (CPU or CUDA) + config: ModelConfig with dimension settings + """ + super().__init__() + triple_type = "ResiduePeakNoeTriple" + + # Instantiate gather operations + self.first_gather = FirstResidueGather(triple_type) + self.second_gather = SecondPeakGather(triple_type) + self.noe_gather = NoeGather(triple_type) + + # Instantiate update operation + self.update = PeakUpdate(triple_type, device, config) + + # Instantiate scatter operations + self.first_scatter = FirstResidueScatter(triple_type) + self.second_scatter = SecondPeakScatter(triple_type) + self.noe_scatter = NoeScatter(triple_type) + + def forward(self, data): + """ + Execute gather → update → scatter pipeline. + + Args: + data: HeteroData graph + + Returns: + Updated HeteroData graph + """ + data = self.first_gather(data) + data = self.second_gather(data) + data = self.noe_gather(data) + data = self.update(data) + data = self.first_scatter(data) + data = self.second_scatter(data) + data = self.noe_scatter(data) + return data + + +class PeakResidueNoeTriple(nn.Module): + """ + Wire together components for PeakResidueNoeTriple. + + This triple type handles (Peak, Residue, Noe) relationships. Updates shifts + only, NO coordinate updates. + + Components: + - FirstPeakGather: Extract from first peak + - SecondResidueGather: Extract from second residue + - NoeGather: Extract from NOE constraint + - PeakUpdate: Compute shift deltas only + - FirstPeakScatter: Propagate to first peak + - SecondResidueScatter: Propagate to second residue + - NoeScatter: Propagate to NOE constraint + + Forward pass sequence: + gather_first → gather_second → gather_noe → update → + scatter_first → scatter_second → scatter_noe + """ + + def __init__(self, device, config): + """ + Initialize PeakResidueNoeTriple. + + Args: + device: torch device (CPU or CUDA) + config: ModelConfig with dimension settings + """ + super().__init__() + triple_type = "PeakResidueNoeTriple" + + # Instantiate gather operations + self.first_gather = FirstPeakGather(triple_type) + self.second_gather = SecondResidueGather(triple_type) + self.noe_gather = NoeGather(triple_type) + + # Instantiate update operation + self.update = PeakUpdate(triple_type, device, config) + + # Instantiate scatter operations + self.first_scatter = FirstPeakScatter(triple_type) + self.second_scatter = SecondResidueScatter(triple_type) + self.noe_scatter = NoeScatter(triple_type) + + def forward(self, data): + """ + Execute gather → update → scatter pipeline. + + Args: + data: HeteroData graph + + Returns: + Updated HeteroData graph + """ + data = self.first_gather(data) + data = self.second_gather(data) + data = self.noe_gather(data) + data = self.update(data) + data = self.first_scatter(data) + data = self.second_scatter(data) + data = self.noe_scatter(data) + return data + + +class PeakPeakNoeTriple(nn.Module): + """ + Wire together components for PeakPeakNoeTriple. + + This triple type handles (Peak, Peak, Noe) relationships. Updates shifts + only, NO coordinate updates. + + Components: + - FirstPeakGather: Extract from first peak + - SecondPeakGather: Extract from second peak + - NoeGather: Extract from NOE constraint + - PeakUpdate: Compute shift deltas only + - FirstPeakScatter: Propagate to first peak + - SecondPeakScatter: Propagate to second peak + - NoeScatter: Propagate to NOE constraint + + Forward pass sequence: + gather_first → gather_second → gather_noe → update → + scatter_first → scatter_second → scatter_noe + """ + + def __init__(self, device, config): + """ + Initialize PeakPeakNoeTriple. + + Args: + device: torch device (CPU or CUDA) + config: ModelConfig with dimension settings + """ + super().__init__() + triple_type = "PeakPeakNoeTriple" + + # Instantiate gather operations + self.first_gather = FirstPeakGather(triple_type) + self.second_gather = SecondPeakGather(triple_type) + self.noe_gather = NoeGather(triple_type) + + # Instantiate update operation + self.update = PeakUpdate(triple_type, device, config) + + # Instantiate scatter operations + self.first_scatter = FirstPeakScatter(triple_type) + self.second_scatter = SecondPeakScatter(triple_type) + self.noe_scatter = NoeScatter(triple_type) + + def forward(self, data): + """ + Execute gather → update → scatter pipeline. + + Args: + data: HeteroData graph + + Returns: + Updated HeteroData graph + """ + data = self.first_gather(data) + data = self.second_gather(data) + data = self.noe_gather(data) + data = self.update(data) + data = self.first_scatter(data) + data = self.second_scatter(data) + data = self.noe_scatter(data) return data + + +# ============================================================================ +# SECTION 5: Helper Functions +# ============================================================================ +# Calculation utilities used by update operations + + +def calc_noe_difference(x1, x2, noe): + """ + Calculate shift differences between NOE and residue/peak shifts. + + Args: + x1: First node shifts [n, 2] (H, N) + x2: Second node shifts [n, 2] (H, N) + noe: NOE shifts [n, 3] (N, H', H") + + Returns: + Tuple of (diff_N, diff_H1, diff_H2) each [n, 1] + """ + NOE_N1 = slice(0, 1) + NOE_H1 = slice(1, 2) + NOE_H2 = slice(2, 3) + SHIFT_H = slice(0, 1) # H is at index 0 + SHIFT_N = slice(1, 2) # N is at index 1 + + diff_N = noe[:, NOE_N1] - x1[:, SHIFT_N] # N [n, 1] + diff_H1 = noe[:, NOE_H1] - x1[:, SHIFT_H] # H' [n, 1] + diff_H2 = noe[:, NOE_H2] - x2[:, SHIFT_H] # H" [n, 1] + return diff_N, diff_H1, diff_H2 + + +def calc_shift_difference(x1, x2): + """ + Calculate shift differences between two nodes. + + Args: + x1: First node shifts [n, 2] (H, N) + x2: Second node shifts [n, 2] (H, N) + + Returns: + Tuple of (diff_N, diff_H) each [n, 1] + """ + SHIFT_H = slice(0, 1) # H is at index 0 + SHIFT_N = slice(1, 2) # N is at index 1 + + diff_N = x1[:, SHIFT_N] - x2[:, SHIFT_N] # N [n, 1] + diff_H = x1[:, SHIFT_H] - x2[:, SHIFT_H] # H [n, 1] + return diff_N, diff_H + + +def calc_res_distance(x1, x2): + """ + Calculate relative distance and squared distance between residues. + + Args: + x1: First residue coordinates [n, 3] + x2: Second residue coordinates [n, 3] + + Returns: + Tuple of (rel_dist, dist_squared) + rel_dist: Relative distance vector [n, 3] + dist_squared: Squared distance scalar [n, 1] + """ + rel_dist = x1 - x2 # [n, 3] + dist_squared = torch.norm(rel_dist, dim=-1, keepdim=True) ** 2 # [n, 1] + return rel_dist, dist_squared diff --git a/scripts/train_gnn.py b/scripts/train_gnn.py index 5bc0bda..605b535 100644 --- a/scripts/train_gnn.py +++ b/scripts/train_gnn.py @@ -8,10 +8,9 @@ sys.path.insert(0, str(Path(__file__).parent.parent)) from nmr.construct import construct_graph -from nmr.models import NMRNet +from nmr.models import NMRNet, ModelConfig from nmr.nmr_gym.io import load_histories from torch_geometric.loader import DataLoader -from matplotlib import pyplot as plt def extract_data(pickle_file): @@ -29,58 +28,77 @@ def extract_data(pickle_file): def preprocess_data(examples, device): - """Convert state dictionaries to graphs and return actions and values separately.""" + """ + Convert state dictionaries to graphs with targets attached. + + Attaches action and value targets as graph-level attributes following + the same pattern as shift_to_assign in construct_graph(). + This allows the DataLoader to handle graphs and targets together. + + Args: + examples: List of (state_dict, action, value) tuples + device: Device to place tensors on + + Returns: + List of HeteroData graphs with .action and .value attributes + """ graphs = [] - actions = [] - values = [] for state_dict, action, value in examples: graph = construct_graph(state_dict, device) + + # Attach targets as graph-level attributes (similar to shift_to_assign) + graph.action = torch.tensor([action], dtype=torch.long, device=device) + graph.value = torch.tensor([value], dtype=torch.float32, device=device) + graphs.append(graph) - actions.append(action) - values.append(value) - return graphs, actions, values + return graphs if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Train GNN model on NMR assignment histories") + parser = argparse.ArgumentParser( + description="Train GNN model on NMR assignment histories" + ) parser.add_argument( "--histories", type=str, required=True, - help="Path to the pickle file containing training histories" + help="Path to the pickle file containing training histories", ) parser.add_argument( "--device", type=str, default="cpu", choices=["cpu", "cuda"], - help="Device to use for training (default: cpu)" + help="Device to use for training (default: cpu)", ) parser.add_argument( "--epochs", type=int, default=50000, - help="Number of training epochs (default: 50000)" + help="Number of training epochs (default: 50000)", ) parser.add_argument( - "--batch-size", - type=int, - default=1, - help="Batch size for training (default: 1)" + "--batch-size", type=int, default=1, help="Batch size for training (default: 1)" ) parser.add_argument( "--learning-rate", type=float, default=1e-4, - help="Learning rate for optimizer (default: 1e-4)" + help="Learning rate for optimizer (default: 1e-4)", ) parser.add_argument( "--eval-interval", type=int, default=100, - help="Evaluate model every N iterations (default: 100)" + help="Evaluate model every N iterations (default: 100)", + ) + parser.add_argument( + "--num-nmr-layers", + type=int, + default=1, + help="Number of NMR layers in the model (default: 1)", ) args = parser.parse_args() @@ -90,75 +108,87 @@ def preprocess_data(examples, device): # Load data examples = extract_data(args.histories) - nmr_graphs, actions, rewards = preprocess_data(examples, device) + nmr_graphs = preprocess_data(examples, device) batch_size = args.batch_size epochs = args.epochs eval_interval = args.eval_interval + # Targets are now attached to graphs, so shuffling is safe data_loader = DataLoader(nmr_graphs, batch_size=batch_size, shuffle=False) + config = ModelConfig(num_nmr_layers=args.num_nmr_layers) # Create our network and optimizer - net = NMRNet(device) + net = NMRNet(device, config) opt = torch.optim.AdamW(net.parameters(), lr=args.learning_rate, weight_decay=0.01) - # Interactive plotting - plt.ion() - fig, ax = plt.subplots() - iteration = 0 for epoch in range(epochs): - for i, xs in enumerate(data_loader): + for xs in data_loader: net.train() iteration += 1 - # Break targets into batches manually (if in same order - no shuffling) - ys = actions[i * batch_size : (i + 1) * batch_size] - - _, _, policies = net(xs) - - # Plots current iteration in batch ########################################## - ax.cla() - x_val1, y_val1 = xs["SHIFT"].x[:, 0].tolist(), xs["SHIFT"].x[:, 1].tolist() - x_val2, y_val2 = xs["RES"].x[:, 3].tolist(), xs["RES"].x[:, 4].tolist() - - ax.scatter(x_val1, y_val1, color="red", label="shift") - ax.scatter(x_val2, y_val2, color="blue", label="resid") - - for j in range(len(x_val1)): - ax.annotate(j, (x_val1[j], y_val1[j] + 0.5), color="red") - ax.annotate(j, (x_val2[j], y_val2[j] + 0.5), color="blue") + _, policies = net(xs) - ax.legend() - ax.set_title(f"Iteration {iteration} - Batch {i}") - plt.pause(0.01) - ############################################################################# + # Unbatch graphs to access per-graph attributes + # Targets (action, value) are attached to each graph + graphs_list = xs.to_data_list() loss = 0 correct = 0 count = 0 - - for policy, y in zip(policies, ys): - y = torch.tensor([y], dtype=torch.long, device=device) - loss += torch.nn.functional.cross_entropy(policy, y) - y_pred = torch.argmax(policy) - - if y_pred == y: + policy_ce_losses = [] + + # Policy training: Supervise the network to predict the correct residue assignment + # - policy: shape [num_peaks, num_residues], where policy[peak_i, res_j] = logit for peak i -> residue j + # - We compute cross-entropy loss for: + # 1. The peak being assigned (shift_to_assign) with action as target + # 2. All previously assigned peaks with their assigned residues as targets + for graph, policy in zip(graphs_list, policies): + action = graph.action # Target residue for current assignment + shift_to_assign = graph.shift_to_assign.item() + + # Cross-entropy loss for current assignment + # Extract row for the peak being assigned: policy[shift_to_assign, :] + current_logits = policy[shift_to_assign].unsqueeze(0) # [1, num_residues] + ce_loss = torch.nn.functional.cross_entropy(current_logits, action) + + # Cross-entropy loss for all previous assignments + edge_index = graph["Peak", "assigned_to", "Residue"].edge_index + if edge_index.shape[1] > 0: # Check if there are any assigned peaks + assigned_peak_ids = edge_index[0] # Indices of assigned peaks + assigned_residue_ids = edge_index[1] # Indices of assigned residues + + # Extract rows for all assigned peaks: policy[assigned_peak_ids, :] + # Shape: [num_assigned, num_residues] + assigned_logits = policy[assigned_peak_ids] + + # Compute cross-entropy for each previous assignment and sum + # Target for each row is the corresponding assigned residue + prev_ce_loss = torch.nn.functional.cross_entropy( + assigned_logits, assigned_residue_ids, reduction='sum' + ) + ce_loss = ce_loss + prev_ce_loss + + # Track components for logging + policy_ce_losses.append(ce_loss.item()) + + loss += ce_loss + + y_pred = torch.argmax(current_logits) + if y_pred == action: correct += 1 count += 1 - loss = loss / count - accuracy = correct / count - - print(f"Iteration {iteration} - Loss: {loss.item()} - Accuracy: {accuracy}") - - # writer.add_scalar("Loss/train", loss.item(), iteration) - # writer.add_scalar("Accuracy/train", accuracy, iteration) + if iteration % eval_interval == 0: + mean_ce = sum(policy_ce_losses) / len(policy_ce_losses) + print( + f"{iteration}, " + f"Policy Loss: {mean_ce:.4f}, " + f"Accuracy: {correct / count:.4f}" + ) opt.zero_grad() loss.backward() opt.step() - - plt.ioff() - plt.show() diff --git a/tests/test_gather_components.py b/tests/test_gather_components.py new file mode 100644 index 0000000..273725b --- /dev/null +++ b/tests/test_gather_components.py @@ -0,0 +1,303 @@ +""" +Unit tests for gather component operations. + +Tests the 5 gather classes that extract features from source nodes to triple nodes: +- FirstResidueGather +- FirstPeakGather +- SecondResidueGather +- SecondPeakGather +- NoeGather +""" + +import sys +import unittest +from pathlib import Path + +import torch +from torch_geometric.data import HeteroData + +# Add parent directory to path to allow imports from nmr package +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.models.triple import ( + FirstResidueGather, + FirstPeakGather, + SecondResidueGather, + SecondPeakGather, + NoeGather, +) + + +class TestFirstResidueGather(unittest.TestCase): + """Test FirstResidueGather extracts coords, shifts, features correctly.""" + + def setUp(self): + """Set up test graph with Residue nodes and triple edges.""" + self.device = torch.device("cpu") + self.data = HeteroData() + + # Create 3 Residue nodes with coords [n, 5] = [x, y, z, H, N] + self.data["Residue"].x = torch.tensor( + [ + [1.0, 2.0, 3.0, 8.5, 120.0], # Residue 0 + [4.0, 5.0, 6.0, 7.5, 115.0], # Residue 1 + [7.0, 8.0, 9.0, 9.0, 125.0], # Residue 2 + ], + dtype=torch.float32, + device=self.device, + ) + + # Features [n, 2] = [is_being_assigned, is_already_assigned] + self.data["Residue"].f = torch.tensor( + [ + [0.0, 1.0], + [1.0, 0.0], + [0.0, 0.0], + ], + dtype=torch.float32, + device=self.device, + ) + + # Create 2 triple nodes (placeholders) + self.data["ResidueResidueNoeTriple"].x = torch.zeros( + (2, 1), dtype=torch.float32, device=self.device + ) + + # Create gather edges: Residue -> ResidueResidueNoeTriple + # Edge index [2, num_edges] where [0] = source indices, [1] = target indices + # Edge: source Residue 0 -> target triple 0, source Residue 1 -> target triple 1 + self.data["Residue", "prop_first", "ResidueResidueNoeTriple"].edge_index = ( + torch.tensor( + [[0, 1], [0, 1]], # [source_indices, target_indices] + dtype=torch.long, + device=self.device, + ) + ) + + def test_gather_extracts_coordinates(self): + """Test that FirstResidueGather extracts coordinates correctly.""" + gather = FirstResidueGather("ResidueResidueNoeTriple") + data = gather(self.data) + + # Check that first_coords attribute is set on triple nodes + self.assertIn("first_coords", data["ResidueResidueNoeTriple"]) + first_coords = data["ResidueResidueNoeTriple"].first_coords + + # Shape should be [2, 3] (2 triple nodes, 3 coords each) + self.assertEqual(first_coords.shape, (2, 3)) + + # Values should match source Residue nodes + torch.testing.assert_close(first_coords[0], torch.tensor([1.0, 2.0, 3.0])) + torch.testing.assert_close(first_coords[1], torch.tensor([4.0, 5.0, 6.0])) + + def test_gather_extracts_shifts(self): + """Test that FirstResidueGather extracts shifts correctly.""" + gather = FirstResidueGather("ResidueResidueNoeTriple") + data = gather(self.data) + + # Check that first_shifts attribute is set + self.assertIn("first_shifts", data["ResidueResidueNoeTriple"]) + first_shifts = data["ResidueResidueNoeTriple"].first_shifts + + # Shape should be [2, 2] (2 triple nodes, 2 shifts each) + self.assertEqual(first_shifts.shape, (2, 2)) + + # Values should match shifts from Residue.x[:, 3:5] + torch.testing.assert_close(first_shifts[0], torch.tensor([8.5, 120.0])) + torch.testing.assert_close(first_shifts[1], torch.tensor([7.5, 115.0])) + + def test_gather_extracts_features(self): + """Test that FirstResidueGather extracts features correctly.""" + gather = FirstResidueGather("ResidueResidueNoeTriple") + data = gather(self.data) + + # Check that first_features attribute is set + self.assertIn("first_features", data["ResidueResidueNoeTriple"]) + first_features = data["ResidueResidueNoeTriple"].first_features + + # Shape should be [2, 2] + self.assertEqual(first_features.shape, (2, 2)) + + # Values should match Residue.f + torch.testing.assert_close(first_features[0], torch.tensor([0.0, 1.0])) + torch.testing.assert_close(first_features[1], torch.tensor([1.0, 0.0])) + + +class TestFirstPeakGather(unittest.TestCase): + """Test FirstPeakGather extracts shifts and features (no coords).""" + + def setUp(self): + """Set up test graph with Peak nodes.""" + self.device = torch.device("cpu") + self.data = HeteroData() + + # Create 3 Peak nodes with shifts [n, 2] = [H, N] + self.data["Peak"].x = torch.tensor( + [ + [8.5, 120.0], + [7.5, 115.0], + [9.0, 125.0], + ], + dtype=torch.float32, + device=self.device, + ) + + # Features [n, 2] + self.data["Peak"].f = torch.tensor( + [ + [1.0, 0.0], + [0.0, 1.0], + [0.0, 0.0], + ], + dtype=torch.float32, + device=self.device, + ) + + # Create triple nodes + self.data["PeakPeakNoeTriple"].x = torch.zeros( + (2, 1), dtype=torch.float32, device=self.device + ) + + # Create gather edges: Peak -> PeakPeakNoeTriple + self.data["Peak", "prop_first", "PeakPeakNoeTriple"].edge_index = ( + torch.tensor( + [[0, 2], [0, 1]], # Peak 0 -> triple 0, Peak 2 -> triple 1 + dtype=torch.long, + device=self.device, + ) + ) + + def test_gather_extracts_shifts_no_coords(self): + """Test that FirstPeakGather extracts shifts but not coordinates.""" + gather = FirstPeakGather("PeakPeakNoeTriple") + data = gather(self.data) + + # Should have first_shifts + self.assertIn("first_shifts", data["PeakPeakNoeTriple"]) + first_shifts = data["PeakPeakNoeTriple"].first_shifts + self.assertEqual(first_shifts.shape, (2, 2)) + + # Should NOT have first_coords (peaks have no coordinates) + self.assertNotIn("first_coords", data["PeakPeakNoeTriple"]) + + def test_gather_extracts_features(self): + """Test that FirstPeakGather extracts features correctly.""" + gather = FirstPeakGather("PeakPeakNoeTriple") + data = gather(self.data) + + self.assertIn("first_features", data["PeakPeakNoeTriple"]) + first_features = data["PeakPeakNoeTriple"].first_features + self.assertEqual(first_features.shape, (2, 2)) + + +class TestNoeGather(unittest.TestCase): + """Test NoeGather extracts NOE shifts and features.""" + + def setUp(self): + """Set up test graph with Noe nodes.""" + self.device = torch.device("cpu") + self.data = HeteroData() + + # Create Noe nodes with shifts [n, 3] = [N, H', H"] + self.data["Noe"].x = torch.tensor( + [ + [120.0, 8.5, 7.5], + [115.0, 9.0, 8.0], + [125.0, 7.0, 9.5], + ], + dtype=torch.float32, + device=self.device, + ) + + # Features [n, 2] + self.data["Noe"].f = torch.tensor( + [ + [1.0, 0.0], + [0.0, 1.0], + [0.5, 0.5], + ], + dtype=torch.float32, + device=self.device, + ) + + # Create triple nodes + self.data["ResidueResidueNoeTriple"].x = torch.zeros( + (2, 1), dtype=torch.float32, device=self.device + ) + + # Create gather edges: Noe -> ResidueResidueNoeTriple + self.data["Noe", "prop_noe", "ResidueResidueNoeTriple"].edge_index = ( + torch.tensor( + [[0, 1], [0, 1]], # Noe 0 -> triple 0, Noe 1 -> triple 1 + dtype=torch.long, + device=self.device, + ) + ) + + def test_gather_extracts_noe_shifts(self): + """Test that NoeGather extracts NOE shifts correctly.""" + gather = NoeGather("ResidueResidueNoeTriple") + data = gather(self.data) + + # Check noe_shifts attribute + self.assertIn("noe_shifts", data["ResidueResidueNoeTriple"]) + noe_shifts = data["ResidueResidueNoeTriple"].noe_shifts + + # Shape should be [2, 3] + self.assertEqual(noe_shifts.shape, (2, 3)) + + # Values should match Noe.x + torch.testing.assert_close(noe_shifts[0], torch.tensor([120.0, 8.5, 7.5])) + torch.testing.assert_close(noe_shifts[1], torch.tensor([115.0, 9.0, 8.0])) + + def test_gather_extracts_noe_features(self): + """Test that NoeGather extracts NOE features correctly.""" + gather = NoeGather("ResidueResidueNoeTriple") + data = gather(self.data) + + self.assertIn("noe_features", data["ResidueResidueNoeTriple"]) + noe_features = data["ResidueResidueNoeTriple"].noe_features + self.assertEqual(noe_features.shape, (2, 2)) + + +class TestGatherEmptyEdges(unittest.TestCase): + """Test that gather operations handle empty edge sets gracefully.""" + + def setUp(self): + """Set up test graph with nodes but no edges.""" + self.device = torch.device("cpu") + self.data = HeteroData() + + # Create nodes + self.data["Residue"].x = torch.tensor( + [[1.0, 2.0, 3.0, 8.5, 120.0]], + dtype=torch.float32, + device=self.device, + ) + self.data["Residue"].f = torch.tensor( + [[0.0, 1.0]], dtype=torch.float32, device=self.device + ) + self.data["ResidueResidueNoeTriple"].x = torch.zeros( + (1, 1), dtype=torch.float32, device=self.device + ) + + # Create EMPTY edge set with proper shape [2, 0] + self.data["Residue", "prop_first", "ResidueResidueNoeTriple"].edge_index = ( + torch.zeros((2, 0), dtype=torch.long, device=self.device) + ) + + def test_gather_handles_empty_edges(self): + """Test that FirstResidueGather handles empty edge sets without error.""" + gather = FirstResidueGather("ResidueResidueNoeTriple") + + # Should not raise an error + data = gather(self.data) + + # Attributes should still be created (will be zero-filled for empty edges) + self.assertIn("first_coords", data["ResidueResidueNoeTriple"]) + self.assertIn("first_shifts", data["ResidueResidueNoeTriple"]) + self.assertIn("first_features", data["ResidueResidueNoeTriple"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_gather_components.py.bak b/tests/test_gather_components.py.bak new file mode 100644 index 0000000..d593a81 --- /dev/null +++ b/tests/test_gather_components.py.bak @@ -0,0 +1,303 @@ +""" +Unit tests for gather component operations. + +Tests the 5 gather classes that extract features from source nodes to triple nodes: +- FirstResidueGather +- FirstPeakGather +- SecondResidueGather +- SecondPeakGather +- NoeGather +""" + +import sys +import unittest +from pathlib import Path + +import torch +from torch_geometric.data import HeteroData + +# Add parent directory to path to allow imports from nmr package +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.models.triple import ( + FirstResidueGather, + FirstPeakGather, + SecondResidueGather, + SecondPeakGather, + NoeGather, +) + + +class TestFirstResidueGather(unittest.TestCase): + """Test FirstResidueGather extracts coords, shifts, features correctly.""" + + def setUp(self): + """Set up test graph with Residue nodes and triple edges.""" + self.device = torch.device("cpu") + self.data = HeteroData() + + # Create 3 Residue nodes with coords [n, 5] = [x, y, z, H, N] + self.data["Residue"].x = torch.tensor( + [ + [1.0, 2.0, 3.0, 8.5, 120.0], # Residue 0 + [4.0, 5.0, 6.0, 7.5, 115.0], # Residue 1 + [7.0, 8.0, 9.0, 9.0, 125.0], # Residue 2 + ], + dtype=torch.float32, + device=self.device, + ) + + # Features [n, 2] = [is_being_assigned, is_already_assigned] + self.data["Residue"].f = torch.tensor( + [ + [0.0, 1.0], + [1.0, 0.0], + [0.0, 0.0], + ], + dtype=torch.float32, + device=self.device, + ) + + # Create 2 triple nodes (placeholders) + self.data["ResidueResidueNoeTriple"].x = torch.zeros( + (2, 1), dtype=torch.float32, device=self.device + ) + + # Create gather edges: Residue -> ResidueResidueNoeTriple + # Edge index [2, num_edges] where [0] = source indices, [1] = target indices + # Edge: source Residue 0 -> target triple 0, source Residue 1 -> target triple 1 + self.data["Residue", "gather_first", "ResidueResidueNoeTriple"].edge_index = ( + torch.tensor( + [[0, 1], [0, 1]], # [source_indices, target_indices] + dtype=torch.long, + device=self.device, + ) + ) + + def test_gather_extracts_coordinates(self): + """Test that FirstResidueGather extracts coordinates correctly.""" + gather = FirstResidueGather("ResidueResidueNoeTriple") + data = gather(self.data) + + # Check that first_coords attribute is set on triple nodes + self.assertIn("first_coords", data["ResidueResidueNoeTriple"]) + first_coords = data["ResidueResidueNoeTriple"].first_coords + + # Shape should be [2, 3] (2 triple nodes, 3 coords each) + self.assertEqual(first_coords.shape, (2, 3)) + + # Values should match source Residue nodes + torch.testing.assert_close(first_coords[0], torch.tensor([1.0, 2.0, 3.0])) + torch.testing.assert_close(first_coords[1], torch.tensor([4.0, 5.0, 6.0])) + + def test_gather_extracts_shifts(self): + """Test that FirstResidueGather extracts shifts correctly.""" + gather = FirstResidueGather("ResidueResidueNoeTriple") + data = gather(self.data) + + # Check that first_shifts attribute is set + self.assertIn("first_shifts", data["ResidueResidueNoeTriple"]) + first_shifts = data["ResidueResidueNoeTriple"].first_shifts + + # Shape should be [2, 2] (2 triple nodes, 2 shifts each) + self.assertEqual(first_shifts.shape, (2, 2)) + + # Values should match shifts from Residue.x[:, 3:5] + torch.testing.assert_close(first_shifts[0], torch.tensor([8.5, 120.0])) + torch.testing.assert_close(first_shifts[1], torch.tensor([7.5, 115.0])) + + def test_gather_extracts_features(self): + """Test that FirstResidueGather extracts features correctly.""" + gather = FirstResidueGather("ResidueResidueNoeTriple") + data = gather(self.data) + + # Check that first_features attribute is set + self.assertIn("first_features", data["ResidueResidueNoeTriple"]) + first_features = data["ResidueResidueNoeTriple"].first_features + + # Shape should be [2, 2] + self.assertEqual(first_features.shape, (2, 2)) + + # Values should match Residue.f + torch.testing.assert_close(first_features[0], torch.tensor([0.0, 1.0])) + torch.testing.assert_close(first_features[1], torch.tensor([1.0, 0.0])) + + +class TestFirstPeakGather(unittest.TestCase): + """Test FirstPeakGather extracts shifts and features (no coords).""" + + def setUp(self): + """Set up test graph with Peak nodes.""" + self.device = torch.device("cpu") + self.data = HeteroData() + + # Create 3 Peak nodes with shifts [n, 2] = [H, N] + self.data["Peak"].x = torch.tensor( + [ + [8.5, 120.0], + [7.5, 115.0], + [9.0, 125.0], + ], + dtype=torch.float32, + device=self.device, + ) + + # Features [n, 2] + self.data["Peak"].f = torch.tensor( + [ + [1.0, 0.0], + [0.0, 1.0], + [0.0, 0.0], + ], + dtype=torch.float32, + device=self.device, + ) + + # Create triple nodes + self.data["PeakPeakNoeTriple"].x = torch.zeros( + (2, 1), dtype=torch.float32, device=self.device + ) + + # Create gather edges: Peak -> PeakPeakNoeTriple + self.data["Peak", "gather_first", "PeakPeakNoeTriple"].edge_index = ( + torch.tensor( + [[0, 2], [0, 1]], # Peak 0 -> triple 0, Peak 2 -> triple 1 + dtype=torch.long, + device=self.device, + ) + ) + + def test_gather_extracts_shifts_no_coords(self): + """Test that FirstPeakGather extracts shifts but not coordinates.""" + gather = FirstPeakGather("PeakPeakNoeTriple") + data = gather(self.data) + + # Should have first_shifts + self.assertIn("first_shifts", data["PeakPeakNoeTriple"]) + first_shifts = data["PeakPeakNoeTriple"].first_shifts + self.assertEqual(first_shifts.shape, (2, 2)) + + # Should NOT have first_coords (peaks have no coordinates) + self.assertNotIn("first_coords", data["PeakPeakNoeTriple"]) + + def test_gather_extracts_features(self): + """Test that FirstPeakGather extracts features correctly.""" + gather = FirstPeakGather("PeakPeakNoeTriple") + data = gather(self.data) + + self.assertIn("first_features", data["PeakPeakNoeTriple"]) + first_features = data["PeakPeakNoeTriple"].first_features + self.assertEqual(first_features.shape, (2, 2)) + + +class TestNoeGather(unittest.TestCase): + """Test NoeGather extracts NOE shifts and features.""" + + def setUp(self): + """Set up test graph with Noe nodes.""" + self.device = torch.device("cpu") + self.data = HeteroData() + + # Create Noe nodes with shifts [n, 3] = [N, H', H"] + self.data["Noe"].x = torch.tensor( + [ + [120.0, 8.5, 7.5], + [115.0, 9.0, 8.0], + [125.0, 7.0, 9.5], + ], + dtype=torch.float32, + device=self.device, + ) + + # Features [n, 2] + self.data["Noe"].f = torch.tensor( + [ + [1.0, 0.0], + [0.0, 1.0], + [0.5, 0.5], + ], + dtype=torch.float32, + device=self.device, + ) + + # Create triple nodes + self.data["ResidueResidueNoeTriple"].x = torch.zeros( + (2, 1), dtype=torch.float32, device=self.device + ) + + # Create gather edges: Noe -> ResidueResidueNoeTriple + self.data["Noe", "gather_noe", "ResidueResidueNoeTriple"].edge_index = ( + torch.tensor( + [[0, 1], [0, 1]], # Noe 0 -> triple 0, Noe 1 -> triple 1 + dtype=torch.long, + device=self.device, + ) + ) + + def test_gather_extracts_noe_shifts(self): + """Test that NoeGather extracts NOE shifts correctly.""" + gather = NoeGather("ResidueResidueNoeTriple") + data = gather(self.data) + + # Check noe_shifts attribute + self.assertIn("noe_shifts", data["ResidueResidueNoeTriple"]) + noe_shifts = data["ResidueResidueNoeTriple"].noe_shifts + + # Shape should be [2, 3] + self.assertEqual(noe_shifts.shape, (2, 3)) + + # Values should match Noe.x + torch.testing.assert_close(noe_shifts[0], torch.tensor([120.0, 8.5, 7.5])) + torch.testing.assert_close(noe_shifts[1], torch.tensor([115.0, 9.0, 8.0])) + + def test_gather_extracts_noe_features(self): + """Test that NoeGather extracts NOE features correctly.""" + gather = NoeGather("ResidueResidueNoeTriple") + data = gather(self.data) + + self.assertIn("noe_features", data["ResidueResidueNoeTriple"]) + noe_features = data["ResidueResidueNoeTriple"].noe_features + self.assertEqual(noe_features.shape, (2, 2)) + + +class TestGatherEmptyEdges(unittest.TestCase): + """Test that gather operations handle empty edge sets gracefully.""" + + def setUp(self): + """Set up test graph with nodes but no edges.""" + self.device = torch.device("cpu") + self.data = HeteroData() + + # Create nodes + self.data["Residue"].x = torch.tensor( + [[1.0, 2.0, 3.0, 8.5, 120.0]], + dtype=torch.float32, + device=self.device, + ) + self.data["Residue"].f = torch.tensor( + [[0.0, 1.0]], dtype=torch.float32, device=self.device + ) + self.data["ResidueResidueNoeTriple"].x = torch.zeros( + (1, 1), dtype=torch.float32, device=self.device + ) + + # Create EMPTY edge set with proper shape [2, 0] + self.data["Residue", "gather_first", "ResidueResidueNoeTriple"].edge_index = ( + torch.zeros((2, 0), dtype=torch.long, device=self.device) + ) + + def test_gather_handles_empty_edges(self): + """Test that FirstResidueGather handles empty edge sets without error.""" + gather = FirstResidueGather("ResidueResidueNoeTriple") + + # Should not raise an error + data = gather(self.data) + + # Attributes should still be created (will be zero-filled for empty edges) + self.assertIn("first_coords", data["ResidueResidueNoeTriple"]) + self.assertIn("first_shifts", data["ResidueResidueNoeTriple"]) + self.assertIn("first_features", data["ResidueResidueNoeTriple"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_network_orchestration.py b/tests/test_network_orchestration.py new file mode 100644 index 0000000..1d25d34 --- /dev/null +++ b/tests/test_network_orchestration.py @@ -0,0 +1,146 @@ +""" +Tests for NMRLayer orchestration with explicit triple classes. + +This test module verifies that NMRLayer correctly instantiates all 4 triple +classes as separate attributes and calls them explicitly in sequence without +conditionals or branching. +""" + +import sys +from pathlib import Path +import unittest + +import torch +import numpy as np + +# Add parent directory to path for imports +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.models.network import NMRLayer, ModelConfig, StandardizeShifts, EmbedFeatures +from nmr.models.triple import ( + ResidueResidueNoeTriple, + ResiduePeakNoeTriple, + PeakResidueNoeTriple, + PeakPeakNoeTriple, +) +from nmr.construct import construct_graph + + +class TestNMRLayerOrchestration(unittest.TestCase): + """Test NMRLayer instantiation and forward pass orchestration.""" + + def setUp(self): + """Set up test fixtures.""" + self.device = "cpu" + self.config = ModelConfig() + + def test_nmr_layer_instantiates_all_four_triples(self): + """Test that NMRLayer instantiates all 4 triple classes as separate attributes.""" + layer = NMRLayer(self.device, self.config) + + # Verify all 4 triple classes are instantiated as attributes + self.assertTrue(hasattr(layer, "residue_residue_noe")) + self.assertTrue(hasattr(layer, "residue_peak_noe")) + self.assertTrue(hasattr(layer, "peak_residue_noe")) + self.assertTrue(hasattr(layer, "peak_peak_noe")) + + # Verify each attribute is the correct type + self.assertIsInstance(layer.residue_residue_noe, ResidueResidueNoeTriple) + self.assertIsInstance(layer.residue_peak_noe, ResiduePeakNoeTriple) + self.assertIsInstance(layer.peak_residue_noe, PeakResidueNoeTriple) + self.assertIsInstance(layer.peak_peak_noe, PeakPeakNoeTriple) + + # Verify old attribute does NOT exist + self.assertFalse(hasattr(layer, "triple_layer")) + + def test_nmr_layer_forward_calls_all_triples_in_sequence(self): + """Test that forward pass calls all 4 triples in sequence.""" + layer = NMRLayer(self.device, self.config) + + # Create a small test graph using the correct history structure + # Coordinates should be [x, y, z, H, N] - 5 values per residue + xyz = np.random.randn(5, 3) + pred_shifts = np.random.randn(5, 2) + coordinates = np.concatenate([xyz, pred_shifts], axis=1).tolist() # 5 residues with [x,y,z,H,N] + obs_shifts = np.random.randn(5, 2).tolist() # 5 observed shifts + noes = [[110.0, 8.0, 7.5], [115.0, 8.5, 8.0]] # 2 NOE constraints + + history = { + "coordinates": coordinates, + "obs_chemical_shifts": obs_shifts, + "noes": noes, + "assignments": {}, # No assignments yet + "shift_to_assign": 0, + } + + data = construct_graph(history=history, device=self.device) + + # Embed the data (NMRLayer expects embedded features) + normalize = StandardizeShifts() + embed = EmbedFeatures(self.device, self.config) + data = normalize(data) + data = embed(data) + + # Store original node features to verify they change + original_residue_x = data["Residue"].x.clone() + original_peak_x = data["Peak"].x.clone() + original_noe_x = data["Noe"].x.clone() + + # Run forward pass + output_data = layer(data) + + # Verify data is returned + self.assertIsNotNone(output_data) + + # Verify node features have been updated (at least one should change) + # Note: We don't know the exact values, but we can verify the shapes are preserved + self.assertEqual(output_data["Residue"].x.shape, original_residue_x.shape) + self.assertEqual(output_data["Peak"].x.shape, original_peak_x.shape) + self.assertEqual(output_data["Noe"].x.shape, original_noe_x.shape) + + def test_nmr_layer_forward_no_conditionals(self): + """Test that forward pass has no conditionals or branching logic.""" + layer = NMRLayer(self.device, self.config) + + # Create a small test graph + # Coordinates should be [x, y, z, H, N] - 5 values per residue + xyz = np.random.randn(3, 3) + pred_shifts = np.random.randn(3, 2) + coordinates = np.concatenate([xyz, pred_shifts], axis=1).tolist() # 3 residues with [x,y,z,H,N] + obs_shifts = np.random.randn(3, 2).tolist() + noes = [[110.0, 8.0, 7.5]] + + history = { + "coordinates": coordinates, + "obs_chemical_shifts": obs_shifts, + "noes": noes, + "assignments": {}, + "shift_to_assign": 0, + } + + data = construct_graph(history=history, device=self.device) + + # Embed the data (NMRLayer expects embedded features) + normalize = StandardizeShifts() + embed = EmbedFeatures(self.device, self.config) + data = normalize(data) + data = embed(data) + + # The forward method should have a straightforward sequence of calls + # We verify this by checking that all triple types can process the data + # without errors, regardless of the specific triple instances present + + # This test primarily validates structure - the actual forward logic + # is tested in the previous test + output_data = layer(data) + + # Verify output is valid HeteroData + self.assertIsNotNone(output_data) + self.assertTrue(hasattr(output_data, "node_types")) + self.assertIn("Residue", output_data.node_types) + self.assertIn("Peak", output_data.node_types) + self.assertIn("Noe", output_data.node_types) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_scatter_operations.py b/tests/test_scatter_operations.py new file mode 100644 index 0000000..4987388 --- /dev/null +++ b/tests/test_scatter_operations.py @@ -0,0 +1,218 @@ +""" +Tests for scatter operations in triple system. + +Tests verify that scatter classes correctly propagate deltas from triple nodes +back to source nodes using mean aggregation. +""" + +import sys +import unittest +from pathlib import Path + +import torch +from torch_geometric.data import HeteroData + +# Add parent directory to path for imports +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.models.triple import ( + FirstResidueScatter, + FirstPeakScatter, + SecondResidueScatter, + SecondPeakScatter, + NoeScatter, +) + + +class TestScatterOperations(unittest.TestCase): + """Test scatter operations that propagate deltas to source nodes.""" + + def setUp(self): + """Set up test fixtures.""" + self.device = torch.device("cpu") + + def test_first_residue_scatter_updates_coordinates_and_shifts(self): + """Test FirstResidueScatter updates Residue.x[:, 0:3], Residue.x[:, 3:5], and Residue.f.""" + # Create minimal HeteroData with Residue nodes and ResidueResidueNoeTriple + data = HeteroData() + + # Create 3 Residue nodes with coords [0:3] and shifts [3:5] + data["Residue"].x = torch.zeros(3, 5) # [n, 5] = [coords(3) + shifts(2)] + data["Residue"].f = torch.zeros(3, 2) # [n, 2] features + + # Create 2 triple nodes with deltas + triple_type = "ResidueResidueNoeTriple" + data[triple_type].x = torch.zeros(2, 1) # Dummy attribute + data[triple_type].delta_first_coords = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + data[triple_type].delta_first_shifts = torch.tensor([[0.1, 0.2], [0.3, 0.4]]) + data[triple_type].delta_first_features = torch.tensor([[0.5, 0.6], [0.7, 0.8]]) + + # Create prop edges: Residue <-> Triple (used with reversed flow for scatter) + # Edge format: [source_indices, triple_indices] = [Residue IDs, Triple IDs] + # Triple 0 connects to Residue 0, Triple 1 connects to Residue 1 + data["Residue", "prop_first", triple_type].edge_index = torch.tensor( + [[0, 1], [0, 1]] + ) + + # Apply scatter operation + scatter = FirstResidueScatter(triple_type) + data = scatter(data) + + # Verify coordinates were updated + self.assertTrue(torch.allclose(data["Residue"].x[0, 0:3], torch.tensor([1.0, 2.0, 3.0]))) + self.assertTrue(torch.allclose(data["Residue"].x[1, 0:3], torch.tensor([4.0, 5.0, 6.0]))) + self.assertTrue(torch.allclose(data["Residue"].x[2, 0:3], torch.tensor([0.0, 0.0, 0.0]))) + + # Verify shifts were updated + self.assertTrue(torch.allclose(data["Residue"].x[0, 3:5], torch.tensor([0.1, 0.2]))) + self.assertTrue(torch.allclose(data["Residue"].x[1, 3:5], torch.tensor([0.3, 0.4]))) + + # Verify features were updated + self.assertTrue(torch.allclose(data["Residue"].f[0], torch.tensor([0.5, 0.6]))) + self.assertTrue(torch.allclose(data["Residue"].f[1], torch.tensor([0.7, 0.8]))) + + def test_first_peak_scatter_updates_shifts_and_features(self): + """Test FirstPeakScatter updates Peak.x and Peak.f (no coordinates).""" + # Create minimal HeteroData with Peak nodes and ResiduePeakNoeTriple + data = HeteroData() + + # Create 3 Peak nodes with shifts only + data["Peak"].x = torch.zeros(3, 2) # [n, 2] shifts (H, N) + data["Peak"].f = torch.zeros(3, 2) # [n, 2] features + + # Create 2 triple nodes with deltas + triple_type = "ResiduePeakNoeTriple" + data[triple_type].x = torch.zeros(2, 1) + data[triple_type].delta_first_shifts = torch.tensor([[0.1, 0.2], [0.3, 0.4]]) + data[triple_type].delta_first_features = torch.tensor([[0.5, 0.6], [0.7, 0.8]]) + + # Create prop edges: Peak <-> Triple (used with reversed flow for scatter) + # Edge format: [source_indices, triple_indices] = [Peak IDs, Triple IDs] + data["Peak", "prop_first", triple_type].edge_index = torch.tensor([[0, 1], [0, 1]]) + + # Apply scatter operation + scatter = FirstPeakScatter(triple_type) + data = scatter(data) + + # Verify shifts were updated + self.assertTrue(torch.allclose(data["Peak"].x[0], torch.tensor([0.1, 0.2]))) + self.assertTrue(torch.allclose(data["Peak"].x[1], torch.tensor([0.3, 0.4]))) + self.assertTrue(torch.allclose(data["Peak"].x[2], torch.tensor([0.0, 0.0]))) + + # Verify features were updated + self.assertTrue(torch.allclose(data["Peak"].f[0], torch.tensor([0.5, 0.6]))) + self.assertTrue(torch.allclose(data["Peak"].f[1], torch.tensor([0.7, 0.8]))) + + def test_noe_scatter_updates_shifts_and_features(self): + """Test NoeScatter updates Noe.x and Noe.f.""" + # Create minimal HeteroData with Noe nodes and ResidueResidueNoeTriple + data = HeteroData() + + # Create 3 Noe nodes with shifts [N, H', H"] + data["Noe"].x = torch.zeros(3, 3) # [n, 3] NOE shifts + data["Noe"].f = torch.zeros(3, 2) # [n, 2] features + + # Create 2 triple nodes with deltas + triple_type = "ResidueResidueNoeTriple" + data[triple_type].x = torch.zeros(2, 1) + data[triple_type].delta_noe_shifts = torch.tensor([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]) + data[triple_type].delta_noe_features = torch.tensor([[0.7, 0.8], [0.9, 1.0]]) + + # Create prop edges: Noe <-> Triple (used with reversed flow for scatter) + # Edge format: [source_indices, triple_indices] = [Noe IDs, Triple IDs] + data["Noe", "prop_noe", triple_type].edge_index = torch.tensor([[0, 1], [0, 1]]) + + # Apply scatter operation + scatter = NoeScatter(triple_type) + data = scatter(data) + + # Verify NOE shifts were updated + self.assertTrue(torch.allclose(data["Noe"].x[0], torch.tensor([0.1, 0.2, 0.3]))) + self.assertTrue(torch.allclose(data["Noe"].x[1], torch.tensor([0.4, 0.5, 0.6]))) + self.assertTrue(torch.allclose(data["Noe"].x[2], torch.tensor([0.0, 0.0, 0.0]))) + + # Verify features were updated + self.assertTrue(torch.allclose(data["Noe"].f[0], torch.tensor([0.7, 0.8]))) + self.assertTrue(torch.allclose(data["Noe"].f[1], torch.tensor([0.9, 1.0]))) + + def test_scatter_aggregates_multiple_updates_with_mean(self): + """Test scatter operations aggregate multiple updates using mean.""" + # Create HeteroData where multiple triples scatter to same Residue node + data = HeteroData() + + # Create 2 Residue nodes + data["Residue"].x = torch.zeros(2, 5) + data["Residue"].f = torch.zeros(2, 2) + + # Create 3 triple nodes with deltas + triple_type = "ResidueResidueNoeTriple" + data[triple_type].x = torch.zeros(3, 1) + data[triple_type].delta_first_coords = torch.tensor( + [[1.0, 2.0, 3.0], [2.0, 4.0, 6.0], [3.0, 6.0, 9.0]] + ) + data[triple_type].delta_first_shifts = torch.tensor( + [[0.1, 0.2], [0.2, 0.4], [0.3, 0.6]] + ) + data[triple_type].delta_first_features = torch.tensor( + [[0.5, 0.6], [1.0, 1.2], [1.5, 1.8]] + ) + + # Create prop edges: Residue <-> Triple (used with reversed flow for scatter) + # Edge format: [source_indices, triple_indices] = [Residue IDs, Triple IDs] + # Triples 0,1,2 all connect to Residue 0 + data["Residue", "prop_first", triple_type].edge_index = torch.tensor( + [[0, 0, 0], [0, 1, 2]] + ) + + # Apply scatter operation + scatter = FirstResidueScatter(triple_type) + data = scatter(data) + + # Verify mean aggregation + # Expected coords mean: (1+2+3)/3 = 2, (2+4+6)/3 = 4, (3+6+9)/3 = 6 + expected_coords = torch.tensor([2.0, 4.0, 6.0]) + self.assertTrue(torch.allclose(data["Residue"].x[0, 0:3], expected_coords)) + + # Expected shifts mean: (0.1+0.2+0.3)/3 = 0.2, (0.2+0.4+0.6)/3 = 0.4 + expected_shifts = torch.tensor([0.2, 0.4]) + self.assertTrue(torch.allclose(data["Residue"].x[0, 3:5], expected_shifts)) + + # Expected features mean: (0.5+1.0+1.5)/3 = 1.0, (0.6+1.2+1.8)/3 = 1.2 + expected_features = torch.tensor([1.0, 1.2]) + self.assertTrue(torch.allclose(data["Residue"].f[0], expected_features)) + + # Residue 1 should remain unchanged (no edges to it) + self.assertTrue(torch.allclose(data["Residue"].x[1], torch.zeros(5))) + self.assertTrue(torch.allclose(data["Residue"].f[1], torch.zeros(2))) + + def test_scatter_edge_directions_correct(self): + """Test scatter operations use bidirectional edges with reversed flow.""" + # This test verifies the edge type is configured correctly + # Edge types are now bidirectional (source, "prop_*", triple) + # Scatter uses the same edges as gather but with flow='target_to_source' + + triple_type = "ResidueResidueNoeTriple" + + # Test FirstResidueScatter - uses ("Residue", "prop_first", triple) edges + scatter = FirstResidueScatter(triple_type) + self.assertEqual(scatter.edge_type, ("Residue", "prop_first", triple_type)) + + # Test FirstPeakScatter - uses ("Peak", "prop_first", triple) edges + scatter = FirstPeakScatter("ResiduePeakNoeTriple") + self.assertEqual(scatter.edge_type, ("Peak", "prop_first", "ResiduePeakNoeTriple")) + + # Test SecondResidueScatter - uses ("Residue", "prop_second", triple) edges + scatter = SecondResidueScatter(triple_type) + self.assertEqual(scatter.edge_type, ("Residue", "prop_second", triple_type)) + + # Test SecondPeakScatter - uses ("Peak", "prop_second", triple) edges + scatter = SecondPeakScatter("PeakPeakNoeTriple") + self.assertEqual(scatter.edge_type, ("Peak", "prop_second", "PeakPeakNoeTriple")) + + # Test NoeScatter - uses ("Noe", "prop_noe", triple) edges + scatter = NoeScatter(triple_type) + self.assertEqual(scatter.edge_type, ("Noe", "prop_noe", triple_type)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_scatter_operations.py.bak b/tests/test_scatter_operations.py.bak new file mode 100644 index 0000000..4c19e68 --- /dev/null +++ b/tests/test_scatter_operations.py.bak @@ -0,0 +1,212 @@ +""" +Tests for scatter operations in triple system. + +Tests verify that scatter classes correctly propagate deltas from triple nodes +back to source nodes using mean aggregation. +""" + +import sys +import unittest +from pathlib import Path + +import torch +from torch_geometric.data import HeteroData + +# Add parent directory to path for imports +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.models.triple import ( + FirstResidueScatter, + FirstPeakScatter, + SecondResidueScatter, + SecondPeakScatter, + NoeScatter, +) + + +class TestScatterOperations(unittest.TestCase): + """Test scatter operations that propagate deltas to source nodes.""" + + def setUp(self): + """Set up test fixtures.""" + self.device = torch.device("cpu") + + def test_first_residue_scatter_updates_coordinates_and_shifts(self): + """Test FirstResidueScatter updates Residue.x[:, 0:3], Residue.x[:, 3:5], and Residue.f.""" + # Create minimal HeteroData with Residue nodes and ResidueResidueNoeTriple + data = HeteroData() + + # Create 3 Residue nodes with coords [0:3] and shifts [3:5] + data["Residue"].x = torch.zeros(3, 5) # [n, 5] = [coords(3) + shifts(2)] + data["Residue"].f = torch.zeros(3, 2) # [n, 2] features + + # Create 2 triple nodes with deltas + triple_type = "ResidueResidueNoeTriple" + data[triple_type].x = torch.zeros(2, 1) # Dummy attribute + data[triple_type].delta_first_coords = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + data[triple_type].delta_first_shifts = torch.tensor([[0.1, 0.2], [0.3, 0.4]]) + data[triple_type].delta_first_features = torch.tensor([[0.5, 0.6], [0.7, 0.8]]) + + # Create scatter edges: triple -> Residue + # Triple 0 -> Residue 0, Triple 1 -> Residue 1 + data[triple_type, "scatter_first", "Residue"].edge_index = torch.tensor( + [[0, 1], [0, 1]] + ) + + # Apply scatter operation + scatter = FirstResidueScatter(triple_type) + data = scatter(data) + + # Verify coordinates were updated + self.assertTrue(torch.allclose(data["Residue"].x[0, 0:3], torch.tensor([1.0, 2.0, 3.0]))) + self.assertTrue(torch.allclose(data["Residue"].x[1, 0:3], torch.tensor([4.0, 5.0, 6.0]))) + self.assertTrue(torch.allclose(data["Residue"].x[2, 0:3], torch.tensor([0.0, 0.0, 0.0]))) + + # Verify shifts were updated + self.assertTrue(torch.allclose(data["Residue"].x[0, 3:5], torch.tensor([0.1, 0.2]))) + self.assertTrue(torch.allclose(data["Residue"].x[1, 3:5], torch.tensor([0.3, 0.4]))) + + # Verify features were updated + self.assertTrue(torch.allclose(data["Residue"].f[0], torch.tensor([0.5, 0.6]))) + self.assertTrue(torch.allclose(data["Residue"].f[1], torch.tensor([0.7, 0.8]))) + + def test_first_peak_scatter_updates_shifts_and_features(self): + """Test FirstPeakScatter updates Peak.x and Peak.f (no coordinates).""" + # Create minimal HeteroData with Peak nodes and ResiduePeakNoeTriple + data = HeteroData() + + # Create 3 Peak nodes with shifts only + data["Peak"].x = torch.zeros(3, 2) # [n, 2] shifts (H, N) + data["Peak"].f = torch.zeros(3, 2) # [n, 2] features + + # Create 2 triple nodes with deltas + triple_type = "ResiduePeakNoeTriple" + data[triple_type].x = torch.zeros(2, 1) + data[triple_type].delta_first_shifts = torch.tensor([[0.1, 0.2], [0.3, 0.4]]) + data[triple_type].delta_first_features = torch.tensor([[0.5, 0.6], [0.7, 0.8]]) + + # Create scatter edges: triple -> Peak + data[triple_type, "scatter_first", "Peak"].edge_index = torch.tensor([[0, 1], [0, 1]]) + + # Apply scatter operation + scatter = FirstPeakScatter(triple_type) + data = scatter(data) + + # Verify shifts were updated + self.assertTrue(torch.allclose(data["Peak"].x[0], torch.tensor([0.1, 0.2]))) + self.assertTrue(torch.allclose(data["Peak"].x[1], torch.tensor([0.3, 0.4]))) + self.assertTrue(torch.allclose(data["Peak"].x[2], torch.tensor([0.0, 0.0]))) + + # Verify features were updated + self.assertTrue(torch.allclose(data["Peak"].f[0], torch.tensor([0.5, 0.6]))) + self.assertTrue(torch.allclose(data["Peak"].f[1], torch.tensor([0.7, 0.8]))) + + def test_noe_scatter_updates_shifts_and_features(self): + """Test NoeScatter updates Noe.x and Noe.f.""" + # Create minimal HeteroData with Noe nodes and ResidueResidueNoeTriple + data = HeteroData() + + # Create 3 Noe nodes with shifts [N, H', H"] + data["Noe"].x = torch.zeros(3, 3) # [n, 3] NOE shifts + data["Noe"].f = torch.zeros(3, 2) # [n, 2] features + + # Create 2 triple nodes with deltas + triple_type = "ResidueResidueNoeTriple" + data[triple_type].x = torch.zeros(2, 1) + data[triple_type].delta_noe_shifts = torch.tensor([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]) + data[triple_type].delta_noe_features = torch.tensor([[0.7, 0.8], [0.9, 1.0]]) + + # Create scatter edges: triple -> Noe + data[triple_type, "scatter_noe", "Noe"].edge_index = torch.tensor([[0, 1], [0, 1]]) + + # Apply scatter operation + scatter = NoeScatter(triple_type) + data = scatter(data) + + # Verify NOE shifts were updated + self.assertTrue(torch.allclose(data["Noe"].x[0], torch.tensor([0.1, 0.2, 0.3]))) + self.assertTrue(torch.allclose(data["Noe"].x[1], torch.tensor([0.4, 0.5, 0.6]))) + self.assertTrue(torch.allclose(data["Noe"].x[2], torch.tensor([0.0, 0.0, 0.0]))) + + # Verify features were updated + self.assertTrue(torch.allclose(data["Noe"].f[0], torch.tensor([0.7, 0.8]))) + self.assertTrue(torch.allclose(data["Noe"].f[1], torch.tensor([0.9, 1.0]))) + + def test_scatter_aggregates_multiple_updates_with_mean(self): + """Test scatter operations aggregate multiple updates using mean.""" + # Create HeteroData where multiple triples scatter to same Residue node + data = HeteroData() + + # Create 2 Residue nodes + data["Residue"].x = torch.zeros(2, 5) + data["Residue"].f = torch.zeros(2, 2) + + # Create 3 triple nodes with deltas + triple_type = "ResidueResidueNoeTriple" + data[triple_type].x = torch.zeros(3, 1) + data[triple_type].delta_first_coords = torch.tensor( + [[1.0, 2.0, 3.0], [2.0, 4.0, 6.0], [3.0, 6.0, 9.0]] + ) + data[triple_type].delta_first_shifts = torch.tensor( + [[0.1, 0.2], [0.2, 0.4], [0.3, 0.6]] + ) + data[triple_type].delta_first_features = torch.tensor( + [[0.5, 0.6], [1.0, 1.2], [1.5, 1.8]] + ) + + # Create scatter edges: triples 0,1,2 all scatter to Residue 0 + data[triple_type, "scatter_first", "Residue"].edge_index = torch.tensor( + [[0, 1, 2], [0, 0, 0]] + ) + + # Apply scatter operation + scatter = FirstResidueScatter(triple_type) + data = scatter(data) + + # Verify mean aggregation + # Expected coords mean: (1+2+3)/3 = 2, (2+4+6)/3 = 4, (3+6+9)/3 = 6 + expected_coords = torch.tensor([2.0, 4.0, 6.0]) + self.assertTrue(torch.allclose(data["Residue"].x[0, 0:3], expected_coords)) + + # Expected shifts mean: (0.1+0.2+0.3)/3 = 0.2, (0.2+0.4+0.6)/3 = 0.4 + expected_shifts = torch.tensor([0.2, 0.4]) + self.assertTrue(torch.allclose(data["Residue"].x[0, 3:5], expected_shifts)) + + # Expected features mean: (0.5+1.0+1.5)/3 = 1.0, (0.6+1.2+1.8)/3 = 1.2 + expected_features = torch.tensor([1.0, 1.2]) + self.assertTrue(torch.allclose(data["Residue"].f[0], expected_features)) + + # Residue 1 should remain unchanged (no edges to it) + self.assertTrue(torch.allclose(data["Residue"].x[1], torch.zeros(5))) + self.assertTrue(torch.allclose(data["Residue"].f[1], torch.zeros(2))) + + def test_scatter_edge_directions_correct(self): + """Test scatter operations use correct edge directions (triple -> source).""" + # This test verifies the edge type is configured correctly + # Edge direction should be from triple nodes TO source nodes + + triple_type = "ResidueResidueNoeTriple" + + # Test FirstResidueScatter + scatter = FirstResidueScatter(triple_type) + self.assertEqual(scatter.edge_type, (triple_type, "scatter_first", "Residue")) + + # Test FirstPeakScatter + scatter = FirstPeakScatter("ResiduePeakNoeTriple") + self.assertEqual(scatter.edge_type, ("ResiduePeakNoeTriple", "scatter_first", "Peak")) + + # Test SecondResidueScatter + scatter = SecondResidueScatter(triple_type) + self.assertEqual(scatter.edge_type, (triple_type, "scatter_second", "Residue")) + + # Test SecondPeakScatter + scatter = SecondPeakScatter("PeakPeakNoeTriple") + self.assertEqual(scatter.edge_type, ("PeakPeakNoeTriple", "scatter_second", "Peak")) + + # Test NoeScatter + scatter = NoeScatter(triple_type) + self.assertEqual(scatter.edge_type, (triple_type, "scatter_noe", "Noe")) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_terminology_updates.py b/tests/test_terminology_updates.py new file mode 100644 index 0000000..3dc820c --- /dev/null +++ b/tests/test_terminology_updates.py @@ -0,0 +1,98 @@ +""" +Tests for terminology updates in the Triple System Re-Architecture. + +This test file verifies that the graph construction uses the new naming conventions: +- Node types: Residue, Peak, Noe (instead of RES, SHIFT, NOE) +- Triple types: ResidueResidueNoeTriple, ResiduePeakNoeTriple, etc. (instead of TRIPLE0/1/2/3) +- Edge relations: prop_first/second/noe bidirectional edges (instead of NH1/NH2_extract/add) +""" + +import sys +from pathlib import Path +sys.path.insert(0, str(Path(__file__).parent.parent)) + +import unittest +import torch +from nmr.construct import construct_graph + + +class TestTerminologyUpdates(unittest.TestCase): + """Test that graph construction uses new node and edge naming conventions.""" + + def setUp(self): + """Create a minimal test history for graph construction.""" + self.device = torch.device("cpu") + self.history = { + "coordinates": [[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]], + "obs_chemical_shifts": [[8.5, 120.0], [7.8, 118.0]], + "noes": [[120.0, 8.5, 7.8], [118.0, 7.8, 8.5]], + "assignments": {}, + "shift_to_assign": 0, + } + + def test_node_types_use_new_names(self): + """Test that graph construction creates nodes with new names (Residue, Peak, Noe).""" + data = construct_graph(self.history, self.device) + + # Verify new node types exist + self.assertIn("Residue", data.node_types) + self.assertIn("Peak", data.node_types) + self.assertIn("Noe", data.node_types) + + # Verify old node types do not exist + self.assertNotIn("RES", data.node_types) + self.assertNotIn("SHIFT", data.node_types) + self.assertNotIn("NOE", data.node_types) + + def test_triple_names_use_descriptive_names(self): + """Test that triple node types use full descriptive names.""" + data = construct_graph(self.history, self.device) + + # Verify new triple node types exist + self.assertIn("ResidueResidueNoeTriple", data.node_types) + self.assertIn("ResiduePeakNoeTriple", data.node_types) + self.assertIn("PeakResidueNoeTriple", data.node_types) + self.assertIn("PeakPeakNoeTriple", data.node_types) + + # Verify old triple node types do not exist + self.assertNotIn("ResResNoe", data.node_types) + self.assertNotIn("ResShiftNoe", data.node_types) + self.assertNotIn("ShiftResNoe", data.node_types) + self.assertNotIn("ShiftShiftNoe", data.node_types) + self.assertNotIn("TRIPLE0", data.node_types) + self.assertNotIn("TRIPLE1", data.node_types) + self.assertNotIn("TRIPLE2", data.node_types) + self.assertNotIn("TRIPLE3", data.node_types) + + def test_edge_relations_use_prop_naming(self): + """Test that edge relations use new naming (prop_first, prop_second, etc.).""" + data = construct_graph(self.history, self.device) + + # Get all edge types as strings + edge_types = [str(et) for et in data.edge_types] + + # Verify new edge relation names exist (bidirectional prop_* edges) + prop_edges = [et for et in edge_types if "prop_first" in et or "prop_second" in et or "prop_noe" in et] + + self.assertGreater(len(prop_edges), 0, "Should have prop_* edges") + + # Verify old triple-related edge names do not exist (exclude value aggregation edges) + old_triple_edges = [et for et in edge_types if + ("NH1_extract" in et or "NH2_extract" in et or + "NH1_add" in et or "NH2_add" in et or + "res1_add" in et or "res2_add" in et)] + + self.assertEqual(len(old_triple_edges), 0, + f"Should not have old triple edge names, found: {old_triple_edges}") + + def test_triple_names_dictionary_removed(self): + """Test that TRIPLE_NAMES dictionary has been removed from construct module.""" + import nmr.construct as construct_module + + # TRIPLE_NAMES should not exist as an attribute + self.assertFalse(hasattr(construct_module, "TRIPLE_NAMES"), + "TRIPLE_NAMES dictionary should be removed from construct.py") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_triple_compositions.py b/tests/test_triple_compositions.py new file mode 100644 index 0000000..78d5f02 --- /dev/null +++ b/tests/test_triple_compositions.py @@ -0,0 +1,351 @@ +""" +Unit tests for triple composition classes. + +Tests the 4 triple composition classes that wire together gather, update, and +scatter operations: +- ResidueResidueNoeTriple +- ResiduePeakNoeTriple +- PeakResidueNoeTriple +- PeakPeakNoeTriple +""" + +import sys +import unittest +from pathlib import Path + +import torch +from torch_geometric.data import HeteroData + +# Add parent directory to path to allow imports from nmr package +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.models.triple import ( + ResidueResidueNoeTriple, + ResiduePeakNoeTriple, + PeakResidueNoeTriple, + PeakPeakNoeTriple, +) +from nmr.models import ModelConfig + + +class TestResidueResidueNoeTriple(unittest.TestCase): + """Test ResidueResidueNoeTriple wiring and forward pass.""" + + def setUp(self): + """Set up test graph with all necessary nodes and edges.""" + self.device = torch.device("cpu") + self.config = ModelConfig() + self.shift_dim = self.config.shift_embed.output_dim + self.feature_dim = self.config.feature_embed.output_dim + self.data = HeteroData() + + # Create Residue nodes [n, 3 + shift_dim] = [x, y, z, shifts...] + coords = torch.tensor( + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + dtype=torch.float32, + device=self.device, + ) + shifts = torch.randn((2, self.shift_dim), dtype=torch.float32, device=self.device) + self.data["Residue"].x = torch.cat([coords, shifts], dim=-1) + self.data["Residue"].f = torch.randn( + (2, self.feature_dim), dtype=torch.float32, device=self.device + ) + + # Create Noe nodes [n, shift_dim] = embedded NOE features + self.data["Noe"].x = torch.randn( + (1, self.shift_dim), dtype=torch.float32, device=self.device + ) + self.data["Noe"].f = torch.randn( + (1, self.feature_dim), dtype=torch.float32, device=self.device + ) + + # Create triple nodes + self.data["ResidueResidueNoeTriple"].x = torch.zeros( + (1, 1), dtype=torch.float32, device=self.device + ) + + # Create gather edges + self.data[ + "Residue", "prop_first", "ResidueResidueNoeTriple" + ].edge_index = torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + + self.data[ + "Residue", "prop_second", "ResidueResidueNoeTriple" + ].edge_index = torch.tensor([[1], [0]], dtype=torch.long, device=self.device) + + self.data["Noe", "prop_noe", "ResidueResidueNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + + # Create scatter edges (reverse direction) + self.data[ + "ResidueResidueNoeTriple", "prop_first", "Residue" + ].edge_index = torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + + self.data[ + "ResidueResidueNoeTriple", "prop_second", "Residue" + ].edge_index = torch.tensor([[0], [1]], dtype=torch.long, device=self.device) + + self.data["ResidueResidueNoeTriple", "prop_noe", "Noe"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + + def test_triple_instantiation(self): + """Test that ResidueResidueNoeTriple instantiates correctly.""" + triple = ResidueResidueNoeTriple(self.device, self.config) + + # Check that it has all the required components + self.assertTrue(hasattr(triple, "first_gather")) + self.assertTrue(hasattr(triple, "second_gather")) + self.assertTrue(hasattr(triple, "noe_gather")) + self.assertTrue(hasattr(triple, "update")) + self.assertTrue(hasattr(triple, "first_scatter")) + self.assertTrue(hasattr(triple, "second_scatter")) + self.assertTrue(hasattr(triple, "noe_scatter")) + + def test_forward_pass_executes(self): + """Test that forward pass executes without error.""" + triple = ResidueResidueNoeTriple(self.device, self.config) + + # Should not raise an error + data = triple(self.data) + + # Data should be returned + self.assertIsNotNone(data) + + def test_forward_pass_sequence(self): + """Test that forward pass calls gather -> update -> scatter in sequence.""" + triple = ResidueResidueNoeTriple(self.device, self.config) + data = triple(self.data) + + # After forward pass, check that: + # 1. Gathered attributes exist on triple nodes (from gather operations) + self.assertIn("first_coords", data["ResidueResidueNoeTriple"]) + self.assertIn("second_coords", data["ResidueResidueNoeTriple"]) + self.assertIn("noe_shifts", data["ResidueResidueNoeTriple"]) + + # 2. Delta attributes exist on triple nodes (from update operation) + self.assertIn("delta_first_coords", data["ResidueResidueNoeTriple"]) + self.assertIn("delta_second_shifts", data["ResidueResidueNoeTriple"]) + + # 3. Residue and Noe nodes have been updated (from scatter operations) + # We can't easily check the exact values, but we can verify the tensors exist + self.assertIsNotNone(data["Residue"].x) + self.assertIsNotNone(data["Noe"].x) + + +class TestResiduePeakNoeTriple(unittest.TestCase): + """Test ResiduePeakNoeTriple wiring and forward pass.""" + + def setUp(self): + """Set up test graph with Residue, Peak, and Noe nodes.""" + self.device = torch.device("cpu") + self.config = ModelConfig() + self.shift_dim = self.config.shift_embed.output_dim + self.feature_dim = self.config.feature_embed.output_dim + self.data = HeteroData() + + # Create Residue nodes + coords = torch.tensor([[1.0, 2.0, 3.0]], dtype=torch.float32, device=self.device) + shifts = torch.randn((1, self.shift_dim), dtype=torch.float32, device=self.device) + self.data["Residue"].x = torch.cat([coords, shifts], dim=-1) + self.data["Residue"].f = torch.randn( + (1, self.feature_dim), dtype=torch.float32, device=self.device + ) + + # Create Peak nodes + self.data["Peak"].x = torch.randn((1, self.shift_dim), dtype=torch.float32, device=self.device) + self.data["Peak"].f = torch.randn( + (1, self.feature_dim), dtype=torch.float32, device=self.device + ) + + # Create Noe nodes + self.data["Noe"].x = torch.randn( + (1, self.shift_dim), dtype=torch.float32, device=self.device + ) + self.data["Noe"].f = torch.randn( + (1, self.feature_dim), dtype=torch.float32, device=self.device + ) + + # Create triple nodes + self.data["ResiduePeakNoeTriple"].x = torch.zeros( + (1, 1), dtype=torch.float32, device=self.device + ) + + # Create gather edges + self.data["Residue", "prop_first", "ResiduePeakNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["Peak", "prop_second", "ResiduePeakNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["Noe", "prop_noe", "ResiduePeakNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + + # Create scatter edges + self.data["ResiduePeakNoeTriple", "prop_first", "Residue"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["ResiduePeakNoeTriple", "prop_second", "Peak"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["ResiduePeakNoeTriple", "prop_noe", "Noe"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + + def test_triple_uses_correct_components(self): + """Test that ResiduePeakNoeTriple uses correct gather/scatter types.""" + triple = ResiduePeakNoeTriple(self.device, self.config) + + # Check component types by examining their triple_type attribute + self.assertEqual(triple.first_gather.triple_type, "ResiduePeakNoeTriple") + self.assertEqual(triple.second_gather.triple_type, "ResiduePeakNoeTriple") + self.assertEqual(triple.noe_gather.triple_type, "ResiduePeakNoeTriple") + self.assertEqual(triple.update.triple_type, "ResiduePeakNoeTriple") + + # Check edge types to ensure correct gather/scatter classes + self.assertEqual( + triple.first_gather.edge_type, + ("Residue", "prop_first", "ResiduePeakNoeTriple"), + ) + self.assertEqual( + triple.second_gather.edge_type, + ("Peak", "prop_second", "ResiduePeakNoeTriple"), + ) + + def test_forward_pass_executes(self): + """Test that forward pass executes without error.""" + triple = ResiduePeakNoeTriple(self.device, self.config) + data = triple(self.data) + self.assertIsNotNone(data) + + +class TestPeakResidueNoeTriple(unittest.TestCase): + """Test PeakResidueNoeTriple wiring and forward pass.""" + + def setUp(self): + """Set up test graph with Peak, Residue, and Noe nodes.""" + self.device = torch.device("cpu") + self.config = ModelConfig() + self.shift_dim = self.config.shift_embed.output_dim + self.feature_dim = self.config.feature_embed.output_dim + self.data = HeteroData() + + # Create Peak nodes + self.data["Peak"].x = torch.randn((1, self.shift_dim), dtype=torch.float32, device=self.device) + self.data["Peak"].f = torch.randn( + (1, self.feature_dim), dtype=torch.float32, device=self.device + ) + + # Create Residue nodes + coords = torch.tensor([[1.0, 2.0, 3.0]], dtype=torch.float32, device=self.device) + shifts = torch.randn((1, self.shift_dim), dtype=torch.float32, device=self.device) + self.data["Residue"].x = torch.cat([coords, shifts], dim=-1) + self.data["Residue"].f = torch.randn( + (1, self.feature_dim), dtype=torch.float32, device=self.device + ) + + # Create Noe nodes + self.data["Noe"].x = torch.randn( + (1, self.shift_dim), dtype=torch.float32, device=self.device + ) + self.data["Noe"].f = torch.randn( + (1, self.feature_dim), dtype=torch.float32, device=self.device + ) + + # Create triple nodes + self.data["PeakResidueNoeTriple"].x = torch.zeros( + (1, 1), dtype=torch.float32, device=self.device + ) + + # Create gather edges + self.data["Peak", "prop_first", "PeakResidueNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["Residue", "prop_second", "PeakResidueNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["Noe", "prop_noe", "PeakResidueNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + + # Create scatter edges + self.data["PeakResidueNoeTriple", "prop_first", "Peak"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["PeakResidueNoeTriple", "prop_second", "Residue"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["PeakResidueNoeTriple", "prop_noe", "Noe"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + + def test_forward_pass_executes(self): + """Test that forward pass executes without error.""" + triple = PeakResidueNoeTriple(self.device, self.config) + data = triple(self.data) + self.assertIsNotNone(data) + + +class TestPeakPeakNoeTriple(unittest.TestCase): + """Test PeakPeakNoeTriple wiring and forward pass.""" + + def setUp(self): + """Set up test graph with Peak and Noe nodes.""" + self.device = torch.device("cpu") + self.config = ModelConfig() + self.shift_dim = self.config.shift_embed.output_dim + self.feature_dim = self.config.feature_embed.output_dim + self.data = HeteroData() + + # Create Peak nodes + self.data["Peak"].x = torch.randn((2, self.shift_dim), dtype=torch.float32, device=self.device) + self.data["Peak"].f = torch.randn( + (2, self.feature_dim), dtype=torch.float32, device=self.device + ) + + # Create Noe nodes + self.data["Noe"].x = torch.randn( + (1, self.shift_dim), dtype=torch.float32, device=self.device + ) + self.data["Noe"].f = torch.randn( + (1, self.feature_dim), dtype=torch.float32, device=self.device + ) + + # Create triple nodes + self.data["PeakPeakNoeTriple"].x = torch.zeros( + (1, 1), dtype=torch.float32, device=self.device + ) + + # Create gather edges + self.data["Peak", "prop_first", "PeakPeakNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["Peak", "prop_second", "PeakPeakNoeTriple"].edge_index = ( + torch.tensor([[1], [0]], dtype=torch.long, device=self.device) + ) + self.data["Noe", "prop_noe", "PeakPeakNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + + # Create scatter edges + self.data["PeakPeakNoeTriple", "prop_first", "Peak"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["PeakPeakNoeTriple", "prop_second", "Peak"].edge_index = ( + torch.tensor([[0], [1]], dtype=torch.long, device=self.device) + ) + self.data["PeakPeakNoeTriple", "prop_noe", "Noe"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + + def test_forward_pass_executes(self): + """Test that forward pass executes without error.""" + triple = PeakPeakNoeTriple(self.device, self.config) + data = triple(self.data) + self.assertIsNotNone(data) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_triple_compositions.py.bak b/tests/test_triple_compositions.py.bak new file mode 100644 index 0000000..5eb8c7a --- /dev/null +++ b/tests/test_triple_compositions.py.bak @@ -0,0 +1,345 @@ +""" +Unit tests for triple composition classes. + +Tests the 4 triple composition classes that wire together gather, update, and +scatter operations: +- ResidueResidueNoeTriple +- ResiduePeakNoeTriple +- PeakResidueNoeTriple +- PeakPeakNoeTriple +""" + +import sys +import unittest +from pathlib import Path + +import torch +from torch_geometric.data import HeteroData + +# Add parent directory to path to allow imports from nmr package +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.models.triple import ( + ResidueResidueNoeTriple, + ResiduePeakNoeTriple, + PeakResidueNoeTriple, + PeakPeakNoeTriple, +) + + +class TestResidueResidueNoeTriple(unittest.TestCase): + """Test ResidueResidueNoeTriple wiring and forward pass.""" + + def setUp(self): + """Set up test graph with all necessary nodes and edges.""" + self.device = torch.device("cpu") + self.data = HeteroData() + + # Create Residue nodes [n, 5] = [x, y, z, H, N] + self.data["Residue"].x = torch.tensor( + [ + [1.0, 2.0, 3.0, 8.5, 120.0], + [4.0, 5.0, 6.0, 7.5, 115.0], + ], + dtype=torch.float32, + device=self.device, + ) + self.data["Residue"].f = torch.tensor( + [[0.0, 1.0], [1.0, 0.0]], dtype=torch.float32, device=self.device + ) + + # Create Noe nodes [n, 3] = [N, H', H"] + self.data["Noe"].x = torch.tensor( + [[120.0, 8.5, 7.5]], dtype=torch.float32, device=self.device + ) + self.data["Noe"].f = torch.tensor( + [[1.0, 0.0]], dtype=torch.float32, device=self.device + ) + + # Create triple nodes + self.data["ResidueResidueNoeTriple"].x = torch.zeros( + (1, 1), dtype=torch.float32, device=self.device + ) + + # Create gather edges + self.data[ + "Residue", "gather_first", "ResidueResidueNoeTriple" + ].edge_index = torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + + self.data[ + "Residue", "gather_second", "ResidueResidueNoeTriple" + ].edge_index = torch.tensor([[1], [0]], dtype=torch.long, device=self.device) + + self.data["Noe", "gather_noe", "ResidueResidueNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + + # Create scatter edges (reverse direction) + self.data[ + "ResidueResidueNoeTriple", "scatter_first", "Residue" + ].edge_index = torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + + self.data[ + "ResidueResidueNoeTriple", "scatter_second", "Residue" + ].edge_index = torch.tensor([[0], [1]], dtype=torch.long, device=self.device) + + self.data["ResidueResidueNoeTriple", "scatter_noe", "Noe"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + + def test_triple_instantiation(self): + """Test that ResidueResidueNoeTriple instantiates correctly.""" + triple = ResidueResidueNoeTriple(self.device) + + # Check that it has all the required components + self.assertTrue(hasattr(triple, "first_gather")) + self.assertTrue(hasattr(triple, "second_gather")) + self.assertTrue(hasattr(triple, "noe_gather")) + self.assertTrue(hasattr(triple, "update")) + self.assertTrue(hasattr(triple, "first_scatter")) + self.assertTrue(hasattr(triple, "second_scatter")) + self.assertTrue(hasattr(triple, "noe_scatter")) + + def test_forward_pass_executes(self): + """Test that forward pass executes without error.""" + triple = ResidueResidueNoeTriple(self.device) + + # Should not raise an error + data = triple(self.data) + + # Data should be returned + self.assertIsNotNone(data) + + def test_forward_pass_sequence(self): + """Test that forward pass calls gather -> update -> scatter in sequence.""" + triple = ResidueResidueNoeTriple(self.device) + data = triple(self.data) + + # After forward pass, check that: + # 1. Gathered attributes exist on triple nodes (from gather operations) + self.assertIn("first_coords", data["ResidueResidueNoeTriple"]) + self.assertIn("second_coords", data["ResidueResidueNoeTriple"]) + self.assertIn("noe_shifts", data["ResidueResidueNoeTriple"]) + + # 2. Delta attributes exist on triple nodes (from update operation) + self.assertIn("delta_first_coords", data["ResidueResidueNoeTriple"]) + self.assertIn("delta_second_shifts", data["ResidueResidueNoeTriple"]) + + # 3. Residue and Noe nodes have been updated (from scatter operations) + # We can't easily check the exact values, but we can verify the tensors exist + self.assertIsNotNone(data["Residue"].x) + self.assertIsNotNone(data["Noe"].x) + + +class TestResiduePeakNoeTriple(unittest.TestCase): + """Test ResiduePeakNoeTriple wiring and forward pass.""" + + def setUp(self): + """Set up test graph with Residue, Peak, and Noe nodes.""" + self.device = torch.device("cpu") + self.data = HeteroData() + + # Create Residue nodes + self.data["Residue"].x = torch.tensor( + [[1.0, 2.0, 3.0, 8.5, 120.0]], dtype=torch.float32, device=self.device + ) + self.data["Residue"].f = torch.tensor( + [[0.0, 1.0]], dtype=torch.float32, device=self.device + ) + + # Create Peak nodes + self.data["Peak"].x = torch.tensor( + [[7.5, 115.0]], dtype=torch.float32, device=self.device + ) + self.data["Peak"].f = torch.tensor( + [[1.0, 0.0]], dtype=torch.float32, device=self.device + ) + + # Create Noe nodes + self.data["Noe"].x = torch.tensor( + [[120.0, 8.5, 7.5]], dtype=torch.float32, device=self.device + ) + self.data["Noe"].f = torch.tensor( + [[1.0, 0.0]], dtype=torch.float32, device=self.device + ) + + # Create triple nodes + self.data["ResiduePeakNoeTriple"].x = torch.zeros( + (1, 1), dtype=torch.float32, device=self.device + ) + + # Create gather edges + self.data["Residue", "gather_first", "ResiduePeakNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["Peak", "gather_second", "ResiduePeakNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["Noe", "gather_noe", "ResiduePeakNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + + # Create scatter edges + self.data["ResiduePeakNoeTriple", "scatter_first", "Residue"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["ResiduePeakNoeTriple", "scatter_second", "Peak"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["ResiduePeakNoeTriple", "scatter_noe", "Noe"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + + def test_triple_uses_correct_components(self): + """Test that ResiduePeakNoeTriple uses correct gather/scatter types.""" + triple = ResiduePeakNoeTriple(self.device) + + # Check component types by examining their triple_type attribute + self.assertEqual(triple.first_gather.triple_type, "ResiduePeakNoeTriple") + self.assertEqual(triple.second_gather.triple_type, "ResiduePeakNoeTriple") + self.assertEqual(triple.noe_gather.triple_type, "ResiduePeakNoeTriple") + self.assertEqual(triple.update.triple_type, "ResiduePeakNoeTriple") + + # Check edge types to ensure correct gather/scatter classes + self.assertEqual( + triple.first_gather.edge_type, + ("Residue", "gather_first", "ResiduePeakNoeTriple"), + ) + self.assertEqual( + triple.second_gather.edge_type, + ("Peak", "gather_second", "ResiduePeakNoeTriple"), + ) + + def test_forward_pass_executes(self): + """Test that forward pass executes without error.""" + triple = ResiduePeakNoeTriple(self.device) + data = triple(self.data) + self.assertIsNotNone(data) + + +class TestPeakResidueNoeTriple(unittest.TestCase): + """Test PeakResidueNoeTriple wiring and forward pass.""" + + def setUp(self): + """Set up test graph with Peak, Residue, and Noe nodes.""" + self.device = torch.device("cpu") + self.data = HeteroData() + + # Create Peak nodes + self.data["Peak"].x = torch.tensor( + [[8.5, 120.0]], dtype=torch.float32, device=self.device + ) + self.data["Peak"].f = torch.tensor( + [[0.0, 1.0]], dtype=torch.float32, device=self.device + ) + + # Create Residue nodes + self.data["Residue"].x = torch.tensor( + [[1.0, 2.0, 3.0, 7.5, 115.0]], dtype=torch.float32, device=self.device + ) + self.data["Residue"].f = torch.tensor( + [[1.0, 0.0]], dtype=torch.float32, device=self.device + ) + + # Create Noe nodes + self.data["Noe"].x = torch.tensor( + [[120.0, 8.5, 7.5]], dtype=torch.float32, device=self.device + ) + self.data["Noe"].f = torch.tensor( + [[1.0, 0.0]], dtype=torch.float32, device=self.device + ) + + # Create triple nodes + self.data["PeakResidueNoeTriple"].x = torch.zeros( + (1, 1), dtype=torch.float32, device=self.device + ) + + # Create gather edges + self.data["Peak", "gather_first", "PeakResidueNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["Residue", "gather_second", "PeakResidueNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["Noe", "gather_noe", "PeakResidueNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + + # Create scatter edges + self.data["PeakResidueNoeTriple", "scatter_first", "Peak"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["PeakResidueNoeTriple", "scatter_second", "Residue"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["PeakResidueNoeTriple", "scatter_noe", "Noe"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + + def test_forward_pass_executes(self): + """Test that forward pass executes without error.""" + triple = PeakResidueNoeTriple(self.device) + data = triple(self.data) + self.assertIsNotNone(data) + + +class TestPeakPeakNoeTriple(unittest.TestCase): + """Test PeakPeakNoeTriple wiring and forward pass.""" + + def setUp(self): + """Set up test graph with Peak and Noe nodes.""" + self.device = torch.device("cpu") + self.data = HeteroData() + + # Create Peak nodes + self.data["Peak"].x = torch.tensor( + [[8.5, 120.0], [7.5, 115.0]], dtype=torch.float32, device=self.device + ) + self.data["Peak"].f = torch.tensor( + [[0.0, 1.0], [1.0, 0.0]], dtype=torch.float32, device=self.device + ) + + # Create Noe nodes + self.data["Noe"].x = torch.tensor( + [[120.0, 8.5, 7.5]], dtype=torch.float32, device=self.device + ) + self.data["Noe"].f = torch.tensor( + [[1.0, 0.0]], dtype=torch.float32, device=self.device + ) + + # Create triple nodes + self.data["PeakPeakNoeTriple"].x = torch.zeros( + (1, 1), dtype=torch.float32, device=self.device + ) + + # Create gather edges + self.data["Peak", "gather_first", "PeakPeakNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["Peak", "gather_second", "PeakPeakNoeTriple"].edge_index = ( + torch.tensor([[1], [0]], dtype=torch.long, device=self.device) + ) + self.data["Noe", "gather_noe", "PeakPeakNoeTriple"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + + # Create scatter edges + self.data["PeakPeakNoeTriple", "scatter_first", "Peak"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + self.data["PeakPeakNoeTriple", "scatter_second", "Peak"].edge_index = ( + torch.tensor([[0], [1]], dtype=torch.long, device=self.device) + ) + self.data["PeakPeakNoeTriple", "scatter_noe", "Noe"].edge_index = ( + torch.tensor([[0], [0]], dtype=torch.long, device=self.device) + ) + + def test_forward_pass_executes(self): + """Test that forward pass executes without error.""" + triple = PeakPeakNoeTriple(self.device) + data = triple(self.data) + self.assertIsNotNone(data) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_triple_helpers.py b/tests/test_triple_helpers.py new file mode 100644 index 0000000..755ffa4 --- /dev/null +++ b/tests/test_triple_helpers.py @@ -0,0 +1,224 @@ +""" +Unit tests for triple.py helper functions. + +Tests the three helper functions at the bottom of triple.py: +- calc_noe_difference +- calc_shift_difference +- calc_res_distance +""" + +import sys +from pathlib import Path +import unittest +import torch + +# Add parent directory to path for imports +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.models.triple import calc_noe_difference, calc_shift_difference, calc_res_distance + + +class TestCalcNOEDifference(unittest.TestCase): + """Test calc_noe_difference function.""" + + def test_basic_differences(self): + """Test basic NOE difference calculations.""" + # Create sample data: shifts are [H, N] + x1 = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32) # [2, 2] (H, N) + x2 = torch.tensor([[5.0, 6.0], [7.0, 8.0]], dtype=torch.float32) # [2, 2] (H, N) + noe = torch.tensor([[10.0, 11.0, 12.0], [20.0, 21.0, 22.0]], dtype=torch.float32) # [2, 3] (N, H', H") + + diff_N, diff_H1, diff_H2 = calc_noe_difference(x1, x2, noe) + + # Expected calculations: + # diff_N = noe[:, 0] - x1[:, 1] (NOE_N - first_N) + # diff_H1 = noe[:, 1] - x1[:, 0] (NOE_H' - first_H) + # diff_H2 = noe[:, 2] - x2[:, 0] (NOE_H" - second_H) + + expected_diff_N = torch.tensor([[10.0 - 2.0], [20.0 - 4.0]], dtype=torch.float32) + expected_diff_H1 = torch.tensor([[11.0 - 1.0], [21.0 - 3.0]], dtype=torch.float32) + expected_diff_H2 = torch.tensor([[12.0 - 5.0], [22.0 - 7.0]], dtype=torch.float32) + + torch.testing.assert_close(diff_N, expected_diff_N) + torch.testing.assert_close(diff_H1, expected_diff_H1) + torch.testing.assert_close(diff_H2, expected_diff_H2) + + def test_output_shapes(self): + """Test that output shapes are correct.""" + x1 = torch.randn(10, 2) + x2 = torch.randn(10, 2) + noe = torch.randn(10, 3) + + diff_N, diff_H1, diff_H2 = calc_noe_difference(x1, x2, noe) + + self.assertEqual(diff_N.shape, (10, 1)) + self.assertEqual(diff_H1.shape, (10, 1)) + self.assertEqual(diff_H2.shape, (10, 1)) + + def test_zero_differences(self): + """Test with zero differences.""" + # x1 has shifts [1.0, 10.0], x2 has shifts [12.0, 6.0] + # noe has shifts [10.0, 1.0, 12.0] to match + x1 = torch.tensor([[1.0, 10.0]], dtype=torch.float32) + x2 = torch.tensor([[12.0, 6.0]], dtype=torch.float32) + noe = torch.tensor([[10.0, 1.0, 12.0]], dtype=torch.float32) + + diff_N, diff_H1, diff_H2 = calc_noe_difference(x1, x2, noe) + + # Should all be zero + torch.testing.assert_close(diff_N, torch.zeros((1, 1))) + torch.testing.assert_close(diff_H1, torch.zeros((1, 1))) + torch.testing.assert_close(diff_H2, torch.zeros((1, 1))) + + def test_negative_differences(self): + """Test with negative differences.""" + x1 = torch.tensor([[5.0, 6.0]], dtype=torch.float32) + x2 = torch.tensor([[7.0, 8.0]], dtype=torch.float32) + noe = torch.tensor([[2.0, 3.0, 4.0]], dtype=torch.float32) + + diff_N, diff_H1, diff_H2 = calc_noe_difference(x1, x2, noe) + + expected_diff_N = torch.tensor([[2.0 - 6.0]], dtype=torch.float32) # -4.0 + expected_diff_H1 = torch.tensor([[3.0 - 5.0]], dtype=torch.float32) # -2.0 + expected_diff_H2 = torch.tensor([[4.0 - 7.0]], dtype=torch.float32) # -3.0 + + torch.testing.assert_close(diff_N, expected_diff_N) + torch.testing.assert_close(diff_H1, expected_diff_H1) + torch.testing.assert_close(diff_H2, expected_diff_H2) + + +class TestCalcShiftDifference(unittest.TestCase): + """Test calc_shift_difference function.""" + + def test_basic_differences(self): + """Test basic shift difference calculations.""" + # Create sample data: shifts are [H, N] + x1 = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32) + x2 = torch.tensor([[5.0, 6.0], [7.0, 8.0]], dtype=torch.float32) + + diff_N, diff_H = calc_shift_difference(x1, x2) + + # Expected: x1 - x2 + expected_diff_N = torch.tensor([[2.0 - 6.0], [4.0 - 8.0]], dtype=torch.float32) + expected_diff_H = torch.tensor([[1.0 - 5.0], [3.0 - 7.0]], dtype=torch.float32) + + torch.testing.assert_close(diff_N, expected_diff_N) + torch.testing.assert_close(diff_H, expected_diff_H) + + def test_output_shapes(self): + """Test that output shapes are correct.""" + x1 = torch.randn(10, 2) + x2 = torch.randn(10, 2) + + diff_N, diff_H = calc_shift_difference(x1, x2) + + self.assertEqual(diff_N.shape, (10, 1)) + self.assertEqual(diff_H.shape, (10, 1)) + + def test_zero_differences(self): + """Test with identical shifts.""" + x1 = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32) + x2 = x1.clone() + + diff_N, diff_H = calc_shift_difference(x1, x2) + + torch.testing.assert_close(diff_N, torch.zeros((2, 1))) + torch.testing.assert_close(diff_H, torch.zeros((2, 1))) + + def test_symmetry(self): + """Test that calc_shift_difference(x1, x2) = -calc_shift_difference(x2, x1).""" + x1 = torch.randn(5, 2) + x2 = torch.randn(5, 2) + + diff_N1, diff_H1 = calc_shift_difference(x1, x2) + diff_N2, diff_H2 = calc_shift_difference(x2, x1) + + torch.testing.assert_close(diff_N1, -diff_N2) + torch.testing.assert_close(diff_H1, -diff_H2) + + +class TestCalcResDistance(unittest.TestCase): + """Test calc_res_distance function.""" + + def test_basic_distance(self): + """Test basic distance calculations.""" + # Simple case: points along x-axis + x1 = torch.tensor([[3.0, 0.0, 0.0]], dtype=torch.float32) + x2 = torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float32) + + rel_dist, dist_squared = calc_res_distance(x1, x2) + + expected_rel_dist = torch.tensor([[3.0, 0.0, 0.0]], dtype=torch.float32) + expected_dist_squared = torch.tensor([[9.0]], dtype=torch.float32) + + torch.testing.assert_close(rel_dist, expected_rel_dist) + torch.testing.assert_close(dist_squared, expected_dist_squared) + + def test_output_shapes(self): + """Test that output shapes are correct.""" + x1 = torch.randn(10, 3) + x2 = torch.randn(10, 3) + + rel_dist, dist_squared = calc_res_distance(x1, x2) + + self.assertEqual(rel_dist.shape, (10, 3)) + self.assertEqual(dist_squared.shape, (10, 1)) + + def test_zero_distance(self): + """Test with identical coordinates.""" + x1 = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=torch.float32) + x2 = x1.clone() + + rel_dist, dist_squared = calc_res_distance(x1, x2) + + torch.testing.assert_close(rel_dist, torch.zeros((2, 3))) + torch.testing.assert_close(dist_squared, torch.zeros((2, 1))) + + def test_3d_distance(self): + """Test 3D Euclidean distance calculation.""" + # Use 3-4-5 right triangle in 3D + x1 = torch.tensor([[3.0, 4.0, 0.0]], dtype=torch.float32) + x2 = torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float32) + + rel_dist, dist_squared = calc_res_distance(x1, x2) + + expected_rel_dist = torch.tensor([[3.0, 4.0, 0.0]], dtype=torch.float32) + expected_dist_squared = torch.tensor([[25.0]], dtype=torch.float32) # 3^2 + 4^2 + 0^2 + + torch.testing.assert_close(rel_dist, expected_rel_dist) + torch.testing.assert_close(dist_squared, expected_dist_squared) + + def test_symmetry(self): + """Test that calc_res_distance(x1, x2) = -calc_res_distance(x2, x1) for rel_dist.""" + x1 = torch.randn(5, 3) + x2 = torch.randn(5, 3) + + rel_dist1, dist_squared1 = calc_res_distance(x1, x2) + rel_dist2, dist_squared2 = calc_res_distance(x2, x1) + + # Relative distance should be opposite + torch.testing.assert_close(rel_dist1, -rel_dist2) + + # Squared distance should be the same (distance is symmetric) + torch.testing.assert_close(dist_squared1, dist_squared2) + + def test_batch_processing(self): + """Test batch processing with multiple coordinate pairs.""" + x1 = torch.tensor([ + [1.0, 0.0, 0.0], + [0.0, 2.0, 0.0], + [0.0, 0.0, 3.0], + ], dtype=torch.float32) + x2 = torch.zeros((3, 3), dtype=torch.float32) + + rel_dist, dist_squared = calc_res_distance(x1, x2) + + expected_rel_dist = x1 # x1 - 0 = x1 + expected_dist_squared = torch.tensor([[1.0], [4.0], [9.0]], dtype=torch.float32) + + torch.testing.assert_close(rel_dist, expected_rel_dist) + torch.testing.assert_close(dist_squared, expected_dist_squared) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_triple_integration.py b/tests/test_triple_integration.py new file mode 100644 index 0000000..5e0dff8 --- /dev/null +++ b/tests/test_triple_integration.py @@ -0,0 +1,564 @@ +""" +Integration tests for triple system re-architecture. + +These tests verify the complete pipeline functionality including: +- End-to-end forward pass through entire network +- Graph construction with new naming creates valid graphs +- All 4 triple types process correctly in sequence +- Gradient flow through gather → update → scatter pipeline +- Equivariance of ResidueResidueNoeTriple coordinate updates +- Peak-based triples produce zero coordinate deltas +- Batch processing with multiple graphs +- Edge case handling (empty triple sets, single node graphs) +""" + +import copy +import sys +import unittest +from pathlib import Path + +import numpy as np +import torch +from torch_geometric.loader import DataLoader + +# Add parent directory to path +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.construct import construct_graph +from nmr.models.network import NMRNet, NMRLayer, ModelConfig, StandardizeShifts, EmbedFeatures +from nmr.nmr_gym.fake_data import FakeDataGenerator +from nmr.nmr_gym.gym_env import GymEnv + + +def create_state_dict_from_env(env_state): + """ + Convert gym environment state to construct_graph compatible state. + + The gym_env.reset() returns coordinates as Protein namedtuples with + (x, y, z, H1, N15). construct_graph expects coordinates to contain all 5 values. + """ + # Extract all 5 values from Protein namedtuples: x, y, z, H1, N15 + coordinates = [[p.x, p.y, p.z, p.H1, p.N15] for p in env_state["coordinates"]] + + # Extract obs_shifts from HSQCPeak namedtuples + obs_shifts = [[s.H1, s.N15] for s in env_state["obs_chemical_shifts"]] + + # Extract noes from NOEPeak namedtuples + noes = [[n.N15, n.H1, n.H2] for n in env_state["noes"]] + + return { + "coordinates": coordinates, + "obs_chemical_shifts": obs_shifts, + "noes": noes, + "assignments": env_state["assignments"], + "shift_to_assign": env_state["shift_to_assign"], + } + + +class TestCompleteIntegration(unittest.TestCase): + """Integration tests for the complete triple system pipeline.""" + + def setUp(self): + """Set up test environment and small test data.""" + self.device = torch.device("cpu") + self.num_resid = 5 + + # Generate a small test dataset + generator = FakeDataGenerator(self.num_resid) + coords, obs_shifts, pred_shifts, noes, connectivity = generator.generate_data_arrays( + random_key=False + ) + self.pred_coords, self.obs_shifts, self.noes, self.connectivity = generator.order_data( + coords, obs_shifts, pred_shifts, noes, connectivity + ) + + # Create environment and initial state + self.env = GymEnv(self.num_resid) + env_state = self.env.reset( + self.pred_coords, self.obs_shifts, self.noes, self.connectivity + ) + + # Convert to construct_graph compatible format + self.state = create_state_dict_from_env(env_state) + + def test_end_to_end_forward_pass(self): + """Test complete forward pass through entire network.""" + # Construct graph + data = construct_graph(self.state, self.device) + + # Verify graph structure has new naming + self.assertIn("Residue", data.node_types) + self.assertIn("Peak", data.node_types) + self.assertIn("Noe", data.node_types) + self.assertIn("ResidueResidueNoeTriple", data.node_types) + self.assertIn("ResiduePeakNoeTriple", data.node_types) + self.assertIn("PeakResidueNoeTriple", data.node_types) + self.assertIn("PeakPeakNoeTriple", data.node_types) + + # Create network and run forward pass + net = NMRNet(self.device, ModelConfig()) + + # Run forward pass + value, policy = net(data) + + # Verify outputs have correct shapes + self.assertEqual(value.shape, (1, 1)) + self.assertIsInstance(policy, list) + self.assertEqual(len(policy), 1) + # Policy shape is [num_peaks, num_residues] for full pairwise distance matrix + self.assertEqual(policy[0].shape, (self.num_resid, self.num_resid)) + + # Verify outputs are finite + self.assertTrue(torch.isfinite(value).all()) + for p in policy: + self.assertTrue(torch.isfinite(p).all()) + + def test_graph_construction_with_new_naming(self): + """Test that graph construction creates valid graphs with new naming.""" + # Construct graph + data = construct_graph(self.state, self.device) + + # Verify node types use new naming + node_types = data.node_types + self.assertIn("Residue", node_types) + self.assertIn("Peak", node_types) + self.assertIn("Noe", node_types) + + # Verify old naming is NOT present + self.assertNotIn("RES", node_types) + self.assertNotIn("SHIFT", node_types) + self.assertNotIn("NOE", node_types) + self.assertNotIn("TRIPLE0", node_types) + self.assertNotIn("TRIPLE1", node_types) + self.assertNotIn("TRIPLE2", node_types) + self.assertNotIn("TRIPLE3", node_types) + + # Verify triple types use full descriptive names + self.assertIn("ResidueResidueNoeTriple", node_types) + self.assertIn("ResiduePeakNoeTriple", node_types) + self.assertIn("PeakResidueNoeTriple", node_types) + self.assertIn("PeakPeakNoeTriple", node_types) + + # Verify edge types use new naming + edge_types = data.edge_types + + # Check for prop edges (bidirectional) + prop_edges = [et for et in edge_types if "prop" in et[1]] + self.assertTrue(len(prop_edges) > 0, "No prop edges found") + + # Verify old triple edge naming is NOT present (exclude VALUE edges) + old_edge_names = ["NH1_extract", "NH2_extract", "NH1_add", "NH2_add"] + for edge_type in edge_types: + self.assertNotIn(edge_type[1], old_edge_names, + f"Old edge naming '{edge_type[1]}' still present") + + def test_all_four_triple_types_process_correctly(self): + """Test that all 4 triple types process correctly in sequence.""" + # Construct graph + data = construct_graph(self.state, self.device) + + # Embed features before calling NMRLayer + config = ModelConfig() + standardize = StandardizeShifts() + embed = EmbedFeatures(self.device, config) + data = standardize(data) + data = embed(data) + + # Store original node features + original_residue_x = data["Residue"].x.clone() + original_peak_x = data["Peak"].x.clone() + original_noe_x = data["Noe"].x.clone() + + # Create a single NMRLayer (which calls all 4 triples) + layer = NMRLayer(self.device, config) + + # Run forward pass through layer + data = layer(data) + + # Verify that node features were updated + # (At least some should change after message passing) + residue_changed = not torch.allclose(data["Residue"].x, original_residue_x) + peak_changed = not torch.allclose(data["Peak"].x, original_peak_x) + noe_changed = not torch.allclose(data["Noe"].x, original_noe_x) + + # At least one node type should have changed + self.assertTrue(residue_changed or peak_changed or noe_changed, + "No node features changed after layer forward pass") + + # Verify outputs are finite + self.assertTrue(torch.isfinite(data["Residue"].x).all()) + self.assertTrue(torch.isfinite(data["Peak"].x).all()) + self.assertTrue(torch.isfinite(data["Noe"].x).all()) + + def test_gradient_flow_through_pipeline(self): + """Test gradient flow through gather → update → scatter pipeline.""" + # Construct graph + data = construct_graph(self.state, self.device) + + # Create network + net = NMRNet(self.device, ModelConfig()) + + # Run forward pass + value, policy = net(data) + + # Compute simple loss + loss = value.sum() + sum(p.sum() for p in policy) + + # Run backward pass + loss.backward() + + # Verify gradients exist and are valid for all parameters + num_nonzero_grads = 0 + for name, param in net.named_parameters(): + if param.requires_grad: + self.assertIsNotNone(param.grad, + f"No gradient for parameter {name}") + + # Verify gradients are finite (not NaN/inf) + self.assertTrue(torch.isfinite(param.grad).all(), + f"Non-finite gradients in parameter {name}") + + # Verify gradients have reasonable magnitudes + grad_norm = param.grad.norm().item() + self.assertLess(grad_norm, 1e6, + f"Gradient norm too large for {name}: {grad_norm}") + + # Count non-zero gradients (zero gradients are OK with ReLU) + if grad_norm > 0.0: + num_nonzero_grads += 1 + + # Verify that at least some gradients are non-zero (not all dead) + total_params = sum(1 for p in net.parameters() if p.requires_grad) + self.assertGreater(num_nonzero_grads, 0, + "All gradients are zero - no gradient flow detected") + self.assertGreater(num_nonzero_grads, total_params * 0.5, + f"Too many zero gradients: {num_nonzero_grads}/{total_params}") + + def test_equivariance_of_residue_residue_noe_triple(self): + """Test equivariance of ResidueResidueNoeTriple coordinate updates.""" + # Construct graph + data1 = construct_graph(self.state, self.device) + + # Embed features before calling NMRLayer + config = ModelConfig() + standardize = StandardizeShifts() + embed = EmbedFeatures(self.device, config) + data1 = standardize(data1) + data1 = embed(data1) + + # Store original coordinates + original_coords = data1["Residue"].x[:, 0:3].clone() + + # Run forward pass and capture coordinate updates + layer = NMRLayer(self.device, config) + data1 = layer(data1) + updated_coords1 = data1["Residue"].x[:, 0:3] + coord_delta1 = updated_coords1 - original_coords + + # Create a rotation matrix (90 degrees around z-axis) + theta = np.pi / 2 + rotation_matrix = torch.tensor([ + [np.cos(theta), -np.sin(theta), 0], + [np.sin(theta), np.cos(theta), 0], + [0, 0, 1] + ], dtype=torch.float32, device=self.device) + + # Create a translation vector + translation = torch.tensor([1.0, 2.0, 3.0], device=self.device) + + # Apply rotation and translation to input coordinates + rotated_coords = torch.matmul(original_coords, rotation_matrix.T) + transformed_coords = rotated_coords + translation + + # Create new state with transformed coordinates + state2 = copy.deepcopy(self.state) + + # Update coordinates in state2 - need to preserve all values [x, y, z, shifts...] + # transformed_coords only has [x, y, z], so we need to append the shifts + original_full_coords = torch.tensor(self.state["coordinates"], dtype=torch.float32, device=self.device) + shifts = original_full_coords[:, 3:] # Extract all shift values + full_transformed_coords = torch.cat([transformed_coords, shifts], dim=-1) + state2["coordinates"] = full_transformed_coords.tolist() + + # Construct graph with transformed coordinates + data2 = construct_graph(state2, self.device) + + # Embed features for data2 + data2 = standardize(data2) + data2 = embed(data2) + + # Run forward pass again + layer2 = NMRLayer(self.device, config) + # Copy weights from first layer to ensure same transformation + layer2.load_state_dict(layer.state_dict()) + + data2 = layer2(data2) + updated_coords2 = data2["Residue"].x[:, 0:3] + coord_delta2 = updated_coords2 - transformed_coords + + # Verify coordinate deltas transform correctly (rotation only, not translation) + expected_delta2 = torch.matmul(coord_delta1, rotation_matrix.T) + + # Allow some numerical tolerance + self.assertTrue(torch.allclose(coord_delta2, expected_delta2, atol=1e-4), + "Coordinate deltas did not transform equivariantly") + + # Verify shift and feature updates are invariant + shifts1 = data1["Residue"].x[:, 3:] + shifts2 = data2["Residue"].x[:, 3:] + + # Shifts should be similar (invariant to coordinate transformation) + # Note: They may not be exactly equal due to coordinate-dependent updates + # but should be in the same ballpark + shift_diff = (shifts1 - shifts2).abs().max().item() + self.assertLess(shift_diff, 1.0, + f"Shift updates changed significantly: {shift_diff}") + + def test_peak_based_triples_produce_zero_coordinate_deltas(self): + """Test that Peak-based triples produce zero coordinate deltas.""" + # Construct graph + data = construct_graph(self.state, self.device) + + # Embed features + config = ModelConfig() + standardize = StandardizeShifts() + embed = EmbedFeatures(self.device, config) + data = standardize(data) + data = embed(data) + + # Store original Residue coordinates + original_residue_coords = data["Residue"].x[:, 0:3].clone() + + # Create network and isolate the Peak-based triples + # We'll test by running the full layer and checking that coordinate + # updates come only from ResidueResidueNoeTriple + layer = NMRLayer(self.device, config) + + # Manually call each triple and check coordinate changes + + # First, run only ResidueResidueNoeTriple + data1 = construct_graph(self.state, self.device) + data1 = standardize(data1) + data1 = embed(data1) + data1 = layer.residue_residue_noe(data1) + coords_after_res_res = data1["Residue"].x[:, 0:3] + + # Check if coordinates changed (they should for ResidueResidueNoeTriple) + torch.allclose(coords_after_res_res, original_residue_coords, atol=1e-6) + + # Now test each Peak-based triple + for triple_name, triple_module in [ + ("ResiduePeakNoeTriple", layer.residue_peak_noe), + ("PeakResidueNoeTriple", layer.peak_residue_noe), + ("PeakPeakNoeTriple", layer.peak_peak_noe) + ]: + data_peak = construct_graph(self.state, self.device) + data_peak = standardize(data_peak) + data_peak = embed(data_peak) + coords_before = data_peak["Residue"].x[:, 0:3].clone() + + # Run the Peak-based triple + data_peak = triple_module(data_peak) + coords_after = data_peak["Residue"].x[:, 0:3] + + # Verify coordinates did NOT change (or changed by at most numerical error) + coord_delta = (coords_after - coords_before).abs().max().item() + self.assertLess(coord_delta, 1e-5, + f"{triple_name} changed coordinates by {coord_delta}") + + def test_batch_processing_with_multiple_graphs(self): + """Test batch processing with multiple graphs.""" + # Generate multiple graphs + graphs = [] + for i in range(3): + # Use different random states for variety + np.random.seed(i) + generator = FakeDataGenerator(self.num_resid) + coords, obs_shifts, pred_shifts, noes, connectivity = generator.generate_data_arrays( + random_key=False + ) + coords, obs_shifts, noes, connectivity = generator.order_data( + coords, obs_shifts, pred_shifts, noes, connectivity + ) + + env = GymEnv(self.num_resid) + env_state = env.reset(coords, obs_shifts, noes, connectivity) + + # Convert to construct_graph format + state = create_state_dict_from_env(env_state) + + # Construct graph + data = construct_graph(state, self.device) + graphs.append(data) + + # Create data loader for batching + loader = DataLoader(graphs, batch_size=3, follow_batch=["x"]) + + # Get batched data + batched_data = next(iter(loader)) + + # Create network + net = NMRNet(self.device, ModelConfig()) + + # Run forward pass on batch + value, policy = net(batched_data) + + # Verify output shapes + # Value should have 3 entries (one per graph) + self.assertEqual(value.shape[0], 3) + + # Policy should have 3 entries (one per graph), each [num_peaks, num_residues] + self.assertIsInstance(policy, list) + self.assertEqual(len(policy), 3) + for p in policy: + self.assertEqual(p.shape, (self.num_resid, self.num_resid)) + + # Verify outputs are finite + self.assertTrue(torch.isfinite(value).all()) + for p in policy: + self.assertTrue(torch.isfinite(p).all()) + + def test_edge_case_handling(self): + """Test handling of edge cases (minimum size graphs).""" + # Test: Small graph (minimum viable size) + # Use 3 residues to ensure we get some NOEs + generator_small = FakeDataGenerator(num_resid=3) + coords, obs_shifts, pred_shifts, noes, connectivity = generator_small.generate_data_arrays( + random_key=False + ) + coords, obs_shifts, noes, connectivity = generator_small.order_data( + coords, obs_shifts, pred_shifts, noes, connectivity + ) + + env_small = GymEnv(3) + env_state_small = env_small.reset(coords, obs_shifts, noes, connectivity) + state_small = create_state_dict_from_env(env_state_small) + + # Construct graph + data_small = construct_graph(state_small, self.device) + + # Create network and run forward pass + net = NMRNet(self.device, ModelConfig()) + + try: + value, policy = net(data_small) + + # Verify outputs are valid + self.assertEqual(value.shape, (1, 1)) + self.assertIsInstance(policy, list) + self.assertEqual(len(policy), 1) + # Policy shape is [num_peaks, num_residues] = [3, 3] + self.assertEqual(policy[0].shape, (3, 3)) + self.assertTrue(torch.isfinite(value).all()) + for p in policy: + self.assertTrue(torch.isfinite(p).all()) + + except Exception as e: + self.fail(f"Forward pass failed on small graph: {e}") + + # Count triple nodes in small graph + for triple_type in ["ResidueResidueNoeTriple", "ResiduePeakNoeTriple", + "PeakResidueNoeTriple", "PeakPeakNoeTriple"]: + if triple_type in data_small.node_types: + num_triples = data_small[triple_type].x.size(0) + # Even if some types have zero instances, forward pass should work + self.assertGreaterEqual(num_triples, 0) + + +class TestCoordinateUpdateComparison(unittest.TestCase): + """Test comparison between triple types for coordinate updates.""" + + def setUp(self): + """Set up test environment.""" + self.device = torch.device("cpu") + self.num_resid = 5 + + # Generate test dataset + generator = FakeDataGenerator(self.num_resid) + coords, obs_shifts, pred_shifts, noes, connectivity = generator.generate_data_arrays( + random_key=False + ) + coords, obs_shifts, noes, connectivity = generator.order_data( + coords, obs_shifts, pred_shifts, noes, connectivity + ) + + # Create environment and initial state + env = GymEnv(self.num_resid) + env_state = env.reset(coords, obs_shifts, noes, connectivity) + self.state = create_state_dict_from_env(env_state) + + def test_coordinate_update_comparison(self): + """Verify only ResidueResidueNoeTriple updates coordinates.""" + config = ModelConfig() + layer = NMRLayer(self.device, config) + standardize = StandardizeShifts() + embed = EmbedFeatures(self.device, config) + + # Test ResidueResidueNoeTriple + data_res_res = construct_graph(self.state, self.device) + data_res_res = standardize(data_res_res) + data_res_res = embed(data_res_res) + original_coords = data_res_res["Residue"].x[:, 0:3].clone() + + data_res_res = layer.residue_residue_noe(data_res_res) + updated_coords = data_res_res["Residue"].x[:, 0:3] + + # ResidueResidueNoeTriple SHOULD update coordinates + # (implementation-dependent if updates are actually non-zero) + + # Test ResiduePeakNoeTriple + data_res_peak = construct_graph(self.state, self.device) + data_res_peak = standardize(data_res_peak) + data_res_peak = embed(data_res_peak) + coords_before = data_res_peak["Residue"].x[:, 0:3].clone() + data_res_peak = layer.residue_peak_noe(data_res_peak) + coords_after = data_res_peak["Residue"].x[:, 0:3] + + # ResiduePeakNoeTriple should NOT update coordinates + self.assertTrue(torch.allclose(coords_before, coords_after, atol=1e-6), + "ResiduePeakNoeTriple should not update coordinates") + + # Test PeakResidueNoeTriple + data_peak_res = construct_graph(self.state, self.device) + data_peak_res = standardize(data_peak_res) + data_peak_res = embed(data_peak_res) + coords_before = data_peak_res["Residue"].x[:, 0:3].clone() + data_peak_res = layer.peak_residue_noe(data_peak_res) + coords_after = data_peak_res["Residue"].x[:, 0:3] + + # PeakResidueNoeTriple should NOT update coordinates + self.assertTrue(torch.allclose(coords_before, coords_after, atol=1e-6), + "PeakResidueNoeTriple should not update coordinates") + + # Test PeakPeakNoeTriple + data_peak_peak = construct_graph(self.state, self.device) + data_peak_peak = standardize(data_peak_peak) + data_peak_peak = embed(data_peak_peak) + coords_before = data_peak_peak["Residue"].x[:, 0:3].clone() + data_peak_peak = layer.peak_peak_noe(data_peak_peak) + coords_after = data_peak_peak["Residue"].x[:, 0:3] + + # PeakPeakNoeTriple should NOT update coordinates + self.assertTrue(torch.allclose(coords_before, coords_after, atol=1e-6), + "PeakPeakNoeTriple should not update coordinates") + + # Verify all types update shifts and features + for triple_name, data_obj in [ + ("ResidueResidueNoeTriple", data_res_res), + ("ResiduePeakNoeTriple", data_res_peak), + ("PeakResidueNoeTriple", data_peak_res), + ("PeakPeakNoeTriple", data_peak_peak) + ]: + # Just verify that shifts and features exist and are finite + self.assertTrue(torch.isfinite(data_obj["Residue"].x[:, 3:]).all(), + f"{triple_name}: Residue shifts are not finite") + self.assertTrue(torch.isfinite(data_obj["Residue"].f).all(), + f"{triple_name}: Residue features are not finite") + self.assertTrue(torch.isfinite(data_obj["Peak"].x).all(), + f"{triple_name}: Peak shifts are not finite") + self.assertTrue(torch.isfinite(data_obj["Peak"].f).all(), + f"{triple_name}: Peak features are not finite") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_update_components.py b/tests/test_update_components.py new file mode 100644 index 0000000..bd52ea0 --- /dev/null +++ b/tests/test_update_components.py @@ -0,0 +1,490 @@ +""" +Unit tests for update component operations. + +Tests the 2 update classes that compute deltas on triple nodes: +- ResidueUpdate (for ResidueResidueNoeTriple) +- PeakUpdate (for all other triple types) +""" + +import sys +import unittest +from pathlib import Path + +import torch +from torch_geometric.data import HeteroData + +# Add parent directory to path to allow imports from nmr package +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from nmr.models.triple import ResidueUpdate, PeakUpdate +from nmr.models import ModelConfig + + +class TestResidueUpdateDimensions(unittest.TestCase): + """Test ResidueUpdate MLP input/output dimensions (12 -> 17).""" + + def setUp(self): + """Set up test graph with gathered attributes on triple nodes.""" + self.device = torch.device("cpu") + self.config = ModelConfig() + self.shift_dim = self.config.shift_embed.output_dim + self.feature_dim = self.config.feature_embed.output_dim + self.data = HeteroData() + + # Create 2 triple nodes + num_triples = 2 + self.data["ResidueResidueNoeTriple"].x = torch.zeros( + (num_triples, 1), dtype=torch.float32, device=self.device + ) + + # Set gathered attributes + # first_coords [n, 3] + self.data["ResidueResidueNoeTriple"].first_coords = torch.tensor( + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + dtype=torch.float32, + device=self.device, + ) + # first_shifts [n, shift_dim] + self.data["ResidueResidueNoeTriple"].first_shifts = torch.randn( + (num_triples, self.shift_dim), + dtype=torch.float32, + device=self.device, + ) + # first_features [n, feature_dim] + self.data["ResidueResidueNoeTriple"].first_features = torch.randn( + (num_triples, self.feature_dim), + dtype=torch.float32, + device=self.device, + ) + # second_coords [n, 3] + self.data["ResidueResidueNoeTriple"].second_coords = torch.tensor( + [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]], + dtype=torch.float32, + device=self.device, + ) + # second_shifts [n, shift_dim] + self.data["ResidueResidueNoeTriple"].second_shifts = torch.randn( + (num_triples, self.shift_dim), + dtype=torch.float32, + device=self.device, + ) + # second_features [n, feature_dim] + self.data["ResidueResidueNoeTriple"].second_features = torch.randn( + (num_triples, self.feature_dim), + dtype=torch.float32, + device=self.device, + ) + # noe_shifts [n, shift_dim] (embedded NOE features) + self.data["ResidueResidueNoeTriple"].noe_shifts = torch.randn( + (num_triples, self.shift_dim), + dtype=torch.float32, + device=self.device, + ) + # noe_features [n, feature_dim] + self.data["ResidueResidueNoeTriple"].noe_features = torch.randn( + (num_triples, self.feature_dim), + dtype=torch.float32, + device=self.device, + ) + + def test_residue_update_sets_all_delta_attributes(self): + """Test that ResidueUpdate sets all 8 delta attributes.""" + update = ResidueUpdate("ResidueResidueNoeTriple", self.device, self.config) + data = update(self.data) + + # Check all delta attributes are set + self.assertIn("delta_first_coords", data["ResidueResidueNoeTriple"]) + self.assertIn("delta_first_shifts", data["ResidueResidueNoeTriple"]) + self.assertIn("delta_first_features", data["ResidueResidueNoeTriple"]) + self.assertIn("delta_second_coords", data["ResidueResidueNoeTriple"]) + self.assertIn("delta_second_shifts", data["ResidueResidueNoeTriple"]) + self.assertIn("delta_second_features", data["ResidueResidueNoeTriple"]) + self.assertIn("delta_noe_shifts", data["ResidueResidueNoeTriple"]) + self.assertIn("delta_noe_features", data["ResidueResidueNoeTriple"]) + + def test_residue_update_delta_shapes(self): + """Test that ResidueUpdate produces correct output shapes.""" + update = ResidueUpdate("ResidueResidueNoeTriple", self.device, self.config) + data = update(self.data) + + # Check shapes + self.assertEqual( + data["ResidueResidueNoeTriple"].delta_first_coords.shape, (2, 3) + ) + self.assertEqual( + data["ResidueResidueNoeTriple"].delta_first_shifts.shape, (2, self.shift_dim) + ) + self.assertEqual( + data["ResidueResidueNoeTriple"].delta_first_features.shape, (2, self.feature_dim) + ) + self.assertEqual( + data["ResidueResidueNoeTriple"].delta_second_coords.shape, (2, 3) + ) + self.assertEqual( + data["ResidueResidueNoeTriple"].delta_second_shifts.shape, (2, self.shift_dim) + ) + self.assertEqual( + data["ResidueResidueNoeTriple"].delta_second_features.shape, (2, self.feature_dim) + ) + self.assertEqual( + data["ResidueResidueNoeTriple"].delta_noe_shifts.shape, (2, self.shift_dim) + ) + self.assertEqual( + data["ResidueResidueNoeTriple"].delta_noe_features.shape, (2, self.feature_dim) + ) + + +class TestPeakUpdateDimensions(unittest.TestCase): + """Test PeakUpdate MLP input/output dimensions (11 -> 15).""" + + def setUp(self): + """Set up test graph with gathered attributes on triple nodes.""" + self.device = torch.device("cpu") + self.config = ModelConfig() + self.shift_dim = self.config.shift_embed.output_dim + self.feature_dim = self.config.feature_embed.output_dim + self.data = HeteroData() + + # Create 2 triple nodes + num_triples = 2 + self.data["PeakPeakNoeTriple"].x = torch.zeros( + (num_triples, 1), dtype=torch.float32, device=self.device + ) + + # Set gathered attributes (peaks have NO coordinates) + # first_shifts [n, shift_dim] + self.data["PeakPeakNoeTriple"].first_shifts = torch.randn( + (num_triples, self.shift_dim), + dtype=torch.float32, + device=self.device, + ) + # first_features [n, feature_dim] + self.data["PeakPeakNoeTriple"].first_features = torch.randn( + (num_triples, self.feature_dim), + dtype=torch.float32, + device=self.device, + ) + # second_shifts [n, shift_dim] + self.data["PeakPeakNoeTriple"].second_shifts = torch.randn( + (num_triples, self.shift_dim), + dtype=torch.float32, + device=self.device, + ) + # second_features [n, feature_dim] + self.data["PeakPeakNoeTriple"].second_features = torch.randn( + (num_triples, self.feature_dim), + dtype=torch.float32, + device=self.device, + ) + # noe_shifts [n, shift_dim] (embedded NOE features) + self.data["PeakPeakNoeTriple"].noe_shifts = torch.randn( + (num_triples, self.shift_dim), + dtype=torch.float32, + device=self.device, + ) + # noe_features [n, feature_dim] + self.data["PeakPeakNoeTriple"].noe_features = torch.randn( + (num_triples, self.feature_dim), + dtype=torch.float32, + device=self.device, + ) + + def test_peak_update_sets_all_delta_attributes(self): + """Test that PeakUpdate sets all 8 delta attributes.""" + update = PeakUpdate("PeakPeakNoeTriple", self.device, self.config) + data = update(self.data) + + # Check all delta attributes are set + self.assertIn("delta_first_coords", data["PeakPeakNoeTriple"]) + self.assertIn("delta_first_shifts", data["PeakPeakNoeTriple"]) + self.assertIn("delta_first_features", data["PeakPeakNoeTriple"]) + self.assertIn("delta_second_coords", data["PeakPeakNoeTriple"]) + self.assertIn("delta_second_shifts", data["PeakPeakNoeTriple"]) + self.assertIn("delta_second_features", data["PeakPeakNoeTriple"]) + self.assertIn("delta_noe_shifts", data["PeakPeakNoeTriple"]) + self.assertIn("delta_noe_features", data["PeakPeakNoeTriple"]) + + def test_peak_update_delta_shapes(self): + """Test that PeakUpdate produces correct output shapes.""" + update = PeakUpdate("PeakPeakNoeTriple", self.device, self.config) + data = update(self.data) + + # Check shapes (note: coordinates are zeros but still [n, 3]) + self.assertEqual(data["PeakPeakNoeTriple"].delta_first_coords.shape, (2, 3)) + self.assertEqual(data["PeakPeakNoeTriple"].delta_first_shifts.shape, (2, self.shift_dim)) + self.assertEqual(data["PeakPeakNoeTriple"].delta_first_features.shape, (2, self.feature_dim)) + self.assertEqual(data["PeakPeakNoeTriple"].delta_second_coords.shape, (2, 3)) + self.assertEqual(data["PeakPeakNoeTriple"].delta_second_shifts.shape, (2, self.shift_dim)) + self.assertEqual( + data["PeakPeakNoeTriple"].delta_second_features.shape, (2, self.feature_dim) + ) + self.assertEqual(data["PeakPeakNoeTriple"].delta_noe_shifts.shape, (2, self.shift_dim)) + self.assertEqual(data["PeakPeakNoeTriple"].delta_noe_features.shape, (2, self.feature_dim)) + + +class TestResidueUpdateEquivariance(unittest.TestCase): + """Test ResidueUpdate computes coordinate deltas using relative distances.""" + + def setUp(self): + """Set up test graph with gathered attributes.""" + self.device = torch.device("cpu") + self.config = ModelConfig() + self.shift_dim = self.config.shift_embed.output_dim + self.feature_dim = self.config.feature_embed.output_dim + self.data = HeteroData() + + # Create 1 triple node + self.data["ResidueResidueNoeTriple"].x = torch.zeros( + (1, 1), dtype=torch.float32, device=self.device + ) + + # Set gathered attributes with specific coordinates + self.data["ResidueResidueNoeTriple"].first_coords = torch.tensor( + [[1.0, 2.0, 3.0]], + dtype=torch.float32, + device=self.device, + ) + self.data["ResidueResidueNoeTriple"].first_shifts = torch.randn( + (1, self.shift_dim), + dtype=torch.float32, + device=self.device, + ) + self.data["ResidueResidueNoeTriple"].first_features = torch.randn( + (1, self.feature_dim), + dtype=torch.float32, + device=self.device, + ) + self.data["ResidueResidueNoeTriple"].second_coords = torch.tensor( + [[4.0, 5.0, 6.0]], # Relative distance: [3.0, 3.0, 3.0] + dtype=torch.float32, + device=self.device, + ) + self.data["ResidueResidueNoeTriple"].second_shifts = torch.randn( + (1, self.shift_dim), + dtype=torch.float32, + device=self.device, + ) + self.data["ResidueResidueNoeTriple"].second_features = torch.randn( + (1, self.feature_dim), + dtype=torch.float32, + device=self.device, + ) + self.data["ResidueResidueNoeTriple"].noe_shifts = torch.randn( + (1, self.shift_dim), + dtype=torch.float32, + device=self.device, + ) + self.data["ResidueResidueNoeTriple"].noe_features = torch.randn( + (1, self.feature_dim), + dtype=torch.float32, + device=self.device, + ) + + def test_coordinate_deltas_proportional_to_relative_distance(self): + """Test that coordinate deltas are proportional to relative distance vector.""" + update = ResidueUpdate("ResidueResidueNoeTriple", self.device, self.config) + data = update(self.data) + + # Get coordinate deltas + delta_first = data["ResidueResidueNoeTriple"].delta_first_coords + delta_second = data["ResidueResidueNoeTriple"].delta_second_coords + + # Deltas should be 3D vectors + self.assertEqual(delta_first.shape, (1, 3)) + self.assertEqual(delta_second.shape, (1, 3)) + + # Deltas should not be all zeros (MLP should produce non-zero weights) + # Note: This is a weak test since random initialization could produce + # small weights, but we mainly want to verify the shape and structure + self.assertEqual(delta_first.dtype, torch.float32) + self.assertEqual(delta_second.dtype, torch.float32) + + +class TestPeakUpdateZeroCoordinates(unittest.TestCase): + """Test PeakUpdate outputs zero coordinate deltas.""" + + def setUp(self): + """Set up test graph with gathered attributes.""" + self.device = torch.device("cpu") + self.config = ModelConfig() + self.shift_dim = self.config.shift_embed.output_dim + self.feature_dim = self.config.feature_embed.output_dim + self.data = HeteroData() + + # Create 2 triple nodes + self.data["PeakPeakNoeTriple"].x = torch.zeros( + (2, 1), dtype=torch.float32, device=self.device + ) + + # Set gathered attributes (no coordinates for peaks) + self.data["PeakPeakNoeTriple"].first_shifts = torch.randn( + (2, self.shift_dim), + dtype=torch.float32, + device=self.device, + ) + self.data["PeakPeakNoeTriple"].first_features = torch.randn( + (2, self.feature_dim), + dtype=torch.float32, + device=self.device, + ) + self.data["PeakPeakNoeTriple"].second_shifts = torch.randn( + (2, self.shift_dim), + dtype=torch.float32, + device=self.device, + ) + self.data["PeakPeakNoeTriple"].second_features = torch.randn( + (2, self.feature_dim), + dtype=torch.float32, + device=self.device, + ) + self.data["PeakPeakNoeTriple"].noe_shifts = torch.randn( + (2, self.shift_dim), + dtype=torch.float32, + device=self.device, + ) + self.data["PeakPeakNoeTriple"].noe_features = torch.randn( + (2, self.feature_dim), + dtype=torch.float32, + device=self.device, + ) + + def test_peak_update_produces_zero_coordinate_deltas(self): + """Test that PeakUpdate produces exactly zero coordinate deltas.""" + update = PeakUpdate("PeakPeakNoeTriple", self.device, self.config) + data = update(self.data) + + # Get coordinate deltas + delta_first_coords = data["PeakPeakNoeTriple"].delta_first_coords + delta_second_coords = data["PeakPeakNoeTriple"].delta_second_coords + + # Deltas should be exactly zero + torch.testing.assert_close( + delta_first_coords, torch.zeros(2, 3, device=self.device) + ) + torch.testing.assert_close( + delta_second_coords, torch.zeros(2, 3, device=self.device) + ) + + +class TestUpdateConfigurableArchitecture(unittest.TestCase): + """Test configurable hidden_size and num_layers parameters.""" + + def setUp(self): + """Set up minimal test graph.""" + self.device = torch.device("cpu") + + def test_residue_update_uses_config(self): + """Test that ResidueUpdate uses config for hidden_size and num_layers.""" + config = ModelConfig() + update = ResidueUpdate("ResidueResidueNoeTriple", self.device, config) + self.assertIsNotNone(update.mlp) + + # Verify it uses config values + self.assertEqual(update.config.mlp.hidden_size, config.mlp.hidden_size) + self.assertEqual(update.config.mlp.num_layers, config.mlp.num_layers) + + def test_peak_update_uses_config(self): + """Test that PeakUpdate uses config for hidden_size and num_layers.""" + config = ModelConfig() + update = PeakUpdate("PeakPeakNoeTriple", self.device, config) + self.assertIsNotNone(update.mlp) + + # Verify it uses config values + self.assertEqual(update.config.mlp.hidden_size, config.mlp.hidden_size) + self.assertEqual(update.config.mlp.num_layers, config.mlp.num_layers) + + +class TestUpdateEmptyTripleSets(unittest.TestCase): + """Test update operations handle empty triple sets.""" + + def setUp(self): + """Set up test graph with zero triple nodes.""" + self.device = torch.device("cpu") + self.config = ModelConfig() + self.shift_dim = self.config.shift_embed.output_dim + self.feature_dim = self.config.feature_embed.output_dim + self.data = HeteroData() + + # Create ZERO triple nodes + self.data["ResidueResidueNoeTriple"].x = torch.zeros( + (0, 1), dtype=torch.float32, device=self.device + ) + + # Set gathered attributes as empty tensors + self.data["ResidueResidueNoeTriple"].first_coords = torch.zeros( + (0, 3), dtype=torch.float32, device=self.device + ) + self.data["ResidueResidueNoeTriple"].first_shifts = torch.zeros( + (0, self.shift_dim), dtype=torch.float32, device=self.device + ) + self.data["ResidueResidueNoeTriple"].first_features = torch.zeros( + (0, self.feature_dim), dtype=torch.float32, device=self.device + ) + self.data["ResidueResidueNoeTriple"].second_coords = torch.zeros( + (0, 3), dtype=torch.float32, device=self.device + ) + self.data["ResidueResidueNoeTriple"].second_shifts = torch.zeros( + (0, self.shift_dim), dtype=torch.float32, device=self.device + ) + self.data["ResidueResidueNoeTriple"].second_features = torch.zeros( + (0, self.feature_dim), dtype=torch.float32, device=self.device + ) + self.data["ResidueResidueNoeTriple"].noe_shifts = torch.zeros( + (0, self.shift_dim), dtype=torch.float32, device=self.device + ) + self.data["ResidueResidueNoeTriple"].noe_features = torch.zeros( + (0, self.feature_dim), dtype=torch.float32, device=self.device + ) + + def test_residue_update_handles_empty_triple_set(self): + """Test that ResidueUpdate handles empty triple sets without error.""" + update = ResidueUpdate("ResidueResidueNoeTriple", self.device, self.config) + + # Should not raise an error + data = update(self.data) + + # Delta attributes should be set with shape [0, ...] + self.assertEqual( + data["ResidueResidueNoeTriple"].delta_first_coords.shape, (0, 3) + ) + self.assertEqual( + data["ResidueResidueNoeTriple"].delta_first_shifts.shape, (0, self.shift_dim) + ) + + def test_peak_update_handles_empty_triple_set(self): + """Test that PeakUpdate handles empty triple sets without error.""" + # Create empty PeakPeakNoeTriple + self.data["PeakPeakNoeTriple"].x = torch.zeros( + (0, 1), dtype=torch.float32, device=self.device + ) + self.data["PeakPeakNoeTriple"].first_shifts = torch.zeros( + (0, self.shift_dim), dtype=torch.float32, device=self.device + ) + self.data["PeakPeakNoeTriple"].first_features = torch.zeros( + (0, self.feature_dim), dtype=torch.float32, device=self.device + ) + self.data["PeakPeakNoeTriple"].second_shifts = torch.zeros( + (0, self.shift_dim), dtype=torch.float32, device=self.device + ) + self.data["PeakPeakNoeTriple"].second_features = torch.zeros( + (0, self.feature_dim), dtype=torch.float32, device=self.device + ) + self.data["PeakPeakNoeTriple"].noe_shifts = torch.zeros( + (0, self.shift_dim), dtype=torch.float32, device=self.device + ) + self.data["PeakPeakNoeTriple"].noe_features = torch.zeros( + (0, self.feature_dim), dtype=torch.float32, device=self.device + ) + + update = PeakUpdate("PeakPeakNoeTriple", self.device, self.config) + + # Should not raise an error + data = update(self.data) + + # Delta attributes should be set with shape [0, ...] + self.assertEqual(data["PeakPeakNoeTriple"].delta_first_coords.shape, (0, 3)) + self.assertEqual(data["PeakPeakNoeTriple"].delta_first_shifts.shape, (0, self.shift_dim)) + + +if __name__ == "__main__": + unittest.main()