Skip to content

Commit

Permalink
Minor doc improvements
Browse files Browse the repository at this point in the history
  • Loading branch information
hexane360 committed Nov 23, 2024
1 parent 798c824 commit ce30dcd
Show file tree
Hide file tree
Showing 4 changed files with 95 additions and 49 deletions.
2 changes: 1 addition & 1 deletion atomlib/defect.py
Original file line number Diff line number Diff line change
Expand Up @@ -469,7 +469,7 @@ def _disp(r: NDArray[numpy.float64]) -> NDArray[numpy.float64]:

def _loop_disp_z(pts: NDArray[numpy.float64], b_vec: numpy.ndarray, loop_r: float, *,
poisson: float = 0.25, branch: t.Optional[numpy.ndarray] = None) -> numpy.ndarray:
from scipy.special import ellipk, ellipe
from scipy.special import ellipk, ellipe # type: ignore

rho = numpy.linalg.norm(pts[..., :2], axis=-1)
r = numpy.linalg.norm(pts, axis=-1)
Expand Down
127 changes: 83 additions & 44 deletions atomlib/elem.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,43 +8,6 @@

from .types import ElemLike, ElemsLike

ELEMENTS = {
'h': 1, 'he': 2, 'li': 3, 'be': 4, 'b': 5, 'c': 6, 'n': 7, 'o': 8,
'f': 9, 'ne': 10, 'na': 11, 'mg': 12, 'al': 13, 'si': 14, 'p': 15, 's': 16,
'cl': 17, 'ar': 18, 'k': 19, 'ca': 20, 'sc': 21, 'ti': 22, 'v': 23, 'cr': 24,
'mn': 25, 'fe': 26, 'co': 27, 'ni': 28, 'cu': 29, 'zn': 30, 'ga': 31, 'ge': 32,
'as': 33, 'se': 34, 'br': 35, 'kr': 36, 'rb': 37, 'sr': 38, 'y': 39, 'zr': 40,
'nb': 41, 'mo': 42, 'tc': 43, 'ru': 44, 'rh': 45, 'pd': 46, 'ag': 47, 'cd': 48,
'in': 49, 'sn': 50, 'sb': 51, 'te': 52, 'i': 53, 'xe': 54, 'cs': 55, 'ba': 56,
'la': 57, 'ce': 58, 'pr': 59, 'nd': 60, 'pm': 61, 'sm': 62, 'eu': 63, 'gd': 64,
'tb': 65, 'dy': 66, 'ho': 67, 'er': 68, 'tm': 69, 'yb': 70, 'lu': 71, 'hf': 72,
'ta': 73, 'w': 74, 're': 75, 'os': 76, 'ir': 77, 'pt': 78, 'au': 79, 'hg': 80,
'tl': 81, 'pb': 82, 'bi': 83, 'po': 84, 'at': 85, 'rn': 86, 'fr': 87, 'ra': 88,
'ac': 89, 'th': 90, 'pa': 91, 'u': 92, 'np': 93, 'pu': 94, 'am': 95, 'cm': 96,
'bk': 97, 'cf': 98, 'es': 99, 'fm': 100, 'md': 101, 'no': 102, 'lr': 103, 'rf': 104,
'db': 105, 'sg': 106, 'bh': 107, 'hs': 108, 'mt': 109, 'ds': 110, 'rg': 111, 'cn': 112,
'nh': 113, 'fl': 114, 'mc': 115, 'lv': 116, 'ts': 117, 'og': 118,
}

