# Wulfric - Cell, Atoms, K-path.
# Copyright (C) 2023-2025 Andrey Rybakov
#
# e-mail: anry@uv.es, web: adrybakov.com
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from wulfric._exceptions import FailedToDeduceAtomSpecies
from wulfric.constants._atoms import ATOM_SPECIES
# Save local scope at this moment
old_dir = set(dir())
old_dir.add("old_dir")
[docs]
def get_atom_species(name: str, raise_on_fail=False) -> str:
r"""
Attempts to identify atom's type based on its name (i.e. Cr1 -> Cr, ...).
If no type is identified, then return "X".
Parameters
----------
name : str
Name of the atom.
raise_on_fail : bool, default False
Whether to raise an exception if automatic species deduction fails.
Returns
-------
species : str
Species of the atom.
Raises
------
FailedToDeduceAtomSpecies
If ``raise_on_fail = True`` and automatic species deduction fails.
Warnings
--------
If ``raise_on_fail = True`` and automatic species deduction fails, then
``RuntimeWarning`` is issued, and atom species is set to "X".
Examples
--------
.. doctest::
>>> from wulfric.crystal import get_atom_species
>>> get_atom_species("@%^#$")
'X'
>>> get_atom_species("Cr")
'Cr'
>>> get_atom_species("Cr1")
'Cr'
>>> get_atom_species("_3341Cr")
'Cr'
>>> get_atom_species("cr")
'Cr'
>>> get_atom_species("S")
'S'
>>> get_atom_species("Se")
'Se'
>>> get_atom_species("Sp")
'S'
>>> get_atom_species("123a")
'X'
>>> get_atom_species("CrSBr")
'Cr'
Notes
-----
If ``name`` contains several possible atom types of length 2
as substrings, then the type is equal to the first one found.
"""
atom_type = "X"
for trial_type in ATOM_SPECIES:
if trial_type.lower() in name.lower():
atom_type = trial_type
# Maximum amount of characters in the atom type
# Some 1-character types are parts of some 2-character types (i.e. "Se" and "S")
# If type of two characters is found then it is unique,
# If type of one character is found, then the search must continue
if len(atom_type) == 2:
break
if atom_type == "X":
if raise_on_fail:
raise FailedToDeduceAtomSpecies(name=name)
else:
import warnings
warnings.warn(
f"Atom species deduction failed for '{name}'. Set species to 'X'",
RuntimeWarning,
)
return atom_type
[docs]
def populate_atom_species(atoms, raise_on_fail=False) -> None:
r"""
Populate atom species, based on their names.
If atom species are already present in the ``atoms``, then they will be overwritten.
Parameters
----------
atoms : dict
Dictionary with atoms. Must have a ``names`` with the value of ``list`` of N
``str``.
raise_on_fail : bool, default False
Whether to raise an error if the atom type can not be deduced based on its name.
Raises
------
FailedToDeduceAtomSpecies
If ``raise_on_fail = True`` and automatic species deduction fails.
Warnings
--------
If ``raise_on_fail = True`` and automatic species deduction fails, then
``RuntimeWarning`` is issued, and atom species is set to "X".
Examples
--------
.. doctest::
>>> import wulfric as wulf
>>> atoms = {"names" : ["Cr1", "cr2", "Br3", "S4", "fe5", "Fe6"]}
>>> atoms
{'names': ['Cr1', 'cr2', 'Br3', 'S4', 'fe5', 'Fe6']}
>>> wulf.crystal.populate_atom_species(atoms)
>>> atoms
{'names': ['Cr1', 'cr2', 'Br3', 'S4', 'fe5', 'Fe6'], 'species': ['Cr', 'Cr', 'Br', 'S', 'Fe', 'Fe']}
"""
atoms["species"] = []
for i in range(len(atoms["names"])):
atoms["species"].append(
get_atom_species(atoms["names"][i], raise_on_fail=raise_on_fail)
)
[docs]
def ensure_unique_names(atoms, strategy: str = "all") -> None:
r"""
Ensures that atoms have unique ``"names"``.
If atom names are already unique, then this function does nothing.
.. versionadded:: 0.5.1
Parameters
----------
atoms : dict
Dictionary with atoms. Must have a ``"names"`` keyword with the value of
``list`` of N ``str``.
strategy : str, default "all"
Strategy for the modification of atom names. Supported strategies are
* "all"
Add an index to the end of every atom, starting from 1.
* "repeated-only"
Add an index only to the repeated names, index starts with 1, independently for
each repeated grooup. (See examples)
Case-insensitive.
Raises
------
ValueError
If ``strategy`` is not supported.
Examples
--------
.. doctest::
>>> import wulfric as wulf
>>> atoms = {"names" : ["Cr1", "Cr2", "Br", "Br", "S", "S"]}
>>> # Default strategy is "all"
>>> wulf.crystal.ensure_unique_names(atoms)
>>> atoms
{'names': ['Cr11', 'Cr22', 'Br3', 'Br4', 'S5', 'S6']}
>>> atoms = {"names" : ["Cr1", "Cr2", "Br", "Br", "S", "S"]}
>>> wulf.crystal.ensure_unique_names(atoms, strategy="repeated-only")
>>> atoms
{'names': ['Cr1', 'Cr2', 'Br1', 'Br2', 'S1', 'S2']}
>>> # Nothing happens if atom names are already unique
>>> wulf.crystal.ensure_unique_names(atoms)
>>> atoms
{'names': ['Cr1', 'Cr2', 'Br1', 'Br2', 'S1', 'S2']}
>>> wulf.crystal.ensure_unique_names(atoms, strategy="repeated-only")
>>> atoms
{'names': ['Cr1', 'Cr2', 'Br1', 'Br2', 'S1', 'S2']}
"""
SUPPORTED_STRATEGIES = ["all", "repeated-only"]
strategy = strategy.lower()
if strategy not in SUPPORTED_STRATEGIES:
raise ValueError(
f"{strategy} strategy is not supported. Supported are:\n"
+ ("\n").join([f" * {i}" for i in SUPPORTED_STRATEGIES])
)
names_unique = len(atoms["names"]) == len(set(atoms["names"]))
if not names_unique and strategy == "all":
for i in range(len(atoms["names"])):
atoms["names"][i] += f"{i + 1}"
if not names_unique and strategy == "repeated-only":
counter = {}
for name in atoms["names"]:
if name not in counter:
counter[name] = [1, 1]
else:
counter[name][1] += 1
for i in range(len(atoms["names"])):
name = atoms["names"][i]
total = counter[name][1]
met = counter[name][0]
if total > 1:
atoms["names"][i] += str(counter[name][0])
counter[name][0] += 1
# Populate __all__ with objects defined in this file
__all__ = list(set(dir()) - old_dir)
# Remove all semi-private objects
__all__ = [i for i in __all__ if not i.startswith("_")]
del old_dir
