Add compare module for quantifying embedding differences (#9)

examples/comparisons/compare_embeddings.py

@@ -0,0 +1,160 @@
+"""Compare all embedding schemes and generate visualisations.
+
+Produces:
+- Pairwise correlation heatmap between all embedding schemes
+- 2D UMAP map of embedding schemes (each point = one embedding)
+- Mantel test results table
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import seaborn as sns
+from umap import UMAP
+
+from elementembeddings.compare import (
+    _get_similarity_matrix,
+    _upper_triangle,
+    mantel_test,
+    pairwise_embedding_comparison,
+)
+from elementembeddings.core import Embedding
+
+OUTPUT_DIR = Path(__file__).parent
+OUTPUT_DIR.mkdir(exist_ok=True)
+
+EMBEDDINGS = {
+    "Magpie": "magpie_sc",
+    "Mat2Vec": "mat2vec",
+    "MEGNet": "megnet16",
+    "SkipAtom": "skipatom",
+    "Oliynyk": "oliynyk_sc",
+    "XenonPy": "xenonpy",
+    "CGNF": "cgnf",
+    "CrystaLLM": "crystallm",
+    "MACE-MP-0": "mace_mp0",
+    "SevenNet": "sevennet",
+    "ORB-v2": "orb_v2",
+}
+
+
+def load_all() -> dict[str, Embedding]:
+    """Load all embedding schemes."""
+    embeddings = {}
+    for display_name, code_name in EMBEDDINGS.items():
+        print(f"  Loading {display_name}...")
+        embeddings[display_name] = Embedding.load_data(code_name)
+    return embeddings
+
+
+def plot_comparison_heatmap(comparison_df: pd.DataFrame) -> None:
+    """Plot the pairwise embedding comparison as a heatmap."""
+    fig, ax = plt.subplots(figsize=(10, 8))
+    sns.heatmap(
+        comparison_df,
+        cmap="RdBu_r",
+        vmin=-1,
+        vmax=1,
+        annot=True,
+        fmt=".2f",
+        square=True,
+        linewidths=0.5,
+        ax=ax,
+        cbar_kws={"label": "Pearson correlation", "shrink": 0.8},
+        annot_kws={"fontsize": 7},
+    )
+    ax.set_title("Embedding Scheme Similarity\n(Pearson correlation of cosine similarity matrices)", fontsize=12)
+    ax.tick_params(labelsize=9)
+    fig.tight_layout()
+    fig.savefig(OUTPUT_DIR / "embedding_comparison_heatmap.png", dpi=300, bbox_inches="tight")
+    plt.close(fig)
+    print("Saved embedding_comparison_heatmap.png")
+
+
+def plot_embedding_map(embeddings: dict[str, Embedding]) -> None:
+    """Plot embeddings in 2D using UMAP on flattened similarity vectors."""
+    # Find common elements across ALL embeddings
+    all_elements = None
+    for emb in embeddings.values():
+        els = set(emb.element_list)
+        all_elements = els if all_elements is None else all_elements & els
+    common = sorted(all_elements)
+    print(f"  Common elements for UMAP: {len(common)}")
+
+    # Build feature matrix: each row = one embedding's flattened similarity
+    names = list(embeddings.keys())
+    vectors = []
+    for name in names:
+        mat = _get_similarity_matrix(embeddings[name], common, "cosine_similarity")
+        vectors.append(_upper_triangle(mat))
+
+    X = np.array(vectors)
+    reduced = UMAP(n_components=2, random_state=42, n_neighbors=5).fit_transform(X)
+
+    fig, ax = plt.subplots(figsize=(8, 8))
+    ax.scatter(reduced[:, 0], reduced[:, 1], s=120, c="steelblue", edgecolors="white", linewidths=0.5, zorder=5)
+    for i, name in enumerate(names):
+        ax.annotate(
+            name,
+            (reduced[i, 0], reduced[i, 1]),
+            fontsize=9,
+            fontweight="bold",
+            ha="left",
+            va="bottom",
+            xytext=(5, 5),
+            textcoords="offset points",
+        )
+    ax.set_xlabel("UMAP 1")
+    ax.set_ylabel("UMAP 2")
+    ax.set_title("Embedding Schemes in 2D\n(UMAP of flattened cosine similarity matrices)")
+    fig.tight_layout()
+    fig.savefig(OUTPUT_DIR / "embedding_umap_map.png", dpi=300, bbox_inches="tight")
+    plt.close(fig)
+    print("Saved embedding_umap_map.png")
+
+
+def run_mantel_tests(embeddings: dict[str, Embedding]) -> None:
+    """Run Mantel tests between selected embedding pairs."""
+    pairs = [
+        ("Magpie", "Oliynyk"),
+        ("Mat2Vec", "SkipAtom"),
+        ("MACE-MP-0", "SevenNet"),
+        ("MACE-MP-0", "ORB-v2"),
+        ("SevenNet", "ORB-v2"),
+        ("Mat2Vec", "MACE-MP-0"),
+        ("Magpie", "MACE-MP-0"),
+    ]
+
+    results = []
+    for name1, name2 in pairs:
+        print(f"  Mantel: {name1} vs {name2}...")
+        r, p = mantel_test(embeddings[name1], embeddings[name2], n_permutations=999)
+        results.append({"Embedding 1": name1, "Embedding 2": name2, "r": round(r, 4), "p-value": round(p, 4)})
+
+    df = pd.DataFrame(results)
+    print("\nMantel test results:")
+    print(df.to_string(index=False))
+    df.to_csv(OUTPUT_DIR / "mantel_test_results.csv", index=False)
+    print("Saved mantel_test_results.csv")
+
+
+if __name__ == "__main__":
+    embeddings = load_all()
+
+    print("\nComputing pairwise comparisons...")
+    comparison_df = pairwise_embedding_comparison(embeddings)
+    comparison_df.to_csv(OUTPUT_DIR / "pairwise_comparison.csv")
+    print("Saved pairwise_comparison.csv")
+
+    print("\nPlotting comparison heatmap...")
+    plot_comparison_heatmap(comparison_df)
+
+    print("\nPlotting embedding UMAP map...")
+    plot_embedding_map(embeddings)
+
+    print("\nRunning Mantel tests...")
+    run_mantel_tests(embeddings)

examples/comparisons/mantel_test_results.csv

@@ -0,0 +1,8 @@
+Embedding 1,Embedding 2,r,p-value
+Magpie,Oliynyk,0.8601,0.001
+Mat2Vec,SkipAtom,0.5012,0.001
+MACE-MP-0,SevenNet,0.7378,0.001
+MACE-MP-0,ORB-v2,0.2994,0.001
+SevenNet,ORB-v2,0.2312,0.001
+Mat2Vec,MACE-MP-0,0.4326,0.001
+Magpie,MACE-MP-0,0.4664,0.001

examples/comparisons/pairwise_comparison.csv

@@ -0,0 +1,12 @@
+,Magpie,Mat2Vec,MEGNet,SkipAtom,Oliynyk,XenonPy,CGNF,CrystaLLM,MACE-MP-0,SevenNet,ORB-v2
+Magpie,1.0,0.38444501825844374,0.4553642750628917,0.5642335937898679,0.8600963517236999,0.2748278760128914,0.5590163677849405,0.5203272305810318,0.4664052148840695,0.46174106593553954,0.33631491972375316
+Mat2Vec,0.38444501825844374,1.0,0.3052321154458533,0.501246576386556,0.5579225775837166,0.17178596020210996,0.3994978625459398,0.5095680996699704,0.4326065637306554,0.3786883732276776,0.23676388467278026
+MEGNet,0.4553642750628917,0.3052321154458533,1.0,0.4771113385853821,0.5485634850030142,0.26246780787053914,0.5151846366087817,0.5329223265829559,0.5894036940890679,0.5940620826834662,0.26547618464220046
+SkipAtom,0.5642335937898679,0.501246576386556,0.4771113385853821,1.0,0.6539974257888233,0.29024273573255527,0.5271906811331359,0.5760408513405436,0.5847010884089361,0.48115132958169127,0.3954443021704164
+Oliynyk,0.8600963517236999,0.5579225775837166,0.5485634850030142,0.6539974257888233,1.0,0.37136514035668144,0.6147521163512288,0.6142492810182167,0.5293620589508641,0.4988308453813105,0.36521341226873416
+XenonPy,0.2748278760128914,0.17178596020210996,0.26246780787053914,0.29024273573255527,0.37136514035668144,1.0,0.3274082551055765,0.3703648763220639,0.2994838505037849,0.22809408768655826,0.09763740111715223
+CGNF,0.5590163677849405,0.3994978625459398,0.5151846366087817,0.5271906811331359,0.6147521163512288,0.3274082551055765,1.0,0.5888649700312897,0.5952072589284931,0.5322148922272025,0.40359410047802025
+CrystaLLM,0.5203272305810318,0.5095680996699704,0.5329223265829559,0.5760408513405436,0.6142492810182167,0.3703648763220639,0.5888649700312897,1.0,0.5785319769971016,0.5374742964367392,0.19537583294839733
+MACE-MP-0,0.4664052148840695,0.4326065637306554,0.5894036940890679,0.5847010884089361,0.5293620589508641,0.2994838505037849,0.5952072589284931,0.5785319769971016,1.0,0.7377670281507327,0.29935107497815794
+SevenNet,0.46174106593553954,0.3786883732276776,0.5940620826834662,0.48115132958169127,0.4988308453813105,0.22809408768655826,0.5322148922272025,0.5374742964367392,0.7377670281507327,1.0,0.23121900538461082
+ORB-v2,0.33631491972375316,0.23676388467278026,0.26547618464220046,0.3954443021704164,0.36521341226873416,0.09763740111715223,0.40359410047802025,0.19537583294839733,0.29935107497815794,0.23121900538461082,1.0

pyproject.toml

@@ -148,6 +148,7 @@ reportUnboundVariable = true
 skip = "*.csv,*/site/*,publications/element_similarity/**,*periodic-table*.json"
 check-filenames = true
 ignore-words-list = [
+            "fro",
             "H",
             "He",
             "Li",

src/elementembeddings/__init__.py

@@ -5,4 +5,5 @@
 - `elementembeddings.core`: Provides the `Embedding` class.
 - `elementembeddings.composition`: Tools to featurise compositions.
 - `elementembeddings.plotter`: Tools to plot embeddings.
+- `elementembeddings.compare`: Tools to compare embedding schemes.
 """