New: add PNA

frame/evaluate.py

@@ -41,6 +41,7 @@ def main():
     tune["feat_size"] = data["metadata"]["feat_size"]
     tune["edge_dim"] = data["metadata"]["edge_dim"]
     tune["bce_weight"] = data["metadata"]["bce_weight"]
+    tune["deg"] = data["metadata"].get("deg")
 
     # * Prepare dataloader
     test_data = [data for data in dataset if data.set == "test"]

frame/explain.py

@@ -97,6 +97,7 @@ def _model_config(tune: dict, metadata: dict, task: str, params: dict):
     cfg["feat_size"] = metadata["feat_size"]
     cfg["edge_dim"] = metadata["edge_dim"]
     cfg["bce_weight"] = metadata["bce_weight"]
+    cfg["deg"] = metadata.get("deg")
     cfg["task"] = task
     cfg["regression_loss"] = params["Data"].get("regression_loss", "mse")
     return cfg

frame/generate.py

@@ -49,6 +49,42 @@ def _write_dataset_stats(dataset, loader: str, path_csv: str,
         json.dump(stats, fh, indent=2)
 
 
+def _degree_histogram(dataset, max_bins: int = 512):
+    """Train-split node in-degree histogram for PNA degree scalers.
+
+    Accumulates, over the training graphs (data.set == "train"),
+    how many nodes have each in-degree. The node count comes from
+    data.x (the definitive atom set) rather than num_nodes,
+    and destination indices are clamped into range, so each graph
+    allocates only per-graph-sized tensors. Per-node degrees are
+    capped at max_bins - 1 to bound the histogram length. Edgeless
+    graphs (the (2, 0) fallback) are skipped. Falls back to the
+    full dataset if no graph is labelled "train".
+
+    Args:
+        dataset: The in-memory dataset of PyG Data objects.
+        max_bins: Hard cap on the histogram length; degrees at or
+            above it fall in the final bin. Defaults to 512.
+    """
+    graphs = [d for d in dataset if getattr(d, "set", None) == "train"]
+    if not graphs:
+        graphs = list(dataset)
+    deg = torch.zeros(1, dtype=torch.long)
+    for data in graphs:
+        if data.edge_index.numel() == 0:
+            continue
+        n = int(data.x.size(0))
+        dst = data.edge_index[1].clamp(min=0, max=n - 1)
+        node_deg = torch.bincount(dst, minlength=n).clamp(max=max_bins - 1)
+        counts = torch.bincount(node_deg, minlength=deg.numel())
+        if counts.numel() > deg.numel():
+            counts[:deg.numel()] += deg
+            deg = counts
+        else:
+            deg += counts
+    return deg
+
+
 def main():
     args_parser = argparse.ArgumentParser()
     args_parser.add_argument("-c", "--config", dest="config", required=True)
@@ -87,7 +123,14 @@ def main():
                 "edge_dim": dataset.num_edge_features,
                 "bce_weight": bce_weight,
                 "loader": loader,
-                "project_dir": project_dir}
+                "project_dir": project_dir,
+                "deg": _degree_histogram(dataset)}
+
+    # Iterating the dataset (stats, degree histogram) fills PyG's
+    # InMemoryDataset _data_list cache with per-graph tensor views
+    # into the collated store; persisting that cache would bloat the
+    # joblib ~N-fold. Reset it so only the compact store is serialized.
+    dataset._data_list = None
 
     dump_data = {"dataset": dataset, "metadata": metadata}
     joblib.dump(dump_data, project_dir / "data.joblib")

frame/source/models/__init__.py

@@ -71,6 +71,8 @@ def select_model(model_name, config):
         model = pyg_models.GNN_GIN(config).to(device)
     elif model_name == "gcn":
         model = pyg_models.GNN_GCN(config).to(device)
+    elif model_name == "pna":
+        model = pyg_models.GNN_PNA(config).to(device)
     else:
         raise NotImplementedError("Model not available")
 
@@ -84,6 +86,8 @@ def _cast_value(val):
         return int(val)
     if isinstance(val, float):
         return float(val)
+    if isinstance(val, (list, tuple)):
+        return list(val)
     return str(val)
 
 
@@ -127,4 +131,14 @@ def optuna_suggest(params, trial):
             trial.set_user_attr("hidden_channels_suggested", original)
             trial.set_user_attr("hidden_channels_used", rounded)
 
