Sample and Feature Selection with FPS and CUR

Authors:

Davide Tisi @DavideTisi and Hanna Tuerk @HannaTuerk

In this tutorial we generate descriptors using featomic, then select a subset of structures using both the farthest-point sampling (FPS) and CUR algorithms implemented in scikit-matter. Finally, we also generate a selection of the most important features using the same techniques.

First, import all the necessary packages

import ase.io
import chemiscope
import metatensor
import numpy as np
from featomic import SoapPowerSpectrum
from matplotlib import pyplot as plt
from sklearn.decomposition import PCA
from skmatter import feature_selection, sample_selection


# Note that you will need a the specific new version of skmatter (state: Feb. 2025).
# The link is provided in the environment.yml file that you can find on github.

Load molecular data

Load 500 example BTO structures from file, reading them using ASE. The data is available as a download in the data.zip file below or in the atomistic-cookbook github.

# Load a subset of `structures <data/input-fps.xyz>` of the example dataset
n_frames = 500
frames = ase.io.read("data/input-fps.xyz", f":{n_frames}", format="extxyz")

Compute SOAP descriptors using featomic

First, define the featomic hyperparameters used to compute SOAP.

# featomic hyperparameters
hypers = {
    "cutoff": {"radius": 6.0, "smoothing": {"type": "ShiftedCosine", "width": 0.5}},
    "density": {
        "type": "Gaussian",
        "width": 0.3,
        "scaling": {"type": "Willatt2018", "exponent": 4, "rate": 1, "scale": 3.5},
    },
    "basis": {
        "type": "TensorProduct",
        "max_angular": 6,
        "radial": {"type": "Gto", "max_radial": 7},
    },
}

# Generate a SOAP power spectrum
calculator = SoapPowerSpectrum(**hypers)
rho2i = calculator.compute(frames)


# Makes a dense block
atom_soap = rho2i.keys_to_properties(["neighbor_1_type", "neighbor_2_type"])

atom_soap_single_block = atom_soap.keys_to_samples(keys_to_move=["center_type"])

# Sum over atomic centers to compute structure features
struct_soap = metatensor.sum_over_samples(
    atom_soap_single_block, sample_names=["atom", "center_type"]
)


print("atom feature descriptor shape:", atom_soap.block(0).values.shape)
print(
    "atom feature descriptor (all in one block) shape:",
    atom_soap_single_block.block(0).values.shape,
)
print("structure feature descriptor shape:", struct_soap.block(0).values.shape)

atom feature descriptor shape: (12000, 2688)
atom feature descriptor (all in one block) shape: (20000, 2688)
structure feature descriptor shape: (500, 2688)

Perform atomic environment (i.e. sample) selection

Using FPS and CUR algorithms, we can perform selection of atomic environments. These are implemented in skmatter which uses data stored in the metatensor format.

Suppose we want to select the 10 most diverse environments for each chemical species.

First, we can use the keys_to_properties operation in metatensor to move the neighbor species indices to the properties of the TensorBlocks. The resulting descriptor will be a TensorMap comprised of three blocks, one for each chemical species, where the chemical species indices are solely present in the keys.

print("----Atomic environment selection-----")

print(atom_soap)
print(atom_soap.block(0))

----Atomic environment selection-----
TensorMap with 3 blocks
keys: center_type
           8
          22
          56
TensorBlock
    samples (12000): ['system', 'atom']
    components (): []
    properties (2688): ['neighbor_1_type', 'neighbor_2_type', 'l', 'n_1', 'n_2']
    gradients: None

select 10 atomic environments for each chemical species.

# Define the number of structures *per block* to select using FPS
n_envs = 10

# FPS sample selection
for key, block in atom_soap.items():
    sample_fps = sample_selection.FPS(n_to_select=n_envs, initialize="random").fit(
        atom_soap.block(key).values
    )
    sample_fps_idxs = sample_fps.selected_idx_

    # Print the selected envs for this block
    print("atomic envs selected with FPS:\n")

    selected_structures_idx = atom_soap.block(key).samples.values[sample_fps_idxs]
    newblock = metatensor.slice_block(block, axis="samples", selection=sample_fps_idxs)
    print("center_type:", key, "\n(struct_idx, atom_idx)\n", newblock.samples.values)

# CUR sample selection
for key, block in atom_soap.items():
    sample_cur = sample_selection.CUR(n_to_select=n_envs).fit(
        atom_soap.block(key).values
    )

    print("atomic envs selected with CUR:\n")
    newblock = metatensor.slice_block(
        block, axis="samples", selection=sample_cur.selected_idx_
    )
    print("center_type:", key, "\n(struct_idx, atom_idx)\n", newblock.samples.values)

atomic envs selected with FPS:

center_type: LabelsEntry(center_type=8)
(struct_idx, atom_idx)
 [[113  36]
 [140  21]
 [339  23]
 [ 68  18]
 [339  16]
 [347  37]
 [341  17]
 [339  24]
 [436  22]
 [285  20]]
atomic envs selected with FPS:

