Chimie Théorique » scripts_chimie4psmn » DockOnSurf

import logging

import numpy as np

import ase

logger = logging.getLogger('DockOnSurf')

def assign_prop(atoms: ase.Atoms, prop_name: str, prop_val):  # TODO Needed?

    atoms.info[prop_name] = prop_val

def select_confs(orig_conf_list: list, magns: list, num_sel: int, code: str):

    """Takes a list ase.Atoms and selects the most different magnitude-wise.

    Given a list of ase.Atoms objects and a list of magnitudes, it selects a

    number of the most different conformers according to every magnitude

    specified.

    @param orig_conf_list: list of ase.Atoms objects to select among.

    @param magns: list of str with the names of the magnitudes to use for the

        conformer selection.

        Supported magnitudes: 'energy', 'moi'.

    @param num_sel: number of conformers to select for every of the magnitudes.

    @param code: The code that generated the magnitude information.

         Supported codes: See formats.py

    @return: list of the selected ase.Atoms objects.

    """

    from copy import deepcopy

    from modules.formats import read_energies

    conf_list = deepcopy(orig_conf_list)

    selected_ids = []

    if num_sel >= len(conf_list):

        logger.warning('Number of conformers per magnitude is equal or larger '

                       'than the total number of conformers. Using all '

                       f'available conformers: {len(conf_list)}.')

        return conf_list

    # Read properties

    if 'energy' in magns:

        conf_enrgs = read_energies(code, 'isolated')

    if 'moi' in magns:

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

    # Assign values

    for i, conf in enumerate(conf_list):

        assign_prop(conf, 'idx', i)

        if 'energy' in magns:

            assign_prop(conf, 'energy', conf_enrgs[i])

        if 'moi' in magns:

            assign_prop(conf, 'moi', mois[i, 2])

    # pick ids

    for magn in magns:

        sorted_list = sorted(conf_list, key=lambda conf: abs(conf.info[magn]))

        if sorted_list[-1].info['idx'] not in selected_ids:

            selected_ids.append(sorted_list[-1].info['idx'])

        if num_sel > 1:

            for i in range(0, len(sorted_list) - 1,

                           len(conf_list) // (num_sel - 1)):

                if sorted_list[i].info['idx'] not in selected_ids:

                    selected_ids.append(sorted_list[i].info['idx'])

    logger.info(f'Selected {len(selected_ids)} conformers for adsorption.')

    return [conf_list[idx] for idx in selected_ids]

def get_vect_angle(v1, v2, degrees=False):

    """Computes the angle between two vectors.

    @param v1: The first vector.

    @param v2: The second vector.

    @param degrees: Whether the result should be in radians (True) or in

        degrees (False).

    @return: The angle in radians if degrees = False, or in degrees if

        degrees =True

    """

    v1_u = v1 / np.linalg.norm(v1)

    v2_u = v2 / np.linalg.norm(v2)

    angle = np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))

    return angle if not degrees else angle * 180 / np.pi

