tce-lib 0.0.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,176 @@
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[cod]
4
+ *$py.class
5
+
6
+ # C extensions
7
+ *.so
8
+
9
+ # Distribution / packaging
10
+ .Python
11
+ build/
12
+ develop-eggs/
13
+ dist/
14
+ downloads/
15
+ eggs/
16
+ .eggs/
17
+ lib/
18
+ lib64/
19
+ parts/
20
+ sdist/
21
+ var/
22
+ wheels/
23
+ share/python-wheels/
24
+ *.egg-info/
25
+ .installed.cfg
26
+ *.egg
27
+ MANIFEST
28
+
29
+ # PyInstaller
30
+ # Usually these files are written by a python script from a template
31
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
32
+ *.manifest
33
+ *.spec
34
+
35
+ # Installer logs
36
+ pip-log.txt
37
+ pip-delete-this-directory.txt
38
+
39
+ # Unit test / coverage reports
40
+ htmlcov/
41
+ .tox/
42
+ .nox/
43
+ .coverage
44
+ .coverage.*
45
+ .cache
46
+ nosetests.xml
47
+ coverage.xml
48
+ *.cover
49
+ *.py,cover
50
+ .hypothesis/
51
+ .pytest_cache/
52
+ cover/
53
+
54
+ # Translations
55
+ *.mo
56
+ *.pot
57
+
58
+ # Django stuff:
59
+ *.log
60
+ local_settings.py
61
+ db.sqlite3
62
+ db.sqlite3-journal
63
+
64
+ # Flask stuff:
65
+ instance/
66
+ .webassets-cache
67
+
68
+ # Scrapy stuff:
69
+ .scrapy
70
+
71
+ # documentation
72
+ docs/_build/
73
+ _build/
74
+
75
+ # PyBuilder
76
+ .pybuilder/
77
+ target/
78
+
79
+ # Jupyter Notebook
80
+ .ipynb_checkpoints
81
+
82
+ # IPython
83
+ profile_default/
84
+ ipython_config.py
85
+
86
+ # pyenv
87
+ # For a library or package, you might want to ignore these files since the code is
88
+ # intended to run in multiple environments; otherwise, check them in:
89
+ # .python-version
90
+
91
+ # pipenv
92
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
93
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
94
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
95
+ # install all needed dependencies.
96
+ #Pipfile.lock
97
+
98
+ # UV
99
+ # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
100
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
101
+ # commonly ignored for libraries.
102
+ #uv.lock
103
+
104
+ # poetry
105
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
106
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
107
+ # commonly ignored for libraries.
108
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
109
+ #poetry.lock
110
+
111
+ # pdm
112
+ # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
113
+ #pdm.lock
114
+ # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
115
+ # in version control.
116
+ # https://pdm.fming.dev/latest/usage/project/#working-with-version-control
117
+ .pdm.toml
118
+ .pdm-python
119
+ .pdm-build/
120
+
121
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
122
+ __pypackages__/
123
+
124
+ # Celery stuff
125
+ celerybeat-schedule
126
+ celerybeat.pid
127
+
128
+ # SageMath parsed files
129
+ *.sage.py
130
+
131
+ # Environments
132
+ .env
133
+ .venv
134
+ env/
135
+ venv/
136
+ ENV/
137
+ env.bak/
138
+ venv.bak/
139
+
140
+ # Spyder project settings
141
+ .spyderproject
142
+ .spyproject
143
+
144
+ # Rope project settings
145
+ .ropeproject
146
+
147
+ # mkdocs documentation
148
+ /site
149
+
150
+ # mypy
151
+ .mypy_cache/
152
+ .dmypy.json
153
+ dmypy.json
154
+
155
+ # Pyre type checker
156
+ .pyre/
157
+
158
+ # pytype static type analyzer
159
+ .pytype/
160
+
161
+ # Cython debug symbols
162
+ cython_debug/
163
+
164
+ # PyCharm
165
+ # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
166
+ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
167
+ # and can be added to the global gitignore or merged into this file. For a more nuclear
168
+ # option (not recommended) you can uncomment the following to ignore the entire idea folder.
169
+ .idea/
170
+
171
+ # Ruff stuff:
172
+ .ruff_cache/
173
+
174
+ # PyPI configuration file
175
+ .pypirc
176
+ trash
tce_lib-0.0.1/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2025 MUEXLY
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
tce_lib-0.0.1/PKG-INFO ADDED
@@ -0,0 +1,67 @@
1
+ Metadata-Version: 2.4
2
+ Name: tce-lib
3
+ Version: 0.0.1
4
+ Summary: tensor cluster expansion
5
+ Project-URL: Homepage, https://github.com/MUEXLY/tce-lib
6
+ License-Expression: MIT
7
+ License-File: LICENSE
8
+ Keywords: alloys
9
+ Classifier: Programming Language :: Python :: 3.10
10
+ Classifier: Programming Language :: Python :: 3.11
11
+ Classifier: Programming Language :: Python :: 3.12
12
+ Requires-Python: >=3.10
13
+ Requires-Dist: numpy~=2.0.2
14
+ Requires-Dist: opt-einsum~=3.4.0
15
+ Requires-Dist: scipy~=1.13.1
16
+ Requires-Dist: sparse~=0.16.0
17
+ Provides-Extra: dev
18
+ Requires-Dist: hatchling~=1.27.0; extra == 'dev'
19
+ Requires-Dist: mypy~=1.13.0; extra == 'dev'
20
+ Requires-Dist: pdoc~=15.0.1; extra == 'dev'
21
+ Requires-Dist: pytest~=8.0.2; extra == 'dev'
22
+ Requires-Dist: ruff~=0.9.4; extra == 'dev'
23
+ Requires-Dist: scipy-stubs~=1.15.3.0; extra == 'dev'
24
+ Provides-Extra: examples
25
+ Requires-Dist: ase~=3.26.0; extra == 'examples'
26
+ Description-Content-Type: text/markdown
27
+
28
+ # tce-lib
29
+
30
+ Library for building a tensor cluster expansion
31
+
32
+ [![Custom shields.io](https://img.shields.io/badge/docs-brightgreen?logo=github&logoColor=green&label=gh-pages)](https://muexly.github.io/tce-lib)
33
+
34
+ ## 🔎 What is tce?
35
+
36
+ Placeholder text
37
+
38
+ ## 📩 Installation
39
+
40
+ `tce-lib` is installable via the Python Package Index:
41
+
42
+ ```shell
43
+ pip install tce-lib
44
+ ```
45
+
46
+ or, from source:
47
+
48
+ ```shell
49
+ git clone https://github.com/MUEXLY/tce-lib
50
+ pip install -e tce-lib/
51
+ ```
52
+
53
+ ## 📌 Citation
54
+
55
+ Please cite our work [here](https://google.com/) if you use `tce-lib` in your work.
56
+
57
+ ## 💙 Acknowledgements
58
+
59
+ 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.
60
+
61
+ ## 🐝 Found a bug?
62
+
63
+ 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.
64
+
65
+ ## 📑 License
66
+
67
+ `tce-lib` is released under the MIT license
@@ -0,0 +1,40 @@
1
+ # tce-lib
2
+
3
+ Library for building a tensor cluster expansion
4
+
5
+ [![Custom shields.io](https://img.shields.io/badge/docs-brightgreen?logo=github&logoColor=green&label=gh-pages)](https://muexly.github.io/tce-lib)
6
+
7
+ ## 🔎 What is tce?
8
+
9
+ Placeholder text
10
+
11
+ ## 📩 Installation
12
+
13
+ `tce-lib` is installable via the Python Package Index:
14
+
15
+ ```shell
16
+ pip install tce-lib
17
+ ```
18
+
19
+ or, from source:
20
+
21
+ ```shell
22
+ git clone https://github.com/MUEXLY/tce-lib
23
+ pip install -e tce-lib/
24
+ ```
25
+
26
+ ## 📌 Citation
27
+
28
+ Please cite our work [here](https://google.com/) if you use `tce-lib` in your work.
29
+
30
+ ## 💙 Acknowledgements
31
+
32
+ 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.
33
+
34
+ ## 🐝 Found a bug?
35
+
36
+ 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.
37
+
38
+ ## 📑 License
39
+
40
+ `tce-lib` is released under the MIT license
@@ -0,0 +1,55 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "tce-lib"
7
+ dynamic = ["version"]
8
+ description = "tensor cluster expansion"
9
+ readme = "README.md"
10
+ license = "MIT"
11
+ requires-python = ">=3.10"
12
+ author = { name = "Jacob Jeffries", email = "jwjeffr@clemson.edu" }
13
+ keywords = [
14
+ "alloys"
15
+ ]
16
+ classifiers = [
17
+ "Programming Language :: Python :: 3.10",
18
+ "Programming Language :: Python :: 3.11",
19
+ "Programming Language :: Python :: 3.12",
20
+ ]
21
+ dependencies = [
22
+ "numpy~=2.0.2",
23
+ "sparse~=0.16.0",
24
+ "scipy~=1.13.1",
25
+ "opt-einsum~=3.4.0"
26
+ ]
27
+
28
+ [project.urls]
29
+ Homepage = "https://github.com/MUEXLY/tce-lib"
30
+
31
+ [project.optional-dependencies]
32
+ dev = [
33
+ "hatchling~=1.27.0",
34
+ "pytest~=8.0.2",
35
+ "ruff~=0.9.4",
36
+ "mypy~=1.13.0",
37
+ "pdoc~=15.0.1",
38
+ "scipy-stubs~=1.15.3.0"
39
+ ]
40
+ examples = [
41
+ "ase~=3.26.0"
42
+ ]
43
+
44
+ [tool.hatch.version]
45
+ path = "tce/__init__.py"
46
+
47
+ [tool.hatch.build.targets.sdist]
48
+ include = ["/tce"]
49
+
50
+ [tool.hatch.build.targets.wheel]
51
+ include = ["tce"]
52
+
53
+ [[tool.mypy.overrides]]
54
+ module = ["sparse", "opt_einsum"]
55
+ ignore_missing_imports = true
@@ -0,0 +1,24 @@
1
+ r"""
2
+ .. include:: ../README.md
3
+
4
+ # Examples
5
+
6
+ ## ⚛️ Using Atomic Simulation Environment (ASE)
7
+
8
+ Below is an example of converting an `ase.Atoms` object into a feature vector $\mathbf{t}$. The mapping is not exactly
9
+ one-to-one, since an `ase.Atoms` object sits on a dynamic lattice rather than a static one, but we can regardless
10
+ provide `tce-lib` sufficient information to compute $\mathbf{t}$. The code snippet below uses the version `ase==3.26.0`.
11
+
12
+ ```py
13
+ .. include:: ../examples/using-ase.py
14
+ ```
15
+ """
16
+
17
+ __version__ = "0.0.1"
18
+ __authors__ = ["Jacob Jeffries"]
19
+
20
+ __url__ = "https://github.com/MUEXLY/tce-lib"
21
+
22
+ from . import constants as constants
23
+ from . import structures as structures
24
+ from . import topology as topology
@@ -0,0 +1,135 @@
1
+ from enum import Enum, auto
2
+ from typing import Dict
3
+ from itertools import product, permutations
4
+
5
+ import numpy as np
6
+ from scipy.spatial import KDTree
7
+ import sparse
8
+
9
+
10
+ class LatticeStructure(Enum):
11
+
12
+ r"""
13
+ Enum type defining the most typical lattice structures: simple cubic, body-centered cubic, and face-centered cubic.
14
+
15
+ [<img
16
+ src="https://wisc.pb.unizin.org/app/uploads/sites/293/2019/07/CNX_Chem_10_06_CubUntCll.png"
17
+ width=100%
18
+ alt="SC, BCC, and FCC lattice structures"
19
+ title="Lattice structures. Source: UW-Madison Chemistry 103/104 Resource Book"
20
+ />](https://wisc.pb.unizin.org/app/uploads/sites/293/2019/07/CNX_Chem_10_06_CubUntCll.png)
21
+
22
+ Chiefly, this data type defines mappings between lattice structure and three body labels. This is additionally
23
+ useful for creating a `Supercell` instance, e.g.:
24
+
25
+ ```py
26
+ from tce.structures.supercell import Supercell
27
+
28
+ supercell = Supercell(
29
+ lattice_structure=LatticeStructure.BCC,
30
+ lattice_parameter=2.5,
31
+ size=(4, 4, 4)
32
+ )
33
+ ```
34
+
35
+ which generates a $4\times 4\times 4$ bcc supercell with lattice parameter $a = 2.5$, typically in units of
36
+ $\mathring{\mathrm{A}}$.
37
+ """
38
+
39
+ SC = auto()
40
+ r"""simple cubic lattice structure"""
41
+ BCC = auto()
42
+ r"""body-centered cubic lattice structure"""
43
+ FCC = auto()
44
+ r"""face-centered cubic lattice structure"""
45
+
46
+
47
+ STRUCTURE_TO_ATOMIC_BASIS: Dict[LatticeStructure, np.typing.NDArray[np.floating]] = {
48
+ LatticeStructure.SC: np.array([
49
+ [0.0, 0.0, 0.0]
50
+ ]),
51
+ LatticeStructure.BCC: np.array([
52
+ [0.0, 0.0, 0.0],
53
+ [0.5, 0.5, 0.5]
54
+ ]),
55
+ LatticeStructure.FCC: np.array([
56
+ [0.0, 0.0, 0.0],
57
+ [0.0, 0.5, 0.5],
58
+ [0.5, 0.0, 0.5],
59
+ [0.5, 0.5, 0.0]
60
+ ])
61
+ }
62
+ r"""Mapping from lattice structure to atomic basis, i.e. positions of atoms within a unit cell. Here, we use the
63
+ conventional unit cell"""
64
+
65
+ STRUCTURE_TO_CUTOFF_LISTS: Dict[LatticeStructure, np.typing.NDArray[np.floating]] = {
66
+ LatticeStructure.SC: np.array([1.0, np.sqrt(2.0), np.sqrt(3.0), 2.0]),
67
+ LatticeStructure.BCC: np.array([0.5 * np.sqrt(3.0), 1.0, np.sqrt(2.0), 0.5 * np.sqrt(11.0)]),
68
+ LatticeStructure.FCC: np.array([0.5 * np.sqrt(2.0), 1.0, np.sqrt(1.5), np.sqrt(2.0)])
69
+ }
70
+ r"""Mapping from lattice structure to neighbor cutoffs, in units of the lattice parameter $a$"""
71
+
72
+
73
+ def load_three_body_labels(
74
+ tolerance: float = 0.01,
75
+ min_num_sites: int = 125,
76
+ ) -> Dict[LatticeStructure, np.typing.NDArray[np.integer]]:
77
+
78
+ r"""
79
+ Function to generate three body labels for a given lattice structure. Here, we compute the set of three body labels
80
+ for all lattice structures up to fourth-nearest neighbors, and store them in a mapping allowing for
81
+ $\mathcal{O}(1)$ access. This function is called at import, with the result stored in the module-level constant
82
+ `STRUCTURE_TO_THREE_BODY_LABELS`.
83
+ """
84
+
85
+ label_dict = {}
86
+ for lattice_structure in LatticeStructure:
87
+
88
+ min_num_unit_cells = min_num_sites // len(STRUCTURE_TO_ATOMIC_BASIS[lattice_structure])
89
+ s = np.ceil(np.cbrt(min_num_unit_cells))
90
+ size = (s, s, s)
91
+ i, j, k = (np.arange(s) for s in size)
92
+ unit_cell_positions = np.array(np.meshgrid(i, j, k, indexing='ij')).reshape(3, -1).T
93
+
94
+ cutoffs = STRUCTURE_TO_CUTOFF_LISTS[lattice_structure]
95
+ positions = unit_cell_positions[:, np.newaxis, :] + \
96
+ STRUCTURE_TO_ATOMIC_BASIS[lattice_structure][np.newaxis, :, :]
97
+ positions = positions.reshape(-1, 3)
98
+
99
+ tree = KDTree(positions, boxsize=np.array(size))
100
+ distances = tree.sparse_distance_matrix(tree, max_distance=(1.0 + tolerance) * cutoffs[-1]).tocsr()
101
+ distances.eliminate_zeros()
102
+ distances = sparse.COO.from_scipy_sparse(distances)
103
+
104
+ adjacency_tensors = sparse.stack([
105
+ sparse.where(
106
+ sparse.logical_and(distances > (1.0 - tolerance) * c, distances < (1.0 + tolerance) * c),
107
+ x=True, y=False
108
+ ) for c in cutoffs
109
+ ])
110
+
111
+ max_adj_order = adjacency_tensors.shape[0]
112
+ non_zero_labels = []
113
+ for labels in product(*[range(max_adj_order) for _ in range(3)]):
114
+ if not labels[0] <= labels[1] <= labels[2]:
115
+ continue
116
+ three_body_tensor = sum(
117
+ sparse.einsum(
118
+ "ij,jk,ki->ijk",
119
+ adjacency_tensors[i],
120
+ adjacency_tensors[j],
121
+ adjacency_tensors[k]
122
+ ) for i, j, k in set(permutations(labels))
123
+ )
124
+ if not three_body_tensor.nnz:
125
+ continue
126
+ non_zero_labels.append(list(labels))
127
+
128
+ non_zero_labels.sort(key=lambda x: (max(x), x))
129
+ label_dict[lattice_structure] = np.array(non_zero_labels)
130
+
131
+ return label_dict
132
+
133
+
134
+ STRUCTURE_TO_THREE_BODY_LABELS = load_three_body_labels()
135
+ r"""Mapping from lattice structure to set of three body labels"""
@@ -0,0 +1,123 @@
1
+ from dataclasses import dataclass
2
+ from functools import cached_property, lru_cache
3
+ from typing import Union
4
+
5
+ import numpy as np
6
+ from scipy.spatial import KDTree
7
+ import sparse
8
+
9
+ from .constants import LatticeStructure, STRUCTURE_TO_ATOMIC_BASIS, STRUCTURE_TO_CUTOFF_LISTS, STRUCTURE_TO_THREE_BODY_LABELS
10
+ from . import topology
11
+
12
+
13
+ @dataclass(eq=True, frozen=True)
14
+ class Supercell:
15
+
16
+ r"""
17
+ class representing a simulation supercell. `eq=True` and `frozen=True` ensures we can hash a `Supercell` instance,
18
+ which we need to cache the topology tensors later
19
+ """
20
+
21
+ lattice_structure: LatticeStructure
22
+ lattice_parameter: float
23
+ size: tuple[int, int, int]
24
+
25
+ @cached_property
26
+ def num_sites(self) -> Union[int, np.integer]:
27
+
28
+ r"""number of total lattice sites (NOT number of unit cells!)"""
29
+
30
+ return np.prod(self.size) * STRUCTURE_TO_ATOMIC_BASIS[self.lattice_structure].shape[0]
31
+
32
+ @cached_property
33
+ def positions(self) -> np.typing.NDArray[np.floating]:
34
+
35
+ r"""
36
+ positions of lattice sites
37
+ create a meshgrid of unit cell positions, and add lattice sites at atomic basis positions in each unit cell
38
+ """
39
+
40
+ i, j, k = (np.arange(s) for s in self.size)
41
+
42
+ unit_cell_positions = np.array(np.meshgrid(i, j, k, indexing='ij')).reshape(3, -1).T
43
+ positions = unit_cell_positions[:, np.newaxis, :] + \
44
+ STRUCTURE_TO_ATOMIC_BASIS[self.lattice_structure][np.newaxis, :, :]
45
+ return self.lattice_parameter * positions.reshape(-1, 3)
46
+
47
+ @lru_cache
48
+ def adjacency_tensors(self, max_order: int, tolerance: float = 1.0e-6) -> sparse.COO:
49
+
50
+ r"""
51
+ two-body adjacency tensors $A_{ij}^{(n)}$. computed by binning interatomic distances
52
+ """
53
+
54
+ return topology.get_adjacency_tensors(
55
+ tree=KDTree(self.positions, boxsize=self.lattice_parameter * np.array(self.size)),
56
+ cutoffs=[self.lattice_parameter * c for c in STRUCTURE_TO_CUTOFF_LISTS[self.lattice_structure][:max_order]],
57
+ tolerance=tolerance
58
+ )
59
+
60
+ @lru_cache
61
+ def three_body_tensors(self, max_order: int) -> sparse.COO:
62
+
63
+ r"""
64
+ three-body tensors, computed by summing the two-body tensors
65
+
66
+ a set of labels defines each three-body tensor. e.g., in an fcc solid, the first-order triplet is formed
67
+ by three first-nearest neighbor pairs, so its label is $(0, 0, 0)$. similarly, the second-order triplet in fcc is
68
+ formed by two first-nearest neighbor pairs, and one second-nearest neighbor pair, so its label is $(0, 0, 1)$. we
69
+ sum over the different permutations (which triple-counts triplets), and then stack them over
70
+ the labels
71
+ """
72
+
73
+ three_body_labels = [
74
+ STRUCTURE_TO_THREE_BODY_LABELS[self.lattice_structure][order] for order in range(max_order)
75
+ ]
76
+
77
+ return topology.get_three_body_tensors(
78
+ lattice_structure=self.lattice_structure,
79
+ adjacency_tensors=self.adjacency_tensors(max_order=np.concatenate(three_body_labels).max() + 1),
80
+ max_three_body_order=max_order
81
+ )
82
+
83
+ def feature_vector(
84
+ self,
85
+ state_matrix: sparse.COO,
86
+ max_adjacency_order: int,
87
+ max_triplet_order: int
88
+ ) -> np.typing.NDArray[np.integer]:
89
+
90
+ r"""
91
+ feature vector $\mathbf{t}$ extracting topological features, i.e. number of bonds, and number of triplets
92
+ """
93
+
94
+ return topology.get_feature_vector(
95
+ adjacency_tensors=self.adjacency_tensors(max_order=max_adjacency_order),
96
+ three_body_tensors=self.three_body_tensors(max_order=max_triplet_order),
97
+ state_matrix=state_matrix
98
+ )
99
+
100
+ def clever_feature_diff(
101
+ self,
102
+ initial_state_matrix: sparse.COO,
103
+ final_state_matrix: sparse.COO,
104
+ max_adjacency_order: int,
105
+ max_triplet_order: int,
106
+ ) -> np.typing.NDArray[np.floating]:
107
+
108
+ r"""
109
+ clever shortcut for computing feature vector difference $\Delta\mathbf{t}$ between two nearby states. here, we
110
+ perform a truncated contraction, only caring about "active" sites, or lattice sites that changed
111
+ """
112
+
113
+ if max_triplet_order == 0:
114
+ three_body_tensors = None
115
+ else:
116
+ three_body_tensors = self.three_body_tensors(max_order=max_triplet_order)
117
+
118
+ return topology.get_feature_vector_difference(
119
+ adjacency_tensors=self.adjacency_tensors(max_order=max_adjacency_order),
120
+ three_body_tensors=three_body_tensors,
121
+ initial_state_matrix=initial_state_matrix,
122
+ final_state_matrix=final_state_matrix
123
+ )
@@ -0,0 +1,175 @@
1
+ from itertools import permutations
2
+ from typing import Sequence
3
+
4
+ from scipy.spatial import KDTree
5
+ import numpy as np
6
+ import sparse
7
+ from opt_einsum import contract
8
+
9
+ from .constants import LatticeStructure, STRUCTURE_TO_THREE_BODY_LABELS
10
+
11
+
12
+ def symmetrize(tensor: sparse.COO, axes=None) -> sparse.COO:
13
+ r"""
14
+ symmetrize a tensor $T$:
15
+
16
+ $$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)}$$
17
+
18
+ where $S_n$ is the symmetric group on $n$ elements, so we are summing over the permutations of the indices.
19
+
20
+ e.g. $T_{(12)} = \frac{T_{12} + T_{21}}{2}$, or equivalently $\text{symmetrize}(T) = \frac{T + T^\intercal}{2}$
21
+
22
+ specify the `axes` argument if you only want to symmetrize over a subset of indices
23
+ """
24
+
25
+ if not axes:
26
+ axes = tuple(range(tensor.ndim))
27
+
28
+ perms = list(permutations(axes))
29
+
30
+ return sum(sparse.moveaxis(tensor, axes, perm) for perm in perms) / len(perms)
31
+
32
+
33
+ def get_adjacency_tensors(
34
+ tree: KDTree,
35
+ cutoffs: Sequence[float],
36
+ tolerance: float = 0.01
37
+ ) -> sparse.COO:
38
+
39
+ r"""
40
+ compute adjacency tensors $\mathbf{A}_{ij}^{(n)}$. we first compute the sparse distance matrix using the
41
+ `scipy.spatial.KDTree` data structure, and then convert to a `sparse.COO` tensor. then, we stack according to
42
+ neighbor order, i.e. $A_{ij}^{(n)} = 1$ if sites $i$ and $j$ are $n$'th order neighbors, and $0$ else.
43
+ """
44
+
45
+ distances = tree.sparse_distance_matrix(tree, max_distance=(1.0 + tolerance) * cutoffs[-1]).tocsr()
46
+ distances.eliminate_zeros()
47
+ distances_sp = sparse.COO.from_scipy_sparse(distances)
48
+
49
+ return sparse.stack([
50
+ sparse.where(
51
+ sparse.logical_and(distances_sp > (1.0 - tolerance) * c, distances_sp < (1.0 + tolerance) * c),
52
+ x=True, y=False
53
+ ) for c in cutoffs
54
+ ])
55
+
56
+
57
+ def get_three_body_tensors(
58
+ lattice_structure: LatticeStructure,
59
+ adjacency_tensors: sparse.COO,
60
+ max_three_body_order: int,
61
+ ) -> sparse.COO:
62
+
63
+ r"""
64
+ compute three-body tensors $B_{ijk}^{(n)}$:
65
+
66
+ $$ B_{ijk}^{(n)} = \bigvee_{\sigma\in S_3}
67
+ A_{ij}^{(\sigma(\mathfrak{a}))}A_{jk}^{(\sigma(\mathfrak{b}))}A_{ki}^{(\sigma(\mathfrak{c}))} $$
68
+
69
+ where the mapping between $n$ and $(\mathfrak{a}, \mathfrak{b}, \mathfrak{c})$ is defined by
70
+ `constants.STRUCTURE_TO_THREE_BODY_LABELS`.
71
+ """
72
+
73
+ three_body_labels = [
74
+ STRUCTURE_TO_THREE_BODY_LABELS[lattice_structure][order] for order in range(max_three_body_order)
75
+ ]
76
+
77
+ three_body_tensors = sparse.stack([
78
+ sum(
79
+ sparse.einsum(
80
+ "ij,jk,ki->ijk",
81
+ adjacency_tensors[i],
82
+ adjacency_tensors[j],
83
+ adjacency_tensors[k]
84
+ ) for i, j, k in set(permutations(labels))
85
+ ) for labels in three_body_labels
86
+ ])
87
+
88
+ return three_body_tensors
89
+
90
+
91
+ def get_feature_vector(
92
+ adjacency_tensors: sparse.COO,
93
+ three_body_tensors: sparse.COO,
94
+ state_matrix: np.typing.NDArray
95
+ ) -> np.typing.NDArray:
96
+
97
+ r"""
98
+ topological feature vector $\mathbf{t}$ with components $N_{\alpha\beta}^{(n)} = A_{ij}^{(n)}X_{i\alpha}X_{j\beta}$
99
+ and $M_{\alpha\beta\gamma}^{(n)} = B_{ijk}^{(n)} X_{i\alpha}X_{j\beta}X_{k\gamma}$.
100
+ """
101
+
102
+ return np.concatenate([
103
+ contract(
104
+ "nij,iα,jβ->nαβ",
105
+ adjacency_tensors,
106
+ state_matrix,
107
+ state_matrix
108
+ ).flatten(),
109
+ contract(
110
+ "nijk,iα,jβ,kγ->nαβγ",
111
+ three_body_tensors,
112
+ state_matrix,
113
+ state_matrix,
114
+ state_matrix
115
+ ).flatten()
116
+ ])
117
+
118
+
119
+ def get_feature_vector_difference(
120
+ adjacency_tensors: sparse.COO,
121
+ three_body_tensors: sparse.COO,
122
+ initial_state_matrix: np.typing.NDArray,
123
+ final_state_matrix: np.typing.NDArray
124
+ ) -> np.typing.NDArray:
125
+
126
+ r"""
127
+ shortcut method for computing feature vector difference $\mathbf{t}$ between two nearby states
128
+ """
129
+
130
+ sites, _ = np.where(initial_state_matrix != final_state_matrix)
131
+ sites = np.unique(sites).tolist()
132
+
133
+ truncated_adj = sparse.take(adjacency_tensors, sites, axis=1)
134
+ initial_feature_vec_truncated = 2 * symmetrize(contract(
135
+ "nij,iα,jβ->nαβ",
136
+ truncated_adj,
137
+ initial_state_matrix[sites, :],
138
+ initial_state_matrix
139
+ ), axes=(1, 2)).flatten()
140
+ final_feature_vec_truncated = 2 * symmetrize(contract(
141
+ "nij,iα,jβ->nαβ",
142
+ truncated_adj,
143
+ final_state_matrix[sites, :],
144
+ final_state_matrix
145
+ ), axes=(1, 2)).flatten()
146
+
147
+ if three_body_tensors is None:
148
+ return final_feature_vec_truncated - initial_feature_vec_truncated
149
+
150
+ truncated_thr = sparse.take(three_body_tensors, sites, axis=1)
151
+ initial_feature_vec_truncated = np.concatenate(
152
+ [
153
+ initial_feature_vec_truncated,
154
+ 3 * symmetrize(contract(
155
+ "nijk,iα,jβ,kγ->nαβγ",
156
+ truncated_thr,
157
+ initial_state_matrix[sites, :],
158
+ initial_state_matrix,
159
+ initial_state_matrix
160
+ ), axes=(1, 2, 3)).flatten()
161
+ ]
162
+ )
163
+ final_feature_vec_truncated = np.concatenate(
164
+ [
165
+ final_feature_vec_truncated,
166
+ 3 * symmetrize(contract(
167
+ "nijk,iα,jβ,kγ->nαβγ",
168
+ truncated_thr,
169
+ final_state_matrix[sites, :],
170
+ final_state_matrix,
171
+ final_state_matrix
172
+ ), axes=(1, 2, 3)).flatten()
173
+ ]
174
+ )
175
+ return final_feature_vec_truncated - initial_feature_vec_truncated