Skip to content

Model Module

FineTunerGNN

Bases: Module

The main Graph Neural Network model for structural refinement.

It consists of several GNN layers followed by an output head that predicts delta dihedrals for each residue.

Source code in torsiontuner/model.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
class FineTunerGNN(eqx.Module):
    """
    The main Graph Neural Network model for structural refinement.

    It consists of several GNN layers followed by an output head that predicts
    delta dihedrals for each residue.
    """

    layers: list
    output_head: eqx.nn.Linear

    def __init__(
        self, node_dim: int, hidden_dim: int, out_dim: int, n_layers: int, key: jnp.ndarray
    ) -> None:
        keys = jr.split(key, n_layers + 1)
        self.layers = []
        curr_dim = node_dim
        for i in range(n_layers):
            self.layers.append(GNNLayer(curr_dim, hidden_dim, key=keys[i]))
            curr_dim = hidden_dim

        self.output_head = eqx.nn.Linear(hidden_dim, out_dim, key=keys[-1])

    def __call__(self, x: jnp.ndarray, adj: jnp.ndarray, edge_features: jnp.ndarray) -> jnp.ndarray:
        """
        Predict structural adjustments.

        Args:
            x: Input node features (one-hot residue types).
            adj: Adjacency matrix (sequential + spatial).
            edge_features: Edge features (normalized distances).

        Returns:
            Predicted delta dihedrals (delta_phi, delta_psi) for each residue.
        """
        for layer in self.layers:
            x = layer(x, adj, edge_features)

        # Predict delta dihedrals
        out = jax.vmap(self.output_head)(x)
        return out

__call__(x, adj, edge_features)

Predict structural adjustments.

Parameters:

Name Type Description Default
x ndarray

Input node features (one-hot residue types).

 required
adj ndarray

Adjacency matrix (sequential + spatial).

 required
edge_features ndarray

Edge features (normalized distances).

 required

Returns:

Type Description
ndarray

Predicted delta dihedrals (delta_phi, delta_psi) for each residue.
Source code in torsiontuner/model.py 
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def __call__(self, x: jnp.ndarray, adj: jnp.ndarray, edge_features: jnp.ndarray) -> jnp.ndarray:
    """
    Predict structural adjustments.

    Args:
        x: Input node features (one-hot residue types).
        adj: Adjacency matrix (sequential + spatial).
        edge_features: Edge features (normalized distances).

    Returns:
        Predicted delta dihedrals (delta_phi, delta_psi) for each residue.
    """
    for layer in self.layers:
        x = layer(x, adj, edge_features)

    # Predict delta dihedrals
    out = jax.vmap(self.output_head)(x)
    return out

GNNLayer

Bases: Module

A single Graph Neural Network layer that processes node and edge features.

This layer uses message passing to update node representations by aggregating information from neighbors, weighted by edge features (e.g., distances).

Source code in torsiontuner/model.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
class GNNLayer(eqx.Module):
    """
    A single Graph Neural Network layer that processes node and edge features.

    This layer uses message passing to update node representations by aggregating
    information from neighbors, weighted by edge features (e.g., distances).
    """

    lin_node: eqx.nn.Linear
    lin_edge: eqx.nn.Linear
    lin_out: eqx.nn.Linear

    def __init__(self, node_dim: int, out_dim: int, key: jnp.ndarray) -> None:
        keys = jr.split(key, 3)
        self.lin_node = eqx.nn.Linear(node_dim, out_dim, key=keys[0])
        self.lin_edge = eqx.nn.Linear(1, out_dim, key=keys[1])  # 1D edge feature: distance
        self.lin_out = eqx.nn.Linear(out_dim, out_dim, key=keys[2])

    def __call__(self, x: jnp.ndarray, adj: jnp.ndarray, edge_features: jnp.ndarray) -> jnp.ndarray:
        """
        Forward pass of the GNN layer.

        Args:
            x: Node features of shape (n_nodes, node_dim).
            adj: Adjacency matrix of shape (n_nodes, n_nodes).
            edge_features: Edge features of shape (n_nodes, n_nodes, 1).

        Returns:
            Updated node features of shape (n_nodes, out_dim).
        """
        h = jax.vmap(self.lin_node)(x)

        # Message passing with edge features
        # We can apply a linear transformation to edge features and use them as weights
        # or add them to the node messages.
        # Here: m_i = sum_j (adj_ij * (h_j + lin_edge(e_ij)))

        # Project edge features
        e_proj = jax.vmap(jax.vmap(self.lin_edge))(edge_features)  # (n, n, out_dim)

        # Combine messages
        # h[None, :, :] has shape (1, n, out_dim)
        # e_proj has shape (n, n, out_dim)
        messages = (h[None, :, :] + e_proj) * adj[:, :, None]

        m = jnp.sum(messages, axis=1)

        # Update
        out = jax.nn.relu(h + m)
        return out

__call__(x, adj, edge_features)

Forward pass of the GNN layer.

Parameters:

Name Type Description Default
x ndarray

Node features of shape (n_nodes, node_dim).

 required
adj ndarray

Adjacency matrix of shape (n_nodes, n_nodes).

 required
edge_features ndarray

Edge features of shape (n_nodes, n_nodes, 1).

 required

Returns:

Type Description
ndarray

Updated node features of shape (n_nodes, out_dim).
Source code in torsiontuner/model.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def __call__(self, x: jnp.ndarray, adj: jnp.ndarray, edge_features: jnp.ndarray) -> jnp.ndarray:
    """
    Forward pass of the GNN layer.

    Args:
        x: Node features of shape (n_nodes, node_dim).
        adj: Adjacency matrix of shape (n_nodes, n_nodes).
        edge_features: Edge features of shape (n_nodes, n_nodes, 1).

    Returns:
        Updated node features of shape (n_nodes, out_dim).
    """
    h = jax.vmap(self.lin_node)(x)

    # Message passing with edge features
    # We can apply a linear transformation to edge features and use them as weights
    # or add them to the node messages.
    # Here: m_i = sum_j (adj_ij * (h_j + lin_edge(e_ij)))

    # Project edge features
    e_proj = jax.vmap(jax.vmap(self.lin_edge))(edge_features)  # (n, n, out_dim)

    # Combine messages
    # h[None, :, :] has shape (1, n, out_dim)
    # e_proj has shape (n, n, out_dim)
    messages = (h[None, :, :] + e_proj) * adj[:, :, None]

    m = jnp.sum(messages, axis=1)

    # Update
    out = jax.nn.relu(h + m)
    return out