tce-lib 0.0.1__py3-none-any.whl

tce/__init__.py

@@ -0,0 +1,24 @@
+r"""
+.. include:: ../README.md
+
+# Examples
+
+## ⚛️ Using Atomic Simulation Environment (ASE)
+
+Below is an example of converting an `ase.Atoms` object into a feature vector $\mathbf{t}$. The mapping is not exactly
+one-to-one, since an `ase.Atoms` object sits on a dynamic lattice rather than a static one, but we can regardless
+provide `tce-lib` sufficient information to compute $\mathbf{t}$. The code snippet below uses the version `ase==3.26.0`.
+
+```py
+.. include:: ../examples/using-ase.py
+```
+"""
+
+__version__ = "0.0.1"
+__authors__ = ["Jacob Jeffries"]
+
+__url__ = "https://github.com/MUEXLY/tce-lib"
+
+from . import constants as constants
+from . import structures as structures
+from . import topology as topology

tce/constants.py

@@ -0,0 +1,135 @@
+from enum import Enum, auto
+from typing import Dict
+from itertools import product, permutations
+
+import numpy as np
+from scipy.spatial import KDTree
+import sparse
+
+
+class LatticeStructure(Enum):
+
+    r"""
+    Enum type defining the most typical lattice structures: simple cubic, body-centered cubic, and face-centered cubic.
+
+    [<img
+        src="https://wisc.pb.unizin.org/app/uploads/sites/293/2019/07/CNX_Chem_10_06_CubUntCll.png"
+        width=100%
+        alt="SC, BCC, and FCC lattice structures"
+        title="Lattice structures. Source: UW-Madison Chemistry 103/104 Resource Book"
+    />](https://wisc.pb.unizin.org/app/uploads/sites/293/2019/07/CNX_Chem_10_06_CubUntCll.png)
+
+    Chiefly, this data type defines mappings between lattice structure and three body labels. This is additionally
+    useful for creating a `Supercell` instance, e.g.:
+
+    ```py
+    from tce.structures.supercell import Supercell
+
+    supercell = Supercell(
+        lattice_structure=LatticeStructure.BCC,
+        lattice_parameter=2.5,
+        size=(4, 4, 4)
+    )
+    ```
+
+    which generates a $4\times 4\times 4$ bcc supercell with lattice parameter $a = 2.5$, typically in units of
+    $\mathring{\mathrm{A}}$.
+    """
+
+    SC = auto()
+    r"""simple cubic lattice structure"""
+    BCC = auto()
+    r"""body-centered cubic lattice structure"""
+    FCC = auto()
+    r"""face-centered cubic lattice structure"""
+
+
+STRUCTURE_TO_ATOMIC_BASIS: Dict[LatticeStructure, np.typing.NDArray[np.floating]] = {
+    LatticeStructure.SC: np.array([
+        [0.0, 0.0, 0.0]
+    ]),
+    LatticeStructure.BCC: np.array([
+        [0.0, 0.0, 0.0],
+        [0.5, 0.5, 0.5]
+    ]),
+    LatticeStructure.FCC: np.array([
+        [0.0, 0.0, 0.0],
+        [0.0, 0.5, 0.5],
+        [0.5, 0.0, 0.5],
+        [0.5, 0.5, 0.0]
+    ])
+}
+r"""Mapping from lattice structure to atomic basis, i.e. positions of atoms within a unit cell. Here, we use the
+conventional unit cell"""
+
+STRUCTURE_TO_CUTOFF_LISTS: Dict[LatticeStructure, np.typing.NDArray[np.floating]] = {
+    LatticeStructure.SC: np.array([1.0, np.sqrt(2.0), np.sqrt(3.0), 2.0]),
+    LatticeStructure.BCC: np.array([0.5 * np.sqrt(3.0), 1.0, np.sqrt(2.0), 0.5 * np.sqrt(11.0)]),
+    LatticeStructure.FCC: np.array([0.5 * np.sqrt(2.0), 1.0, np.sqrt(1.5), np.sqrt(2.0)])
+}
+r"""Mapping from lattice structure to neighbor cutoffs, in units of the lattice parameter $a$"""
+
+
+def load_three_body_labels(
+        tolerance: float = 0.01,
+        min_num_sites: int = 125,
+) -> Dict[LatticeStructure, np.typing.NDArray[np.integer]]:
+
+    r"""
+    Function to generate three body labels for a given lattice structure. Here, we compute the set of three body labels
+    for all lattice structures up to fourth-nearest neighbors, and store them in a mapping allowing for
+    $\mathcal{O}(1)$ access. This function is called at import, with the result stored in the module-level constant
+    `STRUCTURE_TO_THREE_BODY_LABELS`.
+    """
+
+    label_dict = {}
+    for lattice_structure in LatticeStructure:
+
+        min_num_unit_cells = min_num_sites // len(STRUCTURE_TO_ATOMIC_BASIS[lattice_structure])
+        s = np.ceil(np.cbrt(min_num_unit_cells))
+        size = (s, s, s)
+        i, j, k = (np.arange(s) for s in size)
+        unit_cell_positions = np.array(np.meshgrid(i, j, k, indexing='ij')).reshape(3, -1).T
+
+        cutoffs = STRUCTURE_TO_CUTOFF_LISTS[lattice_structure]
+        positions = unit_cell_positions[:, np.newaxis, :] + \
+            STRUCTURE_TO_ATOMIC_BASIS[lattice_structure][np.newaxis, :, :]
+        positions = positions.reshape(-1, 3)
+
+        tree = KDTree(positions, boxsize=np.array(size))
+        distances = tree.sparse_distance_matrix(tree, max_distance=(1.0 + tolerance) * cutoffs[-1]).tocsr()
+        distances.eliminate_zeros()
+        distances = sparse.COO.from_scipy_sparse(distances)
+
+        adjacency_tensors = sparse.stack([
+            sparse.where(
+                sparse.logical_and(distances > (1.0 - tolerance) * c, distances < (1.0 + tolerance) * c),
+                x=True, y=False
+            ) for c in cutoffs
+        ])
+
+        max_adj_order = adjacency_tensors.shape[0]
+        non_zero_labels = []
+        for labels in product(*[range(max_adj_order) for _ in range(3)]):
+            if not labels[0] <= labels[1] <= labels[2]:
+                continue
+            three_body_tensor = sum(
+                sparse.einsum(
+                    "ij,jk,ki->ijk",
+                    adjacency_tensors[i],
+                    adjacency_tensors[j],
+                    adjacency_tensors[k]
+                ) for i, j, k in set(permutations(labels))
+            )
+            if not three_body_tensor.nnz:
+                continue
+            non_zero_labels.append(list(labels))
+
+        non_zero_labels.sort(key=lambda x: (max(x), x))
+        label_dict[lattice_structure] = np.array(non_zero_labels)
+
+    return label_dict
+
+
+STRUCTURE_TO_THREE_BODY_LABELS = load_three_body_labels()
+r"""Mapping from lattice structure to set of three body labels"""