center_type: LabelsEntry(center_type=22)
(struct_idx, atom_idx)
 [[341  12]
 [166  15]
 [216   8]
 [285   9]
 [ 19   9]
 [198  12]
 [ 55  13]
 [466   9]
 [324   8]
 [433  13]]
atomic envs selected with FPS:

center_type: LabelsEntry(center_type=56)
(struct_idx, atom_idx)
 [[341   4]
 [407   0]
 [238   3]
 [289   6]
 [ 40   7]
 [407   7]
 [140   2]
 [451   7]
 [339   3]
 [436   6]]
atomic envs selected with CUR:

center_type: LabelsEntry(center_type=8)
(struct_idx, atom_idx)
 [[ 77  30]
 [339  36]
 [267  32]
 [436  19]
 [339  24]
 [341  17]
 [ 55  21]
 [ 68  20]
 [336  33]
 [198  36]]
atomic envs selected with CUR:

center_type: LabelsEntry(center_type=22)
(struct_idx, atom_idx)
 [[166  15]
 [216   8]
 [285   9]
 [ 70  10]
 [326  10]
 [ 10  39]
 [466  10]
 [170  14]
 [130  10]
 [ 40  10]]
atomic envs selected with CUR:

center_type: LabelsEntry(center_type=56)
(struct_idx, atom_idx)
 [[407   0]
 [219   7]
 [172   3]
 [339   5]
 [ 77   3]
 [289   6]
 [339   6]
 [296   1]
 [ 40   7]
 [436   2]]

Selecting from a combined pool of atomic environments

One can also select from a combined pool of atomic environments and structures, instead of selecting an equal number of atomic environments for each chemical species. In this case, we can move the ‘center_type’ key to samples such that our descriptor is a TensorMap consisting of a single block. Upon sample selection, the most diverse atomic environments will be selected, regardless of their chemical species.

print("----All atomic environment selection-----")

atom_soap_single_block = atom_soap.keys_to_samples(keys_to_move=["center_type"])
print("keys", atom_soap_single_block.keys)
print("blocks", atom_soap_single_block[0])
print("samples in first and only block", atom_soap_single_block[0].samples)

# Using the original SOAP descriptor, move all keys to properties.

# Define the number of structures to select using FPS
n_envs = 10

# FPS sample selection
sample_fps = sample_selection.FPS(n_to_select=n_envs, initialize="random").fit(
    atom_soap_single_block.block(0).values
)

sample_fps_idxs = sample_fps.selected_idx_
selected_structures_idxs = atom_soap_single_block.block(0).samples["system"][
    sample_fps_idxs
]
newblock = metatensor.slice_block(
    atom_soap_single_block.block(0), axis="samples", selection=sample_fps.selected_idx_
)

print(
    "atomic envs selected with FPS: \n (struct_idx, atom_idx, center_type) \n",
    newblock.samples.values,
)

----All atomic environment selection-----
keys Labels(
    _
    0
)
blocks TensorBlock
    samples (20000): ['system', 'atom', 'center_type']
    components (): []
    properties (2688): ['neighbor_1_type', 'neighbor_2_type', 'l', 'n_1', 'n_2']
    gradients: None
samples in first and only block Labels(
    system  atom  center_type
      0      0        56
      0      1        56
              ...
     499     38       22
     499     39       22
)
atomic envs selected with FPS:
 (struct_idx, atom_idx, center_type)
 [[ 68  12  22]
 [407   0  56]
 [ 77  30   8]
 [289   6  56]
 [166  15  22]
 [216   8  22]
 [140  21   8]
 [460   4  56]
 [339  18   8]
 [466  34   8]]

Perform structure (i.e. sample) selection with FPS/CUR

Instead of atomic environments, one can also select diverse structures. We can use the sum_over_samples operation in metatensor to define features in the structural basis instead of the atomic basis. This is done by summing over the atomic environments, labeled by the ‘center’ index in the samples of the TensorMap.

Alternatively, one could use the mean_over_samples operation, depending on the specific inhomogeneity of the size of the structures in the training set.

print("----Structure selection-----")

struct_soap = metatensor.sum_over_samples(
    atom_soap_single_block, sample_names=["atom", "center_type"]
)

print("keys", struct_soap.keys)
print("blocks", struct_soap[0])
print("samples in first block", struct_soap[0].samples)

# Define the number of structures to select *per block* using FPS
n_structures = 10

# FPS structure selection
sample_fps = sample_selection.FPS(n_to_select=n_structures, initialize="random").fit(
    struct_soap.block(0).values
)
struct_fps_idxs = sample_fps.selected_idx_
print("structures selected with FPS:\n", sample_fps.selected_idx_)


# CUR structure selection
sample_cur = sample_selection.CUR(n_to_select=n_structures).fit(
    struct_soap.block(0).values
)

struct_cur_idxs = sample_cur.selected_idx_
print("structures selected with CUR:\n", struct_cur_idxs)

# Slice structure descriptor along axis 0 to contain only the selected structures
struct_soap_fps = struct_soap.block(0).values[struct_fps_idxs, :]
struct_soap_cur = struct_soap.block(0).values[struct_cur_idxs, :]
assert struct_soap_fps.shape == struct_soap_cur.shape

