Chimie Théorique » scripts_chimie4psmn » DockOnSurf

import logging

import numpy as np

import ase

logger = logging.getLogger('DockOnSurf')

def vect_avg(vects):

    """Computes the element-wise mean of a set of vectors.

    @param vects: np.ndarray with shape (num_vectors, length_vector).

    @return: vector average computed doing the element-wise mean.

    """

    from utilities import try_command

    err = "vect_avg parameter vects must be a list-like, able to be converted" \

          " np.array"

    vects = try_command(np.array, [(ValueError, err)], vects)

    if len(vects.shape) == 1:

        return vects

    else:

        num_vects = vects.shape[1]

        return np.array([np.average(vects[:, i]) for i in range(num_vects)])

def get_atom_coords(atoms: ase.Atoms, ctrs_list):

    """Gets the coordinates of the specified atoms from a ase.Atoms object.

    Given an ase.Atoms object and a list of atom indices specified in ctrs_list

    it gets the coordinates of the specified atoms. If the element in the

    ctrs_list is not an index but yet a list of indices, it computes the

    element-wise mean of the coordinates of the atoms specified in the inner

    list.

    @param atoms: ase.Atoms object for which to obtain the coordinates of.

    @param ctrs_list: list of (indices/list of indices) of the atoms for which

                      the coordinates should be extracted.

    @return: np.ndarray of atomic coordinates.

    """

    coords = []

    for i, elem in enumerate(ctrs_list):

        if isinstance(elem, list):

            coords.append(vect_avg(np.array([atoms[c].position for c in elem])))

        elif isinstance(elem, int):

            coords.append(atoms[elem].position)

        else:

            err = f"'ctrs_list must be a list of ints or lists, {type(elem)} " \

                  "found."

            logger.error(err)

            raise ValueError

    return np.array(coords)

def add_adsorbate(slab, adsorbate, surf_pos, mol_pos, height, offset=None,

                  norm_vect=(0, 0, 1)):

    """Add an adsorbate to a surface.

    This function extends the functionality of ase.build.add_adsorbate

    (https://wiki.fysik.dtu.dk/ase/ase/build/surface.html#ase.build.add_adsorbate)

    by enabling to change the z coordinate and the axis perpendicular to the

    surface.

    @param slab: ase.Atoms object containing the coordinates of the surface

    @param adsorbate: ase.Atoms object containing the coordinates of the

        adsorbate.

    @param surf_pos: The coordinates of the adsorption site on the surface.

    @param mol_pos: The coordinates of the adsorption center in the molecule.

    @param height: The height above the surface

    @param offset: Offsets the adsorbate by a number of unit cells. Mostly

        useful when adding more than one adsorbate.

    @param norm_vect: The vector perpendicular to the surface.

    """

    info = slab.info.get('adsorbate_info', {})

    pos = np.array([0.0, 0.0, 0.0])  # (x, y, z) part

    spos = np.array([0.0, 0.0, 0.0])  # part relative to unit cell

    norm_vect = np.array(norm_vect) / np.linalg.norm(norm_vect)

    if offset is not None:

        spos += np.asarray(offset, float)

    if isinstance(surf_pos, str):

        # A site-name:

        if 'sites' not in info:

            raise TypeError('If the atoms are not made by an ase.build '

                            'function, position cannot be a name.')

        if surf_pos not in info['sites']:

            raise TypeError('Adsorption site %s not supported.' % surf_pos)

        spos += info['sites'][surf_pos]

    else:

        pos += surf_pos

    if 'cell' in info:

        cell = info['cell']

    else:

        cell = slab.get_cell()

    pos += np.dot(spos, cell)

    # Convert the adsorbate to an Atoms object

    if isinstance(adsorbate, ase.Atoms):

        ads = adsorbate

    elif isinstance(adsorbate, ase.Atom):

        ads = ase.Atoms([adsorbate])

    else:

        # Assume it is a string representing a single Atom

        ads = ase.Atoms([ase.Atom(adsorbate)])

    pos += height * norm_vect

    # Move adsorbate into position

    ads.translate(pos - mol_pos)

    # Attach the adsorbate

    slab.extend(ads)

