diff --git a/tce/__init__.py b/tce/__init__.py index 8855c53..29da033 100644 --- a/tce/__init__.py +++ b/tce/__init__.py @@ -342,6 +342,23 @@ def callback(step_: int, num_steps_: int): all lattice sites are equivalent. This is not a problem - we just likely need to feature reduce later using something like PCA, which is relatively easy with an sklearn.pipeline.Pipeline object. +## 🔬 OVITO Plugin + +We've implemented a plugin for OVITO Pro that allows you to visualize clusters in your system. This is a very useful tool for +debugging your cluster expansion model, and for visualizing the clusters that are being used in your model. You can find +the plugin [here](https://github.com/jwjeffr/tce-modifier). This plugin can used within the OVITO Pro software, or within +a stand-alone Python script using the OVITO Python API. See the README within that repository for more instructions! The stand-alone +Python script is useful for generating images of clusters for use in publications, such as the figure below: + +[WS2 feature grid](https://raw.githubusercontent.com/jwjeffr/tce-modifier/refs/heads/main/examples/ws2-grid/grid.png) + +See the example [here](https://github.com/jwjeffr/tce-modifier/tree/main/examples/ws2-grid) for the full script that creates this grid. + # Sharp Edges `tce-lib` has a couple of sharp edges (or gotcha's) that one needs to look out for. @@ -440,13 +457,13 @@ def callback(step_: int, num_steps_: int): which is invalid unless $\alpha = 0$. To address this, the Monte Carlo function will remove this intercept by probing basis vectors, i.e. by evaluating the intercept and subtracting out that intercept: -$$ f_{\text{new}}(\Delta\mathbf{t}) = f(\Delta\mathbf{t}) - f(\mathbf{I}) $$ +$$ f_{\text{new}}(\Delta\mathbf{t}) = f(\Delta\mathbf{t}) - f(\mathbf{0}) $$ where $\mathbf{I}$ is the identity matrix. """ -__version__ = "1.0.1" +__version__ = "1.0.2" __authors__ = ["Jacob Jeffries"] __url__ = "https://github.com/MUEXLY/tce-lib" @@ -459,6 +476,7 @@ def callback(step_: int, num_steps_: int): from . import monte_carlo as monte_carlo from . import topology as topology from . import training as training +from . import citations as citations if __version__.startswith("0."): diff --git a/tce/calculator.py b/tce/calculator.py index 004fe7e..f50558b 100644 --- a/tce/calculator.py +++ b/tce/calculator.py @@ -19,13 +19,13 @@ import numpy as np from numpy.typing import NDArray import sparse -from scipy.spatial import KDTree from opt_einsum import contract from multiset import Multiset from .training import Model, LimitingRidge from .topology import hash_topology, symmetrize from .topology import get_adjacency_tensors +from .citations import cite, ORIGINAL_PAPER, KMC_PAPER LOGGER = logging.getLogger(__name__) @@ -342,15 +342,10 @@ def get_topological_tensors(self, atoms: Atoms) -> dict[int, sparse.COO]: topological_tensors = self.topological_tensors.get(topology_key) if topological_tensors is None: - - if not np.all(atoms.cell.angles() == 90): - raise ValueError("supercells must be orthogonal (for now)") - - tree = KDTree(data=atoms.positions, boxsize=np.diag(atoms.cell)) # these are boolean, so we can sum corresponding to logical or adjacency_tensors = get_adjacency_tensors( - tree=tree, + atoms=atoms, cutoffs=self.neighbor_cutoffs, tolerance=self.neighbor_tolerance ) @@ -396,6 +391,7 @@ def get_topological_tensors(self, atoms: Atoms) -> dict[int, sparse.COO]: return topological_tensors + @cite(paper_link=ORIGINAL_PAPER) def get_feature_vector( self, atoms: Atoms @@ -417,7 +413,6 @@ def get_feature_vector( topological_tensors = self.get_topological_tensors(atoms) - #symbols = np.array(atoms.get_chemical_symbols()) indicator_tensor = atoms.numbers[:, None] == self.atomic_numbers[None, :] indicator_tensor = indicator_tensor.astype(float) @@ -439,6 +434,119 @@ def get_feature_vector( return feature_vec + def get_normalizer( + self, + atoms: Atoms + ) -> NDArray[np.floating]: + + topological_tensors = self.get_topological_tensors(atoms) + + indicator_tensor = atoms.numbers[:, None] == self.atomic_numbers[None, :] + indicator_tensor = indicator_tensor.astype(float) + + # Pre-allocate feature vector + normalizer = np.zeros(self.feature_vector_size, dtype=np.float64) + pos = 0 + + for body_order, t in topological_tensors.items(): + + num_features = len(self.species) ** body_order * t.shape[0] + normalizer[pos:pos+num_features] = t.sum() + pos += num_features + + return normalizer + + + @cite(paper_link=ORIGINAL_PAPER) + def get_batched_feature_vectors( + self, + atoms_list: list[Atoms] + ) -> NDArray[np.floating]: + + r""" + Compute batched feature vectors for many structures. + + This function is quite similar to `TCECalculator.get_feature_vector`, but replaces the contractions: + + $$ N_{\alpha_1\cdots\alpha_m}^{[\ell]} = T_{i_1\cdots i_m}^{[\ell]}\prod_{n=1}^m X_{i_n\alpha_n} $$ + + with a batched contraction instead: + + $$ N_{S\alpha_1\cdots\alpha_m}^{[\ell]} = T_{i_1\cdots i_m}^{[\ell]}\prod_{n=1}^m X_{Si_n\alpha_n} $$ + + where $S$ indexes configurations, and the new indicator tensor $\mathbf{X}$ is: + + $$ X_{Si\alpha} = [\text{site $i$ in sample $S$ is occupied by type $\alpha$}] $$ + + i.e., the function computes the cluster counts in a list of configurations, + rather than for just one. Alternatively, the two calls are equivalent: + + ```py + configurations: list[Atoms] = ... + calc: TCECalculator = ... + + feature_matrix = np.array([ + calc.get_feature_vector(atoms) for atoms in configurations + ]) + feature_matrix = calc.get_batched_feature_vectors(configurations) + ``` + + Args: + atoms_list (list[Atoms]): + The list of configurations to compute feature vectors for. Every system must have the same geometry + and topology. + """ + + topology_hashes = {hash_topology(atoms) for atoms in atoms_list} + if len(topology_hashes) != 1: + raise ValueError("For the batched calculation, every sample must have the same geometry and topology.") + + num_sites = len(atoms_list[0]) + topological_tensors = self.get_topological_tensors(atoms_list[0]) + + # first modify einsum string to have a sample index + # eg Lij,iα,jβ->Lαβ needs to become Lij,Siα,Sjβ->LSαβ, where S denotes a sample + batch_einsum_strs = {} + for body_order, einsum_str in self.einsum_strs.items(): + + input_indices, output_indices = einsum_str.split("->") + input_indices = input_indices.replace(",", ",S") + output_indices = output_indices.replace("L", "LS") + + batch_einsum_strs[body_order] = f"{input_indices}->{output_indices}" + + indicator_tensors = np.zeros( + (len(atoms_list), num_sites, len(self.species)), + dtype=float + ) + + for i, atoms in enumerate(atoms_list): + indicator_tensors[i, :, :] = ( + atoms.numbers[:, None] == self.atomic_numbers[None, :] + ).astype(float) + + feature_matrix = np.zeros((len(atoms_list), self.feature_vector_size), dtype=np.float64) + pos = 0 + + for body_order, t in topological_tensors.items(): + + einsum_str = batch_einsum_strs[body_order] + cluster_counts = contract( + einsum_str, + t, + *repeat(indicator_tensors, body_order) + ) + cluster_counts = np.moveaxis(cluster_counts, 1, 0) + + # Now flatten each sample the same way the single-structure version does + flattened = cluster_counts.reshape(len(atoms_list), -1) + + feature_matrix[:, pos:pos + flattened.shape[1]] = flattened + pos += flattened.shape[1] + + return feature_matrix + + def _get_feature_vector_difference_for_sites( self, initial: Atoms, @@ -616,6 +724,7 @@ def get_feature_vector_difference_nvt(self, initial: Atoms, final: Atoms) -> NDA return total_feature_diff + @cite(paper_link=ORIGINAL_PAPER) def get_feature_vector_difference(self, initial: Atoms, final: Atoms) -> NDArray[np.floating]: r""" @@ -643,6 +752,7 @@ def get_feature_vector_difference(self, initial: Atoms, final: Atoms) -> NDArray raise NotImplementedError + def calculate( self, atoms: Optional[Atoms] = None, @@ -703,6 +813,7 @@ def train(self, configurations: list[Atoms]): return self + @cite(paper_link=KMC_PAPER) def difference_train(self, configuration_pairs: list[tuple[Atoms, Atoms]]): r""" @@ -752,6 +863,7 @@ def difference_train(self, configuration_pairs: list[tuple[Atoms, Atoms]]): return self + def save(self, path: Union[Path, str]): r""" diff --git a/tce/citations.py b/tce/citations.py new file mode 100644 index 0000000..29ecd39 --- /dev/null +++ b/tce/citations.py @@ -0,0 +1,71 @@ +import logging +from typing import Callable, Optional +from functools import wraps +from string import Template + + +LOGGER = logging.getLogger(__name__) +ORIGINAL_PAPER: str = "https://doi.org/10.1016/j.commatsci.2025.114338" +KMC_PAPER: str = "https://arxiv.org/abs/2605.23612" + + +def cite( + paper_link: str, + msg_template: Optional[Template] = None +) -> Callable[[Callable], Callable]: + + r""" + Function decorator to log a citation message. Example usage: + + ```py + from tce.citations import cite + + @cite(paper_url="https://google.com") + def add(x, y): + return x + y + ``` + + Then, the first call to `add` will log a citation message to the user. + + Args: + paper_link (str): + The url to cite + msg_template (Optional[Template]): + The template for the message. This template must have the following keys: + + - `${fn_name}` representing the name of the function it wraps + - `${url}` representing the url within the message + + If not supplied, defaults to: + ```py + from string import Template + msg_template = Template( + "Function ${fn_name} uses the work at ${url}. Please cite it." + ) + ``` + """ + + if not msg_template: + msg_template = Template( + "Function ${fn_name} uses the work at ${url}. Please cite it." + ) + + name_and_urls: set[tuple[str, str]] = set() + + def decorator(fn: Callable) -> Callable: + + @wraps(fn) + def wrapper(*args, **kwargs): + + key = (fn.__name__, paper_link) + if key not in name_and_urls: + name, url = key + msg = msg_template.substitute({"fn_name": name, "url": url}) + LOGGER.info(msg) + name_and_urls.add(key) + + return fn(*args, **kwargs) + + return wrapper + + return decorator diff --git a/tce/constants.py b/tce/constants.py index 2c5977b..b55e58b 100644 --- a/tce/constants.py +++ b/tce/constants.py @@ -11,7 +11,7 @@ LOGGER = logging.getLogger(__name__) -CUTOFFS: dict[Literal["sc", "bcc", "fcc"], NDArray[np.floating]] = { +CUTOFFS: dict[Literal["sc", "bcc", "fcc", "graphene"], NDArray[np.floating]] = { "sc": np.array([1.0, np.sqrt(2.0), np.sqrt(3.0), 2.0, np.sqrt(5.0), np.sqrt(6.0), 2.0 * np.sqrt(2.0), 3.0, np.sqrt(10.0), np.sqrt(11.0), 2 * np.sqrt(3.0), np.sqrt(13.0), @@ -23,6 +23,9 @@ "fcc": np.array([0.5 * np.sqrt(2.0), 1.0, np.sqrt(1.5), np.sqrt(2.0), np.sqrt(2.5), np.sqrt(3.0), np.sqrt(3.5), 2.0, 1.5 * np.sqrt(2.0), np.sqrt(5.0), np.sqrt(0.5 * 11.0), np.sqrt(6.0), np.sqrt(0.5 * 13.0), np.sqrt(0.5 * 15.0), - 2.0 * np.sqrt(2.0), np.sqrt(0.5 * 17.0), 3.0, np.sqrt(0.5 * 19.0)]) + 2.0 * np.sqrt(2.0), np.sqrt(0.5 * 17.0), 3.0, np.sqrt(0.5 * 19.0)]), + "graphene": np.array([1.0 / np.sqrt(3.0), 1.0, 2.0 / np.sqrt(3.0), np.sqrt(7.0 / 3.0), + np.sqrt(3.0), 2.0, np.sqrt(13.0 / 3.0), 4.0 / np.sqrt(3.0) + ]) } r"""Mapping from lattice structure to neighbor cutoffs, in units of the lattice parameter $a$""" diff --git a/tce/topology.py b/tce/topology.py index 80c694e..05e64fc 100644 --- a/tce/topology.py +++ b/tce/topology.py @@ -7,11 +7,11 @@ import hashlib import logging -from scipy.spatial import KDTree import numpy as np from numpy.typing import NDArray import sparse from ase import Atoms +from ase.neighborlist import neighbor_list LOGGER = logging.getLogger(__name__) @@ -45,19 +45,19 @@ def symmetrize(tensor: sparse.COO, axes: Optional[tuple[int, ...]] =None) -> spa def get_adjacency_tensors( - tree: KDTree, + atoms: Atoms, cutoffs: Union[list[float], NDArray[np.floating]], tolerance: float = 0.01 ) -> sparse.COO: r""" - compute adjacency tensors $A_{ij}^{(n)}$. we first compute the sparse distance matrix using the - `scipy.spatial.KDTree` data structure, and then convert to a `sparse.COO` tensor. then we stack the tensors - according to neighbor order, i.e., $A_{ij}^{(n)} = 1$ if sites $i$ and $j$ are $n$th order neighbors, and $0$ else. + compute adjacency tensors $A_{ij}^{(n)}$. we first compute the sparse distance matrix as a `sparse.COO` tensor. + then we stack the tensors according to neighbor order, i.e., $A_{ij}^{(n)} = 1$ if sites $i$ and $j$ are $n$th order + neighbors, and $0$ else. Args: - tree (scipy.spatial.KDTree): - The KDTree to compute adjacency tensors from. this structure stores lattice positions as well as lattice + atoms (ase.Atoms): + The atoms to compute adjacency tensors from. this structure stores lattice positions as well as lattice vectors to encode periodic boundary conditions. cutoffs (Union[list[float], NDArray[np.floating]]): Distance cutoffs for interatomic distances. @@ -67,9 +67,15 @@ def get_adjacency_tensors( should be a small number. defaults to $0.01$. """ - distances = tree.sparse_distance_matrix(tree, max_distance=(1.0 + tolerance) * cutoffs[-1]).tocsr() - distances.eliminate_zeros() - distances_sp = sparse.COO.from_scipy_sparse(distances) + i, j, d = neighbor_list("ijd", atoms, cutoff=(1.0 + tolerance) * cutoffs[-1]) + + distances_sp = sparse.COO( + np.vstack((i, j)), + d, + shape=(len(atoms), len(atoms)), + fill_value=0.0, + has_duplicates=False + ) return sparse.stack([ sparse.where( diff --git a/test_lib.py b/test_lib.py index 46b3190..4b71d98 100644 --- a/test_lib.py +++ b/test_lib.py @@ -31,6 +31,14 @@ def get_supercell() -> Callable[[], Atoms]: def supercell(lattice_structure: str) -> Atoms: + if lattice_structure == "graphene": + return build.graphene( + formula="X2", + a=1.0, + size=(4, 4, 4), + vacuum=20.0 + ) + size = None if lattice_structure == "sc": size = (5, 5, 5) @@ -56,7 +64,8 @@ def supercell(lattice_structure: str) -> Atoms: [ ("sc", 6), ("bcc", 8), - ("fcc", 12) + ("fcc", 12), + ("graphene", 3) ] ) def test_num_neighbors(lattice_structure: str, num_expected_neighbors: int, get_supercell): @@ -125,7 +134,7 @@ def test_noncubic_cell_raises_value_error(): build.bulk("Cr", crystalstructure="bcc", a=2.7, cubic=False).repeat((3, 3, 3)) ] for configuration in configurations: - configuration.info["energy"] = -1.0 + configuration.calc = SinglePointCalculator(atoms=configuration, energy=-1.0) calc = TCECalculator( neighbor_cutoffs=[0.5 * np.sqrt(3.0) * 2.7], @@ -133,8 +142,7 @@ def test_noncubic_cell_raises_value_error(): species=["Fe", "Cr"] ) - with pytest.raises(ValueError): - calc.train(configurations) + calc.train(configurations) def test_no_energy_computation_raises_attribute_error(): @@ -643,4 +651,55 @@ def test_monte_carlo_generator(): assert isinstance(list_results, list) assert isinstance(generator_results, GeneratorType) # check if generator - assert list_results == list(generator_results) \ No newline at end of file + assert list_results == list(generator_results) + + +def test_batched_feature_calculation(): + + pure_w = build.bulk("W", cubic=True).repeat((3, 3, 3)) + a = np.cbrt(build.bulk("W", cubic=True).get_volume()) + + calc = TCECalculator( + neighbor_cutoffs=[ + 0.5 * np.sqrt(3.0) * a, a + ], + many_body_features=[ + (0, 0, 1) + ], + species=["W", "Ta"] + ) + + rng = np.random.default_rng(seed=0) + samples = [] + for _ in range(30): + alloy = pure_w.copy() + alloy.symbols = rng.choice(["Ta", "W"], size=len(alloy)) + samples.append(alloy) + + X_batched = calc.get_batched_feature_vectors(samples) + + X = np.zeros((len(samples), calc.feature_vector_size)) + for i, sample in enumerate(samples): + X[i, :] = calc.get_feature_vector(sample) + + assert np.all(np.isclose(X, X_batched)) + + +def test_batched_calculation_throws_error_on_size_mismatch(): + + pure_w = build.bulk("W", cubic=True).repeat((3, 3, 3)) + pure_w_larger = pure_w.repeat((2, 1, 1)) + a = np.cbrt(build.bulk("W", cubic=True).get_volume()) + + calc = TCECalculator( + neighbor_cutoffs=[ + 0.5 * np.sqrt(3.0) * a, a + ], + many_body_features=[ + (0, 0, 1) + ], + species=["W"] + ) + + with pytest.raises(ValueError): + calc.get_batched_feature_vectors(atoms_list=[pure_w, pure_w_larger])