tce/structures.py

@@ -0,0 +1,123 @@
+from dataclasses import dataclass
+from functools import cached_property, lru_cache
+from typing import Union
+
+import numpy as np
+from scipy.spatial import KDTree
+import sparse
+
+from .constants import LatticeStructure, STRUCTURE_TO_ATOMIC_BASIS, STRUCTURE_TO_CUTOFF_LISTS, STRUCTURE_TO_THREE_BODY_LABELS
+from . import topology
+
+
+@dataclass(eq=True, frozen=True)
+class Supercell:
+
+    r"""
+    class representing a simulation supercell. `eq=True` and `frozen=True` ensures we can hash a `Supercell` instance,
+    which we need to cache the topology tensors later
+    """
+
+    lattice_structure: LatticeStructure
+    lattice_parameter: float
+    size: tuple[int, int, int]
+
+    @cached_property
+    def num_sites(self) -> Union[int, np.integer]:
+
+        r"""number of total lattice sites (NOT number of unit cells!)"""
+
+        return np.prod(self.size) * STRUCTURE_TO_ATOMIC_BASIS[self.lattice_structure].shape[0]
+
+    @cached_property
+    def positions(self) -> np.typing.NDArray[np.floating]:
+
+        r"""
+        positions of lattice sites
+        create a meshgrid of unit cell positions, and add lattice sites at atomic basis positions in each unit cell
+        """
+
+        i, j, k = (np.arange(s) for s in self.size)
+
+        unit_cell_positions = np.array(np.meshgrid(i, j, k, indexing='ij')).reshape(3, -1).T
+        positions = unit_cell_positions[:, np.newaxis, :] + \
+            STRUCTURE_TO_ATOMIC_BASIS[self.lattice_structure][np.newaxis, :, :]
+        return self.lattice_parameter * positions.reshape(-1, 3)
+
+    @lru_cache
+    def adjacency_tensors(self, max_order: int, tolerance: float = 1.0e-6) -> sparse.COO:
+
+        r"""
+        two-body adjacency tensors $A_{ij}^{(n)}$. computed by binning interatomic distances
+        """
+
+        return topology.get_adjacency_tensors(
+            tree=KDTree(self.positions, boxsize=self.lattice_parameter * np.array(self.size)),
+            cutoffs=[self.lattice_parameter * c for c in STRUCTURE_TO_CUTOFF_LISTS[self.lattice_structure][:max_order]],
+            tolerance=tolerance
+        )
+
+    @lru_cache
+    def three_body_tensors(self, max_order: int) -> sparse.COO:
+
+        r"""
+        three-body tensors, computed by summing the two-body tensors
+
+        a set of labels defines each three-body tensor. e.g., in an fcc solid, the first-order triplet is formed
+        by three first-nearest neighbor pairs, so its label is $(0, 0, 0)$. similarly, the second-order triplet in fcc is
+        formed by two first-nearest neighbor pairs, and one second-nearest neighbor pair, so its label is $(0, 0, 1)$. we
+        sum over the different permutations (which triple-counts triplets), and then stack them over
+        the labels
+        """
+
+        three_body_labels = [
+            STRUCTURE_TO_THREE_BODY_LABELS[self.lattice_structure][order] for order in range(max_order)
+        ]
+
+        return topology.get_three_body_tensors(
+            lattice_structure=self.lattice_structure,
+            adjacency_tensors=self.adjacency_tensors(max_order=np.concatenate(three_body_labels).max() + 1),
+            max_three_body_order=max_order
+        )
+
+    def feature_vector(
+        self,
+        state_matrix: sparse.COO,
+        max_adjacency_order: int,
+        max_triplet_order: int
+    ) -> np.typing.NDArray[np.integer]:
+
+        r"""
+        feature vector $\mathbf{t}$ extracting topological features, i.e. number of bonds, and number of triplets
+        """
+
+        return topology.get_feature_vector(
+            adjacency_tensors=self.adjacency_tensors(max_order=max_adjacency_order),
+            three_body_tensors=self.three_body_tensors(max_order=max_triplet_order),
+            state_matrix=state_matrix
+        )
+
+    def clever_feature_diff(
+        self,
+        initial_state_matrix: sparse.COO,
+        final_state_matrix: sparse.COO,
+        max_adjacency_order: int,
+        max_triplet_order: int,
+    ) -> np.typing.NDArray[np.floating]:
+
+        r"""
+        clever shortcut for computing feature vector difference $\Delta\mathbf{t}$ between two nearby states. here, we
+        perform a truncated contraction, only caring about "active" sites, or lattice sites that changed
+        """
+
+        if max_triplet_order == 0:
+            three_body_tensors = None
+        else:
+            three_body_tensors = self.three_body_tensors(max_order=max_triplet_order)
+
+        return topology.get_feature_vector_difference(
+            adjacency_tensors=self.adjacency_tensors(max_order=max_adjacency_order),
+            three_body_tensors=three_body_tensors,
+            initial_state_matrix=initial_state_matrix,
+            final_state_matrix=final_state_matrix
+        )