def ads_euler(molec, surf, site, ctr, pts_angle, neigh_ctr, norm_vect):

    from random import random

    plane_vect = np.cross(norm_vect, [random() for i in range(3)])

    add_adsorbate()

    return

def ads_chemcat(site, ctr, pts_angle):

    return "TO IMPLEMENT"

def adsorb_confs(conf_list, surf, ads_ctrs, sites, algo, num_pts, neigh_ctrs,

                 norm_vect):

    """Generates a number of adsorbate-surface structure coordinates.

    Given a list of conformers, a surface, a list of atom indices (or list of

    list of indices) of both the surface and the adsorbate, it generates a

    number of adsorbate-surface structures for every possible combination of

    them at different orientations.

    @param conf_list: list of ase.Atoms of the different conformers

    @param surf: the ase.Atoms object of the surface

    @param ads_ctrs: the list atom indices of the adsorbate.

    @param sites: the list of atom indices of the surface.

    @param algo: the algorithm to use for the generation of adsorbates.

    @param num_pts: the number of points per angle orientation to sample

    @param neigh_ctrs: the indices of the neighboring atoms to the adsorption

    atoms.

    @param norm_vect: The vector perpendicular to the surface.

    @return: list of ase.Atoms for the adsorbate-surface structures

    """

    surf_ads_list = []

    sites_coords = get_atom_coords(surf, sites)

    for conf in conf_list:

        molec_ctr_coords = get_atom_coords(conf, ads_ctrs)

        molec_neigh_coords = get_atom_coords(conf, neigh_ctrs)

        for site in sites_coords:

            for i, molec_ctr in enumerate(molec_ctr_coords):

                if algo == 'euler':

                    surf_ads_list.append(ads_euler(conf, surf, site, molec_ctr,

                                                   num_pts, norm_vect,

                                                   molec_neigh_coords[i]))

                elif algo == 'chemcat':

                    surf_ads_list.append(ads_chemcat(site, molec_ctr, num_pts))

    return surf_ads_list

def run_screening(inp_vars):

    """Carry out the screening of adsorbate coordinates on a surface

    @param inp_vars: Calculation parameters from input file.

    """

    import os

    from modules.formats import read_coords, read_energies, \

        rdkit_mol_to_ase_atoms, adapt_format

    from modules.clustering import get_rmsd, clustering

    from modules.isolated import get_moments_of_inertia

    from modules.calculation import run_calc

    if not os.path.isdir("isolated"):

        err = "'isolated' directory not found. It is needed in order to carry "

        "out the screening of structures to be adsorbed"

        logger.error(err)

        raise ValueError(err)

    conf_list = read_coords(inp_vars['code'], 'isolated', 'rdkit')

    # TODO Implement neighbors algorithm

    # neigh_list = get_neighbors(conf_list[0], inp_vars['molec_ads_ctrs'])

    conf_enrgs = read_energies(inp_vars['code'], 'isolated')

    mois = np.array([get_moments_of_inertia(conf)[0] for conf in conf_list])

    rmsd_mtx = get_rmsd(conf_list)

    exemplars = clustering(rmsd_mtx)

    conf_list = [conf_list[idx] for idx in exemplars]

    conf_list = [rdkit_mol_to_ase_atoms(conf) for conf in conf_list]

    surf = adapt_format('ase', inp_vars['surf_file'])

    surf_ads_list = adsorb_confs(conf_list, surf, inp_vars['molec_ads_ctrs'],

                                 inp_vars['sites'], inp_vars['ads_algo'],

                                 inp_vars['sample_points_per_angle'],

                                 inp_vars['molec_neigh_ctrs'],

                                 inp_vars['surf_norm_vect'])

    run_calc('screening', inp_vars, surf_ads_list)