1# coding: utf-8
2# Copyright (c) Pymatgen Development Team.
3# Distributed under the terms of the MIT License.
4
5
6"""
7This module provides classes used to enumerate surface sites
8and to find adsorption sites on slabs
9"""
10
11import itertools
12import os
13
14import numpy as np
15from matplotlib import patches
16from matplotlib.path import Path
17from monty.serialization import loadfn
18from scipy.spatial import Delaunay
19
20from pymatgen import vis
21from pymatgen.core.structure import Structure
22from pymatgen.analysis.local_env import VoronoiNN
23from pymatgen.analysis.structure_matcher import StructureMatcher
24from pymatgen.core.operations import SymmOp
25from pymatgen.core.surface import generate_all_slabs
26from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
27from pymatgen.util.coord import in_coord_list_pbc
28
29__author__ = "Joseph Montoya"
30__copyright__ = "Copyright 2016, The Materials Project"
31__version__ = "0.1"
32__maintainer__ = "Joseph Montoya"
33__credits__ = "Richard Tran"
34__email__ = "montoyjh@lbl.gov"
35__status__ = "Development"
36__date__ = "December 2, 2015"
37
38
39class AdsorbateSiteFinder:
40    """
41    This class finds adsorbate sites on slabs and generates
42    adsorbate structures according to user-defined criteria.
43    The algorithm for finding sites is essentially as follows:
44        1. Determine "surface sites" by finding those within
45            a height threshold along the miller index of the
46            highest site
47        2. Create a network of surface sites using the Delaunay
48            triangulation of the surface sites
49        3. Assign on-top, bridge, and hollow adsorption sites
50            at the nodes, edges, and face centers of the Del.
51            Triangulation
52        4. Generate structures from a molecule positioned at
53            these sites
54    """
55
56    def __init__(self, slab, selective_dynamics=False, height=0.9, mi_vec=None):
57        """
58        Create an AdsorbateSiteFinder object.
59
60        Args:
61            slab (Slab): slab object for which to find adsorbate sites
62            selective_dynamics (bool): flag for whether to assign
63                non-surface sites as fixed for selective dynamics
64            height (float): height criteria for selection of surface sites
65            mi_vec (3-D array-like): vector corresponding to the vector
66                concurrent with the miller index, this enables use with
67                slabs that have been reoriented, but the miller vector
68                must be supplied manually
69
70        """
71        # get surface normal from miller index
72        if mi_vec:
73            self.mvec = mi_vec
74        else:
75            self.mvec = get_mi_vec(slab)
76        slab = self.assign_site_properties(slab, height)
77        if selective_dynamics:
78            slab = self.assign_selective_dynamics(slab)
79        self.slab = slab
80
81    @classmethod
82    def from_bulk_and_miller(
83        cls,
84        structure,
85        miller_index,
86        min_slab_size=8.0,
87        min_vacuum_size=10.0,
88        max_normal_search=None,
89        center_slab=True,
90        selective_dynamics=False,
91        undercoord_threshold=0.09,
92    ):
93        """
94        This method constructs the adsorbate site finder from a bulk
95        structure and a miller index, which allows the surface sites
96        to be determined from the difference in bulk and slab coordination,
97        as opposed to the height threshold.
98
99        Args:
100            structure (Structure): structure from which slab
101                input to the ASF is constructed
102            miller_index (3-tuple or list): miller index to be used
103            min_slab_size (float): min slab size for slab generation
104            min_vacuum_size (float): min vacuum size for slab generation
105            max_normal_search (int): max normal search for slab generation
106            center_slab (bool): whether to center slab in slab generation
107            selective dynamics (bool): whether to assign surface sites
108                to selective dynamics
109            undercoord_threshold (float): threshold of "undercoordation"
110                to use for the assignment of surface sites.  Default is
111                0.1, for which surface sites will be designated if they
112                are 10% less coordinated than their bulk counterpart
113        """
114        # TODO: for some reason this works poorly with primitive cells
115        #       may want to switch the coordination algorithm eventually
116        vnn_bulk = VoronoiNN(tol=0.05)
117        bulk_coords = [len(vnn_bulk.get_nn(structure, n)) for n in range(len(structure))]
118        struct = structure.copy(site_properties={"bulk_coordinations": bulk_coords})
119        slabs = generate_all_slabs(
120            struct,
121            max_index=max(miller_index),
122            min_slab_size=min_slab_size,
123            min_vacuum_size=min_vacuum_size,
124            max_normal_search=max_normal_search,
125            center_slab=center_slab,
126        )
127
128        slab_dict = {slab.miller_index: slab for slab in slabs}
129
130        if miller_index not in slab_dict:
131            raise ValueError("Miller index not in slab dict")
132
133        this_slab = slab_dict[miller_index]
134
135        vnn_surface = VoronoiNN(tol=0.05, allow_pathological=True)
136
137        surf_props, undercoords = [], []
138        this_mi_vec = get_mi_vec(this_slab)
139        mi_mags = [np.dot(this_mi_vec, site.coords) for site in this_slab]
140        average_mi_mag = np.average(mi_mags)
141        for n, site in enumerate(this_slab):
142            bulk_coord = this_slab.site_properties["bulk_coordinations"][n]
143            slab_coord = len(vnn_surface.get_nn(this_slab, n))
144            mi_mag = np.dot(this_mi_vec, site.coords)
145            undercoord = (bulk_coord - slab_coord) / bulk_coord
146            undercoords += [undercoord]
147            if undercoord > undercoord_threshold and mi_mag > average_mi_mag:
148                surf_props += ["surface"]
149            else:
150                surf_props += ["subsurface"]
151        new_site_properties = {
152            "surface_properties": surf_props,
153            "undercoords": undercoords,
154        }
155        new_slab = this_slab.copy(site_properties=new_site_properties)
156        return cls(new_slab, selective_dynamics)
157
158    def find_surface_sites_by_height(self, slab, height=0.9, xy_tol=0.05):
159        """
160        This method finds surface sites by determining which sites are within
161        a threshold value in height from the topmost site in a list of sites
162
163        Args:
164            site_list (list): list of sites from which to select surface sites
165            height (float): threshold in angstroms of distance from topmost
166                site in slab along the slab c-vector to include in surface
167                site determination
168            xy_tol (float): if supplied, will remove any sites which are
169                within a certain distance in the miller plane.
170
171        Returns:
172            list of sites selected to be within a threshold of the highest
173        """
174
175        # Get projection of coordinates along the miller index
176        m_projs = np.array([np.dot(site.coords, self.mvec) for site in slab.sites])
177
178        # Mask based on window threshold along the miller index.
179        mask = (m_projs - np.amax(m_projs)) >= -height
180        surf_sites = [slab.sites[n] for n in np.where(mask)[0]]
181        if xy_tol:
182            # sort surface sites by height
183            surf_sites = [s for (h, s) in zip(m_projs[mask], surf_sites)]
184            surf_sites.reverse()
185            unique_sites, unique_perp_fracs = [], []
186            for site in surf_sites:
187                this_perp = site.coords - np.dot(site.coords, self.mvec)
188                this_perp_frac = slab.lattice.get_fractional_coords(this_perp)
189                if not in_coord_list_pbc(unique_perp_fracs, this_perp_frac):
190                    unique_sites.append(site)
191                    unique_perp_fracs.append(this_perp_frac)
192            surf_sites = unique_sites
193
194        return surf_sites
195
196    def assign_site_properties(self, slab, height=0.9):
197        """
198        Assigns site properties.
199        """
200        if "surface_properties" in slab.site_properties.keys():
201            return slab
202
203        surf_sites = self.find_surface_sites_by_height(slab, height)
204        surf_props = ["surface" if site in surf_sites else "subsurface" for site in slab.sites]
205        return slab.copy(site_properties={"surface_properties": surf_props})
206
207    def get_extended_surface_mesh(self, repeat=(5, 5, 1)):
208        """
209        Gets an extended surface mesh for to use for adsorption
210        site finding by constructing supercell of surface sites
211
212        Args:
213            repeat (3-tuple): repeat for getting extended surface mesh
214        """
215        surf_str = Structure.from_sites(self.surface_sites)
216        surf_str.make_supercell(repeat)
217        return surf_str
218
219    @property
220    def surface_sites(self):
221        """
222        convenience method to return a list of surface sites
223        """
224        return [site for site in self.slab.sites if site.properties["surface_properties"] == "surface"]
225
226    def subsurface_sites(self):
227        """
228        convenience method to return list of subsurface sites
229        """
230        return [site for site in self.slab.sites if site.properties["surface_properties"] == "subsurface"]
231
232    def find_adsorption_sites(
233        self,
234        distance=2.0,
235        put_inside=True,
236        symm_reduce=1e-2,
237        near_reduce=1e-2,
238        positions=["ontop", "bridge", "hollow"],
239        no_obtuse_hollow=True,
240    ):
241        """
242        Finds surface sites according to the above algorithm.  Returns
243        a list of corresponding cartesian coordinates.
244
245        Args:
246            distance (float): distance from the coordinating ensemble
247                of atoms along the miller index for the site (i. e.
248                the distance from the slab itself)
249            put_inside (bool): whether to put the site inside the cell
250            symm_reduce (float): symm reduction threshold
251            near_reduce (float): near reduction threshold
252            positions (list): which positions to include in the site finding
253                "ontop": sites on top of surface sites
254                "bridge": sites at edges between surface sites in Delaunay
255                    triangulation of surface sites in the miller plane
256                "hollow": sites at centers of Delaunay triangulation faces
257                "subsurface": subsurface positions projected into miller plane
258            no_obtuse_hollow (bool): flag to indicate whether to include
259                obtuse triangular ensembles in hollow sites
260        """
261        ads_sites = {k: [] for k in positions}
262        if "ontop" in positions:
263            ads_sites["ontop"] = [s.coords for s in self.surface_sites]
264        if "subsurface" in positions:
265            # Get highest site
266            ref = self.slab.sites[np.argmax(self.slab.cart_coords[:, 2])]
267            # Project diff between highest site and subs site into miller
268            ss_sites = [
269                self.mvec * np.dot(ref.coords - s.coords, self.mvec) + s.coords for s in self.subsurface_sites()
270            ]
271            ads_sites["subsurface"] = ss_sites
272        if "bridge" in positions or "hollow" in positions:
273            mesh = self.get_extended_surface_mesh()
274            sop = get_rot(self.slab)
275            dt = Delaunay([sop.operate(m.coords)[:2] for m in mesh])
276            # TODO: refactor below to properly account for >3-fold
277            for v in dt.simplices:
278                if -1 not in v:
279                    dots = []
280                    for i_corner, i_opp in zip(range(3), ((1, 2), (0, 2), (0, 1))):
281                        corner, opp = v[i_corner], [v[o] for o in i_opp]
282                        vecs = [mesh[d].coords - mesh[corner].coords for d in opp]
283                        vecs = [vec / np.linalg.norm(vec) for vec in vecs]
284                        dots.append(np.dot(*vecs))
285                        # Add bridge sites at midpoints of edges of D. Tri
286                        if "bridge" in positions:
287                            ads_sites["bridge"].append(self.ensemble_center(mesh, opp))
288                    # Prevent addition of hollow sites in obtuse triangles
289                    obtuse = no_obtuse_hollow and (np.array(dots) < 1e-5).any()
290                    # Add hollow sites at centers of D. Tri faces
291                    if "hollow" in positions and not obtuse:
292                        ads_sites["hollow"].append(self.ensemble_center(mesh, v))
293        for key, sites in ads_sites.items():
294            # Pare off outer sites for bridge/hollow
295            if key in ["bridge", "hollow"]:
296                frac_coords = [self.slab.lattice.get_fractional_coords(ads_site) for ads_site in sites]
297                frac_coords = [
298                    frac_coord
299                    for frac_coord in frac_coords
300                    if (frac_coord[0] > 1 and frac_coord[0] < 4 and frac_coord[1] > 1 and frac_coord[1] < 4)
301                ]
302                sites = [self.slab.lattice.get_cartesian_coords(frac_coord) for frac_coord in frac_coords]
303            if near_reduce:
304                sites = self.near_reduce(sites, threshold=near_reduce)
305            if put_inside:
306                sites = [put_coord_inside(self.slab.lattice, coord) for coord in sites]
307            if symm_reduce:
308                sites = self.symm_reduce(sites, threshold=symm_reduce)
309            sites = [site + distance * self.mvec for site in sites]
310
311            ads_sites[key] = sites
312        ads_sites["all"] = sum(ads_sites.values(), [])
313        return ads_sites
314
315    def symm_reduce(self, coords_set, threshold=1e-6):
316        """
317        Reduces the set of adsorbate sites by finding removing
318        symmetrically equivalent duplicates
319
320        Args:
321            coords_set: coordinate set in cartesian coordinates
322            threshold: tolerance for distance equivalence, used
323                as input to in_coord_list_pbc for dupl. checking
324        """
325        surf_sg = SpacegroupAnalyzer(self.slab, 0.1)
326        symm_ops = surf_sg.get_symmetry_operations()
327        unique_coords = []
328        # Convert to fractional
329        coords_set = [self.slab.lattice.get_fractional_coords(coords) for coords in coords_set]
330        for coords in coords_set:
331            incoord = False
332            for op in symm_ops:
333                if in_coord_list_pbc(unique_coords, op.operate(coords), atol=threshold):
334                    incoord = True
335                    break
336            if not incoord:
337                unique_coords += [coords]
338        # convert back to cartesian
339        return [self.slab.lattice.get_cartesian_coords(coords) for coords in unique_coords]
340
341    def near_reduce(self, coords_set, threshold=1e-4):
342        """
343        Prunes coordinate set for coordinates that are within
344        threshold
345
346        Args:
347            coords_set (Nx3 array-like): list or array of coordinates
348            threshold (float): threshold value for distance
349        """
350        unique_coords = []
351        coords_set = [self.slab.lattice.get_fractional_coords(coords) for coords in coords_set]
352        for coord in coords_set:
353            if not in_coord_list_pbc(unique_coords, coord, threshold):
354                unique_coords += [coord]
355        return [self.slab.lattice.get_cartesian_coords(coords) for coords in unique_coords]
356
357    @classmethod
358    def ensemble_center(cls, site_list, indices, cartesian=True):
359        """
360        Finds the center of an ensemble of sites selected from
361        a list of sites.  Helper method for the find_adsorption_sites
362        algorithm.
363
364        Args:
365            site_list (list of sites): list of sites
366            indices (list of ints): list of ints from which to select
367                sites from site list
368            cartesian (bool): whether to get average fractional or
369                cartesian coordinate
370        """
371        if cartesian:
372            return np.average([site_list[i].coords for i in indices], axis=0)
373
374        return np.average([site_list[i].frac_coords for i in indices], axis=0)
375
376    def add_adsorbate(self, molecule, ads_coord, repeat=None, translate=True, reorient=True):
377        """
378        Adds an adsorbate at a particular coordinate.  Adsorbate
379        represented by a Molecule object and is translated to (0, 0, 0) if
380        translate is True, or positioned relative to the input adsorbate
381        coordinate if translate is False.
382
383        Args:
384            molecule (Molecule): molecule object representing the adsorbate
385            ads_coord (array): coordinate of adsorbate position
386            repeat (3-tuple or list): input for making a supercell of slab
387                prior to placing the adsorbate
388            translate (bool): flag on whether to translate the molecule so
389                that its CoM is at the origin prior to adding it to the surface
390            reorient (bool): flag on whether to reorient the molecule to
391                have its z-axis concurrent with miller index
392        """
393        molecule = molecule.copy()
394        if translate:
395            # Translate the molecule so that the center of mass of the atoms
396            # that have the most negative z coordinate is at (0, 0, 0)
397            front_atoms = molecule.copy()
398            front_atoms._sites = [
399                s for s in molecule.sites if s.coords[2] == min([s.coords[2] for s in molecule.sites])
400            ]
401            x, y, z = front_atoms.center_of_mass
402            molecule.translate_sites(vector=[-x, -y, -z])
403        if reorient:
404            # Reorient the molecule along slab m_index
405            sop = get_rot(self.slab)
406            molecule.apply_operation(sop.inverse)
407        struct = self.slab.copy()
408        if repeat:
409            struct.make_supercell(repeat)
410        if "surface_properties" in struct.site_properties.keys():
411            molecule.add_site_property("surface_properties", ["adsorbate"] * molecule.num_sites)
412        if "selective_dynamics" in struct.site_properties.keys():
413            molecule.add_site_property("selective_dynamics", [[True, True, True]] * molecule.num_sites)
414        for site in molecule:
415            struct.append(
416                site.specie,
417                ads_coord + site.coords,
418                coords_are_cartesian=True,
419                properties=site.properties,
420            )
421        return struct
422
423    @classmethod
424    def assign_selective_dynamics(cls, slab):
425        """
426        Helper function to assign selective dynamics site_properties
427        based on surface, subsurface site properties
428
429        Args:
430            slab (Slab): slab for which to assign selective dynamics
431        """
432        sd_list = []
433        sd_list = [
434            [False, False, False] if site.properties["surface_properties"] == "subsurface" else [True, True, True]
435            for site in slab.sites
436        ]
437        new_sp = slab.site_properties
438        new_sp["selective_dynamics"] = sd_list
439        return slab.copy(site_properties=new_sp)
440
441    def generate_adsorption_structures(
442        self,
443        molecule,
444        repeat=None,
445        min_lw=5.0,
446        translate=True,
447        reorient=True,
448        find_args=None,
449    ):
450        """
451        Function that generates all adsorption structures for a given
452        molecular adsorbate.  Can take repeat argument or minimum
453        length/width of precursor slab as an input
454
455        Args:
456            molecule (Molecule): molecule corresponding to adsorbate
457            repeat (3-tuple or list): repeat argument for supercell generation
458            min_lw (float): minimum length and width of the slab, only used
459                if repeat is None
460            translate (bool): flag on whether to translate the molecule so
461                that its CoM is at the origin prior to adding it to the surface
462            reorient (bool): flag on whether or not to reorient adsorbate
463                along the miller index
464            find_args (dict): dictionary of arguments to be passed to the
465                call to self.find_adsorption_sites, e.g. {"distance":2.0}
466        """
467        if repeat is None:
468            xrep = np.ceil(min_lw / np.linalg.norm(self.slab.lattice.matrix[0]))
469            yrep = np.ceil(min_lw / np.linalg.norm(self.slab.lattice.matrix[1]))
470            repeat = [xrep, yrep, 1]
471        structs = []
472
473        find_args = find_args or {}
474        for coords in self.find_adsorption_sites(**find_args)["all"]:
475            structs.append(
476                self.add_adsorbate(
477                    molecule,
478                    coords,
479                    repeat=repeat,
480                    translate=translate,
481                    reorient=reorient,
482                )
483            )
484        return structs
485
486    def adsorb_both_surfaces(
487        self,
488        molecule,
489        repeat=None,
490        min_lw=5.0,
491        translate=True,
492        reorient=True,
493        find_args=None,
494    ):
495        """
496        Function that generates all adsorption structures for a given
497        molecular adsorbate on both surfaces of a slab. This is useful
498        for calculating surface energy where both surfaces need to be
499        equivalent or if we want to calculate nonpolar systems.
500
501        Args:
502            molecule (Molecule): molecule corresponding to adsorbate
503            repeat (3-tuple or list): repeat argument for supercell generation
504            min_lw (float): minimum length and width of the slab, only used
505                if repeat is None
506            reorient (bool): flag on whether or not to reorient adsorbate
507                along the miller index
508            find_args (dict): dictionary of arguments to be passed to the
509                call to self.find_adsorption_sites, e.g. {"distance":2.0}
510        """
511
512        # Get the adsorbed surfaces first
513        find_args = find_args or {}
514        adslabs = self.generate_adsorption_structures(
515            molecule,
516            repeat=repeat,
517            min_lw=min_lw,
518            translate=translate,
519            reorient=reorient,
520            find_args=find_args,
521        )
522
523        new_adslabs = []
524        for adslab in adslabs:
525
526            # Find the adsorbate sites and indices in each slab
527            _, adsorbates, indices = False, [], []
528            for i, site in enumerate(adslab.sites):
529                if site.surface_properties == "adsorbate":
530                    adsorbates.append(site)
531                    indices.append(i)
532
533            # Start with the clean slab
534            adslab.remove_sites(indices)
535            slab = adslab.copy()
536
537            # For each site, we add it back to the slab along with a
538            # symmetrically equivalent position on the other side of
539            # the slab using symmetry operations
540            for adsorbate in adsorbates:
541                p2 = adslab.get_symmetric_site(adsorbate.frac_coords)
542                slab.append(adsorbate.specie, p2, properties={"surface_properties": "adsorbate"})
543                slab.append(
544                    adsorbate.specie,
545                    adsorbate.frac_coords,
546                    properties={"surface_properties": "adsorbate"},
547                )
548            new_adslabs.append(slab)
549
550        return new_adslabs
551
552    def generate_substitution_structures(
553        self,
554        atom,
555        target_species=None,
556        sub_both_sides=False,
557        range_tol=1e-2,
558        dist_from_surf=0,
559    ):
560        """
561        Function that performs substitution-type doping on the surface and
562            returns all possible configurations where one dopant is substituted
563            per surface. Can substitute one surface or both.
564
565        Args:
566            atom (str): atom corresponding to substitutional dopant
567            sub_both_sides (bool): If true, substitute an equivalent
568                site on the other surface
569            target_species (list): List of specific species to substitute
570            range_tol (float): Find viable substitution sites at a specific
571                distance from the surface +- this tolerance
572            dist_from_surf (float): Distance from the surface to find viable
573                substitution sites, defaults to 0 to substitute at the surface
574        """
575
576        target_species = target_species or []
577
578        # Get symmetrized structure in case we want to substitue both sides
579        sym_slab = SpacegroupAnalyzer(self.slab).get_symmetrized_structure()
580
581        # Define a function for substituting a site
582        def substitute(site, i):
583            slab = self.slab.copy()
584            props = self.slab.site_properties
585            if sub_both_sides:
586                # Find an equivalent site on the other surface
587                eq_indices = [indices for indices in sym_slab.equivalent_indices if i in indices][0]
588                for ii in eq_indices:
589                    if "%.6f" % (sym_slab[ii].frac_coords[2]) != "%.6f" % (site.frac_coords[2]):
590                        props["surface_properties"][ii] = "substitute"
591                        slab.replace(ii, atom)
592                        break
593
594            props["surface_properties"][i] = "substitute"
595            slab.replace(i, atom)
596            slab.add_site_property("surface_properties", props["surface_properties"])
597            return slab
598
599        # Get all possible substitution sites
600        substituted_slabs = []
601        # Sort sites so that we can define a range relative to the position of the
602        # surface atoms, i.e. search for sites above (below) the bottom (top) surface
603        sorted_sites = sorted(sym_slab, key=lambda site: site.frac_coords[2])
604        if sorted_sites[0].surface_properties == "surface":
605            d = sorted_sites[0].frac_coords[2] + dist_from_surf
606        else:
607            d = sorted_sites[-1].frac_coords[2] - dist_from_surf
608
609        for i, site in enumerate(sym_slab):
610            if d - range_tol < site.frac_coords[2] < d + range_tol:
611                if target_species and site.species_string in target_species:
612                    substituted_slabs.append(substitute(site, i))
613                elif not target_species:
614                    substituted_slabs.append(substitute(site, i))
615
616        matcher = StructureMatcher()
617        return [s[0] for s in matcher.group_structures(substituted_slabs)]
618
619
620def get_mi_vec(slab):
621    """
622    Convenience function which returns the unit vector aligned
623    with the miller index.
624    """
625    mvec = np.cross(slab.lattice.matrix[0], slab.lattice.matrix[1])
626    return mvec / np.linalg.norm(mvec)
627
628
629def get_rot(slab):
630    """
631    Gets the transformation to rotate the z axis into the miller index
632    """
633    new_z = get_mi_vec(slab)
634    a, b, c = slab.lattice.matrix
635    new_x = a / np.linalg.norm(a)
636    new_y = np.cross(new_z, new_x)
637    x, y, z = np.eye(3)
638    rot_matrix = np.array([np.dot(*el) for el in itertools.product([x, y, z], [new_x, new_y, new_z])]).reshape(3, 3)
639    rot_matrix = np.transpose(rot_matrix)
640    sop = SymmOp.from_rotation_and_translation(rot_matrix)
641    return sop
642
643
644def put_coord_inside(lattice, cart_coordinate):
645    """
646    converts a cartesian coordinate such that it is inside the unit cell.
647    """
648    fc = lattice.get_fractional_coords(cart_coordinate)
649    return lattice.get_cartesian_coords([c - np.floor(c) for c in fc])
650
651
652def reorient_z(structure):
653    """
654    reorients a structure such that the z axis is concurrent with the
655    normal to the A-B plane
656    """
657    struct = structure.copy()
658    sop = get_rot(struct)
659    struct.apply_operation(sop)
660    return struct
661
662
663# Get color dictionary
664colors = loadfn(os.path.join(os.path.dirname(vis.__file__), "ElementColorSchemes.yaml"))
665color_dict = {el: [j / 256.001 for j in colors["Jmol"][el]] for el in colors["Jmol"].keys()}
666
667
668def plot_slab(
669    slab,
670    ax,
671    scale=0.8,
672    repeat=5,
673    window=1.5,
674    draw_unit_cell=True,
675    decay=0.2,
676    adsorption_sites=True,
677    inverse=False,
678):
679    """
680    Function that helps visualize the slab in a 2-D plot, for
681    convenient viewing of output of AdsorbateSiteFinder.
682
683    Args:
684        slab (slab): Slab object to be visualized
685        ax (axes): matplotlib axes with which to visualize
686        scale (float): radius scaling for sites
687        repeat (int): number of repeating unit cells to visualize
688        window (float): window for setting the axes limits, is essentially
689            a fraction of the unit cell limits
690        draw_unit_cell (bool): flag indicating whether or not to draw cell
691        decay (float): how the alpha-value decays along the z-axis
692        inverse (bool): invert z axis to plot opposite surface
693    """
694    orig_slab = slab.copy()
695    slab = reorient_z(slab)
696    orig_cell = slab.lattice.matrix.copy()
697    if repeat:
698        slab.make_supercell([repeat, repeat, 1])
699    coords = np.array(sorted(slab.cart_coords, key=lambda x: x[2]))
700    sites = sorted(slab.sites, key=lambda x: x.coords[2])
701    alphas = 1 - decay * (np.max(coords[:, 2]) - coords[:, 2])
702    alphas = alphas.clip(min=0)
703    corner = [0, 0, slab.lattice.get_fractional_coords(coords[-1])[-1]]
704    corner = slab.lattice.get_cartesian_coords(corner)[:2]
705    verts = orig_cell[:2, :2]
706    lattsum = verts[0] + verts[1]
707    # inverse coords, sites, alphas, to show other side of slab
708    if inverse:
709        alphas = np.array(reversed(alphas))
710        sites = list(reversed(sites))
711        coords = np.array(reversed(coords))
712    # Draw circles at sites and stack them accordingly
713    for n, coord in enumerate(coords):
714        r = sites[n].specie.atomic_radius * scale
715        ax.add_patch(patches.Circle(coord[:2] - lattsum * (repeat // 2), r, color="w", zorder=2 * n))
716        color = color_dict[sites[n].species_string]
717        ax.add_patch(
718            patches.Circle(
719                coord[:2] - lattsum * (repeat // 2),
720                r,
721                facecolor=color,
722                alpha=alphas[n],
723                edgecolor="k",
724                lw=0.3,
725                zorder=2 * n + 1,
726            )
727        )
728    # Adsorption sites
729    if adsorption_sites:
730        asf = AdsorbateSiteFinder(orig_slab)
731        if inverse:
732            inverse_slab = orig_slab.copy()
733            inverse_slab.make_supercell([1, 1, -1])
734            asf = AdsorbateSiteFinder(inverse_slab)
735        ads_sites = asf.find_adsorption_sites()["all"]
736        sop = get_rot(orig_slab)
737        ads_sites = [sop.operate(ads_site)[:2].tolist() for ads_site in ads_sites]
738        ax.plot(*zip(*ads_sites), color="k", marker="x", markersize=10, mew=1, linestyle="", zorder=10000)
739    # Draw unit cell
740    if draw_unit_cell:
741        verts = np.insert(verts, 1, lattsum, axis=0).tolist()
742        verts += [[0.0, 0.0]]
743        verts = [[0.0, 0.0]] + verts
744        codes = [Path.MOVETO, Path.LINETO, Path.LINETO, Path.LINETO, Path.CLOSEPOLY]
745        verts = [(np.array(vert) + corner).tolist() for vert in verts]
746        path = Path(verts, codes)
747        patch = patches.PathPatch(path, facecolor="none", lw=2, alpha=0.5, zorder=2 * n + 2)
748        ax.add_patch(patch)
749    ax.set_aspect("equal")
750    center = corner + lattsum / 2.0
751    extent = np.max(lattsum)
752    lim_array = [center - extent * window, center + extent * window]
753    x_lim = [ele[0] for ele in lim_array]
754    y_lim = [ele[1] for ele in lim_array]
755    ax.set_xlim(x_lim)
756    ax.set_ylim(y_lim)
757    return ax
758