tce/topology.py

@@ -0,0 +1,175 @@
+from itertools import permutations
+from typing import Sequence
+
+from scipy.spatial import KDTree
+import numpy as np
+import sparse
+from opt_einsum import contract
+
+from .constants import LatticeStructure, STRUCTURE_TO_THREE_BODY_LABELS
+
+
+def symmetrize(tensor: sparse.COO, axes=None) -> sparse.COO:
+    r"""
+    symmetrize a tensor $T$:
+
+    $$T_{(i_1 i_2 \cdots i_r)} = \frac{1}{r!}\sum_{\sigma\in S_n} T_{\sigma(i_1) \sigma(i_2) \cdots \sigma(i_r)}$$
+
+    where $S_n$ is the symmetric group on $n$ elements, so we are summing over the permutations of the indices.
+
+    e.g. $T_{(12)} = \frac{T_{12} + T_{21}}{2}$, or equivalently $\text{symmetrize}(T) = \frac{T + T^\intercal}{2}$
+
+    specify the `axes` argument if you only want to symmetrize over a subset of indices
+    """
+
+    if not axes:
+        axes = tuple(range(tensor.ndim))
+
+    perms = list(permutations(axes))
+
+    return sum(sparse.moveaxis(tensor, axes, perm) for perm in perms) / len(perms)
+
+
+def get_adjacency_tensors(
+        tree: KDTree,
+        cutoffs: Sequence[float],
+        tolerance: float = 0.01
+) -> sparse.COO:
+
+    r"""
+    compute adjacency tensors $\mathbf{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 according to
+    neighbor order, i.e. $A_{ij}^{(n)} = 1$ if sites $i$ and $j$ are $n$'th order neighbors, and $0$ else.
+    """
+
+    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)
+
+    return sparse.stack([
+        sparse.where(
+            sparse.logical_and(distances_sp > (1.0 - tolerance) * c, distances_sp < (1.0 + tolerance) * c),
+            x=True, y=False
+        ) for c in cutoffs
+    ])
+
+
+def get_three_body_tensors(
+        lattice_structure: LatticeStructure,
+        adjacency_tensors: sparse.COO,
+        max_three_body_order: int,
+) -> sparse.COO:
+
+    r"""
+    compute three-body tensors $B_{ijk}^{(n)}$:
+
+    $$ B_{ijk}^{(n)} = \bigvee_{\sigma\in S_3}
+        A_{ij}^{(\sigma(\mathfrak{a}))}A_{jk}^{(\sigma(\mathfrak{b}))}A_{ki}^{(\sigma(\mathfrak{c}))} $$
+
+    where the mapping between $n$ and $(\mathfrak{a}, \mathfrak{b}, \mathfrak{c})$ is defined by
+    `constants.STRUCTURE_TO_THREE_BODY_LABELS`.
+    """
+
+    three_body_labels = [
+        STRUCTURE_TO_THREE_BODY_LABELS[lattice_structure][order] for order in range(max_three_body_order)
+    ]
+
+    three_body_tensors = sparse.stack([
+        sum(
+            sparse.einsum(
+                "ij,jk,ki->ijk",
+                adjacency_tensors[i],
+                adjacency_tensors[j],
+                adjacency_tensors[k]
+            ) for i, j, k in set(permutations(labels))
+        ) for labels in three_body_labels
+    ])
+
+    return three_body_tensors
+
+
+def get_feature_vector(
+        adjacency_tensors: sparse.COO,
+        three_body_tensors: sparse.COO,
+        state_matrix: np.typing.NDArray
+) -> np.typing.NDArray:
+
+    r"""
+    topological feature vector $\mathbf{t}$ with components $N_{\alpha\beta}^{(n)} = A_{ij}^{(n)}X_{i\alpha}X_{j\beta}$
+    and $M_{\alpha\beta\gamma}^{(n)} = B_{ijk}^{(n)} X_{i\alpha}X_{j\beta}X_{k\gamma}$.
+    """
+
+    return np.concatenate([
+        contract(
+            "nij,iα,jβ->nαβ",
+            adjacency_tensors,
+            state_matrix,
+            state_matrix
+        ).flatten(),
+        contract(
+            "nijk,iα,jβ,kγ->nαβγ",
+            three_body_tensors,
+            state_matrix,
+            state_matrix,
+            state_matrix
+        ).flatten()
+    ])
+
+
+def get_feature_vector_difference(
+        adjacency_tensors: sparse.COO,
+        three_body_tensors: sparse.COO,
+        initial_state_matrix: np.typing.NDArray,
+        final_state_matrix: np.typing.NDArray
+) -> np.typing.NDArray:
+
+    r"""
+    shortcut method for computing feature vector difference $\mathbf{t}$ between two nearby states
+    """
+
+    sites, _ = np.where(initial_state_matrix != final_state_matrix)
+    sites = np.unique(sites).tolist()
+
+    truncated_adj = sparse.take(adjacency_tensors, sites, axis=1)
+    initial_feature_vec_truncated = 2 * symmetrize(contract(
+        "nij,iα,jβ->nαβ",
+        truncated_adj,
+        initial_state_matrix[sites, :],
+        initial_state_matrix
+    ), axes=(1, 2)).flatten()
+    final_feature_vec_truncated = 2 * symmetrize(contract(
+        "nij,iα,jβ->nαβ",
+        truncated_adj,
+        final_state_matrix[sites, :],
+        final_state_matrix
+    ), axes=(1, 2)).flatten()
+
+    if three_body_tensors is None:
+        return final_feature_vec_truncated - initial_feature_vec_truncated
+
+    truncated_thr = sparse.take(three_body_tensors, sites, axis=1)
+    initial_feature_vec_truncated = np.concatenate(
+        [
+            initial_feature_vec_truncated,
+            3 * symmetrize(contract(
+                "nijk,iα,jβ,kγ->nαβγ",
+                truncated_thr,
+                initial_state_matrix[sites, :],
+                initial_state_matrix,
+                initial_state_matrix
+            ), axes=(1, 2, 3)).flatten()
+        ]
+    )
+    final_feature_vec_truncated = np.concatenate(
+        [
+            final_feature_vec_truncated,
+            3 * symmetrize(contract(
+                "nijk,iα,jβ,kγ->nαβγ",
+                truncated_thr,
+                final_state_matrix[sites, :],
+                final_state_matrix,
+                final_state_matrix
+            ), axes=(1, 2, 3)).flatten()
+        ]
+    )
+    return final_feature_vec_truncated - initial_feature_vec_truncated

