1# coding: utf-8
2# Copyright (c) Pymatgen Development Team.
3# Distributed under the terms of the MIT License.
4
5"""
6This module defines classes representing non-periodic and periodic sites.
7"""
8
9import collections
10import json
11from typing import Optional, Tuple, Union
12
13import numpy as np
14from monty.dev import deprecated
15from monty.json import MontyDecoder, MontyEncoder, MSONable
16
17from pymatgen.core.composition import Composition
18from pymatgen.core.lattice import Lattice
19from pymatgen.core.periodic_table import DummySpecies, Element, Species, get_el_sp
20from pymatgen.util.coord import pbc_diff
21from pymatgen.util.typing import ArrayLike, SpeciesLike, CompositionLike
22
23
24class Site(collections.abc.Hashable, MSONable):
25    """
26    A generalized *non-periodic* site. This is essentially a composition
27    at a point in space, with some optional properties associated with it. A
28    Composition is used to represent the atoms and occupancy, which allows for
29    disordered site representation. Coords are given in standard cartesian
30    coordinates.
31    """
32
33    position_atol = 1e-5
34
35    def __init__(
36        self,
37        species: Union[SpeciesLike, CompositionLike],
38        coords: ArrayLike,
39        properties: dict = None,
40        skip_checks: bool = False,
41    ):
42        """
43        Creates a non-periodic Site.
44
45        :param species: Species on the site. Can be:
46            i.  A Composition-type object (preferred)
47            ii. An  element / species specified either as a string
48                symbols, e.g. "Li", "Fe2+", "P" or atomic numbers,
49                e.g., 3, 56, or actual Element or Species objects.
50            iii.Dict of elements/species and occupancies, e.g.,
51                {"Fe" : 0.5, "Mn":0.5}. This allows the setup of
52                disordered structures.
53        :param coords: Cartesian coordinates of site.
54        :param properties: Properties associated with the site as a dict, e.g.
55            {"magmom": 5}. Defaults to None.
56        :param skip_checks: Whether to ignore all the usual checks and just
57            create the site. Use this if the Site is created in a controlled
58            manner and speed is desired.
59        """
60        if not skip_checks:
61            if not isinstance(species, Composition):
62                try:
63                    species = Composition({get_el_sp(species): 1})
64                except TypeError:
65                    species = Composition(species)
66            totaloccu = species.num_atoms
67            if totaloccu > 1 + Composition.amount_tolerance:
68                raise ValueError("Species occupancies sum to more than 1!")
69            coords = np.array(coords)
70        self._species: Composition = species  # type: ignore
71        self.coords: np.ndarray = coords  # type: ignore
72        self.properties: dict = properties or {}
73
74    def __getattr__(self, a):
75        # overriding getattr doens't play nice with pickle, so we
76        # can't use self._properties
77        p = object.__getattribute__(self, "properties")
78        if a in p:
79            return p[a]
80        raise AttributeError(a)
81
82    @property
83    def species(self) -> Composition:
84        """
85        :return: The species on the site as a composition, e.g., Fe0.5Mn0.5.
86        """
87        return self._species  # type: ignore
88
89    @species.setter
90    def species(self, species: Union[SpeciesLike, CompositionLike]):
91        if not isinstance(species, Composition):
92            try:
93                species = Composition({get_el_sp(species): 1})
94            except TypeError:
95                species = Composition(species)
96        totaloccu = species.num_atoms
97        if totaloccu > 1 + Composition.amount_tolerance:
98            raise ValueError("Species occupancies sum to more than 1!")
99        self._species = species
100
101    @property
102    def x(self) -> float:
103        """
104        Cartesian x coordinate
105        """
106        return self.coords[0]  # type: ignore
107
108    @x.setter
109    def x(self, x: float):
110        self.coords[0] = x  # type: ignore
111
112    @property
113    def y(self) -> float:
114        """
115        Cartesian y coordinate
116        """
117        return self.coords[1]  # type: ignore
118
119    @y.setter
120    def y(self, y: float):
121        self.coords[1] = y  # type: ignore
122
123    @property
124    def z(self) -> float:
125        """
126        Cartesian z coordinate
127        """
128        return self.coords[2]  # type: ignore
129
130    @z.setter
131    def z(self, z: float):
132        self.coords[2] = z  # type: ignore
133
134    def distance(self, other) -> float:
135        """
136        Get distance between two sites.
137
138        Args:
139            other: Other site.
140
141        Returns:
142            Distance (float)
143        """
144        return np.linalg.norm(other.coords - self.coords)
145
146    def distance_from_point(self, pt) -> float:
147        """
148        Returns distance between the site and a point in space.
149
150        Args:
151            pt: Cartesian coordinates of point.
152
153        Returns:
154            Distance (float)
155        """
156        return np.linalg.norm(np.array(pt) - self.coords)
157
158    @property
159    def species_string(self) -> str:
160        """
161        String representation of species on the site.
162        """
163        if self.is_ordered:
164            return list(self.species.keys())[0].__str__()
165        sorted_species = sorted(self.species.keys())
166        return ", ".join(["{}:{:.3f}".format(sp, self.species[sp]) for sp in sorted_species])
167
168    @property  # type: ignore
169    @deprecated(message="Use site.species instead. This will be deprecated with effect from pymatgen 2020.")
170    def species_and_occu(self):
171        """
172        The species at the site, i.e., a Composition mapping type of
173        element/species to occupancy.
174        """
175        return self.species
176
177    @property
178    def specie(self) -> Union[Element, Species, DummySpecies]:
179        """
180        The Species/Element at the site. Only works for ordered sites. Otherwise
181        an AttributeError is raised. Use this property sparingly.  Robust
182        design should make use of the property species instead. Note that the
183        singular of species is also species. So the choice of this variable
184        name is governed by programmatic concerns as opposed to grammar.
185
186        Raises:
187            AttributeError if Site is not ordered.
188        """
189        if not self.is_ordered:
190            raise AttributeError("specie property only works for ordered " "sites!")
191        return list(self.species.keys())[0]
192
193    @property
194    def is_ordered(self) -> bool:
195        """
196        True if site is an ordered site, i.e., with a single species with
197        occupancy 1.
198        """
199        totaloccu = self.species.num_atoms
200        return totaloccu == 1 and len(self.species) == 1
201
202    def __getitem__(self, el):
203        """
204        Get the occupancy for element
205        """
206        return self.species[el]
207
208    def __eq__(self, other):
209        """
210        Site is equal to another site if the species and occupancies are the
211        same, and the coordinates are the same to some tolerance.  numpy
212        function `allclose` is used to determine if coordinates are close.
213        """
214        if other is None:
215            return False
216        return (
217            self.species == other.species
218            and np.allclose(self.coords, other.coords, atol=Site.position_atol)
219            and self.properties == other.properties
220        )
221
222    def __ne__(self, other):
223        return not self.__eq__(other)
224
225    def __hash__(self):
226        """
227        Minimally effective hash function that just distinguishes between Sites
228        with different elements.
229        """
230        return sum([el.Z for el in self.species.keys()])
231
232    def __contains__(self, el):
233        return el in self.species
234
235    def __repr__(self):
236        return "Site: {} ({:.4f}, {:.4f}, {:.4f})".format(self.species_string, *self.coords)
237
238    def __lt__(self, other):
239        """
240        Sets a default sort order for atomic species by electronegativity. Very
241        useful for getting correct formulas.  For example, FeO4PLi is
242        automatically sorted in LiFePO4.
243        """
244        if self.species.average_electroneg < other.species.average_electroneg:
245            return True
246        if self.species.average_electroneg > other.species.average_electroneg:
247            return False
248        if self.species_string < other.species_string:
249            return True
250        if self.species_string > other.species_string:
251            return False
252        return False
253
254    def __str__(self):
255        return "{} {}".format(self.coords, self.species_string)
256
257    def as_dict(self) -> dict:
258        """
259        Json-serializable dict representation for Site.
260        """
261        species_list = []
262        for spec, occu in self.species.items():
263            d = spec.as_dict()
264            del d["@module"]
265            del d["@class"]
266            d["occu"] = occu
267            species_list.append(d)
268        d = {
269            "name": self.species_string,
270            "species": species_list,
271            "xyz": [float(c) for c in self.coords],  # type: ignore
272            "properties": self.properties,
273            "@module": self.__class__.__module__,
274            "@class": self.__class__.__name__,
275        }
276        if self.properties:
277            d["properties"] = self.properties
278        return d
279
280    @classmethod
281    def from_dict(cls, d: dict) -> "Site":
282        """
283        Create Site from dict representation
284        """
285        atoms_n_occu = {}
286        for sp_occu in d["species"]:
287            if "oxidation_state" in sp_occu and Element.is_valid_symbol(sp_occu["element"]):
288                sp = Species.from_dict(sp_occu)
289            elif "oxidation_state" in sp_occu:
290                sp = DummySpecies.from_dict(sp_occu)
291            else:
292                sp = Element(sp_occu["element"])  # type: ignore
293            atoms_n_occu[sp] = sp_occu["occu"]
294        props = d.get("properties", None)
295        if props is not None:
296            for key in props.keys():
297                props[key] = json.loads(json.dumps(props[key], cls=MontyEncoder), cls=MontyDecoder)
298        return cls(atoms_n_occu, d["xyz"], properties=props)
299
300
301class PeriodicSite(Site, MSONable):
302    """
303    Extension of generic Site object to periodic systems.
304    PeriodicSite includes a lattice system.
305    """
306
307    def __init__(
308        self,
309        species: Union[SpeciesLike, CompositionLike],
310        coords: ArrayLike,
311        lattice: Lattice,
312        to_unit_cell: bool = False,
313        coords_are_cartesian: bool = False,
314        properties: dict = None,
315        skip_checks: bool = False,
316    ):
317        """
318        Create a periodic site.
319
320        :param species: Species on the site. Can be:
321            i.  A Composition-type object (preferred)
322            ii. An  element / species specified either as a string
323                symbols, e.g. "Li", "Fe2+", "P" or atomic numbers,
324                e.g., 3, 56, or actual Element or Species objects.
325            iii.Dict of elements/species and occupancies, e.g.,
326                {"Fe" : 0.5, "Mn":0.5}. This allows the setup of
327                disordered structures.
328        :param coords: Cartesian coordinates of site.
329        :param lattice: Lattice associated with the site.
330        :param to_unit_cell: Translates fractional coordinate to the
331            basic unit cell, i.e. all fractional coordinates satisfy 0
332            <= a < 1. Defaults to False.
333        :param coords_are_cartesian: Set to True if you are providing
334            cartesian coordinates. Defaults to False.
335        :param properties: Properties associated with the site as a dict, e.g.
336            {"magmom": 5}. Defaults to None.
337        :param skip_checks: Whether to ignore all the usual checks and just
338            create the site. Use this if the PeriodicSite is created in a
339            controlled manner and speed is desired.
340        """
341
342        if coords_are_cartesian:
343            frac_coords = lattice.get_fractional_coords(coords)
344        else:
345            frac_coords = coords  # type: ignore
346
347        if to_unit_cell:
348            frac_coords = np.mod(frac_coords, 1)
349
350        if not skip_checks:
351            frac_coords = np.array(frac_coords)
352            if not isinstance(species, Composition):
353                try:
354                    species = Composition({get_el_sp(species): 1})
355                except TypeError:
356                    species = Composition(species)
357
358            totaloccu = species.num_atoms
359            if totaloccu > 1 + Composition.amount_tolerance:
360                raise ValueError("Species occupancies sum to more than 1!")
361
362        self._lattice: Lattice = lattice
363        self._frac_coords: ArrayLike = frac_coords
364        self._species: Composition = species  # type: ignore
365        self._coords: Optional[np.ndarray] = None
366        self.properties: dict = properties or {}
367
368    def __hash__(self):
369        """
370        Minimally effective hash function that just distinguishes between Sites
371        with different elements.
372        """
373        return sum([el.Z for el in self.species.keys()])
374
375    @property
376    def lattice(self) -> Lattice:
377        """
378        Lattice associated with PeriodicSite
379        """
380        return self._lattice
381
382    @lattice.setter
383    def lattice(self, lattice: Lattice):
384        """
385        Sets Lattice associated with PeriodicSite
386        """
387        self._lattice = lattice
388        self._coords = self._lattice.get_cartesian_coords(self._frac_coords)
389
390    @property  # type: ignore
391    def coords(self) -> np.ndarray:  # type: ignore
392        """
393        Cartesian coordinates
394        """
395        if self._coords is None:
396            self._coords = self._lattice.get_cartesian_coords(self._frac_coords)
397        return self._coords
398
399    @coords.setter
400    def coords(self, coords):
401        """
402        Set Cartesian coordinates
403        """
404        self._coords = np.array(coords)
405        self._frac_coords = self._lattice.get_fractional_coords(self._coords)
406
407    @property
408    def frac_coords(self) -> np.ndarray:
409        """
410        Fractional coordinates
411        """
412        return self._frac_coords  # type: ignore
413
414    @frac_coords.setter
415    def frac_coords(self, frac_coords):
416        """
417        Set fractional coordinates
418        """
419        self._frac_coords = np.array(frac_coords)
420        self._coords = self._lattice.get_cartesian_coords(self._frac_coords)
421
422    @property
423    def a(self) -> float:
424        """
425        Fractional a coordinate
426        """
427        return self._frac_coords[0]  # type: ignore
428
429    @a.setter
430    def a(self, a: float):
431        self._frac_coords[0] = a  # type: ignore
432        self._coords = self._lattice.get_cartesian_coords(self._frac_coords)
433
434    @property
435    def b(self) -> float:
436        """
437        Fractional b coordinate
438        """
439        return self._frac_coords[1]  # type: ignore
440
441    @b.setter
442    def b(self, b: float):
443        self._frac_coords[1] = b  # type: ignore
444        self._coords = self._lattice.get_cartesian_coords(self._frac_coords)
445
446    @property
447    def c(self) -> float:
448        """
449        Fractional c coordinate
450        """
451        return self._frac_coords[2]  # type: ignore
452
453    @c.setter
454    def c(self, c: float):
455        self._frac_coords[2] = c  # type: ignore
456        self._coords = self._lattice.get_cartesian_coords(self._frac_coords)
457
458    @property
459    def x(self) -> float:
460        """
461        Cartesian x coordinate
462        """
463        return self.coords[0]
464
465    @x.setter
466    def x(self, x: float):
467        self.coords[0] = x
468        self._frac_coords = self._lattice.get_fractional_coords(self.coords)
469
470    @property
471    def y(self) -> float:
472        """
473        Cartesian y coordinate
474        """
475        return self.coords[1]
476
477    @y.setter
478    def y(self, y: float):
479        self.coords[1] = y
480        self._frac_coords = self._lattice.get_fractional_coords(self.coords)
481
482    @property
483    def z(self) -> float:
484        """
485        Cartesian z coordinate
486        """
487        return self.coords[2]
488
489    @z.setter
490    def z(self, z: float):
491        self.coords[2] = z
492        self._frac_coords = self._lattice.get_fractional_coords(self.coords)
493
494    def to_unit_cell(self, in_place=False) -> Optional["PeriodicSite"]:
495        """
496        Move frac coords to within the unit cell cell.
497        """
498        frac_coords = np.mod(self.frac_coords, 1)
499        if in_place:
500            self.frac_coords = frac_coords
501            return None
502        return PeriodicSite(self.species, frac_coords, self.lattice, properties=self.properties)
503
504    def is_periodic_image(self, other: "PeriodicSite", tolerance: float = 1e-8, check_lattice: bool = True) -> bool:
505        """
506        Returns True if sites are periodic images of each other.
507
508        Args:
509            other (PeriodicSite): Other site
510            tolerance (float): Tolerance to compare fractional coordinates
511            check_lattice (bool): Whether to check if the two sites have the
512                same lattice.
513
514        Returns:
515            bool: True if sites are periodic images of each other.
516        """
517        if check_lattice and self.lattice != other.lattice:
518            return False
519        if self.species != other.species:
520            return False
521
522        frac_diff = pbc_diff(self.frac_coords, other.frac_coords)
523        return np.allclose(frac_diff, [0, 0, 0], atol=tolerance)
524
525    def __eq__(self, other):
526        return (
527            self.species == other.species
528            and self.lattice == other.lattice
529            and np.allclose(self.coords, other.coords, atol=Site.position_atol)
530            and self.properties == other.properties
531        )
532
533    def __ne__(self, other):
534        return not self.__eq__(other)
535
536    def distance_and_image_from_frac_coords(
537        self, fcoords: ArrayLike, jimage: Optional[ArrayLike] = None
538    ) -> Tuple[float, np.ndarray]:
539        """
540        Gets distance between site and a fractional coordinate assuming
541        periodic boundary conditions. If the index jimage of two sites atom j
542        is not specified it selects the j image nearest to the i atom and
543        returns the distance and jimage indices in terms of lattice vector
544        translations. If the index jimage of atom j is specified it returns the
545        distance between the i atom and the specified jimage atom, the given
546        jimage is also returned.
547
548        Args:
549            fcoords (3x1 array): fcoords to get distance from.
550            jimage (3x1 array): Specific periodic image in terms of
551                lattice translations, e.g., [1,0,0] implies to take periodic
552                image that is one a-lattice vector away. If jimage is None,
553                the image that is nearest to the site is found.
554
555        Returns:
556            (distance, jimage): distance and periodic lattice translations
557            of the other site for which the distance applies.
558        """
559        return self.lattice.get_distance_and_image(self.frac_coords, fcoords, jimage=jimage)
560
561    def distance_and_image(self, other: "PeriodicSite", jimage: Optional[ArrayLike] = None) -> Tuple[float, np.ndarray]:
562        """
563        Gets distance and instance between two sites assuming periodic boundary
564        conditions. If the index jimage of two sites atom j is not specified it
565        selects the j image nearest to the i atom and returns the distance and
566        jimage indices in terms of lattice vector translations. If the index
567        jimage of atom j is specified it returns the distance between the ith
568        atom and the specified jimage atom, the given jimage is also returned.
569
570        Args:
571            other (PeriodicSite): Other site to get distance from.
572            jimage (3x1 array): Specific periodic image in terms of lattice
573                translations, e.g., [1,0,0] implies to take periodic image
574                that is one a-lattice vector away. If jimage is None,
575                the image that is nearest to the site is found.
576
577        Returns:
578            (distance, jimage): distance and periodic lattice translations
579            of the other site for which the distance applies.
580        """
581        return self.distance_and_image_from_frac_coords(other.frac_coords, jimage)
582
583    def distance(self, other: "PeriodicSite", jimage: Optional[ArrayLike] = None):
584        """
585        Get distance between two sites assuming periodic boundary conditions.
586
587        Args:
588            other (PeriodicSite): Other site to get distance from.
589            jimage (3x1 array): Specific periodic image in terms of lattice
590                translations, e.g., [1,0,0] implies to take periodic image
591                that is one a-lattice vector away. If jimage is None,
592                the image that is nearest to the site is found.
593
594        Returns:
595            distance (float): Distance between the two sites
596        """
597        return self.distance_and_image(other, jimage)[0]
598
599    def __repr__(self):
600        return "PeriodicSite: {} ({:.4f}, {:.4f}, {:.4f}) [{:.4f}, {:.4f}, " "{:.4f}]".format(
601            self.species_string, self.coords[0], self.coords[1], self.coords[2], *self._frac_coords
602        )
603
604    def as_dict(self, verbosity: int = 0) -> dict:
605        """
606        Json-serializable dict representation of PeriodicSite.
607
608        Args:
609            verbosity (int): Verbosity level. Default of 0 only includes the
610                matrix representation. Set to 1 for more details such as
611                cartesian coordinates, etc.
612        """
613        species_list = []
614        for spec, occu in self._species.items():
615            d = spec.as_dict()
616            del d["@module"]
617            del d["@class"]
618            d["occu"] = occu
619            species_list.append(d)
620
621        d = {
622            "species": species_list,
623            "abc": [float(c) for c in self._frac_coords],  # type: ignore
624            "lattice": self._lattice.as_dict(verbosity=verbosity),
625            "@module": self.__class__.__module__,
626            "@class": self.__class__.__name__,
627        }
628
629        if verbosity > 0:
630            d["xyz"] = [float(c) for c in self.coords]
631            d["label"] = self.species_string
632
633        d["properties"] = self.properties
634
635        return d
636
637    @classmethod
638    def from_dict(cls, d, lattice=None) -> "PeriodicSite":
639        """
640        Create PeriodicSite from dict representation.
641
642        Args:
643            d (dict): dict representation of PeriodicSite
644            lattice: Optional lattice to override lattice specified in d.
645                Useful for ensuring all sites in a structure share the same
646                lattice.
647
648        Returns:
649            PeriodicSite
650        """
651        species = {}
652        for sp_occu in d["species"]:
653            if "oxidation_state" in sp_occu and Element.is_valid_symbol(sp_occu["element"]):
654                sp = Species.from_dict(sp_occu)
655            elif "oxidation_state" in sp_occu:
656                sp = DummySpecies.from_dict(sp_occu)
657            else:
658                sp = Element(sp_occu["element"])  # type: ignore
659            species[sp] = sp_occu["occu"]
660        props = d.get("properties", None)
661        if props is not None:
662            for key in props.keys():
663                props[key] = json.loads(json.dumps(props[key], cls=MontyEncoder), cls=MontyDecoder)
664        lattice = lattice if lattice else Lattice.from_dict(d["lattice"])
665        return cls(species, d["abc"], lattice, properties=props)
666