fix: harden abacuslite ASE interface

Author: QuantumMisaka

interfaces/ASE_interface/abacuslite/core.py

@@ -54,6 +54,7 @@
     read_input,
     read_stru,
     read_kpt,
+    species_group_indices,
     write_input,
     write_stru,
     write_kpt
@@ -142,7 +143,7 @@ def version(self) -> str:
 class AbacusTemplate(CalculatorTemplate):
 
     implemented_properties = [
-        'energy', 'forces', 'stress', 'free_energy', 'magmom', 'dipole'
+        'energy', 'forces', 'stress', 'free_energy', 'magmom'
     ]
     _label = 'abacus'
 
@@ -184,10 +185,6 @@ def get_free_energy_keywords(self) -> Dict[str, str]:
     @staticmethod
     def get_magmom_keywords(self) -> Dict[str, str]:
         return {'nspin': '2'}
-    
-    @staticmethod
-    def get_dipole_keywords(self) -> Dict[str, str]:
-        return {'esolver_type': 'tddft', 'out_dipole': '1'}
 
     def get_property_keywords(self,
                               parameters: Dict[str, str],
@@ -204,17 +201,28 @@ def get_property_keywords(self,
         properties : list of str
             The list of properties to calculate
         '''
-        # update the parameters with the keywords for the properties
-        # however, one should also consider that there may be the case that
-        # contradictory keywords are needed. In this kind of cases, 
-        # we should raise a ValueError
-        param_cache_ = {}
+        def normalize_keyword_value(value):
+            if isinstance(value, (list, tuple, set)):
+                return ' '.join(str(i) for i in value)
+            return str(value)
+
+        param_cache_ = {
+            key: normalize_keyword_value(value)
+            for key, value in parameters.items()
+            if value is not None
+        }
+
         def counter(param_new: Dict[str, str]) -> Dict[str, str]:
-            info = 'desired properties required contradictory keywords'
+            info = 'desired properties or explicit parameters required contradictory keywords'
+            staged = {}
             for k, v in param_new.items():
-                if k in param_cache_ and param_cache_[k] != v:
+                if v is None:
+                    continue
+                normalized_value = normalize_keyword_value(v)
+                if k in param_cache_ and param_cache_[k] != normalized_value:
                     raise ValueError(f'{info}: {k}={v} (now), {param_cache_[k]} (before)')
-            # if it is alright, pass through
+                staged[k] = normalized_value
+            param_cache_.update(staged)
             return param_new
 
         # update the parameters with the keywords for the properties
@@ -260,9 +268,9 @@ def write_input(self,
 
         # STRU
         _ = file_safe_backup(directory / parameters.get('stru_file', 'STRU'))
-        # reorder the atoms according to the alphabet. Keep the reverse map
+        # group atoms by first-occurrence species order. Keep the reverse map
         # so that we will recover the order in function read_results()
-        ind = sorted(range(len(atoms)), key=lambda i: atoms[i].symbol)
+        ind = species_group_indices(atoms.get_chemical_symbols())
         self.atomorder = sorted(range(len(atoms)), key=lambda i: ind[i]) # revmap
         # then we write
         _ = write_stru(atoms[ind], 
@@ -663,5 +671,28 @@ def test_version_number_check(self):
         self.assertFalse(switch_io_backend_version('v3.11.0-beta.2'))
         self.assertFalse(switch_io_backend_version('v3.11.0'))
 
+    def test_property_keywords_reject_conflicting_user_parameters(self):
+        template = AbacusTemplate()
+        with self.assertRaises(ValueError):
+            template.get_property_keywords({'nspin': 1}, ['magmom'])
+
+        parameters = template.get_property_keywords({'nspin': 2}, ['magmom'])
+        self.assertEqual(str(parameters['nspin']), '2')
+
+    def test_property_keywords_reject_conflicting_properties(self):
+        template = AbacusTemplate()
+        template.implemented_properties = ['prop_a', 'prop_b']
+        template.get_prop_a_keywords = lambda parameters: {'calculation': 'scf'}
+        template.get_prop_b_keywords = lambda parameters: {'calculation': 'md'}
+
+        with self.assertRaises(ValueError):
+            template.get_property_keywords({}, ['prop_a', 'prop_b'])
+
+    def test_dipole_property_is_not_implemented(self):
+        template = AbacusTemplate()
+        self.assertNotIn('dipole', template.implemented_properties)
+        with self.assertRaises(AssertionError):
+            template.get_property_keywords({}, ['dipole'])
+
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()

interfaces/ASE_interface/abacuslite/io/generalio.py

@@ -11,6 +11,7 @@
 import numpy as np
 from ase.atoms import Atoms
 from ase.build import bulk
+from ase.constraints import FixAtoms, FixCartesian
 from ase.data import chemical_symbols, atomic_masses
 ATOM_MASS = dict(zip(chemical_symbols, atomic_masses.tolist()))
 
@@ -84,16 +85,17 @@ def file_safe_backup(fn: Path, suffix: str = 'bak'):
     '''
     assert isinstance(fn, Path)
     where = fn.parent
+    prefix = f'{fn.name}.{suffix}.'
+    indexed_backups = []
 
-    # get the backup files
-    fbak = sorted(list(where.glob(f'{fn.name}.{suffix}.*')),
-                  key=lambda p: int(p.name.split('.')[-1]))
-    if fbak:
-        # rename the elder by adding 1 to the suffix
-        for i, f in enumerate(fbak[::-1]): # reverse order, to avoid overwrite
-            j = len(fbak) - i + 1 #: STRU.bak.i -> STRU.bak.i+1
-            fname = f.name.replace(f'.{j}', f'.{j+1}')
-            f.rename(f.parent / fname)
+    for backup in where.glob(f'{fn.name}.{suffix}.*'):
+        index_text = backup.name.removeprefix(prefix)
+        if not re.fullmatch(r'0|[1-9][0-9]*', index_text):
+            continue
+        indexed_backups.append((int(index_text), backup))
+
+    for backup_index, backup in sorted(indexed_backups, key=lambda item: item[0], reverse=True):
+        backup.rename(backup.parent / f'{fn.name}.{suffix}.{backup_index + 1}')
 
     # backup the latest file, if there is one
     if fn.exists():
@@ -177,6 +179,29 @@ def _write_stru(job_dir, stru, fname='STRU'):
 
                 f.write('\n')
 
+def species_group_indices(symbols: List[str]) -> List[int]:
+    """Return indices grouped by first-occurrence species order."""
+    species_order = list(dict.fromkeys(symbols))
+    return [i for species in species_order for i, symbol in enumerate(symbols) if symbol == species]
+
+
+def _constraint_mobility(atoms: Atoms) -> np.ndarray:
+    """Return ABACUS mobility flags derived from ASE constraints."""
+    mobility = np.ones((len(atoms), 3), dtype=int)
+    for constraint in atoms.constraints:
+        if isinstance(constraint, FixAtoms):
+            mobility[constraint.get_indices()] = 0
+        elif isinstance(constraint, FixCartesian):
+            indices = np.asarray(constraint.get_indices(), dtype=int)
+            mask = np.asarray(constraint.mask, dtype=bool)
+            if mask.ndim == 1:
+                mobility[np.ix_(indices, np.where(mask)[0])] = 0
+            else:
+                for atom_index, atom_mask in zip(indices, mask):
+                    mobility[atom_index, atom_mask] = 0
+    return mobility
+
+
 def write_stru(stru: Atoms,
                outdir: str,
                pp_file: Optional[Dict[str, str]],
@@ -216,17 +241,20 @@ def write_stru(stru: Atoms,
 
     elem = stru.get_chemical_symbols()    
     # ABACUS requires the atoms ranged species-by-species, therefore
-    # we need to sort the atoms by species
-    ind = np.argsort(elem)
+    # we need to group atoms by species. Preserve first-occurrence species
+    # order from ASE instead of forcing alphabetical order.
+    ind = species_group_indices(elem)
     coords = stru.get_positions()[ind]
+    mobility = _constraint_mobility(stru)[ind]
     elem = [elem[i] for i in ind]
 
     # handle the atomic magnetic moment (issue #6516)
     magmoms = np.array([stru[i].magmom for i in ind]).reshape(len(stru), -1) # ncol in [1, 3]
     magmoms = [{} if abs(np.linalg.norm(m)) <= 1e-10 
                else {'mag': m[0] if len(m) == 1 else ('Cartesian', m.tolist())}
                for m in magmoms]
-    elem_uniq, nat = np.unique(elem, return_counts=True)
+    elem_uniq = list(dict.fromkeys(elem))
+    nat = np.array([elem.count(e) for e in elem_uniq])
     stru_dict = {
         'coord_type': 'Cartesian',
         'lat': {
@@ -244,7 +272,7 @@ def write_stru(stru: Atoms,
                 'atom': [
                     magmoms[j] | {
                         'coord': coords[j].tolist(), # coordinate
-                        'm': [1, 1, 1], # mobility
+                        'm': mobility[j].tolist(), # mobility
                         'v': [0.0, 0.0, 0.0], # velocity
                     } for j in range(np.sum(nat[:i]), np.sum(nat[:i+1]))
                 ]
@@ -632,6 +660,25 @@ def test_input_io(self):
             self.assertDictEqual(data, data_)
         # will automatically delete the file after the context manager
 
+    def test_file_safe_backup_rotates_numbered_backups(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            workdir = Path(tmpdir)
+            live = workdir / 'STRU'
+            live.write_text('live')
+            (workdir / 'STRU.bak.0').write_text('bak0')
+            (workdir / 'STRU.bak.1').write_text('bak1')
+            (workdir / 'STRU.bak.01').write_text('bak01')
+            (workdir / 'STRU.bak.note').write_text('note')
+
+            file_safe_backup(live)
+
+            self.assertFalse(live.exists())
+            self.assertEqual((workdir / 'STRU.bak.0').read_text(), 'live')
+            self.assertEqual((workdir / 'STRU.bak.1').read_text(), 'bak0')
+            self.assertEqual((workdir / 'STRU.bak.2').read_text(), 'bak1')
+            self.assertEqual((workdir / 'STRU.bak.01').read_text(), 'bak01')
+            self.assertEqual((workdir / 'STRU.bak.note').read_text(), 'note')
+
     def test_stru_io(self):
         from ase.units import Bohr, Angstrom
         nacl = bulk('NaCl', 'rocksalt', a=5.64)
@@ -676,14 +723,56 @@ def test_stru_io(self):
                 self.assertEqual(a['m'], [1, 1, 1])
                 self.assertEqual(a['v'], [0.0, 0.0, 0.0])
 
-        self.assertEqual(stru_['species'][0]['symbol'], 'Cl')
-        self.assertEqual(stru_['species'][1]['symbol'], 'Na')
-        self.assertEqual(stru_['species'][0]['mass'], ATOM_MASS['Cl'])
-        self.assertEqual(stru_['species'][1]['mass'], ATOM_MASS['Na'])
-        self.assertEqual(stru_['species'][0]['pp_file'], 'Cl.pz-bhs.UPF')
-        self.assertEqual(stru_['species'][1]['pp_file'], 'Na.pz-bhs.UPF')
-        self.assertEqual(stru_['species'][0]['orb_file'], 'Cl_gga_6au_100Ry_2s2p1d.orb')
-        self.assertEqual(stru_['species'][1]['orb_file'], 'Na_gga_6au_100Ry_2s2p1d.orb')
+        self.assertEqual(stru_['species'][0]['symbol'], 'Na')
+        self.assertEqual(stru_['species'][1]['symbol'], 'Cl')
+        self.assertEqual(stru_['species'][0]['mass'], ATOM_MASS['Na'])
+        self.assertEqual(stru_['species'][1]['mass'], ATOM_MASS['Cl'])
+        self.assertEqual(stru_['species'][0]['pp_file'], 'Na.pz-bhs.UPF')
+        self.assertEqual(stru_['species'][1]['pp_file'], 'Cl.pz-bhs.UPF')
+        self.assertEqual(stru_['species'][0]['orb_file'], 'Na_gga_6au_100Ry_2s2p1d.orb')
+        self.assertEqual(stru_['species'][1]['orb_file'], 'Cl_gga_6au_100Ry_2s2p1d.orb')
+
+    def test_write_stru_preserves_first_occurrence_species_order(self):
+        atoms = Atoms(
+            symbols=['C', 'C', 'Pt', 'H', 'H'],
+            positions=np.zeros((5, 3)),
+            cell=np.eye(3),
+        )
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            write_stru(
+                atoms,
+                outdir=tmpdir,
+                pp_file={'C': 'C.upf', 'Pt': 'Pt.upf', 'H': 'H.upf'},
+            )
+            stru = read_stru(Path(tmpdir) / 'STRU')
+
+        self.assertEqual([s['symbol'] for s in stru['species']], ['C', 'Pt', 'H'])
+        self.assertEqual([s['natom'] for s in stru['species']], [2, 1, 2])
+
+    def test_write_stru_uses_ase_constraints_for_mobility(self):
+        atoms = Atoms(
+            symbols=['C', 'C', 'Pt', 'H'],
+            positions=np.zeros((4, 3)),
+            cell=np.eye(3),
+        )
+        atoms.set_constraint([
+            FixAtoms(indices=[0]),
+            FixCartesian(2, mask=[True, False, True]),
+        ])
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            write_stru(
+                atoms,
+                outdir=tmpdir,
+                pp_file={'C': 'C.upf', 'Pt': 'Pt.upf', 'H': 'H.upf'},
+            )
+            stru = read_stru(Path(tmpdir) / 'STRU')
+
+        self.assertEqual(stru['species'][0]['atom'][0]['m'], [0, 0, 0])
+        self.assertEqual(stru['species'][0]['atom'][1]['m'], [1, 1, 1])
+        self.assertEqual(stru['species'][1]['atom'][0]['m'], [0, 1, 0])
+        self.assertEqual(stru['species'][2]['atom'][0]['m'], [1, 1, 1])
 
     def test_kpt_io(self):
         kpt = {
@@ -818,4 +907,4 @@ def test_write_stru_with_magmom(self):
         self.assertEqual(stru['species'][2]['atom'][0]['mag'], ('Cartesian', [0.0, 0.0, 3.0]))
 
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()

interfaces/ASE_interface/abacuslite/io/latestio.py

@@ -507,7 +507,7 @@ def read_abacus_out(fileobj,
         calc = SinglePointDFTCalculator(atoms=atoms, energy=ener['E_KohnSham'],
                                         free_energy=ener['E_KohnSham'],
                                         forces=frs, stress=strs,
-                                        magmoms=mag, efermi=ener['E_Fermi'],
+                                        magmoms=mag[ind], efermi=ener['E_Fermi'],
                                         ibzkpts=kvecd, dipole=None)
         # import the eigenvalues and occupations kpoint-by-kpoint
         calc.kpts = []
@@ -531,6 +531,44 @@ class TestLatestIO(unittest.TestCase):
     here = Path(__file__).parent
     testfiles = here / 'testfiles'
 
+    def test_read_abacus_out_reorders_calculator_magmoms(self):
+        import tempfile
+        from unittest.mock import patch
+
+        frame = {
+            'elem': ['Na', 'Na', 'Cl'],
+            'coords': np.array([[0.0, 0.0, 0.0],
+                                [1.0, 0.0, 0.0],
+                                [2.0, 0.0, 0.0]]),
+            'cell': np.eye(3),
+        }
+        elecstate = [{
+            'k': np.zeros((1, 1, 3)),
+            'e': np.zeros((1, 1, 1)),
+            'occ': np.ones((1, 1, 1)),
+        }]
+        energies = [{'E_KohnSham': -1.0, 'E_Fermi': 0.0}]
+        kpoints = ((np.zeros((1, 3)), np.ones(1), None), None, None, None, None)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            running_log = Path(tmpdir) / 'running_scf.log'
+            running_log.write_text('')
+            (Path(tmpdir) / 'eig_occ.txt').write_text('')
+            with patch(__name__ + '.read_esolver_type_from_running_log', return_value='ksdft'), \
+                 patch(__name__ + '.read_traj_from_running_log', return_value=[frame]), \
+                 patch(__name__ + '.read_band_from_eig_occ', return_value=elecstate), \
+                 patch(__name__ + '.read_forces_from_running_log', return_value=[]), \
+                 patch(__name__ + '.read_stress_from_running_log', return_value=[]), \
+                 patch(__name__ + '.read_kpoints_from_running_log', return_value=kpoints), \
+                 patch(__name__ + '.read_energies_from_running_log', return_value=([], [])), \
+                 patch(__name__ + '.read_iter_header_from_running_log', return_value=[]), \
+                 patch(__name__ + '.find_final_info_with_iter_header', return_value=energies), \
+                 patch(__name__ + '.read_magmom_from_running_log', return_value=[np.array([10.0, 20.0, 30.0])]):
+                atoms = read_abacus_out(running_log, sort_atoms_with=[0, 2, 1])[0]
+
+        self.assertEqual(atoms.get_chemical_symbols(), ['Na', 'Cl', 'Na'])
+        self.assertTrue(np.allclose(atoms.calc.results['magmoms'], [10.0, 30.0, 20.0]))
+
     def test_read_esolver_type_from_running_log(self):
         self.assertEqual(
             read_esolver_type_from_running_log(
@@ -674,4 +712,4 @@ def test_read_iter_header_from_running_log(self):
         self.assertTupleEqual(header[2], (2, 1))
 
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()