Skip to content

Minimising Lennard-Jones Atom Cluster Potentials with Evolution

In this example, we'll demonstrate the application of evolution to a challenging, multi-modal black-box optimisation problem, the Lennard-Jones atom cluster potential minimisation task [1]. This problem has previously been studied with SNES [2]. The task is defined for \(N\) atoms, each with a position \(P_i = (x_i, y_i, z_i)\) in 3D space. The distance from atom \(i\) to atom \(j\) is,

\[r_{i,j} = |P_i - P_j|\]

The total atom cluster potential is,

\[E = 4 \epsilon \sum_{i = 1}^N \sum_{j = 1}^{i - 1} \left(\frac{\sigma}{r_{i,j}}\right)^{12} - \left(\frac{\sigma}{r_{i,j}}\right)^{6}\]

where here we will use the reduced units [1], \(\epsilon = \sigma = 1\). This function can be implemented in a vectorised fashion in PyTorch. Let's assume that the \(N\) positions are described be a vector \(x \in \mathbb{R}^{3N}\). Then we can define a vectorised computation of the distances:

import torch

def pairwise_distances(positions: torch.Tensor) -> torch.Tensor:
    # Assume positions has shape [B, 3N] where B is the batch size and N is the number of atoms
    # Reshaping to get individual atoms' positions of shape [B, N, 3]
    positions = positions.view(positions.shape[0], -1, 3)
    # Compute the pairwise differences
    # Subtracting [B, 1, N, 3] from [B, N, 1, 3] gives [B, N, N, 3] 
    deltas = (positions.unsqueeze(2) - positions.unsqueeze(1))
    # Norm the differences gives [B, N, N]
    distances = torch.norm(deltas, dim = -1)
    return distances

Which gives a straightforward definition of the cluster potential

def cluster_potential(positions: torch.Tensor) -> torch.Tensor:
    # Compute the pairwise distances of atoms
    distances = pairwise_distances(positions)

    # Compute the pairwise cost (1 / dist)^12 - (1 / dist)^ 6
    pairwise_cost = (1 / distances).pow(12) - (1 / distances).pow(6.)

    # Get the upper triangle matrix (ignoring the diagonal)
    ut_pairwise_cost = torch.triu(pairwise_cost, diagonal = 1)

    # 4 * Summutation of the upper triangle of pairwise costs gives potential
    potential = 4 * ut_pairwise_cost.sum(dim = (1, 2))
    return potential

Obtaining Reference Solutions

To measure our performance, its helpful to refer to a reference point. For this, we will use a publicly available database of known global optima. First, let's download the tar of optima and extract it. 

import requests
import tarfile

# Url of tar containing known global minima
url = 'http://doye.chem.ox.ac.uk/jon/structures/LJ/LJ.tar'
# Where to save the tar -- modify as desired
target_path = 'LJ_data.tar'

# Download
response = requests.get(url, stream=True)
if response.status_code == 200:
    with open(target_path, 'wb') as f:
        f.write(response.raw.read())

# Open file
file = tarfile.open(target_path)

#  By default save the data to the 'LJ_data' folder in the local directory
data_path = f'./{target_path.replace(".tar", "")}'
file.extractall(data_path) 

file.close()

Now we can plot the computed potential of each solution obtained from the data base. 

import matplotlib.pyplot as plt
import pandas

# Lists to track atom counts and potentials
atom_counts = []
global_potentials = []

# Iterate from 3 to 67 atoms -- the number visited in the paper
for n_atoms in range(3, 68):
    # File path is simply the nuimber of atoms
    file_path = f'{data_path}/{n_atoms}'
    # Get the positions as a dataframe
    dataframe = pandas.read_csv(file_path, header=None, delim_whitespace=True)
    # Make a positions tensor -- note that we add an initial dimension as the objective function is vectorised
    positions = torch.Tensor(dataframe.to_numpy()).unsqueeze(0)
    # Get the potential
    potential = cluster_potential(positions)

    # Update lists of atom counts and potentials
    atom_counts.append(n_atoms)
    global_potentials.append(potential.item())

# Simple plot
plt.plot(atom_counts, global_potentials)
plt.xlabel('Number of Atoms $N$')
plt.ylabel('Cluster potential $E$')
plt.show()

# Sanity check on the last one
print(f'Potential of N={n_atoms} is computed as {potential.item()} vs. published value -347.252007')

Benchmarking SNES

Let's consider benchmarking SNES on this problem. It is worth noting that evaluation can be done on the GPU due to how we've designed the function to minimise, so feel free to use device = 'cuda:0', as we have done, if you have a cuda-capable device available. Otherwise, you should still see some speedup with device = 'cpu' due to PyTorch's efficient implementation of batched operators even on the CPU. In any case, we'll use a population size substantially higher than the default. We'll run the code below for only the first 17 atom clusters, but you can freely push this value higher for your own interest.

from evotorch import Problem
from evotorch.algorithms import SNES

snes_potentials = []

for n_atoms in range(3, 20):

    print(f'Solving case N={n_atoms} with SNES')

    # Set up the problem in vectorised mode
    problem  = Problem(
        'min',
        cluster_potential,
        vectorized = True,
        device = 'cuda:0' if torch.cuda.is_available() else 'cpu',
        dtype = torch.float64, # Higher precision could be helpful 
        solution_length = 3 * n_atoms,
        initial_bounds = (-1e-12, 1e-12), # Taken directly from [2]
        store_solution_stats = True,  # Make sure the problem tracks the best discovered solution, even on GPU
    )

    searcher = SNES(
        problem,
        stdev_init = 0.01,  # Taken directly from [2]
        popsize = 10 * n_atoms,  # Significantly higher than [2]
        center_learning_rate = 0.5,  # Halving value from [2] slows learning
        stdev_learning_rate = 0.5,  # Halving value from [2] slows learning
        scale_learning_rate = True,  # Boolean flag means modifying the above learning rates rescales the default
    )
    searcher.run(1000 * problem.solution_length)   # 2x value used in [2], adjusted for half learning rates

    # Best solution found
    best_potential = problem.status['best_eval']
    # Center is also a good estimate
    center_potential = cluster_potential(searcher.status['center'].cpu().unsqueeze(0))[0].item()
    if center_potential < best_potential:
        best_potential = center_potential

    print(f'Found potential {best_potential}')

    snes_potentials.append(best_potential)

Finally let's take a look at how we did. We should see that for most cases, particularly for smaller atom clusters, SNES was either exactly recovering or was finding a solution very close to the known global optima of atom positions.

# Simple plot
plt.plot(atom_counts[0:len(snes_potentials)], global_potentials[0:len(snes_potentials)], label = 'Known Optima')
plt.plot(atom_counts[0:len(snes_potentials)], snes_potentials[0:len(snes_potentials)], label = 'SNES-discovered Solutions')
plt.legend()
plt.xlabel('Number of Atoms $N$')
plt.ylabel('Atom Cluster Potential $E$')
plt.show()

References

[1] Wales and Doye. "Global optimization by basin-hopping and the lowest energy structures of Lennard-Jones clusters containing up to 110 atoms." The Journal of Physical Chemistry A 101.28 (1997)

[2] Schaul et. al. "High dimensions and heavy tails for natural evolution strategies." Proceedings of the 13th annual conference on Genetic and evolutionary computation. 2011.

See this notebook on GitHub