Add compare module for quantifying embedding differences (#9)#192
Merged
Conversation
…mes (#9) New module `elementembeddings.compare` provides: - embedding_similarity: Pearson/Spearman/Kendall correlation between flattened element similarity matrices - mantel_test: permutation-based significance test for matrix correlation - kl_divergence: KL divergence between normalised similarity distributions - frobenius_distance: Frobenius norm of similarity matrix difference - pairwise_embedding_comparison: compare all pairs, return comparison matrix Includes example script with comparison heatmap, UMAP map of embedding schemes, and Mantel test results. Closes #9
- Validate `method` argument in `mantel_test` (raise ValueError on unknown methods instead of silently using Spearman). - Switch Mantel p-value to two-sided so strong negative correlations are reported as significant rather than near-1. - Drop unused `_get_common_elements` import in the example script. - Whitelist "fro" (Frobenius norm) in codespell to unblock CI. - Add unit tests covering self-similarity, invalid arguments, KL/Frobenius zero-on-self, and pairwise comparison shape/symmetry.
📝 WalkthroughWalkthroughThis PR introduces a new ChangesEmbedding Comparison Framework
Sequence DiagramsequenceDiagram
participant User
participant Script as Example Script
participant Embeddings as Embedding<br/>System
participant Compare as Comparison<br/>Functions
participant Viz as Visualization<br/>Libs
participant FS as File System
User->>Script: Execute compare_embeddings.py
Script->>Embeddings: load_all() for each<br/>(magpie, mat2vec, megnet16)
Embeddings-->>Script: Return Embedding objects
Script->>Compare: pairwise_embedding_comparison()
Compare->>Embeddings: compute_correlation_metric()<br/>for all pairs
Embeddings-->>Compare: similarity matrices
Compare-->>Script: correlation DataFrame
Script->>FS: Save pairwise_comparison.csv
Script->>Compare: Compute heatmap data
Compare-->>Script: correlation matrix
Script->>Viz: plot_comparison_heatmap()
Viz->>FS: Save embedding_comparison_heatmap.png
Script->>Compare: Compute UMAP features
Compare->>Embeddings: compute_correlation_metric()<br/>for common elements
Embeddings-->>Compare: similarity matrices
Compare-->>Script: flattened feature vectors
Script->>Viz: plot_embedding_map() with UMAP
Viz->>FS: Save embedding_umap_map.png
Script->>Compare: run_mantel_tests() on<br/>fixed embedding pairs
Compare->>Embeddings: compute_correlation_metric()<br/>for each pair
Embeddings-->>Compare: similarity matrices
Compare-->>Script: Mantel r & p-values
Script->>FS: Save mantel_test_results.csv
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Review rate limit: 0/1 reviews remaining, refill in 60 minutes.Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
pyproject.toml (1)
147-152: ⚡ Quick winConsider narrowing the codespell ignore entry for
"fro"Ignoring the short token
"fro"can be a bit broad if codespell ever flags other legitimate “fro” occurrences as misspellings. If the underlying warning was specifically about the Frobenius metric shorthand, consider ignoring the more specific token that codespell reports (e.g.,frobenius,fro_, or whatever exact misspelling appears) rather than the generic substring token.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@pyproject.toml` around lines 147 - 152, The codespell ignore entry currently contains the broad token "fro" in the [tool.codespell] ignore-words-list; narrow this to the specific token reported (e.g., replace "fro" with the exact identifier such as "frobenius", "fro_", or the precise misspelling string codespell flagged) so you only suppress that specific false positive; update the ignore-words-list entry accordingly in the tool.codespell section.examples/comparisons/compare_embeddings.py (1)
19-24: Avoid teaching private helpers in a user-facing example.This sample imports
_get_similarity_matrixand_upper_trianglefrom the module's internals, which can change without a public API guarantee. The example uses these to build a feature matrix for UMAP dimensionality reduction. Either expose a public helper for this similarity-to-vector workflow or keep the matrix-building logic local to the example.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/comparisons/compare_embeddings.py` around lines 19 - 24, The example imports private helpers _get_similarity_matrix and _upper_triangle which are internal and unstable; remove those imports and stop relying on them in the example. Fix by either (A) moving the similarity-to-vector matrix-building logic into the example file itself (implement a local function that computes the full similarity matrix from embeddings and extracts the upper-triangle vector) or (B) add a new public helper in the library (e.g., get_similarity_vector or get_similarity_matrix) and import that instead of _get_similarity_matrix/_upper_triangle; update references to use the new local function or the new public API and keep usage of public functions like pairwise_embedding_comparison, mantel_test unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/elementembeddings/compare.py`:
- Around line 125-126: Add upfront validation for the n_permutations parameter:
if n_permutations < 1 raise a ValueError with a clear message (e.g.
"n_permutations must be >= 1") at the start of the function that declares
n_permutations and returns tuple[float, float], and apply the same check in the
other function that accepts n_permutations (the second occurrence referenced in
the review). This prevents division-by-zero or invalid p-value computation by
rejecting invalid inputs before any permutation logic runs.
- Around line 41-47: Change _get_common_elements to accept a min_common argument
(default 3) and use that threshold instead of the hardcoded 3; update all call
sites to pass the appropriate minimum for their metric: pass min_common=3 from
embedding_similarity() and mantel_test(), and pass min_common=2 from
kl_divergence() and frobenius_distance(). Ensure the ValueError message reflects
the passed min_common and keep the function name _get_common_elements to locate
and modify it.
- Around line 241-246: The Frobenius norm is computed over the full matrix
(diff) but the divisor n_pairs counts only unique pairs, causing a sqrt(2)
mismatch; fix by computing the norm only over the unique off-diagonal entries
(e.g., use the upper triangle of diff) and then divide by sqrt(n_pairs).
Concretely, replace the frob calculation that uses np.linalg.norm(diff, "fro")
with a norm over np.triu(diff, k=1) (or equivalent selection) so that frob and
the normalise divisor n_pairs (based on elements) measure the same set of
entries; adjust no other logic around normalise, diff, np.linalg.norm, n_pairs,
or elements.
---
Nitpick comments:
In `@examples/comparisons/compare_embeddings.py`:
- Around line 19-24: The example imports private helpers _get_similarity_matrix
and _upper_triangle which are internal and unstable; remove those imports and
stop relying on them in the example. Fix by either (A) moving the
similarity-to-vector matrix-building logic into the example file itself
(implement a local function that computes the full similarity matrix from
embeddings and extracts the upper-triangle vector) or (B) add a new public
helper in the library (e.g., get_similarity_vector or get_similarity_matrix) and
import that instead of _get_similarity_matrix/_upper_triangle; update references
to use the new local function or the new public API and keep usage of public
functions like pairwise_embedding_comparison, mantel_test unchanged.
In `@pyproject.toml`:
- Around line 147-152: The codespell ignore entry currently contains the broad
token "fro" in the [tool.codespell] ignore-words-list; narrow this to the
specific token reported (e.g., replace "fro" with the exact identifier such as
"frobenius", "fro_", or the precise misspelling string codespell flagged) so you
only suppress that specific false positive; update the ignore-words-list entry
accordingly in the tool.codespell section.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: baa0f2d1-a0de-40ed-b6c8-6fc4da573c98
⛔ Files ignored due to path filters (4)
examples/comparisons/embedding_comparison_heatmap.pngis excluded by!**/*.pngexamples/comparisons/embedding_umap_map.pngis excluded by!**/*.pngexamples/comparisons/mantel_test_results.csvis excluded by!**/*.csvexamples/comparisons/pairwise_comparison.csvis excluded by!**/*.csv
📒 Files selected for processing (5)
examples/comparisons/compare_embeddings.pypyproject.tomlsrc/elementembeddings/__init__.pysrc/elementembeddings/compare.pysrc/elementembeddings/tests/test_compare.py
Comment on lines
+41
to
+47
| def _get_common_elements(emb1: Embedding, emb2: Embedding) -> list[str]: | ||
| """Return sorted list of elements common to both embeddings.""" | ||
| common = sorted(set(emb1.element_list) & set(emb2.element_list)) | ||
| if len(common) < 3: | ||
| msg = f"Only {len(common)} common elements between embeddings. Need at least 3." | ||
| raise ValueError(msg) | ||
| return common |
There was a problem hiding this comment.
Make the minimum-overlap check metric-specific.
This helper is shared by all public APIs, but the 3-element floor is only needed for the correlation-based ones. kl_divergence() and frobenius_distance() are still well-defined with exactly 2 common elements, so they currently reject valid inputs from custom embeddings.
Suggested fix
-def _get_common_elements(emb1: Embedding, emb2: Embedding) -> list[str]:
+def _get_common_elements(
+ emb1: Embedding,
+ emb2: Embedding,
+ *,
+ min_common: int = 3,
+) -> list[str]:
"""Return sorted list of elements common to both embeddings."""
common = sorted(set(emb1.element_list) & set(emb2.element_list))
- if len(common) < 3:
- msg = f"Only {len(common)} common elements between embeddings. Need at least 3."
+ if len(common) < min_common:
+ msg = f"Only {len(common)} common elements between embeddings. Need at least {min_common}."
raise ValueError(msg)
return commonUse min_common=3 for embedding_similarity() / mantel_test(), and min_common=2 for kl_divergence() / frobenius_distance().
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/elementembeddings/compare.py` around lines 41 - 47, Change
_get_common_elements to accept a min_common argument (default 3) and use that
threshold instead of the hardcoded 3; update all call sites to pass the
appropriate minimum for their metric: pass min_common=3 from
embedding_similarity() and mantel_test(), and pass min_common=2 from
kl_divergence() and frobenius_distance(). Ensure the ValueError message reflects
the passed min_common and keep the function name _get_common_elements to locate
and modify it.
Comment on lines
+125
to
+126
| n_permutations: int = 999, | ||
| ) -> tuple[float, float]: |
There was a problem hiding this comment.
Validate n_permutations before computing the p-value.
Negative values currently fall through the loop and can either divide by zero (-1) or return an invalid p-value. Reject < 1 up front.
Suggested fix
def mantel_test(
emb1: Embedding,
emb2: Embedding,
metric: str = "cosine_similarity",
method: str = "pearson",
n_permutations: int = 999,
) -> tuple[float, float]:
@@
+ if n_permutations < 1:
+ msg = "n_permutations must be at least 1."
+ raise ValueError(msg)
+
if method == "pearson":
corr_fn = stats.pearsonrAlso applies to: 174-175
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/elementembeddings/compare.py` around lines 125 - 126, Add upfront
validation for the n_permutations parameter: if n_permutations < 1 raise a
ValueError with a clear message (e.g. "n_permutations must be >= 1") at the
start of the function that declares n_permutations and returns tuple[float,
float], and apply the same check in the other function that accepts
n_permutations (the second occurrence referenced in the review). This prevents
division-by-zero or invalid p-value computation by rejecting invalid inputs
before any permutation logic runs.
Comment on lines
+241
to
+246
| diff = mat1 - mat2 | ||
| frob = float(np.linalg.norm(diff, "fro")) | ||
|
|
||
| if normalise: | ||
| n_pairs = len(elements) * (len(elements) - 1) / 2 | ||
| frob /= np.sqrt(n_pairs) |
There was a problem hiding this comment.
Normalize the same entries you measure.
np.linalg.norm(diff, "fro") includes both symmetric halves of the matrix, but the normalisation divisor is based on unique element pairs only. That inflates the "normalise=True" result by sqrt(2) relative to the stated per-pair scaling.
Suggested fix
- diff = mat1 - mat2
- frob = float(np.linalg.norm(diff, "fro"))
+ diff = _upper_triangle(mat1) - _upper_triangle(mat2)
+ frob = float(np.linalg.norm(diff))
if normalise:
n_pairs = len(elements) * (len(elements) - 1) / 2
frob /= np.sqrt(n_pairs)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/elementembeddings/compare.py` around lines 241 - 246, The Frobenius norm
is computed over the full matrix (diff) but the divisor n_pairs counts only
unique pairs, causing a sqrt(2) mismatch; fix by computing the norm only over
the unique off-diagonal entries (e.g., use the upper triangle of diff) and then
divide by sqrt(n_pairs). Concretely, replace the frob calculation that uses
np.linalg.norm(diff, "fro") with a norm over np.triu(diff, k=1) (or equivalent
selection) so that frob and the normalise divisor n_pairs (based on elements)
measure the same set of entries; adjust no other logic around normalise, diff,
np.linalg.norm, n_pairs, or elements.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Re-opens the compare module work originally proposed in #182. The previous PR was closed pending review; this version rebases onto current
main(post-v0.7.2) and addresses all four CodeRabbit findings from the earlier review.Closes #9.
What's in
elementembeddings.compareembedding_similaritymantel_testkl_divergencefrobenius_distancepairwise_embedding_comparisonembedding_similarityover all pairs and return a DataFrameIncludes an example script (
examples/comparisons/compare_embeddings.py) that produces a comparison heatmap, UMAP map of embedding schemes, and Mantel test results across the bundled embeddings.Changes vs PR #182
mantel_testnow raisesValueErroron unknownmethodinstead of silently falling back to Spearman.mantel_testp-value is two-sided (abs(perm_corr) >= abs(observed)) so strong negative correlations report as significant rather than ~1.0._get_common_elementsimport in the example script."fro"(Frobenius) to the codespell ignore list.tests/test_compare.pycovers self-similarity (= 1), invalid-argument paths, zero-on-self for KL and Frobenius, and pairwise-comparison shape/symmetry.Test plan
pre-commit run --all-filesclean (ruff, ruff-format, pyright, codespell, prettier)pytest src/elementembeddings/tests— 61 passed (52 existing + 9 new)magpie/mat2vec/megnet16Summary by CodeRabbit
New Features
Tests