<jupyter_start><jupyter_text>Adversarial attacks on ammonia using the pretrained modelsIn this notebook, we perform an adversarial attack on zeolites using the [SchNet NN potential](https://github.com/learningmatter-mit/NeuralForceField). We will be using the third generation of ammonia models, as shown in [our paper](https://arxiv.org/abs/2101.11588).The utilities at `nff` will be used to perform the energy/force predictions. `nglview` will be used to visualize the generated trajectories. A few utility functions from this repo, `robust`, will be used as well. For the sake of generality, all steps for performing the adversarial attack are shown in this notebook.<jupyter_code>import os
import sys
sys.path.append('..')

import robust as rb
import torch as ch
import numpy as np
import pandas as pd

import matplotlib.pyplot as plt

from nff.io import NeuralFF, AtomsBatch, EnsembleNFF
from nff.data import Dataset
from nff.train import load_model

from ase.io import Trajectory, read

import nglview as nv<jupyter_output><empty_output><jupyter_text>Loading the dataset, models and initial geometry<jupyter_code>DEVICE = 2

dset = Dataset.from_file('../data/ammonia.pth.tar')

PATH = '../models/ammonia'
models = []
for model_name in sorted(os.listdir(PATH)):
    m = NeuralFF.from_file(os.path.join(PATH, model_name), device=DEVICE).model
    models.append(m)

ensemble = EnsembleNFF(models, device=DEVICE)
CUTOFF = 5

def get_atoms(props):
    atoms = AtomsBatch(
        positions=props['nxyz'][:, 1:],
        numbers=props['nxyz'][:, 0],
        cutoff=CUTOFF,
        props={'energy': 0, 'energy_grad': []},
        calculator=ensemble,
        nbr_torch=False,
        device=DEVICE,
    )
    _ = atoms.update_nbr_list()
    
    return atoms

initial = get_atoms(dset[np.argmin(dset.props['energy'])])<jupyter_output><empty_output><jupyter_text>Defining the adversarial attackThe `Attacker` class allows one to perform an adversarial attack using an ensemble of SchNet models.<jupyter_code>energy_dset = rb.PotentialDataset(
    ch.zeros_like(dset.props['energy']),
    dset.props['energy'],
    ch.zeros_like(dset.props['energy']),
)

loss_fn = rb.loss.AdvLoss(
    train=energy_dset,
    temperature=20,
)
attacker = rb.schnet.Attacker(
    initial,
    ensemble,
    loss_fn,
    device=DEVICE,
)
results = attacker.attack()<jupyter_output>100%|██████████| 60/60 [00:08<00:00,  7.05it/s]<jupyter_text>Post-processing the results: visualizing and plottingAfter the adversarial attack is performed, we can now visualize it and post-process the results. We start by recalculating the variance in forces and energies.<jupyter_code>df = pd.DataFrame(results)
df['forces_var'] = [f.var(-1).mean() for f in df['forces']]
df['energy_var'] = [e.var() for e in df['energy']]
df['energy_avg'] = [e.mean() for e in df['energy']]<jupyter_output><empty_output><jupyter_text>Then, we reconstruct the trajectory of the adversarial attack using the values of `delta` along the attack.<jupyter_code>newatoms = []
for transl in df.delta:
    at = initial.copy()
    at.translate(transl)
    newatoms.append(at)
view = nv.show_asetraj(newatoms)
view.add_unitcell()

view<jupyter_output><empty_output><jupyter_text>Finally, we can plot the trajectory of the adversarial attack based on the sampled properties shown before (see Fig. S10 of the paper). Notice how the geometry that maximizes the adversarial loss is not necessarily the one with highest energy.<jupyter_code>fig, ax = plt.subplots(figsize=(4, 5))
ax.spines['right'].set_visible(True)

ax.plot(-df.loss, color='k')

ax.set_ylabel('Adversarial Loss')

COLOR_TAX = '#d62728'
tax = ax.twinx()
tax.plot(df.energy_avg, color=COLOR_TAX)
tax.set_yticklabels(
    tax.get_yticks(),
    color=COLOR_TAX
)
tax.set_ylabel('Mean energy (kcal/mol)', color=COLOR_TAX)

ax.set_xlabel('Attack step')

plt.show()<jupyter_output>/home/dskoda/.conda/envs/nff/lib/python3.7/site-packages/ipykernel_launcher.py:13: UserWarning: FixedFormatter should only be used together with FixedLocator
  del sys.path[0]
{'fraction_non_alphanumeric': 0.1166296021744502, 'fraction_numerical': 0.0121077341240425, 'mean_word_length': 5.005934718100891, 'pattern_counts': {'":': 0, '<': 24, '<?xml version=': 0, '>': 23, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '45179', 'license': 'MIT', 'n_tokens_mistral': 1420, 'n_tokens_neox': 1341, 'n_words': 367, 'path': 'examples/Ammonia_attack.ipynb', 'repo_name': 'learningmatter-mit/Atomistic-Adversarial-Attacks', 'size': '4047'}