[Feature] Added Pure Random Algo to OrderDisorderedStructureTransformation (#4236)

* occ_tol

* random transformation

* pre-commit auto-fixes

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: jmmshn

src/pymatgen/transformations/standard_transformations.py

@@ -23,6 +23,7 @@
 from pymatgen.transformations.transformation_abc import AbstractTransformation
 
 if TYPE_CHECKING:
+    from numpy.random import Generator
     from typing_extensions import Self
 
     from pymatgen.core.sites import PeriodicSite
@@ -451,6 +452,7 @@ class OrderDisorderedStructureTransformation(AbstractTransformation):
     ALGO_FAST = 0
     ALGO_COMPLETE = 1
     ALGO_BEST_FIRST = 2
+    ALGO_RANDOM = -1
 
     def __init__(self, algo=ALGO_FAST, symmetrized_structures=False, no_oxi_states=False):
         """
@@ -467,7 +469,9 @@ def __init__(self, algo=ALGO_FAST, symmetrized_structures=False, no_oxi_states=F
         self.no_oxi_states = no_oxi_states
         self.symmetrized_structures = symmetrized_structures
 
-    def apply_transformation(self, structure: Structure, return_ranked_list: bool | int = False) -> Structure:
+    def apply_transformation(
+        self, structure: Structure, return_ranked_list: bool | int = False, occ_tol=0.25
+    ) -> Structure:
         """For this transformation, the apply_transformation method will return
         only the ordered structure with the lowest Ewald energy, to be
         consistent with the method signature of the other transformations.
@@ -478,6 +482,9 @@ def apply_transformation(self, structure: Structure, return_ranked_list: bool |
             structure: Oxidation state decorated disordered structure to order
             return_ranked_list (bool | int, optional): If return_ranked_list is int, that number of structures
                 is returned. If False, only the single lowest energy structure is returned. Defaults to False.
+            occ_tol (float): Occupancy tolerance. If the total occupancy of a group is within this value
+                of an integer, it will be rounded to that integer otherwise raise a ValueError.
+                Defaults to 0.25.
 
         Returns:
             Depending on returned_ranked list, either a transformed structure
@@ -529,14 +536,25 @@ def apply_transformation(self, structure: Structure, return_ranked_list: bool |
         # generate the list of manipulations and input structure
         struct = Structure.from_sites(structure)
 
+        # We will first create an initial ordered structure by filling all sites
+        # with the species that has the highest oxidation state (initial_sp)
+        # replacing all other species on a given site.
+        # then, we process a list of manipulations to get the final structure.
+        # The manipulations are of the format:
+        # [oxi_ratio, 1, [0,1,2,3], Li+]
+        # which means -- Place 1 Li+ in any of these 4 sites
+        # the oxi_ratio is the ratio of the oxidation state of the species to
+        # the initial species. This is used to determine the energy of the
+        # manipulation in the EwaldMinimizer, but is not used in the purely random
+        # algorithm.
         manipulations = []
         for group in equivalent_sites:
             total_occupancy = dict(
                 sum((structure[idx].species for idx in group), Composition()).items()  # type: ignore[attr-defined]
             )
             # round total occupancy to possible values
             for key, val in total_occupancy.items():
-                if abs(val - round(val)) > 0.25:
+                if abs(val - round(val)) > occ_tol:
                     raise ValueError("Occupancy fractions not consistent with size of unit cell")
                 total_occupancy[key] = round(val)
             # start with an ordered structure
@@ -555,6 +573,16 @@ def apply_transformation(self, structure: Structure, return_ranked_list: bool |
             if empty > 0.5:
                 manipulations.append([0, empty, list(group), None])
 
+        if self.algo == self.ALGO_RANDOM:
+            rand_structures = get_randomly_manipulated_structures(
+                struct=struct, manipulations=manipulations, n_return=n_to_return
+            )
+            if return_ranked_list:
+                return [
+                    {"energy": 0.0, "energy_above_minimum": 0.0, "structure": s} for s in rand_structures[:n_to_return]
+                ]
+            return rand_structures[0]
+
         matrix = EwaldSummation(struct).total_energy_matrix
         ewald_m = EwaldMinimizer(matrix, manipulations, n_to_return, self.algo)
 
@@ -891,3 +919,82 @@ def apply_transformation(self, structure):
 
     def __repr__(self):
         return "ScaleToRelaxedTransformation"
+
+
+def _sample_random_manipulation(manipulation, rng, manipulated) -> list[tuple[int, SpeciesLike]]:
+    """Sample a single random manipulation.
+
+    Each manipulation is given in the form of a tuple
+    `(oxi_ratio, nsites, indices, sp)` where:
+    Which means choose nsites from the list of indices and replace them
+    With the species `sp`.
+    """
+    _, nsites, indices, sp = manipulation
+    maniped_indices = [i for i, _ in manipulated]
+    allowed_sites = [i for i in indices if i not in maniped_indices]
+    if len(allowed_sites) < nsites:
+        raise RuntimeError(
+            "No valid manipulations possible. "
+            f" You have already applied a manipulation to each site in this group {indices}"
+        )
+    sampled_sites = rng.choice(allowed_sites, nsites, replace=False).tolist()
+    sampled_sites.sort()
+    return [(i, sp) for i in sampled_sites]
+
+
+def _get_manipulation(manipulations: list, rng: Generator, max_attempts, seen: set[tuple]) -> tuple:
+    """Apply each manipulation."""
+    for _ in range(max_attempts):
+        manipulated: list[tuple] = []
+        for manip_ in manipulations:
+            new_manips = _sample_random_manipulation(manip_, rng, manipulated)
+            manipulated += new_manips
+        tm_ = tuple(manipulated)
+        if tm_ not in seen:
+            return tm_
+    raise RuntimeError(
+        "Could not apply manipulations to structure"
+        "this is likely because you have already applied all the possible manipulations"
+    )
+
+
+def _apply_manip(struct, manipulations) -> Structure:
+    """Apply manipulations to a structure."""
+    struct_copy = struct.copy()
+    rm_indices = []
+    for manip in manipulations:
+        idx, sp = manip
+        if sp is None:
+            rm_indices.append(idx)
+        else:
+            struct_copy.replace(idx, sp)
+    struct_copy.remove_sites(rm_indices)
+    return struct_copy
+
+
+def get_randomly_manipulated_structures(
+    struct: Structure, manipulations: list, seed=None, n_return: int = 1
+) -> list[Structure]:
+    """Get a structure with random manipulations applied.
+
+    Args:
+        struct: Input structure
+        manipulations: List of manipulations to apply
+        seed: Seed for random number generator
+        n_return: Number of structures to return
+
+    Returns:
+        List of structures with manipulations applied.
+    """
+    rng = np.random.default_rng(seed)
+    seen: set[tuple] = set()
+    sampled_manips = []
+
+    for _ in range(n_return):
+        manip_ = _get_manipulation(manipulations, rng, 1000, seen)
+        seen.add(manip_)
+        sampled_manips.append(manip_)
+    output_structs = []
+    for manip_ in sampled_manips:
+        output_structs.append(_apply_manip(struct, manip_))
+    return output_structs

tests/transformations/test_standard_transformations.py

@@ -401,6 +401,32 @@ def test_best_first(self):
         output = trafo.apply_transformation(struct, return_ranked_list=3)
         assert output[0]["energy"] == approx(-234.57813667648315, abs=1e-4)
 
+    def test_random_sample(self):
+        struc_str = (
+            "3.333573 0.000000 1.924639\n"
+            "1.111191 3.142924 1.924639\n"
+            "0.000000 0.000000 3.849278\n"
+            "1.0 0.0 0.0\n"
+            "0.0 1.0 0.0\n"
+            "0.0 0.0 1.0\n"
+            "0.875000 0.875000 0.875000 Si=1\n"
+            "0.125000 0.125000 0.125000 Si=1"
+        )
+        si = Structure.from_str(struc_str, fmt="mcsqs")
+        struct = si * [3, 2, 1]
+        struct.replace(0, {"Fe": 0.5, "Ni": 0.5})
+        struct.replace(1, {"Fe": 0.5, "Ni": 0.5})
+        trafo = OrderDisorderedStructureTransformation(
+            algo=OrderDisorderedStructureTransformation.ALGO_RANDOM, no_oxi_states=True
+        )
+        output = trafo.apply_transformation(struct * [2, 2, 2], return_ranked_list=3)
+        assert len(output) == 3
+        for entry in output:
+            assert set(entry.keys()) == {"structure", "energy", "energy_above_minimum"}
+
+        output = trafo.apply_transformation(struct * [2, 2, 2], return_ranked_list=False)
+        assert output.composition.reduced_formula == struct.composition.reduced_formula
+
 
 class TestPrimitiveCellTransformation:
     def test_apply_transformation(self):