# coding: utf-8
#
# This code is part of lattpy.
#
# Copyright (c) 2022, Dylan Jones
#
# This code is licensed under the MIT License. The copyright notice in the
# LICENSE file in the root directory and this permission notice shall
# be included in all copies or substantial portions of the Software.
"""Objects for representing atoms and the unitcell of a lattice."""
import itertools
from collections import abc
from typing import Union, Any, Iterator, Dict
__all__ = ["Atom"]
[docs]class Atom(abc.MutableMapping):
"""Object representing an atom of a Bravais lattice.
Parameters
----------
name : str, optional
The name of the atom. The default is 'A'.
radius : float, optional
The radius of the atom in real space.
color : str or float or array_like, optional
The color used to visualize the atom.
weight : float, optional
The weight of the atom.
**kwargs
Additional attributes of the atom.
"""
_counter = itertools.count()
__slots__ = ["_index", "_name", "_weight", "_params"]
def __init__(
self,
name: str = None,
radius: float = 0.2,
color: str = None,
weight: float = 1.0,
**kwargs,
):
super().__init__()
index = next(Atom._counter)
self._index = index
self._name = name or "A"
self._weight = weight
self._params = dict(color=color, radius=radius, **kwargs)
@property
def id(self):
"""The id of the atom."""
return id(self)
@property
def index(self) -> int:
"""Return the index of the ``Atom`` instance."""
return self._index
@property
def name(self) -> str:
"""Return the name of the ``Atom`` instance."""
return self._name
@property
def weight(self):
"""Return the weight or the ``Atom`` instance."""
return self._weight
[docs] def dict(self) -> Dict[str, Any]:
"""Returns the data of the ``Atom`` instance as a dictionary."""
data = dict(index=self._index, name=self._name)
data.update(self._params)
return data
[docs] def copy(self) -> "Atom":
"""Creates a deep copy of the ``Atom`` instance."""
return Atom(self.name, weight=self.weight, **self._params.copy())
[docs] def get(self, key: str, default=None) -> Any:
try:
return self.__getitem__(key)
except KeyError:
return default
[docs] def is_identical(self, other: "Atom") -> bool:
"""Checks if the other ``Atom`` is identical to this one."""
return self._name == other.name
def __len__(self) -> int:
"""Return the length of the ``Atom`` attributes."""
return len(self._params)
def __iter__(self) -> Iterator[str]:
"""Iterate over the keys of the ``Atom`` attributes."""
return iter(self._params)
def __getitem__(self, key: str) -> Any:
"""Make ``Atom`` attributes accessable as dictionary items."""
return self._params[key]
def __setitem__(self, key: str, value: Any) -> None:
"""Make ``Atom`` attributes accessable as dictionary items."""
self._params[key] = value
def __delitem__(self, key: str) -> None:
"""Make ``Atom`` attributes accessable as dictionary items."""
del self._params[key]
def __getattribute__(self, key: str) -> Any:
"""Make ``Atom`` attributes accessable as attributes."""
key = str(key)
if not key.startswith("_") and key in self._params.keys():
return self._params[key]
else:
return super().__getattribute__(key)
def __setattr__(self, key: str, value: Any) -> None:
"""Make ``Atom`` attributes accessable as attributes."""
key = str(key)
if not key.startswith("_") and key in self._params.keys():
self._params[key] = value
else:
super().__setattr__(key, value)
def __hash__(self) -> hash:
"""Make ``Atom`` instance hashable."""
return hash(self._name)
def __dict__(self) -> Dict[str, Any]: # pragma: no cover
"""Return the information of the atom as a dictionary"""
return self.dict()
def __copy__(self) -> "Atom": # pragma: no cover
"""Creates a deep copy of the ``Atom`` instance."""
return self.copy()
def __eq__(self, other: Union["Atom", str]) -> bool:
if isinstance(other, Atom):
return self.is_identical(other)
else:
return self._name == other
def __repr__(self) -> str:
argstr = f"{self._name}"
paramstr = ", ".join(f"{k}={v}" for k, v in self._params.items() if v)
if paramstr:
argstr += ", " + paramstr
return f"Atom({argstr}, {self.index})"