Converting interchangeably between crystal structure and a grayscale PNG image (basic xtal2png tutorial)

In this notebook, we will install the xtal2png package, encode/decode two example pymatgen Structure objects, and show some visualizations of the intermediate PNG representations and before/after crystal structure plots. Finally, we comment on how you can use xtal2png with state-of-the-art machine learning image models.

Installation

Install the xtal2png package. Optionally install ase and nglview which can be used to visualize crystal structures. You may need to restart the runtime via Ctrl+M, . or Runtime --> Restart Runtime (via menubar).

!pip install xtal2png
!pip install ase nglview # optional, for visualization of crystal structures

Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/
Collecting xtal2png
  Downloading xtal2png-0.1.6-py3-none-any.whl (20 kB)
Requirement already satisfied: pillow in /usr/local/lib/python3.7/dist-packages (from xtal2png) (7.1.2)
Collecting pymatgen
  Downloading pymatgen-2022.0.17.tar.gz (40.6 MB)
     |████████████████████████████████| 40.6 MB 2.0 MB/s
  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Installing backend dependencies ... done
    Preparing wheel metadata ... done
Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from xtal2png) (1.21.6)
Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from xtal2png) (4.11.3)
Requirement already satisfied: typing-extensions>=3.6.4 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->xtal2png) (4.2.0)
Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->xtal2png) (3.8.0)
Requirement already satisfied: palettable>=3.1.1 in /usr/local/lib/python3.7/dist-packages (from pymatgen->xtal2png) (3.3.0)
Collecting monty>=3.0.2
  Downloading monty-2022.4.26-py3-none-any.whl (65 kB)
     |████████████████████████████████| 65 kB 3.9 MB/s
Requirement already satisfied: plotly>=4.5.0 in /usr/local/lib/python3.7/dist-packages (from pymatgen->xtal2png) (5.5.0)
Collecting uncertainties>=3.1.4
  Downloading uncertainties-3.1.6-py2.py3-none-any.whl (98 kB)
     |████████████████████████████████| 98 kB 6.4 MB/s
Collecting ruamel.yaml>=0.15.6
  Downloading ruamel.yaml-0.17.21-py3-none-any.whl (109 kB)
     |████████████████████████████████| 109 kB 50.8 MB/s
Collecting scipy>=1.5.0
  Downloading scipy-1.7.3-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (38.1 MB)
     |████████████████████████████████| 38.1 MB 1.2 MB/s
Requirement already satisfied: matplotlib>=1.5 in /usr/local/lib/python3.7/dist-packages (from pymatgen->xtal2png) (3.2.2)
Requirement already satisfied: sympy in /usr/local/lib/python3.7/dist-packages (from pymatgen->xtal2png) (1.7.1)
Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from pymatgen->xtal2png) (2.23.0)
Requirement already satisfied: pandas in /usr/local/lib/python3.7/dist-packages (from pymatgen->xtal2png) (1.3.5)
Requirement already satisfied: tabulate in /usr/local/lib/python3.7/dist-packages (from pymatgen->xtal2png) (0.8.9)
Requirement already satisfied: networkx>=2.2 in /usr/local/lib/python3.7/dist-packages (from pymatgen->xtal2png) (2.6.3)
Collecting spglib>=1.9.9.44
  Downloading spglib-1.16.5-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (325 kB)
     |████████████████████████████████| 325 kB 64.7 MB/s
Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib>=1.5->pymatgen->xtal2png) (1.4.2)
Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib>=1.5->pymatgen->xtal2png) (3.0.9)
Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib>=1.5->pymatgen->xtal2png) (2.8.2)
Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib>=1.5->pymatgen->xtal2png) (0.11.0)
Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from plotly>=4.5.0->pymatgen->xtal2png) (1.15.0)
Requirement already satisfied: tenacity>=6.2.0 in /usr/local/lib/python3.7/dist-packages (from plotly>=4.5.0->pymatgen->xtal2png) (8.0.1)
Collecting ruamel.yaml.clib>=0.2.6
  Downloading ruamel.yaml.clib-0.2.6-cp37-cp37m-manylinux1_x86_64.whl (546 kB)
     |████████████████████████████████| 546 kB 52.2 MB/s
Requirement already satisfied: future in /usr/local/lib/python3.7/dist-packages (from uncertainties>=3.1.4->pymatgen->xtal2png) (0.16.0)
Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.7/dist-packages (from pandas->pymatgen->xtal2png) (2022.1)
Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->pymatgen->xtal2png) (2022.5.18.1)
Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->pymatgen->xtal2png) (2.10)
Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->pymatgen->xtal2png) (3.0.4)
Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->pymatgen->xtal2png) (1.24.3)
Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.7/dist-packages (from sympy->pymatgen->xtal2png) (1.2.1)
Building wheels for collected packages: pymatgen

Encode(/decode) two pymatgen Structure objects

