import unittest import numpy as np import os import tempfile import shutil from pyabacus import Cell class TestCell(unittest.TestCase): """Test suite for the Cell class.""" def setUp(self): """Set up test fixtures.""" # Get the absolute path to the test directory self.test_dir = os.path.dirname(os.path.abspath(__file__)) self.test_cell_dir = os.path.join(self.test_dir, 'test_cell') # Set paths for test files self.lcao_dir = os.path.join(self.test_cell_dir, 'lcao_ZnO') self.stru_file = os.path.join(self.lcao_dir, 'STRU') self.xyz_file = os.path.join(self.test_cell_dir, 'h2o.xyz') # Create a temporary directory for output files self.temp_dir = tempfile.mkdtemp() # Verify test files exist if not os.path.exists(self.stru_file): raise FileNotFoundError(f"STRU file not found at {self.stru_file}") if not os.path.exists(self.xyz_file): raise FileNotFoundError(f"XYZ file not found at {self.xyz_file}") def tearDown(self): """Clean up test fixtures.""" shutil.rmtree(self.temp_dir) def test_initialization(self): """Test cell initialization.""" cell = Cell() self.assertFalse(cell._built) self.assertEqual(cell.unit, 'Angstrom') self.assertEqual(cell.ecutwfc, 100.0) self.assertEqual(cell.precision, 1e-8) np.testing.assert_array_equal(cell.lattice, np.eye(3)) def test_from_stru_file(self): """Test loading structure from STRU file.""" cell = Cell.from_file(self.stru_file) self.assertTrue(cell._built) # Verify basic structure information self.assertEqual(len(cell.species), 2) self.assertEqual(cell.species, ['Zn', 'O']) # Check lattice parameters and vectors from the actual STRU file self.assertEqual(cell.lattice_constant, 6.1416) # From STRU file scaled_lattice = cell.lattice * cell.lattice_constant expected_lattice = np.array([ [6.1416, 0.0, 0.0], [-3.0708, 5.3186256, 0.0], [0.0, 0.0, 9.82656] ]) np.testing.assert_array_almost_equal(scaled_lattice, expected_lattice) def test_from_xyz_file(self): """Test loading structure from XYZ file.""" cell = Cell.from_file(self.xyz_file) self.assertTrue(cell._built) # Check basic structure self.assertEqual(len(cell.species), 3) self.assertEqual(cell.species, ['O', 'H', 'H']) # Check positions match the h2o.xyz file content positions = cell.positions self.assertEqual(len(positions), 3) def test_build_status(self): """Test build status management.""" cell = Cell() self.assertFalse(cell._built) # Adding atoms should keep cell unbuilt cell.add_atom('H', [0, 0, 0]) self.assertFalse(cell._built) # Building should set the flag cell.build() self.assertTrue(cell._built) def test_auto_parameters(self): """Test automatic parameter setting.""" cell = Cell() cell.add_atom('H', [0, 0, 0]) cell.lattice = [[5.0, 0.0, 0.0], [0.0, 5.0, 0.0], [0.0, 0.0, 5.0]] cell.build() # Check mesh generation self.assertIsNotNone(cell._mesh) self.assertEqual(len(cell._mesh), 3) def test_stru_pseudopotentials(self): """Test pseudopotential information from STRU file.""" cell = Cell.from_file(self.stru_file) # Check pseudopotential information self.assertIn('Zn', cell.pseudo_potentials) self.assertIn('O', cell.pseudo_potentials) # Verify specific pseudopotential files from STRU self.assertEqual(cell.pseudo_potentials['Zn']['pseudo_file'], 'Zn.LDA.UPF') self.assertEqual(cell.pseudo_potentials['O']['pseudo_file'], 'O.LDA.100.UPF') def test_stru_orbital_files(self): """Test orbital information from STRU file.""" cell = Cell.from_file(self.stru_file) # Check orbital files self.assertEqual(len(cell.orbitals), 2) self.assertIn('Zn_lda_8.0au_120Ry_2s2p2d', cell.orbitals) self.assertIn('O_lda_7.0au_50Ry_2s2p1d', cell.orbitals) def test_stru_coordinates(self): """Test coordinate system and positions from STRU file.""" cell = Cell.from_file(self.stru_file) # Check coordinate type self.assertEqual(cell._coord_type, 'Direct') # Get scaled positions scaled_positions = cell.get_scaled_positions() # Expected positions from STRU file expected_scaled = np.array([ [0.00, 0.00, 0.00], # Zn [0.33333, 0.66667, 0.50] # O ]) np.testing.assert_array_almost_equal(scaled_positions, expected_scaled, decimal=4) def test_file_operations(self): """Test file reading and writing operations.""" # Load from STRU original_cell = Cell.from_file(self.stru_file) # Save to new STRU new_stru = os.path.join(self.temp_dir, 'STRU') original_cell.to_file(new_stru, 'stru') # Load back and compare new_cell = Cell.from_file(new_stru) # Compare structures np.testing.assert_array_almost_equal(original_cell.positions, new_cell.positions) np.testing.assert_array_almost_equal(original_cell.lattice, new_cell.lattice) self.assertEqual(original_cell.species, new_cell.species) self.assertEqual(original_cell.pseudo_potentials, new_cell.pseudo_potentials) self.assertEqual(original_cell.orbitals, new_cell.orbitals) # Test XYZ file operations xyz_cell = Cell.from_file(self.xyz_file) new_xyz = os.path.join(self.temp_dir, 'test.xyz') xyz_cell.to_file(new_xyz, 'xyz') new_xyz_cell = Cell.from_file(new_xyz) np.testing.assert_array_almost_equal(xyz_cell.positions, new_xyz_cell.positions) self.assertEqual(xyz_cell.species, new_xyz_cell.species) def test_add_atom(self): """Test atom addition functionality.""" cell = Cell() # Add atom with properties properties = {"mag": 0.5, "constraint": [1, 1, 1]} cell.add_atom("Fe", [0, 0, 0], properties) self.assertEqual(len(cell.atoms), 1) self.assertEqual(cell.species[0], "Fe") atom_data = cell.atoms[0] self.assertEqual(atom_data[2], properties) # Check if adding atom unsets built status cell.build() self.assertTrue(cell._built) cell.add_atom("Fe", [1, 1, 1]) self.assertFalse(cell._built) def test_k_points(self): """Test k-points generation.""" cell = Cell() cell.add_atom('H', [0, 0, 0]) cell.lattice = [[5.0, 0.0, 0.0], [0.0, 5.0, 0.0], [0.0, 0.0, 5.0]] cell.build() # Test k-points generation kpts = cell.make_kpts([2, 2, 2]) self.assertEqual(kpts.shape, (8, 3)) # Test k-points with non-gamma kpts_nogamma = cell.make_kpts([2, 2, 2], with_gamma_point=False) self.assertEqual(kpts_nogamma.shape, (8, 3)) if __name__ == '__main__': unittest.main()