Magnetic materials embedding and dimensionality reduction

Here we have an HTML file generated from Python and Plotly that displays ~5,000 magnetic materials in 3 dimensions. To generate this visual, I took each material and "embedded" it by running it through the Orb v2 MLIP model and extracted the latent space part of the GNN.

I then reduce 256 dimensional output into just 3 with UMAP and visualized each, colored by their magnetic density.

magnetic_materials_umap_visualization

.html file

A 3D interactive scatter plot of magnetic materials from Materials Project. Points are colored by the materials estimated magnetic density and poitioned by UMAP reduction of Orb v2 model latent space

6mo

The visualization is interactive. You can zoom around, find clusters, and see what Materials Project material each is. I've found some interesting clusters of high density while still seemingly different chemical systems.

The idea is to use these dimensions and the "location" of known permanent magnets and find similar materials. It's another layer of candidate filtering. We're basically looking for ranges on the values of the latent features that may lead to good magnets. Ideally, it will encode anisotropy or coercivity, two important features we haven't been able to build a model for yet.

If would could find a material like NdFeB but without rare-earths, that would be amazing.

I'll work on a dataset similar to the recent one I shared of magnetic materials with Curie temperature and magnetic density, with an added column of cosine similarity or L2 distance to a couple known permanent magnets.

I tried the same thing with superconductors to see if we could find clusters of high-Tc materials. Check that out here:

Interactive 3D latent space exploration with Tc

.html file

Using the 256 dimensional latent space output from the Orb model, we visualize the 3DSC(MP) dataset using t-SNE and UMAP. The UMAP projection has been given the target for learning a manifold that keeps similar Tc materials close together.

9mo

Code

python

#!/usr/bin/env python3
import pandas as pd
import numpy as np
import plotly.express as px
from umap import UMAP
import plotly.io as pio

# Set plotly to render in browser for better interaction
pio.renderers.default = "browser"

# Load the data
print("Loading data...")
df = pd.read_csv("magnetic_features_with_density.csv")

# Separate features and target
target = "magnetic_density"
# total_magnetic_moment
feature_cols = [col for col in df.columns if col.startswith("orb_feat_")]
X = df[feature_cols].values
y = df[target].values

print(f"Performing UMAP on {len(feature_cols)} features...")
# Initialize and fit UMAP
umap_3d = UMAP(
    n_components=3,
    n_neighbors=15,
    min_dist=0.1,
    metric="euclidean",
    target_metric="euclidean",
    target_weight=0.5,
    random_state=42,
)

# Fit and transform the data
embedding = umap_3d.fit_transform(
    X,
)

# Create a new dataframe with the embeddings
viz_df = pd.DataFrame(
    {
        "UMAP1": embedding[:, 0],
        "UMAP2": embedding[:, 1],
        "UMAP3": embedding[:, 2],
        target: y,
        "MP ID": df["mp_id"],
        "formula": df["cif_path"].str.replace("cif_path/", ""),
    }
)

print("Creating 3D visualization...")
# Create an interactive 3D scatter plot
fig = px.scatter_3d(
    viz_df,
    x="UMAP1",
    y="UMAP2",
    z="UMAP3",
    color=target,
    hover_data=["formula"],
    color_continuous_scale="Viridis",
    title=f"UMAP visualization colored by {target}",
)

# Update the layout for better visualization
fig.update_layout(
    scene=dict(xaxis_title="UMAP1", yaxis_title="UMAP2", zaxis_title="UMAP3"),
    width=1000,
    height=800,
)

# Show the plot
fig.show()

# Save the plot as HTML for later viewing
fig.write_html("magnetic_umap_visualization.html")
print("Visualization saved as 'magnetic_umap_visualization.html'")

Dataset

magnetic_features_with_density.csv

.csv file

6mo

posts