tce_lib-0.0.1.dist-info/METADATA

@@ -0,0 +1,67 @@
+Metadata-Version: 2.4
+Name: tce-lib
+Version: 0.0.1
+Summary: tensor cluster expansion
+Project-URL: Homepage, https://github.com/MUEXLY/tce-lib
+License-Expression: MIT
+License-File: LICENSE
+Keywords: alloys
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.10
+Requires-Dist: numpy~=2.0.2
+Requires-Dist: opt-einsum~=3.4.0
+Requires-Dist: scipy~=1.13.1
+Requires-Dist: sparse~=0.16.0
+Provides-Extra: dev
+Requires-Dist: hatchling~=1.27.0; extra == 'dev'
+Requires-Dist: mypy~=1.13.0; extra == 'dev'
+Requires-Dist: pdoc~=15.0.1; extra == 'dev'
+Requires-Dist: pytest~=8.0.2; extra == 'dev'
+Requires-Dist: ruff~=0.9.4; extra == 'dev'
+Requires-Dist: scipy-stubs~=1.15.3.0; extra == 'dev'
+Provides-Extra: examples
+Requires-Dist: ase~=3.26.0; extra == 'examples'
+Description-Content-Type: text/markdown
+
+# tce-lib
+
+Library for building a tensor cluster expansion
+
+[![Custom shields.io](https://img.shields.io/badge/docs-brightgreen?logo=github&logoColor=green&label=gh-pages)](https://muexly.github.io/tce-lib)
+
+## 🔎 What is tce?
+
+Placeholder text
+
+## 📩 Installation
+
+`tce-lib` is installable via the Python Package Index:
+
+```shell
+pip install tce-lib
+```
+
+or, from source:
+
+```shell
+git clone https://github.com/MUEXLY/tce-lib
+pip install -e tce-lib/
+```
+
+## 📌 Citation
+
+Please cite our work [here](https://google.com/) if you use `tce-lib` in your work.
+
+## 💙 Acknowledgements
+
+Authors acknowledge support from the U.S. Department of Energy, Office of Basic Energy Sciences, Materials Science and Engineering Division under Award No. DE-SC0022980.
+
+## 🐝 Found a bug?
+
+Please open an issue [here](https://github.com/MUEXLY/tce/issues), with a description of the issue and a [minimal, reproducible example](https://stackoverflow.com/help/minimal-reproducible-example) of the issue.
+
+## 📑 License
+
+`tce-lib` is released under the MIT license

tce_lib-0.0.1.dist-info/RECORD

@@ -0,0 +1,8 @@
+tce/__init__.py,sha256=SSZiFKsC3UJhjyD1AKB_GkcJ7DiLG-bcjLWxbUr_QQg,719
+tce/constants.py,sha256=zR_LROepzRSH8KA-mQuOkDPl_UdbEn8o02TKy3Y1W9E,5003
+tce/structures.py,sha256=HUfxzwgW_VjIzQPmolEpCHF-vhwqOJZTa396s8NylGE,4663
+tce/topology.py,sha256=lIiM3u0wlPIqy3cvFoZcMKY1INjIA3pvYd7a-NcpK8A,5645
+tce_lib-0.0.1.dist-info/METADATA,sha256=NFZckWmJzsL9vWSLFFmeB1BuQFZ3-VBRKBVM6hBMkmU,1970
+tce_lib-0.0.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+tce_lib-0.0.1.dist-info/licenses/LICENSE,sha256=O-iwtzySLgmu-Vcy0RUAgNKBQtj1EcUvP6dIBkes9jY,1063
+tce_lib-0.0.1.dist-info/RECORD,,

tce_lib-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

tce_lib-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 MUEXLY
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.