diff --git a/evaluate_inductive_ood.py b/evaluate_inductive_ood.py new file mode 100644 index 0000000..f6a351f --- /dev/null +++ b/evaluate_inductive_ood.py @@ -0,0 +1,236 @@ +"""Evaluate whether POME's inductive (transform) embeddings of unseen samples +land in the same distribution as the transductive embeddings of training samples. + +For each hancock train/test split we: + 1. fit() an Embedder on the training split, + 2. read off transductive training embeddings via get_embeddings(), + 3. compute inductive embeddings for the test split via transform(), + 4. compute inductive embeddings for the *training* samples via transform() + (the method-effect control), + 5. run scxmatch's Rosenbaum cross-match test (squared euclidean, full + distance matrix) on: + - MAIN: train (transductive) vs test (inductive) + - CONTROL: train (transductive) vs train (inductive) + +Interpretation (per split): + - high p-value (>~0.05) => the two groups are indistinguishable + (in-distribution). + - low p-value + negative z => groups are separable (out-of-distribution). + - If the CONTROL is also low, the MAIN separation is at least partly an + inductive-vs-transductive *method artifact* rather than genuine data OOD; + the gap between MAIN and CONTROL isolates the genuine data-OOD component. + +Note: transform()'s KNN initialisation for a training sample includes that +sample's own neighbours, biasing inductive train embeddings toward their +transductive counterparts. The CONTROL is therefore a conservative +(lower-bound) estimate of the method effect. + +Run (per CLAUDE.md): + conda activate torch + python evaluate_inductive_ood.py # full run, 10 splits, 1000 epochs + python evaluate_inductive_ood.py --splits 1 --epochs 50 # quick smoke test +""" + +import argparse +from pathlib import Path + +import anndata as ad +import numpy as np +import pandas as pd +import torch + +import scxmatch +from pome.gnn_embedding import Embedder, make_deterministic + +# --- Configuration ----------------------------------------------------------- +SPLITS_DIR = Path(__file__).resolve().parent / "data" / "splits" / "hancock" +RESULTS_CSV = SPLITS_DIR / "ood_evaluation_results.csv" +SUMMARY_PNG = SPLITS_DIR / "ood_evaluation_summary.png" + +N_SPLITS = 10 +SEED = 42 + +EMBEDDING_DIMENSION = 32 +EPOCHS = 1000 +BINS_PER_CONTINUOUS = 15 +DISCRETIZATION_TYPE = "z" +NA_ENCODING = -99.0 + +METRIC = "sqeuclidean" +K = None # full distance matrix (exact Rosenbaum cross-match) + + +def load_graph(path: Path) -> pd.DataFrame: + """Load a graph-format split: rows = variables, cols = samples + 'type'.""" + return pd.read_csv(path, sep="\t", index_col=0) + + +def run_xmatch(emb_a: pd.DataFrame, label_a: str, + emb_b: pd.DataFrame, label_b: str) -> dict: + """Cross-match test between two embedding groups. + + emb_a / emb_b are (n_samples, dim) DataFrames. Returns the scxmatch result + dict augmented with the two group sizes. + """ + # Make obs names globally unique (the control reuses the same sample names + # in both groups, which AnnData would otherwise reject / silently collide). + names_a = [f"{label_a}::{s}" for s in emb_a.index] + names_b = [f"{label_b}::{s}" for s in emb_b.index] + + X = np.vstack([emb_a.to_numpy(), emb_b.to_numpy()]).astype(np.float64) + obs = pd.DataFrame( + {"group": [label_a] * len(emb_a) + [label_b] * len(emb_b)}, + index=names_a + names_b, + ) + adata = ad.AnnData(X=X, obs=obs) + + result = scxmatch.test( + adata, + group_by="group", + test_group=label_b, + reference=label_a, + metric=METRIC, + k=K, + ) + result = dict(result) + result["n_reference"] = len(emb_a) + result["n_test"] = len(emb_b) + return result + + +def evaluate_split(split_id: int, epochs: int, device: str) -> list[dict]: + """Fit on one split and run the MAIN and CONTROL cross-match tests. + + Embeddings for unseen samples use the (directed) inductive transform(). + """ + tag = f"split_{split_id:02d}" + train_path = SPLITS_DIR / f"{tag}_train_graph.tsv" + test_path = SPLITS_DIR / f"{tag}_test_graph.tsv" + + make_deterministic(SEED) + + train_df = load_graph(train_path) + test_df = load_graph(test_path) + + embedder = Embedder( + embedding_dimension=EMBEDDING_DIMENSION, + epochs=epochs, + bins_per_continuous=BINS_PER_CONTINUOUS, + discretization_type=DISCRETIZATION_TYPE, + na_encoding=NA_ENCODING, + device=device, + ) + embedder.fit(train_df) + + train_emb, *_ = embedder.get_embeddings() # transductive (train) + test_emb = embedder.transform(test_df) # inductive (test) + train_ind_emb = embedder.transform(train_df) # inductive (train) -> control + + print( + f" {tag}: train_emb {train_emb.shape}, " + f"test_emb {test_emb.shape}, train_ind_emb {train_ind_emb.shape}" + ) + + main = run_xmatch(train_emb, "train", test_emb, "test") + control = run_xmatch(train_emb, "train_transductive", + train_ind_emb, "train_inductive") + + rows = [] + for comparison, res in (("main", main), ("control", control)): + rows.append({ + "split": split_id, + "comparison": comparison, + "n_reference": res["n_reference"], + "n_test": res["n_test"], + "p_value": res["p_value"], + "z_score": res["z_score"], + "coverage": res["coverage"], + "effect_strength_ratio": res["effect_strength_ratio"], + }) + print( + f" {comparison:8s} p={res['p_value']:.4g} " + f"z={res['z_score']:.3f} coverage={res['coverage']:.2%}" + ) + return rows + + +def summarize(df: pd.DataFrame) -> None: + print("\n=== Summary (mean +/- std across splits) ===") + print(" Higher p-value / less-negative z = more in-distribution (better).") + for comparison, sub in df.groupby("comparison"): + n_sig = int((sub["p_value"] < 0.05).sum()) + print( + f" {comparison:8s} " + f"p_value = {sub['p_value'].mean():.4g} +/- {sub['p_value'].std():.4g} " + f"z_score = {sub['z_score'].mean():.3f} +/- {sub['z_score'].std():.3f} " + f"sig(p<0.05) {n_sig}/{len(sub)}" + ) + main = df[df["comparison"] == "main"] + n_sig = int((main["p_value"] < 0.05).sum()) + print( + f"\n MAIN comparison significant (p<0.05, i.e. test embeddings OOD) " + f"in {n_sig}/{len(main)} splits." + ) + print( + " If the CONTROL is not significant, the inductive transform introduces no " + "method artifact, so any MAIN separation reflects genuine data OOD." + ) + + +def make_plot(df: pd.DataFrame) -> None: + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + splits = sorted(df["split"].unique()) + main = df[df["comparison"] == "main"].set_index("split").reindex(splits) + control = df[df["comparison"] == "control"].set_index("split").reindex(splits) + + x = np.arange(len(splits)) + width = 0.4 + + fig, ax = plt.subplots(figsize=(max(6, len(splits) * 0.9), 4.5)) + ax.bar(x - width / 2, main["p_value"], width, label="main (train vs test)") + ax.bar(x + width / 2, control["p_value"], width, + label="control (train transductive vs inductive)") + ax.axhline(0.05, color="red", linestyle="--", linewidth=1, label="p = 0.05") + ax.set_xticks(x) + ax.set_xticklabels([f"{s:02d}" for s in splits]) + ax.set_xlabel("split") + ax.set_ylabel("cross-match p-value") + ax.set_title("POME inductive-embedding OOD test (directed, scxmatch sqeuclidean)") + ax.legend(fontsize=8) + fig.tight_layout() + fig.savefig(SUMMARY_PNG, dpi=150) + print(f"\nSaved plot to {SUMMARY_PNG}") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--splits", type=int, default=N_SPLITS, + help="number of splits to evaluate (default: all 10)") + parser.add_argument("--epochs", type=int, default=EPOCHS, + help="training epochs per split (default: 1000)") + args = parser.parse_args() + + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"Device: {device} | epochs: {args.epochs} | splits: {args.splits}") + + all_rows = [] + for split_id in range(args.splits): + print(f"\n[split {split_id:02d}] fitting Embedder ...") + all_rows.extend(evaluate_split(split_id, args.epochs, device)) + + df = pd.DataFrame(all_rows) + df.to_csv(RESULTS_CSV, index=False) + print(f"\nWrote per-split results to {RESULTS_CSV}") + + summarize(df) + try: + make_plot(df) + except Exception as exc: # plotting is non-essential + print(f"(plot skipped: {exc})") + + +if __name__ == "__main__": + main() diff --git a/src/pome/gnn_embedding.py b/src/pome/gnn_embedding.py index 54636f7..b08b172 100644 --- a/src/pome/gnn_embedding.py +++ b/src/pome/gnn_embedding.py @@ -10,7 +10,7 @@ from sklearn.base import BaseEstimator, ClassifierMixin import joblib from pome.models import GraphAutoencoder -from pome.utils import compute_roc, repeat_pad_to_max_cols, bin_column_non_linear, bin_column_with_na_adjusted +from pome.utils import compute_roc, repeat_pad_to_max_cols, bin_column_non_linear, bin_column_with_na_adjusted, signed_power_bins, get_zscore_bins def make_deterministic(seed=42): # 1. Basic seeding @@ -49,7 +49,6 @@ def __init__(self, self.device = device self.lr = lr self.layer_type = layer_type - self.num_nodes = None self.model = None self.type_column = type_column self.non_informative_na = na_encoding @@ -72,7 +71,7 @@ def fit(self, X, y=None): """ # Translate input data into graph based format. input_data = X.copy() - discrete_data, graph_data, _, sample_node_dict, value_node_dict, var_value_dict, neg_edges_per_pair, variable_names, variable_embedding_ids, cont_bin_names, cont_bin_embedding_ids = self.data_to_graph( + discrete_data, graph_data, _, sample_node_dict, value_node_dict, var_value_dict, neg_edges_per_pair, variable_names, variable_embedding_ids, cont_bin_names, cont_bin_embedding_ids, bin_stats = self.data_to_graph( input_data, dtype_col=self.type_column, ignore_na=self.non_informative_na, @@ -85,21 +84,23 @@ def fit(self, X, y=None): self._variable_names = variable_names self._cont_bin_names = cont_bin_names self._X = input_data - + self._value_node_dict = value_node_dict + self._bin_stats = bin_stats + if self.enable_imputation: self._discretized_data = discrete_data self._var_value_dict = var_value_dict - self._value_node_dict = value_node_dict - self._cont_vars = set(self._X.index[self._X[self.type_column].isin(['numerical', 'cont'])]) self._cat_vars = set(self._X.index[self._X[self.type_column].isin(['nominal', 'cat', 'ordinal'])]) - + # Store actual dataframe without type column for potential imputation later. + self._cont_vars = set(self._X.index[self._X[self.type_column].isin(['numerical', 'cont'])]) self._X.drop(columns=[self.type_column], inplace=True) # Compute node-based embeddings. node_embeddings, variable_embeddings, bin_embeddings, auc, ap = self.compute_node_embeddings(graph_data, variable_embedding_ids, cont_bin_embedding_ids) + self._graph_data = graph_data # graph_data is on CPU after compute_node_embeddings self._auc = auc self._ap = ap self._all_embeddings = node_embeddings @@ -222,6 +223,191 @@ def impute_sample(self, sample_colum : str, na_value : float): return imputed_df + def _remap_to_populated_bin(self, var, bin_id): + """Return the nearest bin index to bin_id that has a node in the training graph.""" + for delta in range(1, self.bins_per_continuous): + for candidate in (bin_id - delta, bin_id + delta): + if 0 <= candidate < self.bins_per_continuous: + if f"{var}={float(candidate)}" in self._value_node_dict: + return candidate + return None + + def _get_bin_id(self, variable, value): + """Map a single continuous value to its training-time bin ID.""" + stats = self._bin_stats[variable] + if stats['type'] == 'z': + std = stats['std'] + z = 0.0 if std == 0 else (value - stats['mean']) / std + result = pd.cut([z], bins=get_zscore_bins(self.bins_per_continuous), labels=False) + bin_val = result[0] + return None if pd.isna(bin_val) else int(bin_val) + elif stats['type'] == 'nonlinear': + return int(np.digitize(value, stats['edges'][1:-1], right=True)) + return None + + def _sample_value_node_pairs(self, input_data, sample): + """Map one new sample to the training value nodes it connects to. + + Reproduces the edge-construction logic used during fit(): continuous values + are mapped to their training-time bin (with empty-bin remapping), NA sentinels + are skipped, and unseen categories are dropped. + + Returns: + list[tuple[str, int]]: (variable, value_node_index) pairs for each observed, + in-vocabulary value of the sample. + """ + pairs = [] + for var in input_data.index: + raw = input_data.loc[var, sample] + if pd.isna(raw): + raise ValueError("Actual NA entry detected in input dataframe.") + value = float(raw) + if value == self.non_informative_na: + continue + + if var in self._cont_vars and value not in self.informative_nas: + bin_id = self._get_bin_id(var, value) + if bin_id is None: + continue + value_label = f"{var}={float(bin_id)}" + if value_label not in self._value_node_dict: + remapped = self._remap_to_populated_bin(var, bin_id) + if remapped is None: + continue + value_label = f"{var}={float(remapped)}" + else: + value_label = f"{var}={value}" + if value_label not in self._value_node_dict: + continue + + pairs.append((var, self._value_node_dict[value_label])) + return pairs + + def _augmented_edge_index(self, pairs_per_sample, base_idx, device): + """Build the augmented edge_index that injects new sample nodes into the training graph. + + Only ``value_node -> new_node`` edges are added: each new sample *receives* messages + from the frozen value nodes but never sends any back. This keeps the shared value-node + representations identical to training and lets every new sample be encoded in a single + pass without perturbing the training nodes or interfering with each other. + + Args: + pairs_per_sample (list[list[tuple[str, int]]]): per-new-sample (var, value_node) pairs. + base_idx (int): node index assigned to the first new sample (subsequent samples + are base_idx+1, base_idx+2, ...). + device: torch device. + + Returns: + torch.Tensor | None: augmented edge_index, or None if no new edges could be built. + """ + value_src, new_dst = [], [] + for offset, pairs in enumerate(pairs_per_sample): + node = base_idx + offset + for _, value_node in pairs: + # PyG convention: a message flows src -> dst, so value_node -> new_node lets + # the new sample aggregate from its value nodes (and never the reverse). + value_src.append(value_node) + new_dst.append(node) + if not value_src: + return None + recv_edges = torch.tensor([value_src, new_dst], dtype=torch.long, device=device) + training_edge_index = self._graph_data.edge_index.to(device) + return torch.cat([training_edge_index, recv_edges], dim=1) + + def _frozen_encode(self, augmented_embeds, augmented_edge_index): + """Normalize node inputs, run the frozen encoder, and L2-normalize the latents. + + Mirrors GraphAutoencoder.get_embeddings exactly so training (transductive) and + inductive nodes are processed identically. + """ + norms = augmented_embeds.norm(p=2, dim=1, keepdim=True).clamp(min=1e-8) + augmented_embeds = augmented_embeds / norms + latent = self.model.encoder(augmented_embeds, augmented_edge_index) + norms = latent.norm(p=2, dim=1, keepdim=True).clamp(min=1e-8) + return latent / norms + + def transform(self, X): + """Generate embeddings for new, unseen samples using the frozen trained encoder. + + The trained model weights are never updated (no retraining). Each new sample is added + to the training graph as a node connected to the training value nodes it observes, with + edges pointing *only* from value nodes to the new sample. Because the new samples never + send messages back, the shared value-node representations stay identical to training + and all new samples are encoded in a single forward pass without affecting the training + nodes or one another. The new node's own ("self-loop") input embedding is set to zeros: + the trained encoder assigns negligible weight to a sample node's self-loop, so the + embedding is determined by which value nodes the sample connects to (its data), and a + data-derived initialization was found to have no measurable effect. + + Empirically this places inductive embeddings in the same distribution as the + transductive training embeddings (Rosenbaum cross-match z ~ 0). Adding the new samples + with bidirectional edges instead perturbs the value nodes and is markedly + out-of-distribution. + + Every sample must share at least one observed, in-vocabulary value with the training + data; an edgeless sample (all values missing / unseen) has no signal and raises. + + Args: + X (pd.DataFrame): New samples in the same format as the training data + (rows=variables, columns=new_samples+type_column). + + Returns: + pd.DataFrame: Embedding matrix with shape (num_new_samples, embedding_dimension), + indexed by sample name. + """ + if not hasattr(self, '_graph_data'): + raise ValueError("Run fit() first before calling transform()!") + + input_data = X.copy() + new_samples = [c for c in input_data.columns if c != self.type_column] + input_data.drop(columns=[self.type_column], inplace=True) + + new_vars = set(input_data.index) + train_vars = set(self._variable_names) + if new_vars != train_vars: + missing = sorted(train_vars - new_vars) + extra = sorted(new_vars - train_vars) + parts = [] + if missing: + parts.append(f"missing from input: {missing}") + if extra: + parts.append(f"not seen during training: {extra}") + raise ValueError("Variable mismatch between transform input and training data — " + "; ".join(parts)) + + # Resolve which training value nodes each new sample connects to. Every sample must + # have at least one valid edge, otherwise it carries no signal for the encoder. + pairs_per_sample = [self._sample_value_node_pairs(input_data, s) for s in new_samples] + edgeless = [s for s, pairs in zip(new_samples, pairs_per_sample) if not pairs] + if edgeless: + raise ValueError( + f"No valid edges for sample(s) {edgeless}: none of their values are observed, " + "in-vocabulary values shared with the training data, so they have no signal. " + "Every sample must share at least one value with the training data.") + + num_existing_nodes = self._graph_data.num_nodes + device = next(self.model.parameters()).device + self.model.eval() + + with torch.no_grad(): + # Frozen training-node input embeddings (component_1 + component_2). New sample + # nodes get a zero self-loop feature (the encoder ignores it; the embedding comes + # entirely from the value nodes the sample connects to). + c1 = self.model.node_embeddings(self.model.node_to_embeddings[:, 0].to(device)) + c2 = self.model.node_embeddings(self.model.node_to_embeddings[:, 1].to(device)) + training_embeds = c1 + c2 + + new_sample_embeds = torch.zeros( + len(new_samples), self.embedding_dimension, + device=device, dtype=training_embeds.dtype) + augmented_embeds = torch.cat([training_embeds, new_sample_embeds], dim=0) + augmented_edge_index = self._augmented_edge_index( + pairs_per_sample, num_existing_nodes, device) + latent = self._frozen_encode(augmented_embeds, augmented_edge_index) + new_latent = latent[num_existing_nodes:].cpu() + + embedding_cols = [f'dim_{i}' for i in range(self.embedding_dimension)] + return pd.DataFrame(new_latent.numpy(), index=new_samples, columns=embedding_cols) + def _move_self_tensors(self, device): """Move all tensor-like attributes of self to given device.""" for name, value in self.__dict__.items(): @@ -349,6 +535,25 @@ def data_to_graph(self, all_data, dtype_col: str, ignore_na, keep_na): variable_types = data[dtype_col].to_list() data.drop(columns=[dtype_col], inplace=True) + # Compute per-variable bin stats for later inductive inference on new samples. + bin_stats = {} + for var in continuous_vars: + col = data.loc[var].copy().astype(float) + valid = ~((col == ignore_na) | col.isin(keep_na) | col.isna()) + non_na = col[valid] + if self.discretization_type == "z": + if len(set(non_na)) == 1: + bin_stats[var] = {'type': 'z', 'mean': float(non_na.iloc[0]), 'std': 1.0} + else: + bin_stats[var] = {'type': 'z', 'mean': float(non_na.mean()), 'std': float(non_na.std(ddof=0))} + elif self.discretization_type == "nonlinear": + if len(set(non_na)) == 1: + val = float(non_na.iloc[0]) + bin_stats[var] = {'type': 'nonlinear', 'edges': np.array([-np.inf, val, np.inf])} + else: + edges, _ = signed_power_bins(non_na.values, n_bins=self.bins_per_continuous) + bin_stats[var] = {'type': 'nonlinear', 'edges': edges} + if self.discretization_type == "z": data.loc[continuous_vars] = data.loc[continuous_vars].apply( lambda row: bin_column_with_na_adjusted(row, self.bins_per_continuous, ignore_na, keep_na), axis=1) @@ -452,6 +657,6 @@ def data_to_graph(self, all_data, dtype_col: str, ignore_na, keep_na): neg_edges_per_pair[(sample, var, pos_value)] = (torch.tensor(negative_edges, dtype=torch.long, device=self.device)) data.drop(columns=[dtype_col], inplace=True) - return (data, graph_data, num_variables, sample_node_ids, value_node_ids, + return (data, graph_data, num_variables, sample_node_ids, value_node_ids, var_value_dict, neg_edges_per_pair, variable_names, variable_embedding_ids, - cont_bin_names, cont_bin_embedding_ids) + cont_bin_names, cont_bin_embedding_ids, bin_stats) diff --git a/src/pome/models.py b/src/pome/models.py index c8bcbc0..f4e25fd 100644 --- a/src/pome/models.py +++ b/src/pome/models.py @@ -1,4 +1,4 @@ -from torch_geometric.nn import GATConv, GCNConv, SAGEConv +from torch_geometric.nn import GATConv import torch import torch.nn as nn import torch.nn.functional as F @@ -13,7 +13,7 @@ def __init__(self, input_dim, hidden_dims, output_dim, layer_type): if layer_type == "GAT": Layer = GATConv else: - raise ValueError(f"Unknown layer_type: {layer_type}. Choose from ['GCN', 'GAT', 'SAGE']") + raise ValueError(f"Unknown layer_type: {layer_type}. Choose from ['GAT']") self.convs = torch.nn.ModuleList() diff --git a/src/pome/utils.py b/src/pome/utils.py index f16dbaf..78986bd 100644 --- a/src/pome/utils.py +++ b/src/pome/utils.py @@ -5,9 +5,7 @@ import numpy as np def link_similarity(embeddings, edge_index, decoder_model=None): - u, v = edge_index # Extract node pairs - similarities = decoder_model(embeddings, edge_index) - return similarities + return decoder_model(embeddings, edge_index) def compute_roc(graph_data, neg_edges_per_pair, node_embeddings, decoder_model): @@ -32,6 +30,19 @@ def compute_roc(graph_data, neg_edges_per_pair, node_embeddings, decoder_model): return roc_auc, avg_precision +def get_zscore_bins(K): + if K == 15: + return [-np.inf, -3.5, -3.0, -2.5, -2.0, -1.5, -1.0, -0.5, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, np.inf] + elif K == 11: + return [-np.inf, -3.5, -2.5, -1.5, -1.0, -0.5, 0.5, 1.0, 1.5, 2.5, 3.5, np.inf] + elif K == 7: + return [-np.inf, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, np.inf] + elif K == 3: + return [-np.inf, -0.5, 0.5, np.inf] + else: + raise ValueError(f"Invalid number of z-score bins: {K}") + + def bin_column_with_na_adjusted(column, K, true_missing, keep_nas): column = column.copy() # Avoid modifying the original data column = column.astype(float) @@ -47,21 +58,12 @@ def bin_column_with_na_adjusted(column, K, true_missing, keep_nas): # Compute z-scores of non-NA values. if len(set(non_na))==1: print(f"Warning: cont variable {column.name} contains only one unique value. Setting zscore to 0.") - non_na_zscores = [0.0] * len(non_na) + non_na_zscores = np.zeros(len(non_na)) else: non_na_zscores = zscore(non_na, nan_policy='raise') # Perform binning on valid values - if K==15: - zscore_bins = [-np.inf, -3.5, -3.0, -2.5, -2.0, -1.5, -1.0, -0.5, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, np.inf] - elif K==11: - zscore_bins = [-np.inf, -3.5, -2.5, -1.5, -1.0, -0.5, 0.5, 1.0, 1.5, 2.5, 3.5, np.inf] - elif K==7: - zscore_bins = [-np.inf, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, np.inf] - elif K==3: - zscore_bins = [-np.inf, -0.5, 0.5, np.inf] - else: - raise ValueError(f"Invalid number of z-score bins: {K}") + zscore_bins = get_zscore_bins(K) binned_non_na = pd.cut(non_na_zscores, bins=zscore_bins, labels=False) @@ -114,7 +116,7 @@ def bin_column_non_linear(column, K, true_missing, keep_nas): # Compute non-linear bins of non-NA values. if len(set(non_na))==1: print(f"Warning: cont variable {column.name} contains only one unique value. Setting all bins to 0.") - binned_non_na = [int((K-1)/2)] * len(non_na) + binned_non_na = np.full(len(non_na), int((K-1)/2)) else: _, binned_non_na = signed_power_bins(non_na, n_bins=K) diff --git a/tests/test_gnn_embedding.py b/tests/test_gnn_embedding.py index 491f4ab..be58ed5 100644 --- a/tests/test_gnn_embedding.py +++ b/tests/test_gnn_embedding.py @@ -159,11 +159,191 @@ def test_layer_failure(): embedder.fit(example_df) def test_discretization_failure(): - embedder = Embedder(epochs=50, - na_encoding=NA_ENCODING, + embedder = Embedder(epochs=50, + na_encoding=NA_ENCODING, embedding_dimension=DIMENSION, device=DEVICE, enable_imputation=False, discretization_type="rubbish") with pytest.raises(ValueError): embedder.fit(example_df) + +def _fit_embedder(epochs=50): + """Helper: fit on the full dataset, return (embedder, last_sample_col).""" + TYPE_COL = "type" + sample_cols = [c for c in example_df.columns if c != TYPE_COL] + embedder = Embedder(epochs=epochs, + na_encoding=NA_ENCODING, + embedding_dimension=DIMENSION, + device=DEVICE) + embedder.fit(example_df) + return embedder, sample_cols[-1] + +def test_transform_new_samples(): + # Fit on full dataset; all categorical values are guaranteed to be seen. + TYPE_COL = "type" + sample_cols = [c for c in example_df.columns if c != TYPE_COL] + embedder, _ = _fit_embedder(epochs=50) + new_df = example_df[sample_cols[-2:] + [TYPE_COL]] + + new_embeddings = embedder.transform(new_df) + + assert isinstance(new_embeddings, pd.DataFrame) + assert new_embeddings.shape == (2, DIMENSION) + assert list(new_embeddings.index) == sample_cols[-2:] + +def test_transform_before_fit(): + embedder = Embedder(epochs=10, + na_encoding=NA_ENCODING, + embedding_dimension=DIMENSION, + device=DEVICE) + with pytest.raises(ValueError): + embedder.transform(example_df) + +def test_transform_variable_mismatch(): + embedder, held_out = _fit_embedder(epochs=10) + + new_df = example_df[[held_out, "type"]].drop(index="cont.4") + with pytest.raises(ValueError, match="Variable mismatch"): + embedder.transform(new_df) + +def test_transform_unseen_category(): + embedder, held_out = _fit_embedder(epochs=10) + + # Inject a categorical value never seen during training (valid cats are 0–3). + # The unseen value is treated as missing: no edge is added for it, but + # transform() succeeds using the remaining observed features. + new_df = example_df[[held_out, "type"]].copy() + new_df.loc["cat", held_out] = 99.0 + result = embedder.transform(new_df) + assert result.shape == (1, embedder.embedding_dimension) + +def test_transform_empty_bin_remapping(): + embedder, held_out = _fit_embedder(epochs=10) + + # An extreme value lands in the highest z-score bin, which is empty with only + # 10 training samples (|z| > 3.5 requires an outlier beyond 3.5 standard deviations). + extreme_bin = embedder._get_bin_id("cont", 1000.0) + assert f"cont={float(extreme_bin)}" not in embedder._value_node_dict + + remapped = embedder._remap_to_populated_bin("cont", extreme_bin) + assert remapped is not None + assert f"cont={float(remapped)}" in embedder._value_node_dict + + # transform() must succeed via remapping rather than raising or dropping. + new_df = example_df[[held_out, "type"]].copy() + new_df.loc["cont", held_out] = 1000.0 + new_embeddings = embedder.transform(new_df) + assert new_embeddings.shape == (1, DIMENSION) + +def test_impute_sample_failure(): + embedder = Embedder(epochs=10, + na_encoding=NA_ENCODING, + embedding_dimension=DIMENSION, + device=DEVICE, + enable_imputation=False) + embedder.fit(example_df) + with pytest.raises(ValueError): + embedder.impute_sample("sample0", NA_ENCODING) + +def test_transform_nonlinear_discretization(): + TYPE_COL = "type" + sample_cols = [c for c in example_df.columns if c != TYPE_COL] + embedder = Embedder(epochs=50, + na_encoding=NA_ENCODING, + embedding_dimension=DIMENSION, + device=DEVICE, + discretization_type="nonlinear") + embedder.fit(example_df) + new_df = example_df[[sample_cols[-1], TYPE_COL]] + new_embeddings = embedder.transform(new_df) + assert new_embeddings.shape == (1, DIMENSION) + +def test_transform_extra_variable(): + embedder, held_out = _fit_embedder(epochs=10) + new_df = example_df[[held_out, "type"]].copy() + new_df.loc["extra_var"] = [0.0, "numerical"] + with pytest.raises(ValueError, match="Variable mismatch"): + embedder.transform(new_df) + +def test_transform_nan_in_input(): + embedder, held_out = _fit_embedder(epochs=10) + new_df = example_df[[held_out, "type"]].copy() + new_df.loc["cont", held_out] = np.nan + with pytest.raises(ValueError, match="Actual NA entry"): + embedder.transform(new_df) + +def test_transform_all_na_input(): + embedder, held_out = _fit_embedder(epochs=10) + new_df = example_df[[held_out, "type"]].copy() + new_df.loc[:, held_out] = NA_ENCODING + with pytest.raises(ValueError, match="No valid edges"): + embedder.transform(new_df) + +def test_transform_one_edgeless_sample_raises(): + # Even a single edgeless sample (all values missing) among otherwise valid samples + # must raise: it carries no signal for the encoder. + TYPE_COL = "type" + sample_cols = [c for c in example_df.columns if c != TYPE_COL] + embedder, _ = _fit_embedder(epochs=10) + new_df = example_df[sample_cols[:2] + [TYPE_COL]].copy() + new_df.loc[:, sample_cols[0]] = NA_ENCODING # first sample edgeless, second valid + with pytest.raises(ValueError, match="No valid edges"): + embedder.transform(new_df) + +def test_constant_continuous_variable_z(): + # A continuous variable with a single unique value triggers the zscore constant-var + # fallback in both data_to_graph (line 503) and bin_column_with_na_adjusted (utils 62-63). + const_df = example_df.copy() + sample_cols = [c for c in const_df.columns if c != "type"] + const_df.loc["cont", sample_cols] = 0.5 + embedder = Embedder(epochs=10, + na_encoding=NA_ENCODING, + embedding_dimension=DIMENSION, + device=DEVICE, + discretization_type="z") + embedder.fit(const_df) + assert embedder._ap >= 0 + +def test_constant_continuous_variable_nonlinear(): + # Same scenario with nonlinear discretization covers lines 508-509 and utils 120-121. + const_df = example_df.copy() + sample_cols = [c for c in const_df.columns if c != "type"] + const_df.loc["cont", sample_cols] = 0.5 + embedder = Embedder(epochs=10, + na_encoding=NA_ENCODING, + embedding_dimension=DIMENSION, + device=DEVICE, + discretization_type="nonlinear") + embedder.fit(const_df) + assert embedder._ap >= 0 + +def test_informative_na_z(): + # Informative NAs get their own embedding slot (gnn line 572, utils line 80). + INFORMATIVE_NA = -88.0 + inf_df = example_df.copy() + sample_cols = [c for c in inf_df.columns if c != "type"] + inf_df.loc["cont", sample_cols[0]] = INFORMATIVE_NA + embedder = Embedder(epochs=10, + na_encoding=NA_ENCODING, + embedding_dimension=DIMENSION, + device=DEVICE, + informative_nas=[INFORMATIVE_NA], + discretization_type="z") + embedder.fit(inf_df) + assert embedder._ap >= 0 + +def test_informative_na_nonlinear(): + # Same with nonlinear discretization covers utils line 133. + INFORMATIVE_NA = -88.0 + inf_df = example_df.copy() + sample_cols = [c for c in inf_df.columns if c != "type"] + inf_df.loc["cont", sample_cols[0]] = INFORMATIVE_NA + embedder = Embedder(epochs=10, + na_encoding=NA_ENCODING, + embedding_dimension=DIMENSION, + device=DEVICE, + informative_nas=[INFORMATIVE_NA], + discretization_type="nonlinear") + embedder.fit(inf_df) + assert embedder._ap >= 0