Features 072826

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:
+
+[<img
+    src="https://raw.githubusercontent.com/jwjeffr/tce-modifier/refs/heads/main/examples/ws2-grid/grid.png"
+    width=100%
+    alt="WS2 feature grid"
+    title="WS2"
+/>](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."):

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"""

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

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$"""

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(