abacus-develop/python/pyabacus/tests/test_cell.py at develop · pplab/abacus-develop · GitHub

History

207 lines (164 loc) · 7.38 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

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()

abacus-develop/python/pyabacus/tests/test_cell.py at develop · pplab/abacus-develop · GitHub

Expand file tree

Breadcrumbs

test_cell.py

Latest commit

History

Breadcrumbs

test_cell.py

File metadata and controls