updated with max_ent

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()

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]