Source code for atom_hifi.soap_ML.soap

import logging
from typing import List
from scipy.spatial.distance import pdist, squareform
import numpy as np
import decaf

logger = logging.getLogger(__name__)





def get_dist(soap_1, soap_2):

    return squareform(pdist(np.vstack([soap_1, soap_2])))




def generate_ED_SOAP(lib: List,
                lengthscale=1.2,
                SOAP_S_cut=(1+np.sqrt(3))/2,
                SOAP_S_nmax=8,
                SOAP_S_lmax=4,
                SOAP_L_cut=(2+3)/2,
                SOAP_L_nmax=4,
                SOAP_L_lmax=3,
                species = None,
                periodic=True,
                atom_ids_per_structure=None,
    ):
    """Embedded Doubled SOAP (ED-SOAP): concatenates a short-range and a
    long-range SOAP per atom (hence 'doubled'), scaled by `lengthscale`.
    Designed for use with DECAF's PCA embedding + MSC clustering
    (`decaf.embed_cluster`).

    ``atom_ids_per_structure`` optionally selects SOAP evaluation centers for
    each structure. It does not remove atoms from the structures, so every
    selected center retains the same neighbours and descriptor row as a full
    evaluation.
    """

    # Gaussian widths and cutoffs are in units of the physical lengthscale.
    SOAP_S_sigma = SOAP_S_cut / 8
    SOAP_L_sigma = SOAP_L_cut / 8
    if species is None:
        species = np.unique(np.concatenate([struct.numbers for struct in lib])).astype(int)

    # Obtain the full SOAP of the library.  Periodicity is decided per structure:
    # the `periodic` master switch is honoured only where a structure actually has
    # periodic boundary conditions, so non-PBC systems (e.g. organic molecules,
    # pbc=F F F) and mixed libraries are handled correctly without a cell error.
    if atom_ids_per_structure is None:
        atom_ids_per_structure = [np.array([], dtype=int) for _ in lib]
    elif len(atom_ids_per_structure) != len(lib):
        raise ValueError(
            'atom_ids_per_structure must contain one index array per structure')

    training_set = np.concatenate([
        decaf.get_SOAP(
            struct, lengthscale, species=species,
            SOAP_S_cut=SOAP_S_cut, SOAP_S_sigma=SOAP_S_sigma,
            SOAP_S_nmax=SOAP_S_nmax, SOAP_S_lmax=SOAP_S_lmax,
            SOAP_L_cut=SOAP_L_cut, SOAP_L_sigma=SOAP_L_sigma,
            SOAP_L_nmax=SOAP_L_nmax, SOAP_L_lmax=SOAP_L_lmax,
            periodic=bool(periodic and struct.pbc.any()),
            atomID=np.asarray(atom_ids, dtype=int),
        )
        for struct, atom_ids in zip(lib, atom_ids_per_structure)
    ])

    logger.debug(f'The shape of the training set is [atoms,soap_shape]: {training_set.shape}')

    return training_set