diff --git a/.gitignore b/.gitignore index 82f9275..beed6a1 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,13 @@ __pycache__/ *.py[cod] *$py.class +# Pickled data and visualization +*.pkl +!fake_histories_r4_0.pkl +*.png +runs/ +*.pdb + # C extensions *.so 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/fake_data.py b/fake_data.py deleted file mode 100644 index 97fc841..0000000 --- a/fake_data.py +++ /dev/null @@ -1,161 +0,0 @@ -import numpy as np -from itertools import product -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 - -class NOEPeak(NamedTuple): - H1: float - N15: float - H2: float - -class Protein(NamedTuple): - 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)) - -# 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 calc_dist(p1, p2): - return np.sqrt((p2[0] - p1[0])**2 + (p2[1] - p1[1])**2 + (p2[2] - p1[2])**2) - -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. - - **Can end up with no NOEs depending on the cutoff** - """ - 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]) - val = len(noe) # Append both shifts to correct noe index - noe.append(list(shifts[i][:])) # H1, N1 - noe[val].append(shifts[j][0]) # H2 - - noisy_noe = add_noise(np.array(noe), scale=0.01) - - 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 - -# 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 - -# 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) 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..e3a35ce --- /dev/null +++ b/nmr/construct.py @@ -0,0 +1,415 @@ +""" +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 typing import Any, Dict, Tuple + +import torch +from torch_geometric.data import HeteroData +from time import time + + +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. + + 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') + + 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. + + 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') + + Returns: + HeteroData graph with all edges constructed + """ + num_noe = len(data["Noe"].x) + num_peak = len(data["Peak"].x) + num_residue = len(data["Residue"].x) + + # Add all edge types for each triple configuration + _add_triple_edges( + data, ("Residue", "Residue", "Noe"), num_noe, num_peak, num_residue, device + ) + _add_triple_edges( + data, ("Residue", "Peak", "Noe"), num_noe, num_peak, num_residue, device + ) + _add_triple_edges( + data, ("Peak", "Residue", "Noe"), num_noe, num_peak, num_residue, device + ) + _add_triple_edges( + data, ("Peak", "Peak", "Noe"), num_noe, num_peak, num_residue, 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, 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["Peak"].x = torch.tensor( + histories["obs_chemical_shifts"], dtype=torch.float32, device=device + ) + data["Peak"].f = torch.zeros( + (len(histories["obs_chemical_shifts"]), 2), + dtype=torch.float32, + device=device, + ) + # 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["Residue"].f = torch.zeros( + (len(histories["coordinates"]), 2), dtype=torch.float32, device=device + ) + return data + + +def _construct_triple_nodes(data: HeteroData, device: torch.device | str) -> HeteroData: + """ + 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 + """ + 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 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 _construct_node_features(data: HeteroData, histories: Dict[str, Any]) -> HeteroData: + """ + 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: + data: HeteroData graph with nodes + histories: Dictionary containing assignment state + + Returns: + HeteroData with node features initialized + """ + # 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 data + + +def _get_triple_edges( + num_noe: int, source1: int, source2: int, device: torch.device | str +) -> torch.Tensor: + """ + 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: + 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: + 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. + """ + total_edges = num_noe * source1 * source2 + + # 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_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_peak: Number of peak nodes + num_residue: Number of residue nodes + device: Device to place tensors on + """ + peak = torch.arange(0, num_peak, device=device) + noe = torch.arange(0, num_noe, device=device) + resid = torch.arange(0, num_residue, device=device) + + peak_repeats = torch.zeros(num_peak, device=device).long() + noe_repeats = torch.zeros(num_noe, device=device).long() + resid_repeats = torch.zeros(num_residue, device=device).long() + + data["Peak", "SHIFT_extract", "VALUE_SHIFT"].edge_index = torch.stack( + (peak, peak_repeats), dim=0 + ) + data["Noe", "aggregate", "VALUE_NOE"].edge_index = torch.stack( + (noe, noe_repeats), dim=0 + ) + data["Residue", "RES_extract", "VALUE_RES"].edge_index = torch.stack( + (resid, resid_repeats), dim=0 + ) + + +def _add_triple_edges( + data: HeteroData, + triple_type: Tuple[str, str, str], + num_noe: int, + num_peak: int, + num_residue: int, + device: torch.device | str, +) -> None: + """ + 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_type: Tuple of (source1, source2, source3) node type names + num_noe: Number of NOE nodes + num_peak: Number of peak nodes + num_residue: Number of residue nodes + device: Device to place tensors on + """ + 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 + ) + + # 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 new file mode 100644 index 0000000..6fbdd53 --- /dev/null +++ b/nmr/models/__init__.py @@ -0,0 +1,86 @@ +""" +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, + ModelConfig, + FeatureEmbedConfig, + ShiftEmbedConfig, + MLPConfig, +) + +# Triple message passing - new modular architecture +from .triple import ( + # 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 +from .heads import BatchMessagePass, PolicyCalc, ValueCalc + +__all__ = [ + # Network + "NMRLayer", + "NMRNet", + # 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", + "PolicyCalc", +] diff --git a/nmr/models/heads.py b/nmr/models/heads.py new file mode 100644 index 0000000..c986522 --- /dev/null +++ b/nmr/models/heads.py @@ -0,0 +1,188 @@ +""" +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, config): + super().__init__(aggr=aggr) + self.device = device + self.config = config + + # 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) + 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): + 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) + 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 + + +class ValueCalc(nn.Module): + """ + Value function estimation head. + + Aggregates information from Noe, Peak, and Residue nodes to predict + a scalar value representing the quality of the current state. + """ + + def __init__(self, device, config): + super().__init__() + self.device = device + self.config = config + + 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), + ) + + 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): + """ + 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["VALUE_NOE"].x, + data["Noe", "aggregate", "VALUE_NOE"].edge_index, + ) + shift = self.get_aggr( + data["Peak"].x, + data["VALUE_SHIFT"].x, + data["Peak", "SHIFT_extract", "VALUE_SHIFT"].edge_index, + ) + resid = self.get_aggr( + data["Residue"].x, + data["VALUE_RES"].x, + data["Residue", "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, config): + super().__init__() + 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 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: 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] + """ + # 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: + # Get all peak shift embeddings: shape [num_peaks, shift_dim] + peak_features = item["Peak"].x + + # Get all residue shift embeddings: shape [num_residues, shift_dim] + residue_features = item["Residue"].x[:, 3:] # Extract shift dimensions starting at index 3 + + # 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) + + return policies diff --git a/nmr/models/network.py b/nmr/models/network.py new file mode 100644 index 0000000..83601db --- /dev/null +++ b/nmr/models/network.py @@ -0,0 +1,489 @@ +""" +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. +""" + +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 + + +@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 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 + + 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, config: ModelConfig): + super(NMRLayer, self).__init__() + # 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) + + def forward(self, data): + """ + Process all four triple types and assigned pairs in sequence. + + Args: + data: HeteroData graph to process + + 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 StandardizeShifts(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 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 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) + # 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["Residue"].x = res_x + data["Peak"].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 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, config: ModelConfig): + super().__init__() + self.config = config + self.device = 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): + # 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 new file mode 100644 index 0000000..2e69767 --- /dev/null +++ b/nmr/models/triple.py @@ -0,0 +1,1511 @@ +""" +Modular triple-based graph message passing components. + +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 +import torch.nn as nn +from torch_geometric.nn import MessagePassing + + +# ============================================================================ +# 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): + """ + 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. + + 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 __init__(self, triple_type: str): + """ + Initialize FirstResidueGather. + + Args: + triple_type: Name of target triple node type + """ + super().__init__(aggr="mean") + self.triple_type = triple_type + self.edge_type = ("Residue", "prop_first", triple_type) + + def forward(self, data): + """ + 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 + """ + # 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) + ) + + # 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): + """ + Initialize FirstPeakGather. + + Args: + triple_type: Name of target triple node type + """ + 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. + + Args: + data: HeteroData graph with Peak nodes and gather edges + + Returns: + Updated HeteroData with first_shifts, first_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].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): + """ + Initialize SecondResidueGather. + + 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 forward(self, data): + """ + 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 + """ + # 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) + ) + + # 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 + + +class SecondPeakGather(MessagePassing): + """ + 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, 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 + + +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) + + # Propagate shifts with explicit size + data[self.triple_type].noe_shifts = self.propagate( + edge_index, x=shifts, size=(shifts.size(0), num_triples) + ) + + # Propagate features + data[self.triple_type].noe_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 + + +# ============================================================================ +# 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, + ) + + # 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, + ) + + # 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 + + +class FirstPeakScatter(MessagePassing): + """ + Propagate deltas from triple nodes to peaks in first position. + + 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 + + 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_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 + """ + # 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] + + # 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)) + ) + + # 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/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/nmr/nmr_gym/assignment_order.py b/nmr/nmr_gym/assignment_order.py new file mode 100644 index 0000000..a9ff933 --- /dev/null +++ b/nmr/nmr_gym/assignment_order.py @@ -0,0 +1,98 @@ +import numpy as np +from itertools import chain + + + +class FractionalActivation(): + + def __init__(self, num_resid, restraints): + self.num_resid = num_resid + self.restraints = restraints + + def activation_loop(self, target, noe, assign_order): + """ + 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 = [] + # missing_shifts = [] + + # while len(assign_order) < self.num_resid: + # temp2 = [] + # for i in range(self.num_resid): + # 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) + # # print(assign_order, temp) + # # print(i, temp2) + + # # 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) + # # 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 + + + + 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/nmr/nmr_gym/energy.py b/nmr/nmr_gym/energy.py new file mode 100644 index 0000000..2e99608 --- /dev/null +++ b/nmr/nmr_gym/energy.py @@ -0,0 +1,158 @@ +import numpy as np +from typing import NamedTuple +from itertools import product +import math +from scipy.spatial.distance import pdist, squareform + + + +class Energy(): + + 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 + + # 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): + """ + 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 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. + """ + restraint_energy = math.inf + for i, j in restraint: + k, l = assignments.get(i), assignments.get(j) + # 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, 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. + """ + 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, dist_grid, 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, dist_grid, tolerance=self.tolerance_dist) + 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)) + + 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/nmr_gym/gym_env.py b/nmr/nmr_gym/gym_env.py new file mode 100644 index 0000000..16da01d --- /dev/null +++ b/nmr/nmr_gym/gym_env.py @@ -0,0 +1,161 @@ +import gymnasium as gym +from gymnasium import error, spaces +import numpy as np +import pickle +import math + +from .energy import Energy +from .assignment_order import FractionalActivation +from scripts.visualize_data import Visualization + +import cProfile +import pstats +import time + + + +class GymEnv(gym.Env): + + def __init__(self, num_resid): + self.state = None + self.num_resid = num_resid + + 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({ + "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), + "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 + # 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 + + # 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 + + # 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.intermediate_energy = 0.0 + + # 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.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, 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) + + # Get restraints + self.energy = Energy(coordinates, obs_chemical_shifts, noes) + self.restraints = self.energy.noe_combinations() + + self.dist_grid = self.energy.calc_pdist() + + self.assignments = {} + self.assign_step = 0 + self.intermediate_energy = 0.0 + + # 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.shift_to_assign}') + + self.state = { + "coordinates": coordinates, + "obs_chemical_shifts": obs_chemical_shifts, + "noes": noes, + "connectivity": connectivity, + "assignments": self.assignments, + "assign_order": self.assign_order, + "shift_to_assign": self.shift_to_assign, + "total_energy": 0.0, + "reward": 0.0 + } + + return self.state + + 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) + + + 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/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..605b535 --- /dev/null +++ b/scripts/train_gnn.py @@ -0,0 +1,194 @@ +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, ModelConfig +from nmr.nmr_gym.io import load_histories +from torch_geometric.loader import DataLoader + + +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 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 = [] + + 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) + + return graphs + + +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)", + ) + parser.add_argument( + "--num-nmr-layers", + type=int, + default=1, + help="Number of NMR layers in the model (default: 1)", + ) + + args = parser.parse_args() + + # Set device + device = args.device + + # Load data + examples = extract_data(args.histories) + 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, config) + opt = torch.optim.AdamW(net.parameters(), lr=args.learning_rate, weight_decay=0.01) + + iteration = 0 + for epoch in range(epochs): + + for xs in data_loader: + net.train() + iteration += 1 + + _, policies = net(xs) + + # 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 + 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 + + 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() diff --git a/scripts/visualize_data.py b/scripts/visualize_data.py new file mode 100644 index 0000000..9c8bde4 --- /dev/null +++ b/scripts/visualize_data.py @@ -0,0 +1,149 @@ +from typing import NamedTuple +import matplotlib.pyplot as plt +import numpy as np +from mpl_toolkits.mplot3d import Axes3D +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] + + return actual_x, actual_y, predict_x, predict_y, contacts + + def plot_shifts(self, named_tuple_used=True): + + actual_x, actual_y, predict_x, predict_y, contacts = self.get_data(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(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)) + + # 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/scripts/write_pdb.py b/scripts/write_pdb.py new file mode 100644 index 0000000..22cff6d --- /dev/null +++ b/scripts/write_pdb.py @@ -0,0 +1,69 @@ +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 + + + +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) diff --git a/tests/test_energy.py b/tests/test_energy.py new file mode 100644 index 0000000..74409ed --- /dev/null +++ b/tests/test_energy.py @@ -0,0 +1,95 @@ +import unittest +import numpy as np +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.energy import Energy + + + +class TestEnergyMethods(unittest.TestCase): + + def setUp(self): + 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 NOE activation - if activated, the expected value will be the lowest energy calculated + def test_energy_activated(self): + restraint = [(2, 1), (0, 3)] + assignments = {2: 1, 0: 3, 1: 2, 3: 4} + + result = self.energy.calc_restraint_energy(assignments, restraint, self.dist_grid, 0.5) + expected = 0.0 + + self.assertEqual(result, expected) + + # Test if NOE not activated - if not activated, the expected value will be 0 (no energy calculated) + def test_energy_not_activated(self): + restraints = [[(2, 1), (0, 3)]] + assignments = {0: 3} + + result = self.energy.get_total_energy(restraints, assignments, self.dist_grid) + expected = 0 + + self.assertEqual(result, expected) + + # 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_total_energy(restraints, assignments, self.dist_grid) + expected = 0.01019237886466845 + + self.assertEqual(result, expected) + + # 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 = 0.5 + x = 5 # distance + result = self.energy.flat_bottom(x, tolerance) + expected = 22 # x^2 - tolerance*x + self.assertEqual(result, expected) + + + +if __name__ == '__main__': + unittest.main() + + ''' + 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) + + ''' 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/tests/test_fractional_activation.py b/tests/test_fractional_activation.py new file mode 100644 index 0000000..bcb1ac2 --- /dev/null +++ b/tests/test_fractional_activation.py @@ -0,0 +1,54 @@ +import unittest +from itertools import chain +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.assignment_order import FractionalActivation + + + +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): + + 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'} + ] + + 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']) + + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file 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_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/tests/test_gym_env.py b/tests/test_gym_env.py new file mode 100644 index 0000000..020c6d1 --- /dev/null +++ b/tests/test_gym_env.py @@ -0,0 +1,92 @@ +import unittest +import numpy as np +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.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 = { + "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) + + for i, action in enumerate(actions): + observation, reward, terminated, total_energy = 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) + + for i, action in enumerate(actions): + observation, reward, terminated, total_energy = 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/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/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()