ELEMENT_SYMBOLS = [
'H', 'He', 'Li', 'Be', 'B', 'C', 'N', 'O', 'F', 'Ne', 'Na', 'Mg', 'Al', 'Si', 'P', 'S',
'Cl', 'Ar', 'K', 'Ca', 'Sc', 'Ti', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Ga', 'Ge',
'As', 'Se', 'Br', 'Kr', 'Rb', 'Sr', 'Y', 'Zr', 'Nb', 'Mo', 'Tc', 'Ru', 'Rh', 'Pd', 'Ag', 'Cd',
'In', 'Sn', 'Sb', 'Te', 'I', 'Xe', 'Cs', 'Ba', 'La', 'Ce', 'Pr', 'Nd', 'Pm', 'Sm', 'Eu', 'Gd',
'Tb', 'Dy', 'Ho', 'Er', 'Tm', 'Yb', 'Lu', 'Hf', 'Ta', 'W', 'Re', 'Os', 'Ir', 'Pt', 'Au', 'Hg',
'Tl', 'Pb', 'Bi', 'Po', 'At', 'Rn', 'Fr', 'Ra', 'Ac', 'Th', 'Pa', 'U', 'Np', 'Pu', 'Am', 'Cm',
'Bk', 'Cf', 'Es', 'Fm', 'Md', 'No', 'Lr', 'Rf', 'Db', 'Sg', 'Bh', 'Hs', 'Mt', 'Ds', 'Rg', 'Cn',
'Nh', 'Fl', 'Mc', 'Lv', 'Ts', 'Og',
]
assert len(ELEMENTS) == len(ELEMENT_SYMBOLS)
ELEMENT_SYMBOLS_POLARS = polars.Series([ELEMENT_SYMBOLS], dtype=polars.List(polars.Utf8))

DATA_PATH = files('atomlib.data')
_ELEMENT_MASSES: t.Optional[numpy.ndarray] = None
_ION_RADII: t.Optional[t.Dict[str, float]] = None
_ELEMENT_RADII: t.Optional[numpy.ndarray] = None


def _open_binary_data(filename: str) -> t.ContextManager[t.BinaryIO]:
return t.cast(t.ContextManager[t.BinaryIO], (DATA_PATH / filename).open('rb'))

Expand All @@ -60,9 +23,6 @@ def _get_sym(elem: int) -> str:
raise ValueError(f"Invalid atomic number {elem}") from None


_SYM_RE = r'[a-zA-Z]{1,3}'


@t.overload
def get_elem(sym: ElemLike) -> int:
...
Expand All @@ -72,6 +32,24 @@ def get_elem(sym: polars.Series) -> polars.Series:
...

def get_elem(sym: t.Union[int, str, polars.Series]):
"""
Get the atomic number corresponding to a given symbol.
# Examples
```python
>>> get_elem("Gd")
62
>>> get_elem(polars.Series(["Gd", "Ce", "O"]))
shape: (3,)
Series: 'elem' [i8]
[
64
58
8
]
```
"""

if isinstance(sym, int):
if not 0 < sym < len(ELEMENTS):
raise ValueError(f"Invalid atomic number {sym}")
Expand All @@ -98,6 +76,18 @@ def get_elem(sym: t.Union[int, str, polars.Series]):


def get_elems(sym: ElemsLike) -> t.List[t.Tuple[int, float]]:
"""
Get the elements and quantities corresponding to a formula unit.
# Examples
```python
>>> get_elems("AlN")
[(13, 1.0), (7, 1.0)]
>>> get_elems("Al0.93Sc0.07N")
[(13, 0.93), (21, 0.07), (7, 1.0)]
```
"""

if not isinstance(sym, str):
if isinstance(sym, int):
return [(sym, 1.0)]
Expand Down Expand Up @@ -142,6 +132,15 @@ def get_sym(elem: polars.Series) -> polars.Series:
...