print("Structure descriptor shape before selection ", struct_soap.block(0).values.shape)
print("Structure descriptor shape after selection (FPS)", struct_soap_fps.shape)
print("Structure descriptor shape after selection (CUR)", struct_soap_cur.shape)

----Structure selection-----
keys Labels(
    _
    0
)
blocks TensorBlock
    samples (500): ['system']
    components (): []
    properties (2688): ['neighbor_1_type', 'neighbor_2_type', 'l', 'n_1', 'n_2']
    gradients: None
samples in first block Labels(
    system
      0
      1
     ...
     498
     499
)
structures selected with FPS:
 [172  15 356 257 326 317 167  71 140 496]
structures selected with CUR:
 [438 140 398 289 110 386 326  40  68  39]
Structure descriptor shape before selection  (500, 2688)
Structure descriptor shape after selection (FPS) (10, 2688)
Structure descriptor shape after selection (CUR) (10, 2688)

Visualize selected structures

sklearn can be used to perform PCA dimensionality reduction on the SOAP descriptors. The resulting PC coordinates can be used to visualize the the data alongside their structures in a chemiscope widget.

# Generate a structure PCA
struct_soap_pca = PCA(n_components=2).fit_transform(struct_soap.block(0).values)
assert struct_soap_pca.shape == (n_frames, 2)

Plot the PCA map

Notice how the selected points avoid the densely-sampled area, and cover the periphery of the dataset

# Matplotlib plot
fig, ax = plt.subplots(1, 1, figsize=(6, 4))
scatter = ax.scatter(struct_soap_pca[:, 0], struct_soap_pca[:, 1], c="red")
ax.plot(
    struct_soap_pca[struct_cur_idxs, 0],
    struct_soap_pca[struct_cur_idxs, 1],
    "ko",
    fillstyle="none",
    label="FPS selection",
)
ax.set_xlabel("PCA[1]")
ax.set_ylabel("PCA[2]")
ax.legend()
fig.show()

Creates a chemiscope viewer

# Selected level
selection_levels = []
for i in range(len(frames)):
    level = 0
    if i in struct_cur_idxs:
        level += 1
    if i in struct_fps_idxs:
        level += 2
    if level == 0:
        level = "Not selected"
    elif level == 1:
        level = "CUR"
    elif level == 2:
        level = "FPS"
    else:
        level = "FPS+CUR"
    selection_levels.append(level)

properties = chemiscope.extract_properties(frames)

properties.update(
    {
        "PC1": struct_soap_pca[:, 0],
        "PC2": struct_soap_pca[:, 1],
        "selection": np.array(selection_levels),
    }
)

widget = chemiscope.show(
    frames,
    properties=properties,
    settings={
        "map": {
            "x": {"property": "PC1"},
            "y": {"property": "PC2"},
            "color": {"property": "energy"},
            "symbol": "selection",
            "size": {"factor": 50},
        },
        "structure": [{"unitCell": True}],
    },
)
widget.save("sample-selection.json.gz")

# display, if in notebook or sphinx
widget

×

/home/runner/work/atomistic-cookbook/atomistic-cookbook/.nox/sample-selection/lib/python3.11/site-packages/chemiscope/structures/_ase.py:147: UserWarning: the following structure properties are only defined for a subset of frames: ['stress']; they will be ignored
  all_properties = _ase_get_structure_properties(frames)

Perform feature selection

Now perform feature selection to reduce the size of the features. In this example we will go back to using the descriptor decomposed into atomic environments, as opposed to the one decomposed into structure environments, but only use FPS for brevity.

print("----Feature selection-----")
print("keys", atom_soap_single_block.keys)
print("blocks", atom_soap_single_block[0])
print("samples in first block", atom_soap_single_block[0].properties)

# Define the number of features to select
n_features = 200

# FPS feature selection
feat_fps = feature_selection.FPS(n_to_select=n_features, initialize="random").fit(
    atom_soap_single_block.block(0).values
)
feat_fps_idxs = feat_fps.selected_idx_
atom_soap_single_block_fps = metatensor.slice_block(
    atom_soap_single_block.block(0), axis="properties", selection=feat_fps_idxs
)

# Slice atomic descriptor along axis 1 to contain only the selected features
print(
    "atomic descriptor shape before selection ",
    atom_soap_single_block.block(0).values.shape,
)
print(
    "atomic descriptor shape after selection ",
    atom_soap_single_block_fps.values.shape,
)

----Feature selection-----
keys Labels(
    _
    0
)
blocks TensorBlock
    samples (20000): ['system', 'atom', 'center_type']
    components (): []
    properties (2688): ['neighbor_1_type', 'neighbor_2_type', 'l', 'n_1', 'n_2']
    gradients: None
samples in first block Labels(
    neighbor_1_type  neighbor_2_type  l  n_1  n_2
           8                8         0   0    0
           8                8         0   0    1
                        ...
          56               56         6   7    6
          56               56         6   7    7
)
atomic descriptor shape before selection  (20000, 2688)
atomic descriptor shape after selection  (20000, 200)