Add tests and tighten helper behavior

.github/workflows/tests.yml

@@ -0,0 +1,24 @@
+name: Tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  pytest:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install package
+        run: python -m pip install -e ".[dev]"
+
+      - name: Run tests
+        run: python -m pytest -q

calculation/CEMC.py

@@ -1,12 +1,9 @@
-from pymatgen.core import Structure, Element, Site, Species
-from smol.io import load_work
 import random
+
 import numpy as np
+from smol.io import load_work
 from smol.moca import Ensemble
 from smol.moca import Sampler
-import os
-import shutil
-import re
 
 from .reorder_structure import reorder_atoms_flexible
 
@@ -151,4 +148,4 @@ def run_cemc(li_content,
     # Reorder atoms in the structure
     ordered_structure = reorder_atoms_flexible(structure, ["Li", "Mn", "Ti", "O"])
 
-    return ordered_structure
+    return ordered_structure

calculation/SROS.py

@@ -74,9 +74,9 @@ def alpha(self):
     def alpha_new(self, a_neighbor: int, b_neighbor: int):
         changed_anion_ls = []
         for i in range(6):
-            if self.cnn.get_nn(self.structure, a_neighbor)[i].specie == Element('F'):
+            if self.cnn.get_nn(self.structure, a_neighbor)[i].specie == Element(self.anion_a):
                 changed_anion_ls.append(self.cnn.get_nn(self.structure, a_neighbor)[i].index)
-            if self.cnn.get_nn(self.structure, b_neighbor)[i].specie == Element('F'):
+            if self.cnn.get_nn(self.structure, b_neighbor)[i].specie == Element(self.anion_a):
                 changed_anion_ls.append(self.cnn.get_nn(self.structure, b_neighbor)[i].index)
 
         new_dict = self.a_dict.copy()
@@ -142,9 +142,9 @@ def alpha_LiLi_new(self, c_neighbor: int, d_neighbor: int, c_is_Li: bool):
 
         changed_cation_ls = []
         for i in range(12):
-            if self.bnn.get_nn(structrue_dup, c_neighbor)[i].specie == Element('Li'):
+            if self.bnn.get_nn(structrue_dup, c_neighbor)[i].specie == Element(self.cation):
                 changed_cation_ls.append(self.bnn.get_nn(structrue_dup, c_neighbor)[i].index)
-            if self.bnn.get_nn(structrue_dup, d_neighbor)[i].specie == Element('Li'):
+            if self.bnn.get_nn(structrue_dup, d_neighbor)[i].specie == Element(self.cation):
                 changed_cation_ls.append(self.bnn.get_nn(structrue_dup, d_neighbor)[i].index)
 
         new_dict_LiLi = self.a_LiLi_dict.copy()
@@ -307,7 +307,7 @@ def exchange_LiLi(
         diff = self.a_LiLi - target_alpha_LiLi
         prob = self.sigmoid(diff * rate)
 
-        c_idxs = self.get_idxs("Li")
+        c_idxs = self.get_idxs(self.cation)
         d_idxs = list(set(self.all_cation_idxs) - set(c_idxs))
         c_site = c_idxs[random.randrange(len(c_idxs))]  # "c_site" represents the position of any arbitrary Li atom.
         d_site = d_idxs[random.randrange(len(d_idxs))]  # "d_site" represents the position of any arbitrary TM atom.

calculation/reorder_structure.py

@@ -1,6 +1,7 @@
-from pymatgen.core import Structure
 from typing import Iterable, List
 
+from pymatgen.core import Structure
+
 
 def reorder_atoms(structure: Structure, input_order: List[str]) -> Structure:
     """
@@ -28,15 +29,15 @@ def reorder_atoms(structure: Structure, input_order: List[str]) -> Structure:
             f"{missing} atoms ({', '.join(missing_elements)}) not in order list. "
             f"Check if input structure matches the specified order."
         )
-        raise RuntimeError(error_message)
-
-    # Process site properties (convert list to dictionary)
-    site_properties = {}
-    if ordered_sites and hasattr(ordered_sites[0], 'properties'):
-        # Get all property keys
-        keys = ordered_sites[0].properties.keys()
-        for key in keys:
-            site_properties[key] = [site.properties.get(key) for site in ordered_sites]
+        raise ValueError(error_message)
+
+    site_property_keys = sorted(
+        {key for site in ordered_sites for key in site.properties}
+    )
+    site_properties = {
+        key: [site.properties.get(key) for site in ordered_sites]
+        for key in site_property_keys
+    }
 
     # Create new structure (preserving lattice parameters)
     return Structure(

tests/test_cemc.py

@@ -0,0 +1,36 @@
+import random
+
+import numpy as np
+from pymatgen.core import Lattice, Structure
+
+from calculation.CEMC import assign_element_numbers, calculate_mn_ratios
+
+
+def test_calculate_mn_ratios_for_charge_balanced_composition():
+    assert calculate_mn_ratios(1.2, 0.4) == (0, 0.4, 0)
+
+
+def test_calculate_mn_ratios_handles_mixed_mn_valence():
+    assert calculate_mn_ratios(1.0, 0.5) == (0.5, 0.0, 0)
+    assert calculate_mn_ratios(1.4, 0.5) == (0, 0.0, 0.5)
+
+
+def test_assign_element_numbers_preserves_site_count_and_symbols():
+    random.seed(1)
+    structure = Structure(
+        Lattice.cubic(4.2),
+        ["Li", "Mn", "Mn", "Ti", "O"],
+        [
+            [0, 0, 0],
+            [0.2, 0.2, 0.2],
+            [0.4, 0.4, 0.4],
+            [0.6, 0.6, 0.6],
+            [0.8, 0.8, 0.8],
+        ],
+    )
+
+    occupancies = assign_element_numbers(structure, 0.5, 0.5, 0)
+
+    assert isinstance(occupancies, np.ndarray)
+    assert occupancies.tolist().count(0) == 2
+    assert len(occupancies) == structure.num_sites

tests/test_generate_random.py

@@ -0,0 +1,41 @@
+from collections import Counter
+
+import pytest
+
+from calculation.generate_random import (
+    create_original_structure,
+    make_supercell_matrix,
+    modify_structure_HEDRX,
+)
+
+
+def test_make_supercell_matrix_validates_shape():
+    with pytest.raises(ValueError, match="3x3"):
+        make_supercell_matrix([[1, 0], [0, 1]])
+
+
+def test_modify_structure_hedrx_is_reproducible_with_seed():
+    structure = create_original_structure()
+    structure.make_supercell([[2, 0, 0], [0, 2, 0], [0, 0, 2]])
+
+    first = modify_structure_HEDRX(structure, "TM4", seed=42)
+    second = modify_structure_HEDRX(structure, "TM4", seed=42)
+
+    assert [site.specie.symbol for site in first] == [site.specie.symbol for site in second]
+
+
+def test_modify_structure_hedrx_orders_tm4_species_and_counts_sites():
+    structure = create_original_structure()
+    structure.make_supercell([[2, 0, 0], [0, 2, 0], [0, 0, 2]])
+
+    modified = modify_structure_HEDRX(structure, "TM4", seed=42)
+    symbols = [site.specie.symbol for site in modified]
+
+    assert symbols == sorted(symbols, key=["Li", "Mn", "Ti", "Nb", "O", "F"].index)
+    assert Counter(symbols) == {
+        "Li": 5,
+        "Mn": 2,
+        "Nb": 1,
+        "O": 7,
+        "F": 1,
+    }

tests/test_reorder_structure.py

@@ -0,0 +1,51 @@
+import pytest
+from pymatgen.core import Lattice, Structure
+
+from calculation.reorder_structure import reorder_atoms, reorder_atoms_flexible
+
+
+def test_reorder_atoms_preserves_all_site_property_keys():
+    structure = Structure(
+        Lattice.cubic(4.2),
+        ["O", "Li", "Mn"],
+        [[0, 0, 0], [0.25, 0.25, 0.25], [0.5, 0.5, 0.5]],
+    )
+    structure[0].properties["label"] = "anion"
+    structure[1].properties["charge"] = 1
+    structure[2].properties["label"] = "transition-metal"
+    structure[2].properties["charge"] = 3
+
+    reordered = reorder_atoms(structure, ["Li", "Mn", "O"])
+
+    assert [site.specie.symbol for site in reordered] == ["Li", "Mn", "O"]
+    assert reordered.site_properties["label"] == [None, "transition-metal", "anion"]
+    assert reordered.site_properties["charge"] == [1, 3, None]
+
+
+def test_reorder_atoms_raises_for_missing_elements():
+    structure = Structure(
+        Lattice.cubic(4.2),
+        ["Li", "O", "F"],
+        [[0, 0, 0], [0.25, 0.25, 0.25], [0.5, 0.5, 0.5]],
+    )
+
+    with pytest.raises(ValueError, match="not in order list"):
+        reorder_atoms(structure, ["Li", "O"])
+
+
+def test_reorder_atoms_flexible_appends_unlisted_elements_alphabetically():
+    structure = Structure(
+        Lattice.cubic(4.2),
+        ["F", "Nb", "Li", "O", "Mn"],
+        [
+            [0, 0, 0],
+            [0.2, 0.2, 0.2],
+            [0.4, 0.4, 0.4],
+            [0.6, 0.6, 0.6],
+            [0.8, 0.8, 0.8],
+        ],
+    )
+
+    reordered = reorder_atoms_flexible(structure, ["Li", "Mn", "O"])
+
+    assert [site.specie.symbol for site in reordered] == ["Li", "Mn", "O", "F", "Nb"]