pack molecules

pack_molecules.py

import ase.build
import ase.data
import ase.geometry
import numpy as np
from ase import Atoms


def compute_radius(atoms: Atoms):
    """Compute the radius of an atom, assuming sphere."""
    positions = atoms.get_positions()
    centroid = np.mean(positions, axis=0)
    if np.linalg.norm(centroid) > 1e-6:
        raise ValueError("Centroid is not at the origin.")
    dists = np.linalg.norm(positions - centroid, axis=1)
    max_dist_idx = np.argmax(dists)
    max_dist = dists[max_dist_idx]
    return max_dist


def distance_sphere(pos1, pos2, radius_1, radius_2, cell=None):
    """Calculate the distance between two spheres."""
    pbc = True if cell is not None else None
    d = ase.geometry.get_distances(pos1, pos2, cell=cell, pbc=pbc)[1]
    return d - radius_1 - radius_2


def pack_molecules(
    molecule,
    num_molecules,
    cell=None,
    base_atoms=None,
    tol=1.0,
    n_iter=10000,
):
    radius = compute_radius(molecule)
    print(f"Radius of molecule: {radius:.2f} Å")
    points = []
    ref_points = base_atoms.get_positions() if base_atoms is not None else []
    if base_atoms is not None:
        cell = base_atoms.cell
    for i in range(n_iter):
        point_frac = np.random.random(3)
        point = (cell @ point_frac.reshape(3, 1)).reshape(3)

        distances = (
            distance_sphere(
                np.stack(points), point, radius, radius, cell=cell
            ).squeeze()
            if len(points) > 0
            else None
        )
        distances_from_base = (
            distance_sphere(ref_points, point, 0.3, radius, cell=cell).squeeze()
            if len(ref_points) > 0
            else None
        )

        min_dist = np.min(distances) if len(points) > 0 else 1e100
        min_dist_from_base = (
            np.min(distances_from_base) if len(ref_points) > 0 else 1e100
        )
        min_dist = min(min_dist, min_dist_from_base)
        if min_dist > tol:
            points.append(point)
        if len(points) == num_molecules:
            break
    if len(points) < num_molecules:
        raise ValueError(f"Could not pack all molecules. Current: {len(points)}")

    # construct box
    atoms = ase.Atoms() if base_atoms is None else base_atoms.copy()
    for point in points:
        mol = molecule.copy()
        mol.positions -= np.mean(mol.positions, axis=0)
        mol.positions += point
        angles = np.random.uniform(-180, 180, 3)
        mol.rotate(angles[0], (1, 0, 0), point)
        mol.rotate(angles[1], (0, 1, 0), point)
        mol.rotate(angles[2], (0, 0, 1), point)
        atoms.extend(mol)
    if cell is not None:
        atoms.cell = cell
        atoms.pbc = True
    atoms = ase.build.sort(atoms)
    return atoms