From 510118b91ace000c4cd96d1ba42a28f68d372449 Mon Sep 17 00:00:00 2001
From: JosePizarro3 <jose.pizarro@physik.hu-berlin.de>
Date: Tue, 5 Mar 2024 16:06:08 +0100
Subject: [PATCH] Added test for AtomicCell

---
 tests/test_model_system.py | 106 ++++++++++++++++++++++++++++++++++---
 1 file changed, 98 insertions(+), 8 deletions(-)

diff --git a/tests/test_model_system.py b/tests/test_model_system.py
index 14ed993f..c04e3023 100644
--- a/tests/test_model_system.py
+++ b/tests/test_model_system.py
@@ -17,16 +17,106 @@
 #
 
 import pytest
+import numpy as np
+
+from nomad.units import ureg
 
 from . import logger
-from nomad_simulations.model_system import AtomicCell
+
+from nomad_simulations.model_system import (
+    GeometricSpace,
+    Cell,
+    AtomicCell,
+    Symmetry,
+    ChemicalFormula,
+    ModelSystem,
+)
+from nomad_simulations.atoms_state import AtomsState
 
 
 class TestAtomicCell:
-    @pytest.fixture(autouse=True)
-    def atomic_cell(self):
-        return AtomicCell(
-            lattice_vectors=[[1, 0, 0], [0, 1, 0], [0, 0, 1]],
-            positions=[[0, 0, 0], [0.5, 0.5, 0.5]],
-            species=['H', 'O'],
-        )
+    @pytest.mark.parametrize(
+        'atoms, atomic_numbers, formula, lattice_vectors, positions, periodic_boundary_conditions',
+        [
+            (
+                ['H', 'H', 'O'],
+                [1, 1, 8],
+                'H2O',
+                [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
+                [[0, 0, 0], [0.5, 0.5, 0.5], [1, 1, 1]],
+                [False, False, False],
+            ),  # full atomic cell
+            (
+                [],
+                [1, 1, 8],
+                'H2O',
+                [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
+                [[0, 0, 0], [0.5, 0.5, 0.5], [1, 1, 1]],
+                [False, False, False],
+            ),  # missing chemical_symbols
+            (
+                ['H', 'H', 'O'],
+                [1, 1, 8],
+                'H2O',
+                [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
+                [],
+                [False, False, False],
+            ),  # missing positions
+            (
+                ['H', 'H', 'O'],
+                [1, 1, 8],
+                'H2O',
+                [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
+                [[0, 0, 0], [0.5, 0.5, 0.5], [1, 1, 1], [2, 2, 2]],
+                [False, False, False],
+            ),  # chemical_symbols and positions with different lengths
+            (
+                ['H', 'H', 'O'],
+                [1, 1, 8],
+                'H2O',
+                [],
+                [[0, 0, 0], [0.5, 0.5, 0.5], [1, 1, 1]],
+                [False, False, False],
+            ),  # missing lattice_vectors
+        ],
+    )
+    def test_generate_ase_atoms(
+        self,
+        atoms,
+        atomic_numbers,
+        formula,
+        lattice_vectors,
+        positions,
+        periodic_boundary_conditions,
+    ):
+        # Define the atomic cell
+        atomic_cell = AtomicCell()
+        if lattice_vectors:
+            atomic_cell.lattice_vectors = lattice_vectors * ureg('angstrom')
+        if positions:
+            atomic_cell.positions = positions * ureg('angstrom')
+        if periodic_boundary_conditions:
+            atomic_cell.periodic_boundary_conditions = periodic_boundary_conditions
+
+        # Add the elements information
+        for index, atom in enumerate(atoms):
+            atom_state = AtomsState()
+            setattr(atom_state, 'chemical_symbol', atom)
+            atomic_number = atom_state.resolve_atomic_number(logger)
+            assert atomic_number == atomic_numbers[index]
+            atom_state.atomic_number = atomic_number
+            atomic_cell.atoms_state.append(atom_state)
+
+        # Test `to_ase_atoms` function
+        ase_atoms = atomic_cell.to_ase_atoms(logger)
+        if not atoms or not positions or len(atoms) != len(positions):
+            assert ase_atoms is None
+        else:
+            if lattice_vectors:
+                assert (ase_atoms.cell == lattice_vectors).all()
+            else:
+                assert (ase_atoms.cell == [0, 0, 0]).all()
+            assert (ase_atoms.positions == positions).all()
+            assert (ase_atoms.pbc == periodic_boundary_conditions).all()
+            assert (ase_atoms.symbols.numbers == atomic_numbers).all()
+            assert ase_atoms.symbols.get_chemical_formula() == formula