def vect_avg(vects):

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

    @param vects: list of lists-like: containing the vectors (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"

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

    if len(array.shape) == 1:

        return array

    else:

        num_vects = array.shape[1]

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

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

    """Gets the coordinates of the specified indices 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 = []

    err = "'ctrs_list' argument must be an integer, a list of integers or a " \

          "list of lists of integers. Every integer must be in the range " \

          "[0, num_atoms)"

    if ctrs_list is None:

        ctrs_list = range(len(atoms))

    elif isinstance(ctrs_list, int):

        if ctrs_list not in range(len(atoms)):

            logger.error(err)

            raise ValueError(err)

        return atoms[ctrs_list].position

    for elem in ctrs_list:

        if isinstance(elem, list):

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

        elif isinstance(elem, int):

            coords.append(atoms[elem].position)

        else:

            logger.error(err)

            raise ValueError

    return np.array(coords)

def add_adsorbate(slab, adsorbate, site_coord, ctr_coord, 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 site_coord: The coordinates of the adsorption site on the surface.

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

    @param height: The height above the surface where to adsorb.

    @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.

    """

    from copy import deepcopy

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

    pos = np.array([0.0, 0.0, 0.0])  # part of absolute coordinates

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

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

    if offset is not None:

        spos += np.asarray(offset, float)

    if isinstance(site_coord, 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 site_coord not in info['sites']:

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

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

    else:

        pos += site_coord

    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 = deepcopy(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_u

    # Move adsorbate into position

    ads.translate(pos - ctr_coord)

    # Attach the adsorbate

    slab.extend(ads)

def check_collision(slab_molec, slab_num_atoms, min_height, vect, nn_slab=0,

                    nn_molec=0, coll_coeff=0.9):

    """Checks whether a slab and a molecule collide or not.

    @param slab_molec: The system of adsorbate-slab for which to detect if there

        are collisions.

    @param nn_slab: Number of neigbors in the surface.

    @param nn_molec: Number of neighbors in the molecule.

    @param coll_coeff: The coefficient that multiplies the covalent radius of

        atoms resulting in a distance that two atoms being closer to that is

        considered as atomic collision.

    @param slab_num_atoms: Number of atoms of the bare slab.

    @param min_height: The minimum height atoms can have to not be considered as

        colliding.

    @param vect: The vector perpendicular to the slab.

    @return: bool, whether the surface and the molecule collide.

    """

    from ase.neighborlist import natural_cutoffs, neighbor_list

    if min_height is not False:

        cart_axes = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0],

                     [-1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, -1.0]]

        if vect.tolist() in cart_axes:

            for atom in slab_molec[slab_num_atoms:]:

                for i, coord in enumerate(vect):

                    if coord == 0:

                        continue

                    if atom.position[i] * coord < min_height:

                        return True

            return False

    else:

        slab_molec_cutoffs = natural_cutoffs(slab_molec, mult=coll_coeff)

        slab_molec_nghbs = len(neighbor_list("i", slab_molec, slab_molec_cutoffs))

        if slab_molec_nghbs > nn_slab + nn_molec:

            return True

        else:

            return False

def correct_coll(molec, slab, ctr_coord, site_coord, num_pts,

                 min_coll_height, norm_vect, slab_nghbs, molec_nghbs,

                 coll_coeff):

    # TODO Rethink this function

    """Tries to adsorb a molecule on a slab trying to avoid collisions by doing

    small rotations.

    @param molec: ase.Atoms object of the molecule to adsorb

    @param slab: ase.Atoms object of the surface on which to adsorb the

        molecule

    @param ctr_coord: The coordinates of the molecule to use as adsorption

        center.

    @param site_coord: The coordinates of the surface on which to adsorb the

        molecule

    @param num_pts: Number on which to sample Euler angles.

    @param min_coll_height: The lowermost height for which to detect a collision

    @param norm_vect: The vector perpendicular to the surface.

    @param slab_nghbs: Number of neigbors in the surface.

    @param molec_nghbs: Number of neighbors in the molecule.

    @param coll_coeff: The coefficient that multiplies the covalent radius of

        atoms resulting in a distance that two atoms being closer to that is

        considered as atomic collision.

    @return collision: bool, whether the structure generated has collisions

        between slab and adsorbate.

    """

    from copy import deepcopy

    slab_num_atoms = len(slab)

    collision = True

    max_corr = 6  # Should be an even number

    d_angle = 180 / ((max_corr / 2.0) * num_pts)

    num_corr = 0

    while collision and num_corr <= max_corr:

        k = num_corr * (-1) ** num_corr

        slab_molec = deepcopy(slab)

        molec.euler_rotate(k * d_angle, k * d_angle / 2, k * d_angle,

                           center=ctr_coord)

        add_adsorbate(slab_molec, molec, site_coord, ctr_coord, 2.5,

                      norm_vect=norm_vect)

        collision = check_collision(slab_molec, slab_num_atoms, min_coll_height,

                                    norm_vect, slab_nghbs, molec_nghbs,

                                    coll_coeff)

        num_corr += 1

    return slab_molec, collision

def ads_euler(orig_molec, slab, ctr_coord, site_coord, num_pts,

              min_coll_height, coll_coeff, norm_vect, slab_nghbs, molec_nghbs):

    """Generates adsorbate-surface structures by sampling over Euler angles.

    This function generates a number of adsorbate-surface structures at

    different orientations of the adsorbate sampled at multiple Euler (zxz)

    angles.

    @param orig_molec: ase.Atoms object of the molecule to adsorb

    @param slab: ase.Atoms object of the surface on which to adsorb the molecule

    @param ctr_coord: The coordinates of the molecule to use as adsorption

        center.

    @param site_coord: The coordinates of the surface on which to adsorb the

        molecule

    @param num_pts: Number on which to sample Euler angles.

    @param min_coll_height: The lowermost height for which to detect a collision

    @param coll_coeff: The coefficient that multiplies the covalent radius of

        atoms resulting in a distance that two atoms being closer to that is

        considered as atomic collision.

    @param norm_vect: The vector perpendicular to the surface.

    @param slab_nghbs: Number of neigbors in the surface.

    @param molec_nghbs: Number of neighbors in the molecule.

    @return: list of ase.Atoms object conatining all the orientations of a given

        conformer

    """

    from copy import deepcopy

    slab_ads_list = []

    # rotation around z

    for alpha in np.arange(0, 360, 360 / num_pts):

        # rotation around x'

        for beta in np.arange(0, 180, 180 / num_pts):

            # rotation around z'

            for gamma in np.arange(0, 360, 360 / num_pts):

                molec = deepcopy(orig_molec)

                molec.euler_rotate(alpha, beta, gamma, center=ctr_coord)

                slab_molec, collision = correct_coll(molec, slab,

                                                     ctr_coord, site_coord,

                                                     num_pts, min_coll_height,

                                                     norm_vect,

                                                     slab_nghbs, molec_nghbs,

                                                     coll_coeff)

                if not collision:

                    slab_ads_list.append(slab_molec)

    return slab_ads_list

def ads_chemcat(site, ctr, pts_angle):

    return "TO IMPLEMENT"

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

                 norm_vect, min_coll_height, coll_coeff):

    """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 molec_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.

    @param min_coll_height: The lowermost height for which to detect a collision

    @param coll_coeff: The coefficient that multiplies the covalent radius of

        atoms resulting in a distance that two atoms being closer to that is

        considered as atomic collision.

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

    """

    from ase.neighborlist import natural_cutoffs, neighbor_list

    surf_ads_list = []

    sites_coords = get_atom_coords(surf, sites)

    if min_coll_height is not False:

        surf_cutoffs = natural_cutoffs(surf, mult=coll_coeff)

        surf_nghbs = len(neighbor_list("i", surf, surf_cutoffs))

    else:

        surf_nghbs = 0

    for conf in conf_list:

        molec_ctr_coords = get_atom_coords(conf, molec_ctrs)

        molec_neigh_coords = get_atom_coords(conf, neigh_ctrs)

        if min_coll_height is not False:

            conf_cutoffs = natural_cutoffs(conf, mult=coll_coeff)

            molec_nghbs = len(neighbor_list("i", conf, conf_cutoffs))

        else:

            molec_nghbs = 0

        for site in sites_coords:

            for ctr in molec_ctr_coords:

                if algo == 'euler':

                    surf_ads_list.extend(ads_euler(conf, surf, ctr, site,

                                                   num_pts, min_coll_height,

                                                   coll_coeff, norm_vect,

                                                   surf_nghbs, molec_nghbs))

                elif algo == 'chemcat':

                    surf_ads_list.extend(ads_chemcat(site, 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, adapt_format

    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', 'ase')

    logger.info(f"Found {len(conf_list)} structures of isolated conformers.")

    selected_confs = select_confs(conf_list, inp_vars['select_magns'],

                                  inp_vars['confs_per_magn'],

                                  inp_vars['code'])

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

    surf_ads_list = adsorb_confs(selected_confs, 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'],

                                 inp_vars['min_coll_height'],

                                 inp_vars['collision_threshold'])

    logger.info(f'Generated {len(surf_ads_list)} adsorbate-surface atomic '

                f'configurations, to carry out a calculation of.')

    run_calc('screening', inp_vars, surf_ads_list)