From beb686483aa31b458f9bca5b1beb173051881f5c Mon Sep 17 00:00:00 2001 From: Henry Plante Date: Sun, 30 Aug 2026 03:24:50 +0000 Subject: [PATCH 1/3] updated with max_ent --- main.py | 62 +++++++++++++++ tce/calculator.py | 197 +++++++++++++++++++++++----------------------- 2 files changed, 160 insertions(+), 99 deletions(-) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..446d458 --- /dev/null +++ b/main.py @@ -0,0 +1,62 @@ +# ref from max_ent - test file edited for confirming the max_ent function works as expected + +from ase.build import bulk +import numpy as np +from tce.calculator import TCECalculator, maximum_entropy_subset_up_to_size_k #changed to import from the calculator +from tce.constants import CUTOFFS + +RNG: np.random.Generator = np.random.default_rng(seed=0) +SUPERSET_SIZE: int = 100 +COPPER_LATTICE_PARAMETER: float = 3.61 + + +def main(): + + pure_cu = bulk("Cu", a=COPPER_LATTICE_PARAMETER, cubic=True).repeat((4, 4, 4)) + + alloys = [] + for sample in range(SUPERSET_SIZE): + + alloy = pure_cu.copy() + nickel_fraction = RNG.uniform(low=0.0, high=1.0) + alloy.symbols = RNG.choice( + a=["Cu", "Ni"], + p=[1.0 - nickel_fraction, nickel_fraction], + size=len(alloy) + ) + alloys.append(alloy) + + calc = TCECalculator( + neighbor_cutoffs=COPPER_LATTICE_PARAMETER * CUTOFFS["fcc"][:4], + many_body_features=[(0, 0, 0)], + species=["Cu", "Ni"] + ) + + # the two types of lists generated: + + #1: generation through max_ent method + subsets_method_1 = list(calc.select_maximum_entropy_subsets(alloys, k=30)) + + #2: normalized method used prior to max_ent + X = calc.get_batched_feature_vectors(alloys) + normalizer = calc.get_normalizer(alloys[0]) + X /= normalizer + + subsets_method_2 = [] + for index_subset in maximum_entropy_subset_up_to_size_k(X, k=30): + alloy_indices = sorted(index_subset) + subsets_method_2.append([alloys[i] for i in alloy_indices]) + + # compare and identify defects, if none, print success message + assert len(subsets_method_1) == len(subsets_method_2) + for s1, s2 in zip(subsets_method_1, subsets_method_2): + assert len(s1) == len(s2) + for a1, a2 in zip(s1, s2): + np.testing.assert_array_equal(a1.numbers, a2.numbers) + np.testing.assert_allclose(a1.positions, a2.positions) + + print("Test passed: Method 1 and Method 2 yield identical results!") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/tce/calculator.py b/tce/calculator.py index f50558b..6bb27af 100644 --- a/tce/calculator.py +++ b/tce/calculator.py @@ -19,19 +19,74 @@ 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 +# new import for optimized sorting: +from typing import Generator, Sequence + 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__) GREEK_ALPHABET = "αβγδεζηθικλμνξοπρστυφχψω" LATIN_ALPHABET = "ijklmnopqrstuvwxyz" +#max_ent funct for optimized matrix sorting +def maximum_entropy_subset_up_to_size_k( + X: np.typing.NDArray, + k: int, + epsilon: float = 1.0e-3 +) -> Generator[set[int], None, None]: + """ + Compute maximum entropy subsets up to size k. + Expects input X of shape (n_samples, n_features). + """ + + #standardizes the matrix's scale to be compared to original calculated matrix + X_proc = X.astype(np.float64, copy=True) + X_proc -= X_proc.mean(axis=0) + std = X_proc.std(axis=0) + std[std == 0.0] = 1.0 # Prevent division by zero for constant features + X_proc /= std + + # Transpose so each column represents a sample vector w_y of dimension d + X_proc = X_proc.T + d, n = X_proc.shape + + # initialize the subset with one sample, that sample being the one closest to the center + subset = [int(np.argmin(np.linalg.norm(X_proc, axis=0)))] + + # initialize a mask that will make sure we only select new samples to add to the subset + mask = np.zeros(n, dtype=bool) + mask[subset[0]] = True + + # compute the covariance matrix for the current subset + covariance = epsilon * np.eye(d) + covariance += X_proc[:, subset[0]:subset[0]+1] @ X_proc[:, subset[0]:subset[0]+1].T + + while len(subset) <= k: + yield set(subset) + + if len(subset) == k or len(subset) == n: + break + + Y = np.linalg.solve(covariance, X_proc) # shape: (d, n) + quad = np.sum(X_proc * Y, axis=0) # shape: (n,) + quad[mask] = 0.0 + + gains = np.log1p(np.maximum(quad, 0.0)) + + j_new = int(np.argmax(gains)) + subset.append(j_new) + mask[j_new] = True + + x_new = X_proc[:, j_new:j_new+1] + covariance += x_new @ x_new.T + @dataclass class TCECalculator(Calculator): @@ -342,10 +397,15 @@ 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( - atoms=atoms, + tree=tree, cutoffs=self.neighbor_cutoffs, tolerance=self.neighbor_tolerance ) @@ -391,7 +451,6 @@ 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 @@ -413,6 +472,7 @@ 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) @@ -433,7 +493,6 @@ def get_feature_vector( return feature_vec - def get_normalizer( self, atoms: Atoms @@ -456,97 +515,6 @@ def get_normalizer( 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, @@ -724,7 +692,6 @@ 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""" @@ -752,7 +719,6 @@ def get_feature_vector_difference(self, initial: Atoms, final: Atoms) -> NDArray raise NotImplementedError - def calculate( self, atoms: Optional[Atoms] = None, @@ -813,7 +779,6 @@ def train(self, configurations: list[Atoms]): return self - @cite(paper_link=KMC_PAPER) def difference_train(self, configuration_pairs: list[tuple[Atoms, Atoms]]): r""" @@ -863,7 +828,6 @@ def difference_train(self, configuration_pairs: list[tuple[Atoms, Atoms]]): return self - def save(self, path: Union[Path, str]): r""" @@ -910,3 +874,38 @@ def load(cls, path: Union[Path, str]) -> "TCECalculator": if not isinstance(obj, cls): raise ValueError(f"loaded object is not of type {cls.__name__}") return obj + +# err handling and center controlling + def get_batched_feature_vectors( + self, + atoms_list: Sequence[Atoms] + ) -> NDArray[np.floating]: + r""" + Compute feature vectors for a sequence of Atoms configurations + and return them as a 2D matrix of shape (n_samples, n_features). + """ + return np.array([self.get_feature_vector(atoms) for atoms in atoms_list]) + + def select_maximum_entropy_subsets( + self, + atoms_list: Sequence[Atoms], + k: int, + epsilon: float = 1.0e-3 + ) -> Generator[list[Atoms], None, None]: + + if not atoms_list: + raise ValueError("atoms_list cannot be empty") + + feature_matrix = self.get_batched_feature_vectors(atoms_list) + + normalizer = self.get_normalizer(atoms_list[0]) + feature_matrix = feature_matrix / normalizer + + for index_subset in maximum_entropy_subset_up_to_size_k( + X=feature_matrix, + k=k, + epsilon=epsilon + ): + # Sort indices to guarantee consistent ordering + alloy_indices = sorted(index_subset) + yield [atoms_list[i] for i in alloy_indices] From 8a8b5d529230903a946b6297c863228e8e432511 Mon Sep 17 00:00:00 2001 From: Henry Plante Date: Mon, 31 Aug 2026 14:37:47 +0000 Subject: [PATCH 2/3] changed test file to not be separate --- main.py | 62 ----------------------------------------------------- test_lib.py | 56 ++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 55 insertions(+), 63 deletions(-) delete mode 100644 main.py diff --git a/main.py b/main.py deleted file mode 100644 index 446d458..0000000 --- a/main.py +++ /dev/null @@ -1,62 +0,0 @@ -# ref from max_ent - test file edited for confirming the max_ent function works as expected - -from ase.build import bulk -import numpy as np -from tce.calculator import TCECalculator, maximum_entropy_subset_up_to_size_k #changed to import from the calculator -from tce.constants import CUTOFFS - -RNG: np.random.Generator = np.random.default_rng(seed=0) -SUPERSET_SIZE: int = 100 -COPPER_LATTICE_PARAMETER: float = 3.61 - - -def main(): - - pure_cu = bulk("Cu", a=COPPER_LATTICE_PARAMETER, cubic=True).repeat((4, 4, 4)) - - alloys = [] - for sample in range(SUPERSET_SIZE): - - alloy = pure_cu.copy() - nickel_fraction = RNG.uniform(low=0.0, high=1.0) - alloy.symbols = RNG.choice( - a=["Cu", "Ni"], - p=[1.0 - nickel_fraction, nickel_fraction], - size=len(alloy) - ) - alloys.append(alloy) - - calc = TCECalculator( - neighbor_cutoffs=COPPER_LATTICE_PARAMETER * CUTOFFS["fcc"][:4], - many_body_features=[(0, 0, 0)], - species=["Cu", "Ni"] - ) - - # the two types of lists generated: - - #1: generation through max_ent method - subsets_method_1 = list(calc.select_maximum_entropy_subsets(alloys, k=30)) - - #2: normalized method used prior to max_ent - X = calc.get_batched_feature_vectors(alloys) - normalizer = calc.get_normalizer(alloys[0]) - X /= normalizer - - subsets_method_2 = [] - for index_subset in maximum_entropy_subset_up_to_size_k(X, k=30): - alloy_indices = sorted(index_subset) - subsets_method_2.append([alloys[i] for i in alloy_indices]) - - # compare and identify defects, if none, print success message - assert len(subsets_method_1) == len(subsets_method_2) - for s1, s2 in zip(subsets_method_1, subsets_method_2): - assert len(s1) == len(s2) - for a1, a2 in zip(s1, s2): - np.testing.assert_array_equal(a1.numbers, a2.numbers) - np.testing.assert_allclose(a1.positions, a2.positions) - - print("Test passed: Method 1 and Method 2 yield identical results!") - - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/test_lib.py b/test_lib.py index 4b71d98..9198ff7 100644 --- a/test_lib.py +++ b/test_lib.py @@ -15,13 +15,14 @@ import numpy as np from numpy.typing import NDArray from ase import build, Atoms +from ase.build import bulk from ase.calculators.singlepoint import SinglePointCalculator import sparse from tce.training import LimitingRidge from tce.topology import symmetrize from tce.datasets import PresetDataset, Dataset -from tce.calculator import TCECalculator +from tce.calculator import TCECalculator, maximum_entropy_subset_up_to_size_k from tce.monte_carlo import monte_carlo, transform_model from tce.constants import CUTOFFS @@ -703,3 +704,56 @@ def test_batched_calculation_throws_error_on_size_mismatch(): with pytest.raises(ValueError): calc.get_batched_feature_vectors(atoms_list=[pure_w, pure_w_larger]) + + +def test_max_ent(): + RNG: np.random.Generator = np.random.default_rng(seed=0) + SUPERSET_SIZE: int = 100 + COPPER_LATTICE_PARAMETER: float = 3.61 + + pure_cu = bulk("Cu", a=COPPER_LATTICE_PARAMETER, cubic=True).repeat((4, 4, 4)) + + alloys = [] + for sample in range(SUPERSET_SIZE): + + alloy = pure_cu.copy() + nickel_fraction = RNG.uniform(low=0.0, high=1.0) + alloy.symbols = RNG.choice( + a=["Cu", "Ni"], + p=[1.0 - nickel_fraction, nickel_fraction], + size=len(alloy) + ) + alloys.append(alloy) + + calc = TCECalculator( + neighbor_cutoffs=COPPER_LATTICE_PARAMETER * CUTOFFS["fcc"][:4], + many_body_features=[(0, 0, 0)], + species=["Cu", "Ni"] + ) + + # the two types of lists generated: + + #1: generation through max_ent method + subsets_method_1 = list(calc.select_maximum_entropy_subsets(alloys, k=30)) + + #2: normalized method used prior to max_ent + X = calc.get_batched_feature_vectors(alloys) + normalizer = calc.get_normalizer(alloys[0]) + X /= normalizer + + subsets_method_2 = [] + for index_subset in maximum_entropy_subset_up_to_size_k(X, k=30): + alloy_indices = sorted(index_subset) + subsets_method_2.append([alloys[i] for i in alloy_indices]) + + # compare and identify defects, if none, print success message + assert len(subsets_method_1) == len(subsets_method_2) + for s1, s2 in zip(subsets_method_1, subsets_method_2): + assert len(s1) == len(s2) + for a1, a2 in zip(s1, s2): + np.testing.assert_array_equal(a1.numbers, a2.numbers) + np.testing.assert_allclose(a1.positions, a2.positions) + + print("Test passed: Method 1 and Method 2 yield identical results!") + + From 9e1d07360f4edc95df853f932cc49552c2b5a701 Mon Sep 17 00:00:00 2001 From: Henry Plante Date: Mon, 31 Aug 2026 14:48:56 +0000 Subject: [PATCH 3/3] Partially done updating calculator.py to sync it with the current version of tce-lib --- tce/calculator.py | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/tce/calculator.py b/tce/calculator.py index 6bb27af..7edcc47 100644 --- a/tce/calculator.py +++ b/tce/calculator.py @@ -4,7 +4,7 @@ from dataclasses import dataclass, field -from typing import Optional, Union +from typing import Optional, Union, Generator, Sequence # updated import order from itertools import permutations, combinations, repeat, product import logging from collections import defaultdict @@ -23,9 +23,6 @@ from opt_einsum import contract from multiset import Multiset -# new import for optimized sorting: -from typing import Generator, Sequence - from .training import Model, LimitingRidge from .topology import hash_topology, symmetrize from .topology import get_adjacency_tensors @@ -36,7 +33,7 @@ LATIN_ALPHABET = "ijklmnopqrstuvwxyz" #max_ent funct for optimized matrix sorting -def maximum_entropy_subset_up_to_size_k( +def _maximum_entropy_subset_up_to_size_k( X: np.typing.NDArray, k: int, epsilon: float = 1.0e-3 @@ -44,6 +41,7 @@ def maximum_entropy_subset_up_to_size_k( """ Compute maximum entropy subsets up to size k. Expects input X of shape (n_samples, n_features). + @private """ #standardizes the matrix's scale to be compared to original calculated matrix @@ -397,15 +395,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 )