Skip to content

Model

Graph neural network architecture for molecular property regression.

Module

project_name.model

GraphNeuralNetwork

Bases: Module

Graph Neural Network for molecular property regression.

Source code in src/project_name/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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
class GraphNeuralNetwork(nn.Module):
    """Graph Neural Network for molecular property regression."""

    def __init__(
        self,
        num_node_features: int = 11,
        num_edge_features: int = 4,
        hidden_dim: int = 128,
        num_layers: int = 3,
        output_dim: int = 1,
        dropout: float = 0.1,
    ) -> None:
        """Initialize the GNN model.

        Args:
            num_node_features: Number of node (atom) features.
            num_edge_features: Number of edge (bond) features.
            hidden_dim: Number of hidden channels.
            num_layers: Number of GraphConv layers.


            output_dim: Output dimension (1 for single property regression).
            dropout: Dropout rate for regularization.
        """
        super().__init__()
        self.dropout_rate = dropout

        self.initial_embedding = nn.Linear(num_node_features, hidden_dim)

        self.conv_layers = nn.ModuleList([GraphConv(hidden_dim, hidden_dim) for _ in range(num_layers)])

        self.pool = global_mean_pool

        self.mlp = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, output_dim),
        )

    def forward(self, data) -> torch.Tensor:
        """Forward pass through the model.

        Args:
            data: PyTorch Geometric Data object with x, edge_index, and batch attributes.

        Returns:
            Predicted property values.
        """
        x, edge_index, batch = data.x, data.edge_index, data.batch

        x = self.initial_embedding(x)
        x = F.relu(x)

        for conv in self.conv_layers:
            x = conv(x, edge_index)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout_rate, training=self.training)

        x = self.pool(x, batch)

        x = self.mlp(x)

        return x

__init__ 

__init__(num_node_features: int = 11, num_edge_features: int = 4, hidden_dim: int = 128, num_layers: int = 3, output_dim: int = 1, dropout: float = 0.1) -> None

Initialize the GNN model.

Parameters:

Name Type Description Default
num_node_features int

Number of node (atom) features.

 11
num_edge_features int

Number of edge (bond) features.

 4
hidden_dim int

Number of hidden channels.

 128
num_layers int

Number of GraphConv layers.

 3
output_dim int

Output dimension (1 for single property regression).

 1
dropout float

Dropout rate for regularization.

 0.1
Source code in src/project_name/model.py 
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
def __init__(
    self,
    num_node_features: int = 11,
    num_edge_features: int = 4,
    hidden_dim: int = 128,
    num_layers: int = 3,
    output_dim: int = 1,
    dropout: float = 0.1,
) -> None:
    """Initialize the GNN model.

    Args:
        num_node_features: Number of node (atom) features.
        num_edge_features: Number of edge (bond) features.
        hidden_dim: Number of hidden channels.
        num_layers: Number of GraphConv layers.


        output_dim: Output dimension (1 for single property regression).
        dropout: Dropout rate for regularization.
    """
    super().__init__()
    self.dropout_rate = dropout

    self.initial_embedding = nn.Linear(num_node_features, hidden_dim)

    self.conv_layers = nn.ModuleList([GraphConv(hidden_dim, hidden_dim) for _ in range(num_layers)])

    self.pool = global_mean_pool

    self.mlp = nn.Sequential(
        nn.Linear(hidden_dim, hidden_dim),
        nn.ReLU(),
        nn.Dropout(dropout),
        nn.Linear(hidden_dim, hidden_dim // 2),
        nn.ReLU(),
        nn.Dropout(dropout),
        nn.Linear(hidden_dim // 2, output_dim),
    )

forward 

forward(data) -> torch.Tensor

Forward pass through the model.

Parameters:

Name Type Description Default
data

PyTorch Geometric Data object with x, edge_index, and batch attributes.

 required

Returns:

Type Description
Tensor

Predicted property values.
Source code in src/project_name/model.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def forward(self, data) -> torch.Tensor:
    """Forward pass through the model.

    Args:
        data: PyTorch Geometric Data object with x, edge_index, and batch attributes.

    Returns:
        Predicted property values.
    """
    x, edge_index, batch = data.x, data.edge_index, data.batch

    x = self.initial_embedding(x)
    x = F.relu(x)

    for conv in self.conv_layers:
        x = conv(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=self.dropout_rate, training=self.training)

    x = self.pool(x, batch)

    x = self.mlp(x)

    return x