diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..adbd10f --- /dev/null +++ b/.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 diff --git a/calculation/CEMC.py b/calculation/CEMC.py index bd3e960..ee5f144 100644 --- a/calculation/CEMC.py +++ b/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 \ No newline at end of file + return ordered_structure diff --git a/calculation/SROS.py b/calculation/SROS.py index 8d5f0a4..fc53ea9 100644 --- a/calculation/SROS.py +++ b/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. diff --git a/calculation/reorder_structure.py b/calculation/reorder_structure.py index 0a2aafc..8572193 100644 --- a/calculation/reorder_structure.py +++ b/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( diff --git a/tests/test_cemc.py b/tests/test_cemc.py new file mode 100644 index 0000000..0dfe3a9 --- /dev/null +++ b/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 diff --git a/tests/test_generate_random.py b/tests/test_generate_random.py new file mode 100644 index 0000000..317f79b --- /dev/null +++ b/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, + } diff --git a/tests/test_reorder_structure.py b/tests/test_reorder_structure.py new file mode 100644 index 0000000..a0e806e --- /dev/null +++ b/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"]