def get_sym(elem: t.Union[int, polars.Series]):
"""
Get the symbol corresponding to an atomic number.
# Examples
```python
>>> get_sym(5)
"B"
```
"""
if isinstance(elem, polars.Series):
sym = elem.cast(polars.Int64).replace_strict(
list(range(1, len(ELEMENT_SYMBOLS)+1)),
Expand Down Expand Up @@ -173,8 +172,9 @@ def get_mass(elem: t.Union[numpy.ndarray, t.Sequence[int]]) -> numpy.ndarray:
def get_mass(elem: t.Union[int, t.Sequence[int], numpy.ndarray, polars.Series]):
"""
Get the standard atomic mass for the given element.
Follows the 2021 IUPAC definitions [1].
Follows the `2021 table of the IUPAC Commission on Isotopic Abundances and Atomic Weights <https://doi.org/10.1515/pac-2019-0603>`_.
[1] 2021 table of the IUPAC Commission on Isotopic Abundances and Atomic Weights <https://doi.org/10.1515/pac-2019-0603>
"""
global _ELEMENT_MASSES

Expand All @@ -192,9 +192,10 @@ def get_mass(elem: t.Union[int, t.Sequence[int], numpy.ndarray, polars.Series]):

def get_ionic_radius(elem: int, charge: int) -> float:
"""
Get crystal ionic radius in angstroms for ``elem`` in charge state ``charge``.
Get crystal ionic radius in angstroms for `elem` in charge state `charge`.
Follows the values in [2].
Follows `R.D. Shannon, Acta Cryst. A32 (1976) <https://doi.org/10.1107/S0567739476001551>`.
[2] R.D. Shannon, Acta Cryst. A32 (1976) <https://doi.org/10.1107/S0567739476001551>
"""
global _ION_RADII

Expand Down Expand Up @@ -228,8 +229,9 @@ def get_radius(elem: t.Union[numpy.ndarray, t.Sequence[int]]) -> numpy.ndarray:
def get_radius(elem: t.Union[int, t.Sequence[int], numpy.ndarray, polars.Series]):
"""
Get the neutral atomic radius for the given element(s), in angstroms.
Follows the values in [3].
Follows the calculated values in `E. Clementi et. al, J. Chem. Phys. 47 (1967) <https://doi.org/10.1063/1.1712084>`_.
[3] E. Clementi et. al, J. Chem. Phys. 47 (1967) <https://doi.org/10.1063/1.1712084>
"""
global _ELEMENT_RADII

Expand All @@ -243,3 +245,40 @@ def get_radius(elem: t.Union[int, t.Sequence[int], numpy.ndarray, polars.Series]
if isinstance(elem, (int, numpy.ndarray)):
return _ELEMENT_RADII[elem-1] # type: ignore
return _ELEMENT_RADII[[e-1 for e in elem]] # type: ignore


ELEMENTS = {
'h': 1, 'he': 2, 'li': 3, 'be': 4, 'b': 5, 'c': 6, 'n': 7, 'o': 8,
'f': 9, 'ne': 10, 'na': 11, 'mg': 12, 'al': 13, 'si': 14, 'p': 15, 's': 16,
'cl': 17, 'ar': 18, 'k': 19, 'ca': 20, 'sc': 21, 'ti': 22, 'v': 23, 'cr': 24,
'mn': 25, 'fe': 26, 'co': 27, 'ni': 28, 'cu': 29, 'zn': 30, 'ga': 31, 'ge': 32,
'as': 33, 'se': 34, 'br': 35, 'kr': 36, 'rb': 37, 'sr': 38, 'y': 39, 'zr': 40,
'nb': 41, 'mo': 42, 'tc': 43, 'ru': 44, 'rh': 45, 'pd': 46, 'ag': 47, 'cd': 48,
'in': 49, 'sn': 50, 'sb': 51, 'te': 52, 'i': 53, 'xe': 54, 'cs': 55, 'ba': 56,
'la': 57, 'ce': 58, 'pr': 59, 'nd': 60, 'pm': 61, 'sm': 62, 'eu': 63, 'gd': 64,
'tb': 65, 'dy': 66, 'ho': 67, 'er': 68, 'tm': 69, 'yb': 70, 'lu': 71, 'hf': 72,
'ta': 73, 'w': 74, 're': 75, 'os': 76, 'ir': 77, 'pt': 78, 'au': 79, 'hg': 80,
'tl': 81, 'pb': 82, 'bi': 83, 'po': 84, 'at': 85, 'rn': 86, 'fr': 87, 'ra': 88,
'ac': 89, 'th': 90, 'pa': 91, 'u': 92, 'np': 93, 'pu': 94, 'am': 95, 'cm': 96,
'bk': 97, 'cf': 98, 'es': 99, 'fm': 100, 'md': 101, 'no': 102, 'lr': 103, 'rf': 104,
'db': 105, 'sg': 106, 'bh': 107, 'hs': 108, 'mt': 109, 'ds': 110, 'rg': 111, 'cn': 112,
'nh': 113, 'fl': 114, 'mc': 115, 'lv': 116, 'ts': 117, 'og': 118,
}

ELEMENT_SYMBOLS = [
'H', 'He', 'Li', 'Be', 'B', 'C', 'N', 'O', 'F', 'Ne', 'Na', 'Mg', 'Al', 'Si', 'P', 'S',
'Cl', 'Ar', 'K', 'Ca', 'Sc', 'Ti', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Ga', 'Ge',
'As', 'Se', 'Br', 'Kr', 'Rb', 'Sr', 'Y', 'Zr', 'Nb', 'Mo', 'Tc', 'Ru', 'Rh', 'Pd', 'Ag', 'Cd',
'In', 'Sn', 'Sb', 'Te', 'I', 'Xe', 'Cs', 'Ba', 'La', 'Ce', 'Pr', 'Nd', 'Pm', 'Sm', 'Eu', 'Gd',
'Tb', 'Dy', 'Ho', 'Er', 'Tm', 'Yb', 'Lu', 'Hf', 'Ta', 'W', 'Re', 'Os', 'Ir', 'Pt', 'Au', 'Hg',
'Tl', 'Pb', 'Bi', 'Po', 'At', 'Rn', 'Fr', 'Ra', 'Ac', 'Th', 'Pa', 'U', 'Np', 'Pu', 'Am', 'Cm',
'Bk', 'Cf', 'Es', 'Fm', 'Md', 'No', 'Lr', 'Rf', 'Db', 'Sg', 'Bh', 'Hs', 'Mt', 'Ds', 'Rg', 'Cn',
'Nh', 'Fl', 'Mc', 'Lv', 'Ts', 'Og',
]
assert len(ELEMENTS) == len(ELEMENT_SYMBOLS)

DATA_PATH = files('atomlib.data')
_ELEMENT_MASSES: t.Optional[numpy.ndarray] = None
_ION_RADII: t.Optional[t.Dict[str, float]] = None
_ELEMENT_RADII: t.Optional[numpy.ndarray] = None
_SYM_RE = r'[a-zA-Z]{1,3}'
4 changes: 2 additions & 2 deletions atomlib/transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -731,7 +731,7 @@ def conjugate(self, transform: Transform3DT) -> Transform3DT: # type: ignore (s
"""
return self.inverse() @ self.compose(transform)

def compose(self, other: Transform3DT) -> Transform3DT: # type: ignore (spurious)
def compose(self, other: Transform3DT) -> Transform3DT:
"""Compose this transformation with another."""
if isinstance(other, LinearTransform3D):
return other.__class__(other.inner @ self.inner)
Expand Down Expand Up @@ -787,7 +787,7 @@ def __matmul__(self, other: BBox3D) -> BBox3D:
def __matmul__(self, other: ArrayLike) -> NDArray[numpy.floating]:
...

def __matmul__(self, other: t.Union[Transform3DT, ArrayLike, BBox3D]) -> t.Union[Transform3DT, NDArray[numpy.floating], BBox3D]:
def __matmul__(self, other: t.Union[Transform3D, ArrayLike, BBox3D]) -> t.Union[Transform3D, NDArray[numpy.floating], BBox3D]:
"""Compose this transformation, or apply it to a given set of points."""
if isinstance(other, Transform3D):
return other.compose(self)
Expand Down
11 changes: 9 additions & 2 deletions atomlib/vec.py
Original file line number Diff line number Diff line change
Expand Up @@ -191,9 +191,16 @@ def in_polygon(poly: numpy.ndarray, pt: t.Optional[numpy.ndarray] = None, *,
def reduce_vec(arr: ArrayLike, max_denom: int = 10000) -> NDArray[numpy.int64]:
"""
Reduce a crystallographic vector (int or float) to lowest common terms.
Example: reduce_vec([3, 3, 3]) = [1, 1, 1]
reduce_vec([0.25, 0.25, 0.25]) = [1, 1, 1]
# Examples
```python
>>> reduce_vec([3, 6, 9])
[1, 2, 3]
>>> reduce_vec([0.5, 0.25, 0.25])
[2, 1, 1]
```
"""

a = numpy.atleast_1d(arr)
if not numpy.issubdtype(a.dtype, numpy.floating):
return a // numpy.gcd.reduce(a, axis=-1, keepdims=True)
Expand Down

0 comments on commit ce30dcd

Please sign in to comment.