Fix test_moca.

smol/capp/generate/random.py

@@ -81,7 +81,8 @@ def _gen_unconstrained_ordered_occu(sublattices, rng=None):
         rand_occu[sublatt.sites] = rng.choice(
             sublatt.encoding, size=len(sublatt.sites), replace=True
         )
-    return np.ascontiguousarray(rand_occu, dtype=int)
+    # Enforce int32.
+    return np.ascontiguousarray(rand_occu, dtype=np.int32)
 
 
 def _gen_neutral_occu(sublattices, lam=10, num_attempts=10000, rng=None):
@@ -135,7 +136,8 @@ def flip(occu, sublattices, lam=10):
     for _ in range(num_attempts):
         occu, C = flip(occu, sublattices, lam=lam)
         if C == 0:
-            return np.ascontiguousarray(occu, dtype=int)
+            # Enforce int32.
+            return np.ascontiguousarray(occu, dtype=np.int32)
 
     raise TimeoutError(
         f"Can not generate a neutral occupancy in {num_attempts} attempts!"
@@ -174,7 +176,8 @@ def _gen_composition_ordered_occu(sublattices, composition, tol, rng=None):
             occu[sites] = code
             all_sites = [i for i in all_sites if i not in sites]
 
-    return np.ascontiguousarray(occu, dtype=int)
+    # Enforce int32.
+    return np.ascontiguousarray(occu, dtype=np.int32)
 
 
 def _composition_compatiblity(sublattices, composition, tol, rng=None):

smol/moca/kernel/base.py

@@ -697,10 +697,14 @@ def set_aux_state(self, occupancy, *args, **kwargs):
         of occupancies for each kernel, ie dim = (n_kernels, n_sites).
         """
         # set the current and previous values based on given occupancy
+        # Enforce int32.
+        if occupancy.dtype != np.int32:
+            occupancy = occupancy.astype(np.int32)
         new_features = []
         if occupancy.ndim == 2 and occupancy.shape[0] == len(self._kernels):
             for kernel, occupancy in zip(self._kernels, occupancy):
-                occupancy = np.ascontiguousarray(occupancy, dtype=int)
+                # Enforce int32.
+                occupancy = np.ascontiguousarray(occupancy, dtype=np.int32)
                 kernel.trace.occupancy = occupancy
                 kernel.set_aux_state(occupancy, *args, **kwargs)
                 new_features.append(kernel.ensemble.compute_feature_vector(occupancy))

smol/moca/processor/base.py

@@ -101,7 +101,7 @@ def compute_feature_vector(self, occupancy):
         symmetrically distinct bit ordering.
 
         Args:
-            occupancy (ndarray):
+            occupancy (ndarray[np.int32]):
                 encoded occupation array
 
         Returns:
@@ -115,7 +115,7 @@ def compute_feature_vector_change(self, occupancy, flips):
         Compute the change in the feature vector from a list of flips.
 
         Args:
-            occupancy (ndarray):
+            occupancy (ndarray[np.int32]):
                 encoded occupancy array
             flips (list of tuple):
                 list of tuples with two elements. Each tuple represents a
@@ -137,7 +137,7 @@ def compute_feature_vector_distance_change(self, feature_vector, occupancy, flip
         Args:
             feature_vector (ndarray):
                 fixed vector to compute absolute distance differences from.
-            occupancy (ndarray):
+            occupancy (ndarray[np.int32]):
                 encoded occupancy array
             flips (list of tuple):
                 list of tuples with two elements. Each tuple represents a

smol/moca/processor/ewald.py

@@ -142,7 +142,7 @@ def compute_feature_vector_change(self, occupancy, flips):
         Compute the change in the feature vector from a list of flips.
 
         Args:
-            occupancy (ndarray):
+            occupancy (ndarray[np.int32]):
                 encoded occupancy string
             flips (list of tuple):
                 list of tuples with two elements. Each tuple represents a
@@ -158,9 +158,14 @@ def compute_feature_vector_change(self, occupancy, flips):
         for f in flips:
             occu_f = occu_i.copy()
             occu_f[f[0]] = f[1]
-            delta_energy += delta_ewald_single_flip(
-                occu_f, occu_i, self.ewald_matrix, self._ewald_inds, f[0]
-            )
+            try:
+                delta_energy += delta_ewald_single_flip(
+                    occu_f, occu_i, self.ewald_matrix, self._ewald_inds, f[0]
+                )
+            except ValueError:
+                raise ValueError(
+                    f"occupancy dtype is: {occu_i.dtype}, must be an array of np.int32!"
+                )
             occu_i = occu_f
         return delta_energy
 

smol/moca/processor/expansion.py

@@ -166,18 +166,24 @@ def compute_feature_vector(self, occupancy):
         other works it is extensive corresponding to the size of the supercell.
 
         Args:
-            occupancy (ndarray):
+            occupancy (ndarray[np.int32]):
                 encoded occupation array
 
         Returns:
             array: correlation vector
         """
-        return (
-            self._evaluator.correlations_from_occupancy(
-                occupancy, self._indices.container
+        try:
+            corr = (
+                self._evaluator.correlations_from_occupancy(
+                    occupancy, self._indices.container
+                )
+                * self.size
             )
-            * self.size
-        )
+        except ValueError:
+            raise ValueError(
+                f"occupancy dtype is: {occupancy.dtype}, must be an array of np.int32!"
+            )
+        return corr
 
     def compute_feature_vector_change(self, occupancy, flips):
         """
@@ -204,9 +210,14 @@ def compute_feature_vector_change(self, occupancy, flips):
             occu_f = occu_i.copy()
             occu_f[f[0]] = f[1]
             eval_data = self._eval_data_by_sites[f[0]]
-            delta_corr += eval_data.evaluator.delta_correlations_from_occupancies(
-                occu_f, occu_i, eval_data.cluster_ratio, eval_data.indices.container
-            )
+            try:
+                delta_corr += eval_data.evaluator.delta_correlations_from_occupancies(
+                    occu_f, occu_i, eval_data.cluster_ratio, eval_data.indices.container
+                )
+            except ValueError:
+                raise ValueError(
+                    f"occupancy dtype is: {occupancy.dtype}, must be an array of np.int32!"
+                )
             occu_i = occu_f
 
         return delta_corr * self.size
@@ -379,12 +390,18 @@ def compute_feature_vector(self, occupancy):
         Returns:
             array: correlation vector
         """
-        return (
-            self._evaluator.interactions_from_occupancy(
-                occupancy, self._indices.container
+        try:
+            corr = (
+                self._evaluator.interactions_from_occupancy(
+                    occupancy, self._indices.container
+                )
+                * self.size
             )
-            * self.size
-        )
+        except ValueError:
+            raise ValueError(
+                f"occupancy dtype is: {occupancy.dtype}, must be an array of np.int32!"
+            )
+        return corr
 
     def compute_feature_vector_change(self, occupancy, flips):
         """
@@ -410,11 +427,19 @@ def compute_feature_vector_change(self, occupancy, flips):
             occu_f = occu_i.copy()
             occu_f[f[0]] = f[1]
             eval_data = self._eval_data_by_sites[f[0]]
-            delta_interactions += (
-                eval_data.evaluator.delta_interactions_from_occupancies(
-                    occu_f, occu_i, eval_data.cluster_ratio, eval_data.indices.container
+            try:
+                delta_interactions += (
+                    eval_data.evaluator.delta_interactions_from_occupancies(
+                        occu_f,
+                        occu_i,
+                        eval_data.cluster_ratio,
+                        eval_data.indices.container,
+                    )
+                )
+            except ValueError:
+                raise ValueError(
+                    f"occupancy dtype is: {occupancy.dtype}, must be an array of np.int32!"
                 )
-            )
             occu_i = occu_f
 
         return delta_interactions * self.size

smol/moca/sublattice.py

@@ -41,7 +41,7 @@ class Sublattice(MSONable):
         array of site indices for all sites in sublattice
      active_sites (ndarray):
         array of site indices for all unrestricted sites in the sublattice.
-     encoding (ndarray):
+     encoding (ndarray[np.int32]):
         array of species encoding in integer indices. By default,
         will be initialized as range(len(site_space)). Might be different
         if a sub-lattice was created from the split of another sub-lattice.
@@ -60,7 +60,8 @@ def __post_init__(self):
             # A single-species sub-lattice should not be active at all.
             self.restrict_sites(self.sites)
 
-        self.encoding = np.arange(len(self.site_space), dtype=int)
+        # Enforce int32 to prevent dtype problems.
+        self.encoding = np.arange(len(self.site_space), dtype=np.int32)
 
     @property
     def is_active(self):
@@ -172,7 +173,8 @@ def get_index(sp, species):
             part_space = SiteSpace(part_comp)
             part_sites = np.array(part_sites, dtype=int)
             part_actives = np.array(part_actives, dtype=int)
-            part_codes = np.array(part_codes, dtype=int)
+            # Enforce int32.
+            part_codes = np.array(part_codes, dtype=np.int32)
             part_sublatt = Sublattice(part_space, part_sites)
             part_sublatt.active_sites = part_actives
             part_sublatt.encoding = part_codes
@@ -219,5 +221,6 @@ def from_dict(cls, d):
             SiteSpace.from_dict(d["site_space"]), sites=np.array(d["sites"], dtype=int)
         )
         sublattice.active_sites = np.array(d["active_sites"], dtype=int)
-        sublattice.encoding = np.array(d["encoding"], dtype=int)
+        # Enforce int32.
+        sublattice.encoding = np.array(d["encoding"], dtype=np.int32)
         return sublattice

tests/test_moca/test_ensemble.py

@@ -101,7 +101,8 @@ def test_from_cluster_expansion(
             ensemble.natural_parameters[: ensemble.num_energy_coefs],
             np.append(cluster_subspace_ewald.orbit_multiplicities, coefs[-1]),
         )
-    occu = np.zeros(ensemble.num_sites, dtype=int)
+    # Enforce int32.
+    occu = np.zeros(ensemble.num_sites, dtype=np.int32)
     for _ in range(50):  # test a few flips
         sublatt = rng.choice(ensemble.active_sublattices)
         site = rng.choice(sublatt.sites)

tests/test_moca/test_processor.py

@@ -175,6 +175,7 @@ def test_get_average_drift(processor):
 def test_compute_property_change(processor, rng):
     sublattices = processor.get_sublattices()
     occu = _gen_unconstrained_ordered_occu(sublattices, rng=rng)
+    occu_wrong_type = occu.astype(np.int64)
     active_sublattices = [sublatt for sublatt in sublattices if sublatt.is_active]
 
     for _ in range(100):
@@ -186,6 +187,19 @@ def test_compute_property_change(processor, rng):
         prop_f = processor.compute_property(new_occu)
         prop_i = processor.compute_property(occu)
         dprop = processor.compute_property_change(occu, [(site, new_sp)])
+        # Check error raising behavior.
+        if isinstance(
+            processor, (ClusterExpansionProcessor, ClusterDecompositionProcessor)
+        ):
+            with pytest.raises(ValueError):
+                _ = processor.compute_property(occu_wrong_type)
+        if isinstance(
+            processor,
+            (ClusterExpansionProcessor, ClusterDecompositionProcessor, EwaldProcessor),
+        ):
+            with pytest.raises(ValueError):
+                _ = processor.compute_property_change(occu_wrong_type, [(site, new_sp)])
+
         # Check with some tight tolerances.
         npt.assert_allclose(dprop, prop_f - prop_i, rtol=RTOL, atol=ATOL)
         # Test reverse matches forward

tests/test_moca/test_sublattice.py

@@ -14,7 +14,7 @@ def sublattice(rng):
         {
             DummySpecies("A"): 0.3,
             DummySpecies("X"): 0.3,
-            DummySpecies("D"): 0.2,
+            DummySpecies("G"): 0.2,
             DummySpecies("E"): 0.2,
         }
     )