ElementEmbeddings

The Element Embeddings package provides high-level tools for analysing elemental and ionic species embeddings data. This primarily involves visualising the correlation between embedding schemes using different statistical measures.

• Documentation: https://wmd-group.github.io/ElementEmbeddings/
• Examples: https://github.com/WMD-group/ElementEmbeddings/tree/main/examples

Motivation

Machine learning approaches for materials informatics have become increasingly widespread. Some of these involve the use of deep learning techniques where the representation of the elements is learned rather than specified by the user of the model. While an important goal of machine learning training is to minimise the chosen error function to make more accurate predictions, it is also important for material scientists to be able to interpret these models. As such, we aim to evaluate and compare different atomic embedding schemes in a consistent framework.

Getting started

ElementEmbeddings's main feature, the Embedding class, is accessible by importing the class.

Installation

The latest stable release can be installed via pip or uv:

pip install ElementEmbeddings
uv pip install ElementEmbeddings

Alternatively, ElementEmbeddings is available via conda through the conda-forge channel on Anaconda Cloud:

conda install -c conda-forge elementembeddings

For installing the development or documentation dependencies from a release:

pip install "ElementEmbeddings[dev]"
pip install "ElementEmbeddings[docs]"
uv pip install "ElementEmbeddings[dev]"
uv pip install "ElementEmbeddings[docs]"

For development, clone the repository and sync the project environment with uv:

git clone https://github.com/WMD-group/ElementEmbeddings.git
cd ElementEmbeddings
uv sync --extra dev --extra docs

This creates a project virtual environment in .venv and installs the package in editable mode, so local code changes are reflected immediately.

Usage

For simple usage, you can instantiate an Embedding object using one of the embeddings in the data directory. For this example, let's use the magpie elemental representation.

# Import the class
>>> from elementembeddings.core import Embedding

# Load the magpie data
>>> magpie = Embedding.load_data("magpie")

We can access some of the properties of the Embedding class. For example, we can find the dimensions of the elemental representation and the list of elements for which an embedding exists.

# Print out some of the properties of the ElementEmbeddings class
>>> print(f"The magpie representation has embeddings of dimension {magpie.dim}")
>>> print(
...     f"The magpie representation contains these elements: \n {magpie.element_list}"
... )  # prints out all the elements considered for this representation
>>> print(
...     f"The magpie representation contains these features: \n {magpie.feature_labels}"
... )  # Prints out the feature labels of the chosen representation

The magpie representation has embeddings of dimension 22
The magpie representation contains these elements:
['H', 'He', 'Li', 'Be', 'B', 'C', 'N', 'O', 'F', 'Ne', 'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'Cl', 'Ar', 'K', 'Ca', 'Sc', 'Ti', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Ga', 'Ge', 'As', 'Se', 'Br', 'Kr', 'Rb', 'Sr', 'Y', 'Zr', 'Nb', 'Mo', 'Tc', 'Ru', 'Rh', 'Pd', 'Ag', 'Cd', 'In', 'Sn', 'Sb', 'Te', 'I', 'Xe', 'Cs', 'Ba', 'La', 'Ce', 'Pr', 'Nd', 'Pm', 'Sm', 'Eu', 'Gd', 'Tb', 'Dy', 'Ho', 'Er', 'Tm', 'Yb', 'Lu', 'Hf', 'Ta', 'W', 'Re', 'Os', 'Ir', 'Pt', 'Au', 'Hg', 'Tl', 'Pb', 'Bi', 'Po', 'At', 'Rn', 'Fr', 'Ra', 'Ac', 'Th', 'Pa', 'U', 'Np', 'Pu', 'Am', 'Cm', 'Bk']
The magpie representation contains these features:
['Number', 'MendeleevNumber', 'AtomicWeight', 'MeltingT', 'Column', 'Row', 'CovalentRadius', 'Electronegativity', 'NsValence', 'NpValence', 'NdValence', 'NfValence', 'NValence', 'NsUnfilled', 'NpUnfilled', 'NdUnfilled', 'NfUnfilled', 'NUnfilled', 'GSvolume_pa', 'GSbandgap', 'GSmagmom', 'SpaceGroupNumber']

Plotting

We can quickly generate heatmaps of distance/similarity measures between the element vectors using heatmap_plotter and plot the representations in two dimensions using the dimension_plotter from the plotter module. Before we do that, we will standardise the embedding using the standardise method available to the Embedding class

from elementembeddings.plotter import heatmap_plotter, dimension_plotter
import matplotlib.pyplot as plt

magpie.standardise(inplace=True)  # Standardises the representation

fig, ax = plt.subplots(1, 1, figsize=(6, 6))
heatmap_params = {"vmin": -1, "vmax": 1}
heatmap_plotter(
    embedding=magpie,
    metric="cosine_similarity",
    show_axislabels=False,
    cmap="Blues_r",
    ax=ax,
    **heatmap_params
)
ax.set_title("Magpie cosine similarities")
fig.tight_layout()
fig.show()

fig, ax = plt.subplots(1, 1, figsize=(6, 6))

reducer_params = {"n_neighbors": 30, "random_state": 42}
scatter_params = {"s": 100}

dimension_plotter(
    embedding=magpie,
    reducer="umap",
    n_components=2,
    ax=ax,
    adjusttext=True,
    reducer_params=reducer_params,
    scatter_params=scatter_params,
)
ax.set_title("Magpie UMAP (n_neighbours=30)")
ax.legend().remove()
handles, labels = ax1.get_legend_handles_labels()
fig.legend(handles, labels, bbox_to_anchor=(1.25, 0.5), loc="center right", ncol=1)

fig.tight_layout()
fig.show()

Compositions

The package can also be used to featurise compositions. Your data could be a list of formula strings or a pandas dataframe of the following format:

The composition_featuriser function can be used to featurise the data. The compositions can be featurised using different representation schemes and different types of pooling through the embedding and stats arguments respectively.

from elementembeddings.composition import composition_featuriser

df_featurised = composition_featuriser(df, embedding="magpie", stats=["mean", "sum"])

df_featurised

The returned dataframe contains the mean-pooled and sum-pooled features of the magpie representation for the four formulas.

Development notes

Bugs, features and questions

Please use the issue tracker to report bugs and any feature requests. Hopefully, most questions should be solvable through the docs. For any other queries related to the project, please contact Anthony Onwuli by e-mail: aonwuli@proton.me or Aron Walsh by e-mail: a.walsh@imperial.ac.uk.

Code contributions

We welcome new contributions to this project. See the contributing guide for detailed instructions on how to contribute to our project.

Add an embedding scheme

The steps required to add a new representation scheme are:

1. Add data file to data/element_representations.
2. Edit docstring table in core.py.
3. Edit utils/config.py to include the representation in DEFAULT_ELEMENT_EMBEDDINGS and CITATIONS.
4. Update the documentation element.md and README.md.

Developer

• Anthony Onwuli (Department of Materials, Imperial College London)

References

A. Onwuli et al, "Ionic species representations for materials informatics"

H. Park et al, "Mapping inorganic crystal chemical space" Faraday Discuss. (2024)

A. Onwuli et al, "Element similarity in high-dimensional materials representations" Digital Discovery 2, 1558 (2023)