diff --git a/examples/PbSnTe.cif b/examples/PbSnTe.cif index 945679c..e4fe110 100644 --- a/examples/PbSnTe.cif +++ b/examples/PbSnTe.cif @@ -227,6 +227,6 @@ _atom_site_fract_x _atom_site_fract_y _atom_site_fract_z _atom_site_occupancy -Pb1 Pb2+ 4 a 0.0 0.0 0.0 0.75000 -Sn1 Sn2+ 4 a 0.0 0.0 0.0 0.25000 -Te1 Te2- 4 b 0.5 0.5 0.5 1.00000 +Pb1 Pb 4 a 0.0 0.0 0.0 0.75000 +Sn1 Sn 4 a 0.0 0.0 0.0 0.25000 +Te1 Te 4 b 0.5 0.5 0.5 1.00000 diff --git a/examples/README.md b/examples/README.md index 6faf202..294c95c 100644 --- a/examples/README.md +++ b/examples/README.md @@ -3,12 +3,12 @@ What each script does, summarized per example. All required CIF inputs are bundled in this directory. ## Example 1: Basic enumeration and CIF output -- File: [examples/example1.py](examples/example1.py) +- File: [example1.py](example1.py) - What it does: Load `SmFe7Ti.cif`, use `Substitutor` to generate all symmetry-inequivalent structures, and write CIFs with their multiplicity (weight). - Key point: Uses `Substitutor.quantities(("cifwriter", "weight"))` and saves to `output/` as `cif_i{index}w{weight}.cif`. ## Example 2: Comparison with enumlib -- File: [examples/example2.py](examples/example2.py) +- File: [example2.py](example2.py) - What it does: Expand `PbSnTe.cif` to a 2×2×2 supercell and compare enumeration results between Pymatgen's `EnumlibAdaptor` and SHRY. - Comparison modes: - SHRY default (group equivalent sites). @@ -16,34 +16,34 @@ What each script does, summarized per example. All required CIF inputs are bundl - Output: Prints the number of structures from each approach to stdout. ## Example 3: LabeledStructure and label-based substitutions -- File: [examples/example3.py](examples/example3.py) +- File: [example3.py](example3.py) - What it does: Load `SmFe12.cif` as a `LabeledStructure`, replace labels `Fe1`→`Fe3Ti`, `Fe2`→`Fe7Ti`, `Fe3`→`Ti`, then expand to a non-diagonal supercell `((2,0,1),(0,1,0),(1,0,1))`. Generates all structures, writes CIFs and weights to `output/`, and prints configuration letters (e.g., `aba...`). - Key point: Keeps `_atom_site_label`, so substitutions remain label-consistent even after supercell expansion. ## Example 4a: Caching pattern generation and saving it -- File: [examples/example4a.py](examples/example4a.py) +- File: [example4a.py](example4a.py) - What it does: Load `SmFe12.cif` as a `LabeledStructure`, replace `Fe1`→`Fe7Ti`, run `Substitutor(cache=True)`, write CIFs, and pickle the internal `pattern_makers` to `pg.pkl`. - Benefit: Reuse the cached pattern makers to skip expensive recomputation for symmetry-equivalent systems later. ## Example 4b: Reusing cache and applying to different compositions -- File: [examples/example4b.py](examples/example4b.py) +- File: [example4b.py](example4b.py) - What it does: Load `pg.pkl` from Example 4a. Starting from `SmFe12.cif`, evaluate two substitution cases (`Fe`→`Fe3Ti` and `Fe`→`FeTi3`) in sequence. Attach the loaded `pattern_makers` to a `Substitutor`, generate CIFs for `run1` and `run2`, then overwrite `pg.pkl` with the updated cache. - Key point: When symmetry is the same, you can rapidly iterate different substitution patterns using the cached pattern makers. ## Example 5: Generating ASE Atoms -- File: [examples/example5.py](examples/example5.py) +- File: [example5.py](example5.py) - What it does: Expand `PbSnTe.cif` to 2×2×2, enumerate all structures with `Substitutor`, and obtain them as ASE `Atoms` objects. - Output: Prints the number of structures to stdout and shows one sample `Atoms` instance; useful as a bridge to downstream ASE workflows. ## Example 6: Ranking by Ewald energy -- File: [examples/example6.py](examples/example6.py) +- File: [example6.py](example6.py) - What it does: Enumerates substitutions, guesses oxidation states, computes Ewald energies, and writes the 100 lowest-energy structures to output_ewald/. - Note: If oxidation states cannot be guessed, Ewald energies may be zero or skipped. ## Example 7: Highest symmetry structures (lowest weight) -- File: [examples/example7.py](examples/example7.py) +- File: [example7.py](example7.py) - What it does: Enumerates substitutions and writes the top-N structures with the smallest weight (highest symmetry) to output_symm_top/. Optionally writes all structures sorted by weight. ## Example 8: Merge multiple Wyckoff labels for shared concentration -- File: [examples/example8.py](examples/example8.py) +- File: [example8.py](example8.py) - What it does: Loads `SmFe12.cif`, relabels multiple Wyckoff labels (e.g., `Fe1` and `Fe2`) into a shared label (`Target`), then applies a single disordered composition (e.g., `Fe0.5Ti0.5`) across that pooled site set. This enables specifying a global substitution concentration across multiple Wyckoff positions. diff --git a/examples/SmFe12.cif b/examples/SmFe12.cif index e540681..c287466 100644 --- a/examples/SmFe12.cif +++ b/examples/SmFe12.cif @@ -8,7 +8,6 @@ data_SmFe12 -loop_ _chemical_formula_sum 'Fe12 Sm' _space_group_IT_number 139 _symmetry_space_group_name_Hall '-I 4 2' diff --git a/examples/SmFe7Ti.cif b/examples/SmFe7Ti.cif index 43b96bd..194000a 100644 --- a/examples/SmFe7Ti.cif +++ b/examples/SmFe7Ti.cif @@ -7,9 +7,7 @@ data_SmFe12 # http://www.crystallography.net/ # - -loop_ -_chemical_formula_sum 'Fe12 Sm' +_chemical_formula_sum 'Fe10.5 Sm Ti1.5' _space_group_IT_number 139 _symmetry_space_group_name_Hall '-I 4 2' _symmetry_space_group_name_H-M 'I 4/m m m' @@ -21,8 +19,8 @@ _cell_length_a 8.59 _cell_length_b 8.59 _cell_length_c 4.804 _cell_volume 354.478 -_cod_data_source_block 'Fe12 Sm1' -_cod_original_formula_sum 'Fe12 Sm1' +_cod_data_source_block 'Fe10.5 Sm1 Ti1.5' +_cod_original_formula_sum 'Fe10.5 Sm1 Ti1.5' loop_ _symmetry_equiv_pos_as_xyz x,y,z diff --git a/examples/example6.py b/examples/example6.py index cd2f551..d8cf4a6 100644 --- a/examples/example6.py +++ b/examples/example6.py @@ -19,7 +19,7 @@ helper = ScriptHelper( structure_file="PbSnTe.cif", # Replace with your CIF if desired from_species=("Sn",), - to_species=("Sn0.5Sb0.5",), + to_species=("Sn0.5Se0.5",), scaling_matrix=(2, 2, 1), # enlarge to break more symmetry and get multiple patterns ) @@ -29,7 +29,16 @@ for i, structure in enumerate(helper.substitutor.structure_writers()): try: - structure.add_oxidation_state_by_guess() + # Ensure oxidation states are set explicitly for this system. + structure.remove_oxidation_states() + structure.add_oxidation_state_by_element( + { + "Pb": +2, + "Sn": +2, + "Se": +4, + "Te": -2, + } + ) energy = EwaldSummation(structure).total_energy except Exception as exc: # noqa: BLE001 - narrow for clarity in real workflows print(f"Skipping structure {i}: {exc}") diff --git a/shry/core.py b/shry/core.py index 94951f3..c2842b0 100755 --- a/shry/core.py +++ b/shry/core.py @@ -11,16 +11,15 @@ import math import random import sys -from typing import OrderedDict, Tuple +from typing import OrderedDict # python modules import numpy as np import spglib import sympy import tqdm -from monty.fractions import gcd_float from pymatgen.analysis.ewald import EwaldSummation -from pymatgen.core.composition import Composition, reduce_formula +from pymatgen.core.composition import Composition from pymatgen.io.cif import CifBlock, CifParser, CifWriter from pymatgen.symmetry.analyzer import SpacegroupAnalyzer, SpacegroupOperations from pymatgen.symmetry.structure import SymmetrizedStructure @@ -29,11 +28,11 @@ from scipy.special import comb from sympy.utilities.iterables import multiset_permutations from tabulate import tabulate -from pymatgen.core.periodic_table import get_el_sp from pymatgen.core import Structure # shry modules from . import const +from .patches import apply_pymatgen_patches # shry version control try: @@ -45,99 +44,9 @@ np.set_printoptions(linewidth=const.LINEWIDTH, threshold=16) np.set_printoptions(linewidth=1000, threshold=sys.maxsize) +spglib.OLD_ERROR_HANDLING = False -def get_integer_formula_and_factor( - self, - max_denominator: int = int(1 / const.DEFAULT_ATOL), - iupac_ordering: bool = False, -) -> Tuple[str, float]: - """The default Composition groups together different ox states which is not ideal...""" - el_amt = self.as_dict() - g = gcd_float(list(el_amt.values()), 1 / max_denominator) - - d = {k: round(v / g) for k, v in el_amt.items()} - (formula, factor) = reduce_formula(d, iupac_ordering=iupac_ordering) - if formula in Composition.special_formulas: - formula = Composition.special_formulas[formula] - factor /= 2 - return formula, factor * g - - -# Patched extra functionalities and bug fixes on top of Pymatgen's classes. -def to_int_dict(self): - """ - Returns: - Dict with element symbol and integer amount - """ - _, factor = self.get_integer_formula_and_factor(max_denominator=int(1 / const.DEFAULT_ATOL)) - int_dict = {e: int(a) for e, a in (self / factor).as_dict().items()} - - # be safe: Composition groups together different ox states which is not ideal... - for x, y in zip(int_dict.values(), self.as_dict().values()): - if not np.isclose(x * factor, y, atol=const.DEFAULT_ATOL): - raise ValueError( - "Composition (Occupancy) is not rational! Please try to increase significant digits " - "e.g., 1/3 = 0.3333 -> 1/3 = 0.3333333333333." - ) - - return int_dict - - -@property -def inted_composition(self): - """ - Return Composition instance with integer formula - """ - _, factor = self.get_integer_formula_and_factor(max_denominator=int(1 / const.DEFAULT_ATOL)) - int_comp = self / factor - - # be safe - int_dict = {e: int(a) for e, a in int_comp.as_dict().items()} - if not all(np.isclose(x * factor, y, atol=const.DEFAULT_ATOL) for x, y in zip(int_dict.values(), self.as_dict().values())): - raise ValueError( - "Composition (Occupancy) is not rational! Please try to increase significant digits " - "e.g., 1/3 = 0.3333 -> 1/3 = 0.3333333333333." - ) - - return int_comp - - -def formula_double_format_tol(afloat, ignore_ones=True, tol: float = const.DEFAULT_ATOL * 10): - """ - This function is used to make pretty formulas by formatting the amounts. - Instead of Li1.0 Fe1.0 P1.0 O4.0, you get LiFePO4. - - Args: - afloat (float): a float - ignore_ones (bool): if true, floats of 1 are ignored. - tol (float): Tolerance to round to nearest int. i.e. 2.0000000001 -> 2 - - Returns: - A string representation of the float for formulas. - """ - if ignore_ones and afloat == 1: - return "" - if abs(afloat - round(afloat)) < tol: - return round(afloat) - return round(afloat, 8) - - -@property -def formula(self) -> str: - """ - Returns a formula string, with elements sorted by electronegativity, - e.g., Li4 Fe4 P4 O16. - """ - sym_amt = self.get_el_amt_dict() - syms = sorted(sym_amt, key=lambda sym: get_el_sp(sym).X) - formula = [f"{s}{formula_double_format_tol(sym_amt[s], False)}" for s in syms] - return " ".join(formula) - - -Composition.to_int_dict = to_int_dict -Composition.get_integer_formula_and_factor = get_integer_formula_and_factor -Composition.inted_composition = inted_composition -Composition.formula = formula +apply_pymatgen_patches() class PatchedSymmetrizedStructure(SymmetrizedStructure): @@ -198,7 +107,7 @@ def get_symmetrized_structure(self): self.get_space_group_number(), self.get_symmetry_operations(), ) - return PatchedSymmetrizedStructure(self._structure, sg, ds["equivalent_atoms"], ds["wyckoffs"]) + return PatchedSymmetrizedStructure(self._structure, sg, ds.equivalent_atoms, ds.wyckoffs) class AltCifBlock(CifBlock): @@ -302,9 +211,6 @@ def __init__(self, data, loops, header): self.string_cache[k] = s -# I just want distinguishable exceptions - - class NeedSupercellError(Exception): """ Just so that this error will raise a distinguishable exception @@ -321,9 +227,6 @@ class TooBigError(Exception): ... -# Some numerical methods - - def rec_asc(a, n, m, k, length): """ Recursive part of the thing below @@ -408,7 +311,7 @@ class Substitutor: "_atol", "_symmops", "_pms", - "_enumerator_collection", + "_get_polya", "disorder_groups", "_group_dmat", "_group_perms", @@ -463,7 +366,7 @@ def __init__( self.cache = cache self._symmops = None self._pms = dict() - self._enumerator_collection = PolyaCollection() + self._get_polya = get_polya self.disorder_groups = dict() self._group_dmat = dict() @@ -528,7 +431,7 @@ def structure(self, structure): logging.info(f"Space group: {sga.get_hall()} ({sga.get_space_group_number()})") logging.info(f"Total {len(self._symmops)} symmetry operations") logging.info(sga.get_symmetrized_structure()) - equivalent_atoms = sga.get_symmetry_dataset()["equivalent_atoms"] + equivalent_atoms = sga.get_symmetry_dataset().equivalent_atoms # Identify and label disorder sites. # Index is used for later recoloring the sites. @@ -759,7 +662,7 @@ def nocache_get_pm(subperm, dmat): pm = PatternMaker( subperm, invar=dmat, - enumerator_collection=self._enumerator_collection, + enumerator_collection=self._get_polya, t_kind=self._t_kind, shuffle=self._shuffle, rg=self._rg, @@ -777,7 +680,7 @@ def cached_get_pm(subperm, dmat): pm = PatternMaker( subperm, invar=dmat, - enumerator_collection=self._enumerator_collection, + enumerator_collection=self._get_polya, t_kind=self._t_kind, cache=self.cache, shuffle=self._shuffle, @@ -864,13 +767,13 @@ def count(self): logging.info(f"\nCounting unique patterns for {self.structure.formula}") if len(self._symmops): - enumerator = self._enumerator_collection.get( + enumerator = self._get_polya( # orbit_symmetries, list(self._group_perms.values()), len(self._symmops), ) else: # No-permutations-supplied - enumerator = self._enumerator_collection.get([]) + enumerator = self._get_polya([]) count = enumerator.count(tuple(self._disorder_amounts().values())) return count @@ -1032,6 +935,15 @@ def _get_cifwriter(self, p, symprec=None): subpattern = next(pi) for i in subpattern: template_structure.sites[indices[i]].species = e + # NOTE: Ensure unique labels only for template creation to avoid CIF warnings. + # This does not affect actual output labels, which are overwritten below. + labels = [f"{site.species_string}{i}" for i, site in enumerate(template_structure.sites)] + if "_atom_site_label" in template_structure.site_properties: + template_structure.remove_site_property("_atom_site_label") + template_structure.add_site_property("_atom_site_label", labels) + for label, site in zip(labels, template_structure.sites): + if hasattr(site, "label"): + site.label = label cifwriter = CifWriter(template_structure) # Use faster CifBlock implementation @@ -1100,14 +1012,17 @@ def _get_cifwriter(self, p, symprec=None): angle_tolerance=self._angle_tolerance, ) + rotations = space_group_data.rotations + translations = space_group_data.translations + equivalent_atoms = space_group_data.equivalent_atoms + international = space_group_data.international + number = space_group_data.number + ops = [ transformation_to_string(rot, trans, delim=", ") - for rot, trans in zip( - space_group_data["rotations"], - space_group_data["translations"], - ) + for rot, trans in zip(rotations, translations) ] - u, inv = np.unique(space_group_data["equivalent_atoms"], return_inverse=True) + u, inv = np.unique(equivalent_atoms, return_inverse=True) equivalent_indices = [[] for _ in range(len(u))] for j, inv in enumerate(inv): equivalent_indices[inv].append(j) @@ -1132,8 +1047,8 @@ def _get_cifwriter(self, p, symprec=None): ), ) - block["_symmetry_space_group_name_H-M"] = space_group_data["international"] - block["_symmetry_Int_Tables_number"] = space_group_data["number"] + block["_symmetry_space_group_name_H-M"] = international + block["_symmetry_Int_Tables_number"] = number block["_symmetry_equiv_pos_site_id"] = [str(i) for i in range(1, len(ops) + 1)] block["_symmetry_equiv_pos_as_xyz"] = ops @@ -1186,10 +1101,10 @@ def _get_ewald(self, p, symprec): # TODO: less ad hoc implementation. cifwriter = self._get_cifwriter(p, symprec) cifparser = CifParser.from_str(str(cifwriter)) - structure = cifparser.get_structures(primitive=False)[0] + structure = cifparser.parse_structures(primitive=False)[0] try: if not np.isclose(structure.charge, 0.0): - logging.warn(f"Unit cell is charged: (total charge = {structure.charge}).") + logging.warning(f"Unit cell is charged: (total charge = {structure.charge}).") return EwaldSummation(structure).total_energy except TypeError as exc: raise ValueError("Ewald summation required CIFs with defined oxidation states.") from exc @@ -1203,7 +1118,7 @@ class PatternMaker: __slots__ = ( "_search", "ap", - "_enumerator_collection", + "_get_polya", "_patterns", "_auts", "_subobj_ts", @@ -1297,9 +1212,11 @@ def _iterate(x): self._iterate = np.flatnonzero if enumerator_collection is None: - self._enumerator_collection = PolyaCollection() + self._get_polya = get_polya + elif callable(enumerator_collection): + self._get_polya = enumerator_collection else: - self._enumerator_collection = enumerator_collection + self._get_polya = enumerator_collection.get ( column_index, @@ -1515,7 +1432,7 @@ def _invarless_search(self, start=0, stop=None): stop = self._nix // 2 stop = min(stop, self._nix - stop) - enumerator = self._enumerator_collection.get([self._perms]) + enumerator = self._get_polya([self._perms]) n_pred = enumerator.count(((stop, self._nix - stop),)) # Progress bar. @@ -1653,7 +1570,7 @@ def _invar_search(self, start=0, stop=None): stop = self._nix // 2 stop = min(stop, self._nix - stop) - enumerator = self._enumerator_collection.get([self._perms]) + enumerator = self._get_polya([self._perms]) n_pred = enumerator.count(((stop, self._nix - stop),)) # Progress bar. @@ -1945,30 +1862,25 @@ def exmul(arrays): return int(sum(o_counts) / self.group_size) -class PolyaCollection: - """ - Collection of Polya objects. This is useful because identical - instances are often recreated, with heavy sympy.expand() operations. - TODO: No longer important so eventually remove. - """ +def _serialize_perm_list(permutation_list): + return tuple((arr.shape, arr.dtype.str, arr.tobytes()) for arr in permutation_list) - def __init__(self): - self._collection = dict() - def __getitem__(self, key): - return self._collection[key] +def _deserialize_perm_list(serialized): + return [np.frombuffer(data, dtype=np.dtype(dtype)).reshape(shape) for shape, dtype, data in serialized] - def __setitem__(self, key, value): - self._collection[key] = value - def get(self, permutation_list, group_size=None): - """ - Get appropriate Polya instance for the given parameters - """ - polya = Polya(permutation_list, group_size) - label = polya.label() - if label in self._collection: - return self._collection[label] - else: - self._collection[label] = polya - return polya +@functools.lru_cache(None) +def _get_polya_cached(serialized_perms, group_size=None): + perm_list = _deserialize_perm_list(serialized_perms) + return Polya(perm_list, group_size) + + +def get_polya(permutation_list, group_size=None): + """ + Return a cached Polya instance for the given permutation list. + + Uses an lru_cache by serializing permutation arrays into a hashable key. + """ + serialized = _serialize_perm_list(permutation_list) + return _get_polya_cached(serialized, group_size) diff --git a/shry/main.py b/shry/main.py index 3cb4339..b536c4a 100644 --- a/shry/main.py +++ b/shry/main.py @@ -7,8 +7,6 @@ # python modules import ast -import collections -import configparser import datetime import io import itertools @@ -16,26 +14,25 @@ import math import operator import os -import re import shutil import signal import sys from fnmatch import fnmatch +from dataclasses import dataclass # python modules import numpy as np import tqdm -from pymatgen.core import Composition, PeriodicSite, Species, Structure -from pymatgen.core.composition import CompositionError +from pymatgen.core import Composition, PeriodicSite, Structure from pymatgen.core.lattice import Lattice -from pymatgen.io.cif import CifParser, str2float -from pymatgen.io.vasp.inputs import Poscar +from pymatgen.io.cif import CifParser from pymatgen.symmetry.analyzer import SpacegroupAnalyzer, SpacegroupOperations from pymatgen.util.coord import in_coord_list_pbc, lattice_points_in_supercell # shry modules from . import const from .core import Substitutor +from .patches import apply_pymatgen_patches # shry version control try: @@ -43,124 +40,110 @@ except (ModuleNotFoundError, ImportError): shry_version = "unknown (not installed)" -# Runtime patches to pymatgen's Composition and Poscar. -# We want some extra functions like separating different oxidation states, etc. - - -def from_string(composition_string) -> Composition: - """ - Workaround when working with strings including oxication states - """ - - def composition_builder(): - components = re.findall(r"[A-z][a-z]*[0-9.\+\-]*[0-9.]*", composition_string) - amount = re.compile(r"(?<=[\+\-A-Za-z])[0-9.]+(?![\+\-])") - amount_part = [amount.findall(x) for x in components] - amount_part = [x[0] if x else "1" for x in amount_part] - species_part = [x.strip(amount) for x, amount in zip(components, amount_part)] - species_part = [x + "0" if not re.search(r"[0-9\+\-]+", x) else x for x in species_part] - amount_part = [float(x) for x in amount_part] - - return Composition({Species.from_string(species): amount for species, amount in zip(species_part, amount_part)}) - - try: - return Composition(composition_string) - except (CompositionError, ValueError, IndexError): - # - CompositionError: get_sym_dict() error - # when "+" oxidation is used. - # - ValueError: Wrong parse by get_sym_dict() - # if the oxidation negative. - # - IndexError: Sometimes appear when the string gets more complex - return composition_builder() - - -@property -def site_symbols(self): - """ - On Poscar: sometimes we would like to use a separate pseudopotential - for each oxidation states. - - Write oxidation states, if, and only if, for the given element, - there are more than 1 state. - """ - return [a[0] for a in itertools.groupby(self.syms)] - - -@property -def syms(self): - """ - Replaced self.syms in some functions of Poscar - """ - e_oxstates = collections.defaultdict(set) - s_symbols = collections.defaultdict(list) - - for site in self.structure: - e_oxstates[site.specie.symbol].add(str(site.specie)) - for e, oes in e_oxstates.items(): - if len(oes) == 1: - s_symbols[list(oes)[0]] = e - else: - for oe in oes: - s_symbols[oe] = oe - - return [s_symbols[str(site.specie)] for site in self.structure] - - -@property -def natoms(self): - """ - Sequence of number of sites of each type associated with the Poscar. - Similar to 7th line in vasp 5+ POSCAR or the 6th line in vasp 4 POSCAR. - """ - return [len(tuple(a[1])) for a in itertools.groupby(self.syms)] - - -@staticmethod -def parse_oxi_states(data): +apply_pymatgen_patches() + + +@dataclass(frozen=True) +class RunConfig: + structure_file: str + from_species: list + to_species: list + scaling_matrix: np.ndarray + symmetrize: bool + sample: int | str | None + seed: int + symprec: float + atol: float + angle_tolerance: float + dir_size: int + write_symm: bool + write_ewald: bool + max_ewald: float | None + no_write: bool + no_dmat: bool + t_kind: str + no_cache: bool = False + + +def _math_eval(expr): """ - Parse oxidation states from data dictionary + Like eval() but limit to +-/*^** expressions for security """ - ox_state_regex = re.compile(r"\d?[\+,\-]?$") - - try: - oxi_states = { - data["_atom_type_symbol"][i]: str2float(data["_atom_type_oxidation_number"][i]) - for i in range(len(data["_atom_type_symbol"])) + operators = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.Div: operator.truediv, + ast.Pow: operator.pow, + ast.BitXor: operator.xor, + ast.USub: operator.neg, + } + + def _tracer(node): + if isinstance(node, ast.Num): + return node.n + if isinstance(node, ast.BinOp): + return operators[type(node.op)](_tracer(node.left), _tracer(node.right)) + if isinstance(node, ast.UnaryOp): + return operators[type(node.op)](_tracer(node.operand)) + raise TypeError("Invalid characters on math expression") + + return _tracer(ast.parse(expr, mode="eval").body) + + +def _normalize_sample(sample): + if isinstance(sample, str): + if sample == "all": + return None + return int(_math_eval(sample)) + return sample + + +def _normalize_config(config: RunConfig) -> RunConfig: + sample = _normalize_sample(config.sample) + scaling_matrix = np.array(config.scaling_matrix) + write_ewald = config.write_ewald or config.max_ewald is not None + return RunConfig( + **{ + **{k: getattr(config, k) for k in config.__dataclass_fields__}, + "sample": sample, + "scaling_matrix": scaling_matrix, + "write_ewald": write_ewald, } - # attempt to strip oxidation state from _atom_type_symbol - # in case the label does not contain an oxidation state - for i, symbol in enumerate(data["_atom_type_symbol"]): - oxi_states[ox_state_regex.sub("", symbol)] = str2float(data["_atom_type_oxidation_number"][i]) - except (ValueError, KeyError): - # Some CIF (including pymatgen's output) are like this. - try: - oxi_states = dict() - for i, symbol in enumerate(data["_atom_site_type_symbol"]): - _symbol = ox_state_regex.sub("", symbol) - - parsed_oxi_state = ox_state_regex.search(symbol).group(0) - if not parsed_oxi_state: - oxi_states[_symbol] = None - continue - - sign = re.search(r"[-+]", parsed_oxi_state) - if sign is None: - sign = "" - else: - sign = sign.group(0) - parsed_oxi_state = parsed_oxi_state.replace("+", "").replace("-", "") - parsed_oxi_state = str2float(sign + parsed_oxi_state) - oxi_states[ox_state_regex.sub("", symbol)] = parsed_oxi_state - except (ValueError, KeyError): - oxi_states = None - return oxi_states - - -Composition.from_string = from_string -Poscar.site_symbols = site_symbols -Poscar.natoms = natoms -Poscar.syms = syms -CifParser.parse_oxi_states = parse_oxi_states + ) + + +def _validate_config(config: RunConfig) -> None: + if len(config.from_species) != len(config.to_species): + raise RuntimeError("from_species and to_species must have the same length.") + + +def prepare_structures(config: RunConfig): + structure = LabeledStructure.from_file( + config.structure_file, + symmetrize=config.symmetrize, + ) + modified_structure = structure.copy() + # Note: since we don't limit the allowable scaling_matrix, + # changing the order of enlargement vs. replace can have different meaning. + modified_structure.replace_species(dict(zip(config.from_species, config.to_species))) + modified_structure *= config.scaling_matrix + return structure, modified_structure + + +def build_substitutor(config: RunConfig, modified_structure): + cache = False if config.no_cache else None + return Substitutor( + modified_structure, + symprec=config.symprec, + atol=config.atol, + angle_tolerance=config.angle_tolerance, + shuffle=config.sample is not None, + seed=config.seed, + no_dmat=config.no_dmat, + t_kind=config.t_kind, + cache=cache, + ) class ScriptHelper: @@ -189,68 +172,42 @@ def __init__( no_cache=False, t_kind=const.DEFAULT_T_KIND, ): - # TODO: Refactor with descriptors. self._timestamp = datetime.datetime.now().timestamp() - self.no_write = no_write - self.no_dmat = no_dmat - self.t_kind = t_kind - self.structure_file = structure_file - self._seed = seed - - if len(from_species) != len(to_species): - raise RuntimeError("from_species and to_species must have the same length.") - self.from_species = from_species - self.to_species = to_species - self.scaling_matrix = np.array(scaling_matrix) - self.symmetrize = symmetrize - - if isinstance(sample, str): - if sample == "all": - sample = None - else: - sample = int(self._math_eval(sample)) - self.sample = sample - self.symprec = symprec - self.atol = atol - self.angle_tolerance = angle_tolerance - self.dir_size = dir_size - self.write_symm = write_symm - self.write_ewald = write_ewald - self.max_ewald = max_ewald - if self.max_ewald is not None: - self.write_ewald = True + config = RunConfig( + structure_file=structure_file, + from_species=from_species, + to_species=to_species, + scaling_matrix=np.array(scaling_matrix), + symmetrize=symmetrize, + sample=sample, + seed=seed, + symprec=symprec, + atol=atol, + angle_tolerance=angle_tolerance, + dir_size=dir_size, + write_symm=write_symm, + write_ewald=write_ewald, + max_ewald=max_ewald, + no_write=no_write, + no_dmat=no_dmat, + t_kind=t_kind, + no_cache=no_cache, + ) + config = _normalize_config(config) + _validate_config(config) + self.config = config + self._seed = config.seed + + for key in config.__dataclass_fields__: + setattr(self, key, getattr(config, key)) logging.info("\nRun configurations:") logging.info(const.HLINE) logging.info(self) logging.info(const.HLINE) - # TODO: these "derivative attributes" should have a proper setter() - # In fact any that are not simple assignment should. - if no_cache: - cache = False - else: - cache = None - self.structure = LabeledStructure.from_file( - self.structure_file, - symmetrize=symmetrize, - ) - self.modified_structure = self.structure.copy() - # Note: since we don't limit the allowable scaling_matrix, - # changing the order of enlargement vs. replace can have different meaning. - self.modified_structure.replace_species(dict(zip(self.from_species, self.to_species))) - self.modified_structure *= self.scaling_matrix - self.substitutor = Substitutor( - self.modified_structure, - symprec=self.symprec, - atol=self.atol, - angle_tolerance=self.angle_tolerance, - shuffle=self.sample is not None, - seed=seed, - no_dmat=self.no_dmat, - t_kind=self.t_kind, - cache=cache, - ) + self.structure, self.modified_structure = prepare_structures(config) + self.substitutor = build_substitutor(config, self.modified_structure) def __str__(self): string = "" @@ -269,87 +226,6 @@ def __str__(self): string += print_format.format("t-kind", self.t_kind) return string - @staticmethod - def _math_eval(expr): - """ - Like eval() but limit to +-/*^** expressions for security - """ - operators = { - ast.Add: operator.add, - ast.Sub: operator.sub, - ast.Mult: operator.mul, - ast.Div: operator.truediv, - ast.Pow: operator.pow, - ast.BitXor: operator.xor, - ast.USub: operator.neg, - } - - def _tracer(node): - if isinstance(node, ast.Num): - return node.n - elif isinstance(node, ast.BinOp): - return operators[type(node.op)](_tracer(node.left), _tracer(node.right)) - elif isinstance(node, ast.UnaryOp): - return operators[type(node.op)](_tracer(node.operand)) - else: - raise TypeError("Invalid characters on math expression") - - return _tracer(ast.parse(expr, mode="eval").body) - - @classmethod - def from_file(cls, input_file): - """ - Reads `*.ini` file containing command line arguments - """ - # TODO: simplify - parser = configparser.ConfigParser(empty_lines_in_values=False, allow_no_value=False) - encoding = getattr(io, "LOCALE_ENCODING", "utf8") - with open(input_file, "r", encoding=encoding, errors="surrogateescape") as f: - parser.read_file(f) - - structure_file = parser.get("DEFAULT", "structure_file") - from_species = parser.get("DEFAULT", "from_species", fallback=const.DEFAULT_FROM_SPECIES) - to_species = parser.get("DEFAULT", "to_species", fallback=const.DEFAULT_TO_SPECIES) - scaling_matrix = parser.get( - "DEFAULT", - "scaling_matrix", - fallback=const.DEFAULT_SCALING_MATRIX_STR, - ) - # Allow ",", ";", and whiteline as separator (';' not allowed by *.ini) - from_species = const.FLEXIBLE_SEPARATOR.split(from_species) - to_species = const.FLEXIBLE_SEPARATOR.split(to_species) - scaling_matrix = [int(x) for x in const.FLEXIBLE_SEPARATOR.split(scaling_matrix)] - symmetrize = parser.getboolean("DEFAULT", "symmetrize", fallback=const.DEFAULT_SYMMETRIZE) - sample = parser.getint("DEFAULT", "sample", fallback=const.DEFAULT_SAMPLE) - symprec = parser.getfloat("DEFAULT", "symprec", fallback=const.DEFAULT_SYMPREC) - angle_tolerance = parser.getfloat( - "DEFAULT", - "angle_tolerance", - fallback=const.DEFAULT_ANGLE_TOLERANCE, - ) - dir_size = parser.getint("DEFAULT", "dir_size", fallback=const.DEFAULT_DIR_SIZE) - write_symm = parser.getboolean("DEFAULT", "write_symm", fallback=const.DEFAULT_WRITE_SYMM) - - no_write = parser.getboolean("DEFAULT", "no_write", fallback=const.DEFAULT_NO_WRITE) - no_dmat = parser.getboolean("DEFAULT", "no_dmat", fallback=const.DEFAULT_NO_DMAT) - t_kind = parser.getboolean("DEFAULT", "t_kind", fallback=const.DEFAULT_T_KIND) - - return cls( - structure_file=structure_file, - from_species=from_species, - to_species=to_species, - scaling_matrix=scaling_matrix, - symmetrize=symmetrize, - sample=sample, - symprec=symprec, - angle_tolerance=angle_tolerance, - dir_size=dir_size, - write_symm=write_symm, - no_write=no_write, - no_dmat=no_dmat, - t_kind=t_kind, - ) - @property def _outdir(self): """ diff --git a/shry/patches.py b/shry/patches.py new file mode 100644 index 0000000..c705cd9 --- /dev/null +++ b/shry/patches.py @@ -0,0 +1,245 @@ +# Copyright (c) SHRY Development Team. +# Distributed under the terms of the MIT License. + +""" +Runtime patches for pymatgen classes. +""" + +from __future__ import annotations + +import collections +import itertools +import re +from typing import Tuple + +import numpy as np +from monty.fractions import gcd_float +from pymatgen.core import Composition, Species +from pymatgen.core.composition import CompositionError, reduce_formula +from pymatgen.core.periodic_table import get_el_sp +from pymatgen.io.cif import CifParser, str2float +from pymatgen.io.vasp.inputs import Poscar + +from . import const + +_APPLIED = False + +# Patched extra functionalities and bug fixes on top of Pymatgen's classes. + +def get_integer_formula_and_factor( + self, + max_denominator: int = int(1 / const.DEFAULT_ATOL), + iupac_ordering: bool = False, +) -> Tuple[str, float]: + """The default Composition groups together different ox states which is not ideal...""" + el_amt = self.as_dict() + g = gcd_float(list(el_amt.values()), 1 / max_denominator) + + d = {k: round(v / g) for k, v in el_amt.items()} + (formula, factor) = reduce_formula(d, iupac_ordering=iupac_ordering) + if formula in Composition.special_formulas: + formula = Composition.special_formulas[formula] + factor /= 2 + return formula, factor * g + + +def to_int_dict(self): + """ + Returns: + Dict with element symbol and integer amount + """ + _, factor = self.get_integer_formula_and_factor(max_denominator=int(1 / const.DEFAULT_ATOL)) + int_dict = {e: int(a) for e, a in (self / factor).as_dict().items()} + + # be safe: Composition groups together different ox states which is not ideal... + for x, y in zip(int_dict.values(), self.as_dict().values()): + if not np.isclose(x * factor, y, atol=const.DEFAULT_ATOL): + raise ValueError( + "Composition (Occupancy) is not rational! Please try to increase significant digits " + "e.g., 1/3 = 0.3333 -> 1/3 = 0.3333333333333." + ) + + return int_dict + + +@property +def inted_composition(self): + """ + Return Composition instance with integer formula + """ + _, factor = self.get_integer_formula_and_factor(max_denominator=int(1 / const.DEFAULT_ATOL)) + int_comp = self / factor + + # be safe + int_dict = {e: int(a) for e, a in int_comp.as_dict().items()} + if not all(np.isclose(x * factor, y, atol=const.DEFAULT_ATOL) for x, y in zip(int_dict.values(), self.as_dict().values())): + raise ValueError( + "Composition (Occupancy) is not rational! Please try to increase significant digits " + "e.g., 1/3 = 0.3333 -> 1/3 = 0.3333333333333." + ) + + return int_comp + + +def formula_double_format_tol(afloat, ignore_ones=True, tol: float = const.DEFAULT_ATOL * 10): + """ + This function is used to make pretty formulas by formatting the amounts. + Instead of Li1.0 Fe1.0 P1.0 O4.0, you get LiFePO4. + + Args: + afloat (float): a float + ignore_ones (bool): if true, floats of 1 are ignored. + tol (float): Tolerance to round to nearest int. i.e. 2.0000000001 -> 2 + + Returns: + A string representation of the float for formulas. + """ + if ignore_ones and afloat == 1: + return "" + if abs(afloat - round(afloat)) < tol: + return round(afloat) + return round(afloat, 8) + + +@property +def formula(self) -> str: + """ + Returns a formula string, with elements sorted by electronegativity, + e.g., Li4 Fe4 P4 O16. + """ + sym_amt = self.get_el_amt_dict() + syms = sorted(sym_amt, key=lambda sym: get_el_sp(sym).X) + formula = [f"{s}{formula_double_format_tol(sym_amt[s], False)}" for s in syms] + return " ".join(formula) + + +def from_string(composition_string) -> Composition: + """ + Workaround when working with strings including oxication states + """ + + def composition_builder(): + components = re.findall(r"[A-z][a-z]*[0-9.\+\-]*[0-9.]*", composition_string) + amount = re.compile(r"(?<=[\+\-A-Za-z])[0-9.]+(?![\+\-])") + amount_part = [amount.findall(x) for x in components] + amount_part = [x[0] if x else "1" for x in amount_part] + species_part = [x.strip(amount) for x, amount in zip(components, amount_part)] + species_part = [x + "0" if not re.search(r"[0-9\+\-]+", x) else x for x in species_part] + amount_part = [float(x) for x in amount_part] + + return Composition({Species.from_string(species): amount for species, amount in zip(species_part, amount_part)}) + + try: + return Composition(composition_string) + except (CompositionError, ValueError, IndexError): + # - CompositionError: get_sym_dict() error + # when "+" oxidation is used. + # - ValueError: Wrong parse by get_sym_dict() + # if the oxidation negative. + # - IndexError: Sometimes appear when the string gets more complex + return composition_builder() + + +@property +def site_symbols(self): + """ + On Poscar: sometimes we would like to use a separate pseudopotential + for each oxidation states. + + Write oxidation states, if, and only if, for the given element, + there are more than 1 state. + """ + return [a[0] for a in itertools.groupby(self.syms)] + + +@property +def syms(self): + """ + Replaced self.syms in some functions of Poscar + """ + e_oxstates = collections.defaultdict(set) + s_symbols = collections.defaultdict(list) + + for site in self.structure: + e_oxstates[site.specie.symbol].add(str(site.specie)) + for e, oes in e_oxstates.items(): + if len(oes) == 1: + s_symbols[list(oes)[0]] = e + else: + for oe in oes: + s_symbols[oe] = oe + + return [s_symbols[str(site.specie)] for site in self.structure] + + +@property +def natoms(self): + """ + Sequence of number of sites of each type associated with the Poscar. + Similar to 7th line in vasp 5+ POSCAR or the 6th line in vasp 4 POSCAR. + """ + return [len(tuple(a[1])) for a in itertools.groupby(self.syms)] + + +@staticmethod +def parse_oxi_states(data): + """ + Parse oxidation states from data dictionary + """ + ox_state_regex = re.compile(r"\d?[\+,\-]?$") + + try: + oxi_states = { + data["_atom_type_symbol"][i]: str2float(data["_atom_type_oxidation_number"][i]) + for i in range(len(data["_atom_type_symbol"])) + } + # attempt to strip oxidation state from _atom_type_symbol + # in case the label does not contain an oxidation state + for i, symbol in enumerate(data["_atom_type_symbol"]): + oxi_states[ox_state_regex.sub("", symbol)] = str2float(data["_atom_type_oxidation_number"][i]) + except (ValueError, KeyError): + # Some CIF (including pymatgen's output) are like this. + try: + oxi_states = dict() + for i, symbol in enumerate(data["_atom_site_type_symbol"]): + _symbol = ox_state_regex.sub("", symbol) + + parsed_oxi_state = ox_state_regex.search(symbol).group(0) + if not parsed_oxi_state: + oxi_states[_symbol] = None + continue + + sign = re.search(r"[-+]", parsed_oxi_state) + if sign is None: + sign = "" + else: + sign = sign.group(0) + parsed_oxi_state = parsed_oxi_state.replace("+", "").replace("-", "") + parsed_oxi_state = str2float(sign + parsed_oxi_state) + oxi_states[ox_state_regex.sub("", symbol)] = parsed_oxi_state + except (ValueError, KeyError): + oxi_states = None + return oxi_states + + +def apply_pymatgen_patches() -> None: + """ + Apply runtime patches to pymatgen classes. + """ + global _APPLIED + if _APPLIED: + return + + Composition.to_int_dict = to_int_dict + Composition.get_integer_formula_and_factor = get_integer_formula_and_factor + Composition.inted_composition = inted_composition + Composition.formula = formula + Composition.from_string = from_string + + Poscar.site_symbols = site_symbols + Poscar.natoms = natoms + Poscar.syms = syms + + CifParser.parse_oxi_states = parse_oxi_states + + _APPLIED = True diff --git a/shry/script.py b/shry/script.py index dbf30c9..3c58201 100755 --- a/shry/script.py +++ b/shry/script.py @@ -25,20 +25,43 @@ shry_version = "unknown (not installed)" -class TqdmLoggingHandler(logging.Handler): +def setup_tqdm_logging(level=logging.INFO, fmt="%(message)s"): """ - Prevent logging from overwriting tqdm + Configure logging to work with tqdm without overwriting the progress bar. + Returns a tuple needed to restore previous logging handlers. """ - def emit(self, record): - try: - msg = self.format(record) - tqdm.tqdm.write(msg) - self.flush() - except (KeyboardInterrupt, SystemExit): - raise - except Exception: # pylint: disable=bare-except - self.handleError(record) + class _TqdmLoggingHandler(logging.Handler): + def emit(self, record): + try: + msg = self.format(record) + tqdm.tqdm.write(msg) + self.flush() + except (KeyboardInterrupt, SystemExit): + raise + except Exception: # pylint: disable=bare-except + self.handleError(record) + + handler = _TqdmLoggingHandler() + handler.setLevel(level) + handler.setFormatter(logging.Formatter(fmt)) + + root = logging.getLogger() + prev_handlers = list(root.handlers) + prev_level = root.level + root.handlers = [] + root.setLevel(level) + root.addHandler(handler) + return root, prev_handlers, prev_level + + +def restore_logging(state): + """ + Restore logging handlers and level from setup_tqdm_logging. + """ + root, prev_handlers, prev_level = state + root.handlers = prev_handlers + root.setLevel(prev_level) def print_header(): @@ -48,10 +71,6 @@ def print_header(): tz = const.NOW.astimezone().tzname() time_string = const.NOW.strftime("%c ") + tz - # logger.setLevel(logging.INFO) - handler = TqdmLoggingHandler() - handler.setLevel(logging.INFO) - logging.basicConfig(format="%(message)s", level=logging.INFO, handlers=[handler]) logging.info("********************************\n") logging.info("SHRY: Suite for High-throughput generation of modelswith atomic substitutions implemented by Python") logging.info("") @@ -70,10 +89,29 @@ def print_footer(): tz = now.astimezone().tzname() time_string = now.strftime("%c ") + tz logging.info(const.HLINE) + logging.info("When you publish a paper, please cite the following paper!!") + logging.info(const.HLINE) + logging.info("SHRY: Application of Canonical Augmentation to the Atomic Substitution Problem") + logging.info("G.I. Prayogo*, A. Tirelli, K. Utimula, K. Hongo, R. Maezono, and K. Nakano*,") + logging.info("J. Chem. Inf. Model. 62, 2909-2915 (2022)") + logging.info("") + logging.info("@article{doi:10.1021/acs.jcim.2c00389,") + logging.info(" author = {Prayogo, Genki and Tirelli, Andrea and Utimula, Keishu and Hongo, Kenta and Maezono, Ryo and Nakano, Kousuke},") + logging.info(" title = {SHRY: Application of Canonical Augmentation to the Atomic Substitution Problem},") + logging.info(" journal = {J. Chem. Inf. Model.},") + logging.info(" volume = {62},") + logging.info(" number = {12},") + logging.info(" pages = {2909-2915},") + logging.info(" year = {2022},") + logging.info(" doi = {10.1021/acs.jcim.2c00389},") + logging.info("}") + logging.info("") + logging.info("If SHRY helps your work, please consider giving us a star on GitHub:") + logging.info("https://github.com/shry-project/SHRY") + logging.info(const.HLINE) logging.info("Ends " + time_string) logging.info(const.HLINE) - def get_parser(): parser = argparse.ArgumentParser( epilog=f"SHRY version {shry_version}", @@ -85,8 +123,7 @@ def get_parser(): "input", type=str, help=( - "A CIF, or an `*.ini` containing command line options as keys.\n" - "If using `*.ini`, write all keys under `DEFAULT` section " + "A CIF containing command line options as keys.\n" "(See `$SHRY_INSTALLDIR/examples`)" ), ) @@ -237,20 +274,17 @@ def main(): # pylint: disable=missing-function-docstring parser = get_parser() args = parser.parse_args() const.DISABLE_PROGRESSBAR = args.disable_progressbar + logging_state = setup_tqdm_logging() + try: + # Print header first for faster perceived response + print_header() - # Print header first for faster perceived response - print_header() - - # Late patch for count/mod/nowrite - if args.count_only: - args.no_write = True + # Late patch for count/mod/nowrite + if args.count_only: + args.no_write = True - from .main import ScriptHelper # pylint:disable=import-outside-toplevel + from .main import ScriptHelper # pylint:disable=import-outside-toplevel - if fnmatch.fnmatch(args.input, "*.ini"): - # Note: when using *.ini, other arguments will be ignored. - helper = ScriptHelper.from_file(args.input) - else: # Trick to allow ",", ";", and whiteline as separator from_species = const.FLEXIBLE_SEPARATOR.split(",".join(args.from_species)) to_species = const.FLEXIBLE_SEPARATOR.split(",".join(args.to_species)) @@ -289,11 +323,13 @@ def main(): # pylint: disable=missing-function-docstring no_dmat=args.no_dmat, t_kind=args.t_kind, ) - helper.count() - helper.save_modified_structure() - if not args.count_only and not args.mod_only: - helper.write() - print_footer() + helper.count() + helper.save_modified_structure() + if not args.count_only and not args.mod_only: + helper.write() + print_footer() + finally: + restore_logging(logging_state) if __name__ == "__main__": diff --git a/tests/test_core.py b/tests/test_core.py index e8e9c81..664a3ed 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -31,6 +31,15 @@ SHRY_ANGLE_TOLERANCE = const.DEFAULT_ANGLE_TOLERANCE SHRY_ATOL = const.DEFAULT_ATOL +pytestmark = [ + pytest.mark.filterwarnings( + "ignore:No Pauling electronegativity for .*:UserWarning", + ), + pytest.mark.filterwarnings( + r"ignore:Set OLD_ERROR_HANDLING to false and catch the errors directly\.:DeprecationWarning", + ) +] + # PatternMaker basic functions.