Plotting structure-level vs. atomic-level diversity using pet-oam-xl in Chemiscope

[amannsonii3]

Hi everyone,

I am working on extracting the chemically most diverse structures from a dataset of ligand-coated nanoparticles, and I am trying to utilise the PETMADFeaturizer with chemiscope (link to chemiscope examples scripts) to visualise this.

Currently, I am using the following setup:

from ase.io import read
import chemiscope

structures_list = read(
    structures_path,
    index=":",
)

featurizer = chemiscope.metatomic_featurizer(
    model="/path/to/pet-oam-xl-v1.0.0.pt",
    check_consistency=True,
)

chemiscope.explore(
    structures=structures_list,
    featurizer=featurizer,
    environments=chemiscope.all_atomic_environments(structures_list),
)

I have a couple of questions regarding this implementation:

• Atomic vs. Structure Level: Am I correct in understanding that this code will only capture the atomic-level diversity ?

• Structure-Level Implementation: To plot structure-level diversity across the dataset, will I need to additionally implement the concatenation-based (or pooling) formulation described in the MAD paper, or is there a native way to pass global environments?

If anyone could point me toward existing code snippets or examples for plotting structure-based diversity using Chemiscope and these MLIPs, it would be incredibly helpful.

[Luthaf]

I moved this discussion to the chemiscope repository, since it is more about how to use chemiscope than how to use upet!

Atomic vs. Structure Level: Am I correct in understanding that this code will only capture the atomic-level diversity ?

Yes, this will plot one point per atomic environment.

Structure-Level Implementation: To plot structure-level diversity across the dataset, will I need to additionally implement the concatenation-based (or pooling) formulation described in the MAD paper, or is there a native way to pass global environments?

I'm not sure what will happen inside the pet-oam-xl model, but you can try removing the environments argument to chemiscope.explore and see what happens, this function should default to per-structure representations.

It is likely that you will have to implement some kind of pooling yourself here though, as a custom metatomic model that calls pet-oam-xl internally. See https://chemiscope.org/docs/examples/8-explore-with-metatomic.html#defining-a-custom-model for a simple example, and we can help you get started with doing the same here.

---

One more point: the pet-mad featurizer in particular is a bit more complex than using pet-oam-xl as a chemiscope.metatomic_featurizer. PetMadFeaturizer is a custom metatomic model that takes pet-mad last layer features for the energy head (and not the pet features), and then applies a small MLP trained to reproduce the SketchMap dimensionality reduction on these.

You could reproduce this setup with pet-oam-xl, but it would require re-training a custom featurizer. @sofiia-chorna do we have scripts we can share to train something similar to the PetMadFeaturizer?

[sofiia-chorna]

Yes, here is the script to re-train a featurizer: https://github.com/lab-cosmo/upet/blob/mad-explorer-script/src/pet_mad/explore/explorer_train_script.py

It takes high-dimensional last layer features from pet-mad and low-dimensional projections as targets (we used sketchmap to obtain them) and trains a small MLP to learn the mapping. Then the model can be directly used in PETMADFeaturizer:

 featurizer = PETMADFeaturizer(
      checkpoint_path="mad_explorer_checkpoint.ckpt",
      pet_checkpoint_path="pet-oam-xl-v1.0.0.pt",
  )

Or it can be exported as a metatomic model and used with chemiscope.metatomic_featurizer

[amannsonii3]

Commenting out the environments=chemiscope.all_atomic_environments(structures_list) works and gives me the per-structure based diversity analysis.
However I'm not sure if the pooling implemented in chemiscope.explore is the one shared in the paper, or if some other pooling mechanism is used?

But this would anyways be wrong to use because:
What I understood from your last point about the featurizer is that:

• chemiscope.metatomic_featurizer is not model agnostic and was designed for pet-mad model not pet-oam-xl models so inorder to use that I'll have to train my own chemiscope.metatomic_featurizer. (This is probably because the small MLP used to go from energy head atomic latent vector to the new vector outputted by featurizer was trained on pet-mad-v1.0.0 model and not the new pet-oam-xl-v1.0.0 model)
• The way to train it will be to use the script shared by @sofiia-chorna

Can you please correct my understanding?

[Luthaf]

chemiscope.metatomic_featurizer is not model agnostic and was designed for pet-mad model not pet-oam-xl models so inorder to use that I'll have to train my own

Not quite, but almost. chemiscope.metatomic_featurizer is model agnostic and accept any metatomic with a features output.

Both the raw pet-mad and pet-oam-xl MLIP models have a "features" output that corresponds to the internal PET features, before they are used for the various trained heads. These features are not very helpful to visualize because they are high-dimensional (I think 128 or 256 dimensions, but I'm not certain). So you can use these with chemiscope, but it is not going to be useful.

For pet-mad, @sofiia-chorna trained a separate PETMADFeaturizer model, with a separate checkpoint that wrap the raw pet-mad MLIP, and then projects the high dimensional features to 2 or 3 dimensions, and exposes this as a metatomic "features" output. If you want to do the same with pet-oam-xl, you'll need to train this "projection" wrapper with the scripts shared.

[amannsonii3]

Ok,

Both the raw pet-mad and pet-oam-xl MLIP models have a "features" output that corresponds to the internal PET features, before they are used for the various trained heads

Are there two types of features?

1. Raw internal PET features - these are features used for visualisations, containing geometry and atomic-env rich information not directly used for energy predictions, but for the energy head feature vector building later...?
2. Energy head features - these are the last layer features used for energy prediction.

I think there's only one feature vector we are talking about here, which is the last layer feature vector before the energy prediction. (i.e. the "raw internal PET features" mean the same as "ast layer features for the energy head ")

So you can use these with chemiscope, but it is not going to be useful.

Chemiscope can squash down the u-MLIP latent space vector down to 2/3D for visualisations automatically; however, this is not useful because it uses PCA instead of more robust algorithms like Sketchmap.

P.S. The link page is showing me a 404 page not found error features

[Luthaf]

Are there two types of features?

Yes, each model can define features how they want, and potentially multiple features.

What's a bit confusing here is that the PET models offer "features" which are NOT the last layer features for energy. These (raw PET features) corresponds to the internal model features after the different transformer and message passing layers are applied in the PET architecture.

The features used for visualization in the paper are NOT built on top of the internal PET features, but rather on top of the last layer features for energy. I think the scripts linked above handle this (these features are accessible as a custom mtt::energy_last_layer_features output). In general, the internal raw pet features are not very useful for visualization, since they are unlikely to be rotationally invariant.

So in summary we have

• pet-mad, with two outputs (features and mtt::energy_last_layer_features)
• pet-mad-featurizer, with one output, features, build on top of pet-mad's mtt::energy_last_layer_features using custom code and SketchMap
• pet-oam-xl, with two outputs (features and mtt::energy_last_layer_features)

Let me know if this is clear.

Chemiscope can squash down the u-MLIP latent space vector down to 2/3D for visualisations automatically; however, this is not useful because it uses PCA instead of more robust algorithms like Sketchmap.

No, chemiscope will not even do PCA by default (unless it was changed while I was not looking, @sofiia-chorna?). It will instead allow you to look into a slice of the 128-dimensional space (i.e. look at the sub-space defined by say dimension 62 and 103).

P.S. The link page is showing me a 404 page not found error features

Sorry, I updated the link, and here it is again: https://docs.metatensor.org/metatomic/latest/outputs/features.html