+    # Round hidden_channels to match pna_towers, and log
+    if model_name == "pna":
+        towers = int(configs.get("pna_towers", 1))
+        original = int(configs["hidden_channels"])
+        rounded = (original // towers) * towers
+        if rounded != original:
+            configs["hidden_channels"] = rounded
+            trial.set_user_attr("hidden_channels_suggested", original)
+            trial.set_user_attr("hidden_channels_used", rounded)
+
     return configs

frame/source/models/pyg_models.py

@@ -1,6 +1,6 @@
 import torch
 from torch_geometric.nn import (GCN, GraphSAGE, GIN, GAT,
-                                AttentiveFP,
+                                AttentiveFP, PNA,
                                 global_mean_pool,
                                 global_add_pool,
                                 global_max_pool)
@@ -122,6 +122,52 @@ def forward(self, x, edge_index, edge_attr, batch):
         return x_pool
 
 
+class GNN_PNA(torch.nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        in_channels = config.get("feat_size", None)
+        hidden_channels = config.get("hidden_channels", 64)
+        num_layers = config.get("num_layers", 2)
+        dropout = config.get("dropout_rate", 0.4)
+        edge_dim = config.get("edge_dim", None)
+        towers = int(config.get("pna_towers", 1))
+        aggregators = config.get("pna_aggregators",
+                                 ["mean", "min", "max", "std"])
+        scalers = config.get("pna_scalers",
+                             ["identity", "amplification", "attenuation"])
+
+        deg = config.get("deg", None)
+        if deg is None:
+            raise ValueError("PNA requires a degree histogram. Re-run "
+                             "frame_gen so metadata stores 'deg'.")
+        deg = torch.as_tensor(deg, dtype=torch.long)
+
+        rounded = (hidden_channels // towers) * towers
+        if rounded == 0:
+            raise ValueError(f"hidden_channels ({hidden_channels}) must "
+                             f">= pna_towers ({towers}).")
+        hidden_channels = rounded
+
+        self.pool = _resolve_pool(config.get("pool", "mean"))
+        self.model = PNA(in_channels=in_channels,
+                         hidden_channels=hidden_channels,
+                         num_layers=num_layers,
+                         out_channels=hidden_channels,
+                         dropout=dropout,
+                         aggregators=aggregators,
+                         scalers=scalers,
+                         deg=deg,
+                         edge_dim=edge_dim,
+                         towers=towers)
+        self.head = torch.nn.Linear(hidden_channels, 1)
+
+    def forward(self, x, edge_index, edge_attr, batch):
+        x = self.model(x, edge_index, edge_attr=edge_attr)
+        x_pool = self.pool(x, batch)
+
+        return self.head(x_pool)
+
+
 class GNN_AttentiveFP(torch.nn.Module):
     def __init__(self, config):
         super().__init__()

frame/train.py

@@ -75,6 +75,7 @@ def run(params, dataset):
     config["regression_loss"] = params["Data"].get("regression_loss", "mse")
     config["huber_delta"] = params["Data"].get("huber_delta", 1.0)
     config["warmup_epochs"] = int(params["Data"].get("warmup_epochs", 0))
+    config["deg"] = params["Data"].get("deg")
     params["Data"]["trial"] = None
 
     size = int(config.get("batch_size", size))
@@ -139,6 +140,7 @@ def main():
     params["Data"]["feat_size"] = data["metadata"]["feat_size"]
     params["Data"]["edge_dim"] = data["metadata"]["edge_dim"]
     params["Data"]["bce_weight"] = data["metadata"]["bce_weight"]
+    params["Data"]["deg"] = data["metadata"].get("deg")
     params["Data"]["project_dir"] = project_dir
 
     if os.path.isdir(cwd / "???"):

frame/tune.py

@@ -74,6 +74,7 @@ def objective(trial, params, dataset):
     config["regression_loss"] = params["Data"].get("regression_loss", "mse")
     config["huber_delta"] = params["Data"].get("huber_delta", 1.0)
     config["warmup_epochs"] = int(params["Data"].get("warmup_epochs", 0))
+    config["deg"] = params["Data"].get("deg")
     params["Data"]["trial"] = trial
     size = int(config.get("batch_size", size))
 
@@ -127,6 +128,20 @@ def objective(trial, params, dataset):
                 raise optuna.exceptions.TrialPruned()
 
 
+def _build_dimension(col: str, series: pd.Series):
+    if pd.api.types.is_numeric_dtype(series):
+        values = series.values
+        return dict(label=col, values=values,
+                    range=[float(np.nanmin(values)),
+                           float(np.nanmax(values))])
+
+    codes, uniques = pd.factorize(series, sort=True)
+    return dict(label=col, values=codes,
+                range=[0, max(len(uniques) - 1, 1)],
+                tickvals=list(range(len(uniques))),
+                ticktext=[str(u) for u in uniques])
+
+
 def get_dataframe(study, task):
     records = []
     for trial in study.trials:
@@ -189,6 +204,7 @@ def main():
     params["Data"]["feat_size"] = data["metadata"]["feat_size"]
     params["Data"]["edge_dim"] = data["metadata"]["edge_dim"]
     params["Data"]["bce_weight"] = data["metadata"]["bce_weight"]
+    params["Data"]["deg"] = data["metadata"].get("deg")
     params["Data"]["project_dir"] = project_dir
 
     if os.path.isdir(cwd / "???"):
@@ -218,12 +234,7 @@ def main():
     feats = [col for col in list(df.columns) if col not in header]
     feats = feats + ["optim"]
 
-    dimensions = []
-    for col in feats:
-        col_values = df[col].values
-        dim = dict(label=col, values=col_values,
-                   range=[col_values.min(), col_values.max()])
-        dimensions.append(dim)
+    dimensions = [_build_dimension(col, df[col]) for col in feats]
 
     fig = go.Figure(data=go.Parcoords(line=dict(color=df["optim"],
                                                 colorscale="viridis",

parameters.yaml

@@ -7,7 +7,7 @@ Data:
 
   trials: 50                           # number of Optuna trials
   epochs: 50                           # max epochs per run
-  model: gat                           # gcn | gat | gin | sage | attentive
+  model: gat                           # gcn | gat | gin | sage | attentive | pna
   batch_size: 128                      # graphs per mini-batch
   patience: 15                         # early-stopping patience (epochs)
   loader: default                      # default | decompose
@@ -48,6 +48,13 @@ Tune:
     min: 1
     max: 4
     # value: 2
+  pna_aggregators:                     # PNA neighbour aggregators
+    value: [mean, min, max, std]
+  pna_scalers:                         # PNA degree scalers
+    value: [identity, amplification, attenuation]
+  pna_towers:                          # PNA towers (auto-rounded)
+    # choices: [1, 2, 4]
+    value: 1
 
   learning_rate:
     # min: 1.0e-4

pyproject.toml

@@ -5,23 +5,23 @@ build-backend = "hatchling.build"
 [project]
 name = "frame"
 version = "0.2.0"
-requires-python = ">=3.11"
+requires-python = ">=3.12"
 authors = [{name = "Rafael Lopes", email = "rafael.lopes@uef.fi"}]
 readme = "README.md"
 description = ""
-dependencies = ["torch==2.11.0",
+dependencies = ["torch==2.12.1",
+                "torch-scatter==2.1.2",
                 "captum==0.9.0",
                 "joblib==1.5.3",
-                "optuna==4.8.0",
-                "pandas==3.0.2",
-                "plotly==6.7.0",
-                "rdkit==2026.3.1",
-                "optuna==4.8.0",
-                "pandas==3.0.2",
-                "scikit-learn==1.8.0",
-                "scipy==1.17.1",
+                "optuna==4.9.0",
+                "pandas==3.0.3",
+                "plotly==6.8.0",
+                "rdkit==2026.3.3",
+                "scikit-learn==1.9.0",
+                "scipy==1.18.0",
                 "svgutils==0.3.4",
-                "torch-geometric==2.7.0"]
+                "torch-geometric==2.8.0",
+                "torch-scatter==2.1.2+pt212cu132"]
 
 [project.scripts]
 frame_tune = "frame.tune:main"
@@ -41,16 +41,23 @@ packages = ["frame"]
 
 [tool.uv]
 package = true
-find-links = ["https://download.pytorch.org/whl/cu128"]
+find-links = ["https://download.pytorch.org/whl/cu132"]
 
 [tool.uv.sources]
-torch = { index = "pytorch-cu128" }
+torch = { index = "pytorch-cu132" }
+torch-scatter = { index = "pyg-cu132" }
 
 [[tool.uv.index]]
-name = "pytorch-cu128"
-url = "https://download.pytorch.org/whl/cu128"
+name = "pytorch-cu132"
+url = "https://download.pytorch.org/whl/cu132"
 explicit = true
 
+[[tool.uv.index]]
+name = "pyg-cu132"
+url = "https://data.pyg.org/whl/torch-2.12.1+cu132.html"
+explicit = true
+format = "flat"
+
 [tool.bumpversion]
 current_version = "0.2.0"
 commit = true