Import a list of two example pymatgen Structure objects (these correspond to mp-560471/\(Zn_2B_2PbO_6\) and mp-7823/\(V_2NiSe_4\), respectively)

from xtal2png.utils.data import example_structures
from xtal2png.core import XtalConverter

Let’s take a look at the second Structure which has a smaller footprint.

example_structures[1]

We will be using the XtalConverter class, for which more information including its __init__ arguments and functions can be displayed via help(XtalConverter). For just the parameters for class instantiation, try help(XtalConverter.__init__). Note that max_sites is not tested for values other than 52.

help(XtalConverter.__init__)

Let’s specify the save directory (save_dir) for the PNG files as "data", which will be automatically created. In this case, it will be saved to temporary Google Colab storage.

xc = XtalConverter(save_dir="data") # DFT surrogate relaxation via m3gnet by default
data = xc.xtal2png(example_structures, save=True)
relaxed_decoded_structures = xc.png2xtal(data, save=False)

We also take a look at the unrelaxed decoded structures.

xc = XtalConverter(save_dir="data", relax_on_decode=False)
data = xc.xtal2png(example_structures, save=True)
decoded_structures = xc.png2xtal(data, save=False)

Visualization

For visualization, we’ll cover two aspects: the structure-encoded PNG images and visualizing before/after crystal structures.

Structure-encoded PNG images

Note that images won’t show via im.show() command on Google Colab even if you specify xc.xtal2png(..., show=True, ...), so for this Colab example we’ll open the images ad-hoc based on where they were saved in local Colab storage. We display the images stacked one on top of another using display(im) instead of im.show(). Note that the filepaths have the chemical formula, volume, and a randomly generated uid portion to promote uniqueness, especially when dealing with allotropes (same chemical formula, different crystal structure).

import glob, os
from PIL import Image
for fpath in glob.glob("data/*.png"):
    with Image.open(fpath) as im:
      im = im.resize((64*5, 64*5), Image.BOX)
      print(fpath)
      display(im)

As mentioned in the README, the legend key for these images is as follows:

        

Also described in the README, the match between the encoded and decoded versions is within an expected tolerance, given that PNG images are represented as discrete RGB values between 0 and 255 (i.e. there is a round-off error).

Original	Decoded	Relaxed Decoded
Structure Summary Lattice abc : 5.033788 11.523021 10.74117 angles : 90.0 90.0 90.0 volume : 623.0356027127609 A : 5.033788 0.0 3.0823061808931787e-16 B : 1.8530431062799525e-15 11.523021 7.055815392078867e-16 C : 0.0 0.0 10.74117 PeriodicSite: Zn2+ (0.9120, 5.7699, 9.1255) [0.1812, 0.5007, 0.8496] PeriodicSite: Zn2+ (4.1218, 5.7531, 1.6156) [0.8188, 0.4993, 0.1504] ...	Structure Summary Lattice abc : 5.0250980392156865 11.533333333333331 10.8 angles : 90.0 90.0 90.0 volume : 625.9262117647058 A : 5.0250980392156865 0.0 0.0 B : 0.0 11.533333333333331 0.0 C : 0.0 0.0 10.8 PeriodicSite: Zn (0.9016, 5.7780, 3.8012) [0.1794, 0.5010, 0.3520] PeriodicSite: Zn (4.1235, 5.7554, 6.9988) [0.8206, 0.4990, 0.6480] ...	Structure Summary Lattice abc : 5.026834307381214 11.578854613685237 10.724087971087924 angles : 90.0 90.0 90.0 volume : 624.1953646135236 A : 5.026834307381214 0.0 0.0 B : 0.0 11.578854613685237 0.0 C : 0.0 0.0 10.724087971087924 PeriodicSite: Zn (0.9050, 5.7978, 3.7547) [0.1800, 0.5007, 0.3501] PeriodicSite: Zn (4.1218, 5.7810, 6.9693) [0.8200, 0.4993, 0.6499] ...

Before/after Crystal Structure Visualization

To visualize the crystal structures before and after, we can use nglview after a bit of Colab finnagling with external ipywidgets.

from google.colab import output
output.enable_custom_widget_manager()

from pymatgen.io.ase import AseAtomsAdaptor
from ase.visualize import view
aaa = AseAtomsAdaptor()

[display(view(aaa.get_atoms(s), viewer='ngl')) for s in example_structures]

[display(view(aaa.get_atoms(s), viewer='ngl')) for s in decoded_structures]

[display(view(aaa.get_atoms(s), viewer='ngl')) for s in relaxed_decoded_structures]

Undo the Colab finnagling of external ipywidgets.

from google.colab import output
output.disable_custom_widget_manager()

Final Remarks

This tool makes it possible to use state-of-the-art image-based machine learning models with minimal “plumbing” required. Just follow the normal instructions for custom image datasets. For example, this can be used with Palette, an image-to-image guided diffusion model by Google, which has an unofficial implementation here.