From fb7de70f2048578e8f9fd8eba063e9fed17f5881 Mon Sep 17 00:00:00 2001 From: Rafael Lopes Date: Tue, 7 Jul 2026 02:17:36 +0300 Subject: [PATCH 1/2] chg: update dimention creaation --- frame/tune.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/frame/tune.py b/frame/tune.py index 2771e2e..6696cee 100644 --- a/frame/tune.py +++ b/frame/tune.py @@ -127,6 +127,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: @@ -218,12 +232,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", From 7a155b7fe0712f34f7712fda99b1e59b59c2d6a3 Mon Sep 17 00:00:00 2001 From: Rafael Lopes Date: Tue, 7 Jul 2026 04:55:38 +0300 Subject: [PATCH 2/2] chg: add PNA --- frame/evaluate.py | 1 + frame/explain.py | 1 + frame/generate.py | 45 ++++++++++++++++++++++++++++- frame/source/models/__init__.py | 14 +++++++++ frame/source/models/pyg_models.py | 48 ++++++++++++++++++++++++++++++- frame/train.py | 2 ++ frame/tune.py | 2 ++ parameters.yaml | 9 +++++- pyproject.toml | 37 ++++++++++++++---------- 9 files changed, 141 insertions(+), 18 deletions(-) diff --git a/frame/evaluate.py b/frame/evaluate.py index f626e79..fac6f2c 100644 --- a/frame/evaluate.py +++ b/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"] diff --git a/frame/explain.py b/frame/explain.py index 9b20d8e..09ec4ec 100644 --- a/frame/explain.py +++ b/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 diff --git a/frame/generate.py b/frame/generate.py index c61a026..464990a 100644 --- a/frame/generate.py +++ b/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") diff --git a/frame/source/models/__init__.py b/frame/source/models/__init__.py index 4c9c073..d1c97ae 100644 --- a/frame/source/models/__init__.py +++ b/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 diff --git a/frame/source/models/pyg_models.py b/frame/source/models/pyg_models.py index aa5e9ae..532c121 100644 --- a/frame/source/models/pyg_models.py +++ b/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__() diff --git a/frame/train.py b/frame/train.py index 1fe18f2..9d190a1 100644 --- a/frame/train.py +++ b/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 / "???"): diff --git a/frame/tune.py b/frame/tune.py index 6696cee..9253470 100644 --- a/frame/tune.py +++ b/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)) @@ -203,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 / "???"): diff --git a/parameters.yaml b/parameters.yaml index b9972a6..66ca0eb 100755 --- a/parameters.yaml +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 0c81a55..427789a 100644 --- a/pyproject.toml +++ b/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