fixed ORB_INDX reading basis

Now I *hope* all is checked and works

Signed-off-by: Nick Papior <[email protected]>

src/sisl/io/siesta/orb_indx.py

@@ -5,12 +5,23 @@
 
 from typing import Optional, Union
 
-from sisl import Atom, AtomicOrbital, Atoms, Geometry, Orbital, PeriodicTable
+import numpy as np
+
+from sisl import (
+    Atom,
+    AtomicOrbital,
+    Atoms,
+    AtomUnknown,
+    Geometry,
+    Orbital,
+    PeriodicTable,
+)
 from sisl._array import arrayi
 from sisl._internal import set_module
 from sisl.messages import deprecate_argument
 from sisl.unit.siesta import unit_convert
 
+from .._help import _fill_basis_empty
 from ..sile import add_sile, sile_fh_open
 from .sile import SileSiesta
 
@@ -78,59 +89,76 @@ def read_basis(self, atoms: Optional[Union[Atoms, Geometry]] = None) -> Atoms:
 
         pt = PeriodicTable()
 
+        if isinstance(atoms, Geometry):
+            atoms = atoms.atoms
+
         if atoms is None:
 
-            def crt_atom(i_s, spec, orbs):
+            def crt_atom(i_s, tag, orbs):
                 # The user has not specified an atomic basis
-                i = pt.Z(spec)
+                i = pt.Z(tag)
                 if isinstance(i, int):
+                    # we can convert tag name to an atom
+                    # Hence we don't need to add the tag
                     return Atom(i, orbs)
-                else:
-                    return Atom(-1, orbs, tag=spec)
+                return AtomUnknown(1000 + i_s, orbs, tag=tag)
 
         else:
-            if isinstance(atoms, Geometry):
-                atoms = atoms.atoms
 
-            def crt_atom(i_s, spec, orbs):
+            def crt_atom(i_s, tag, orbs):
                 # Get the atom and add the orbitals
-                return atoms[i_s].copy(orbitals=orbs)
+                kwargs = {}
+                if atoms[i_s].tag != tag:
+                    # we know ORB_INDX tag is correct
+                    kwargs["tag"] = tag
+                if len(atoms[i_s]) != len(orbs):
+                    # only overwrite if # of orbitals don't match
+                    kwargs["orbitals"] = orbs
+                if kwargs:
+                    return atoms[i_s].copy(**kwargs)
+                return atoms[i_s]
 
         # Now we begin by reading the atoms
-        atom = []
-        orbs = []
-        specs = []
-        ia = 1
+        atom, orbs = [], []
+        species, order_species = [], []
+        current_ia = 1
+        tag = ""
         i_s = 0
         for _ in range(no):
             line = self.readline().split()
 
-            i_a = int(line[1])
-            spec = line[3]
-            if i_a != ia:
-                if i_s not in specs:
-                    atom.append(crt_atom(i_s, spec, orbs))
-                specs.append(i_s)
-                ia = i_a
+            ia = int(line[1])
+            if ia != current_ia:
+                if i_s not in species:
+                    order_species.append(i_s)
+                    atom.append(crt_atom(i_s, tag, orbs))
+                species.append(i_s)
+                current_ia = ia
                 orbs = []
 
+            # Get tag for atom
+            tag = line[3]
+            # and species number
             i_s = int(line[2]) - 1
 
-            if i_s in specs:
+            if i_s in order_species:
+                # no need to collect information for the same orbital
                 continue
+
             nlmz = list(map(int, line[5:9]))
             P = line[9] == "T"
             rc = float(line[11]) * Bohr2Ang
             # Create the orbital
             o = AtomicOrbital(n=nlmz[0], l=nlmz[1], m=nlmz[2], zeta=nlmz[3], P=P, R=rc)
             orbs.append(o)
 
-        if i_s not in specs:
-            atom.append(crt_atom(i_s, spec, orbs))
-        specs.append(i_s)
+        if i_s not in species:
+            order_species.append(i_s)
+            atom.append(crt_atom(i_s, tag, orbs))
+        species.append(i_s)
+        atom = Atoms([atom[i] for i in np.argsort(order_species)])
 
-        # Now re-arrange the atoms and create the Atoms object
-        return Atoms([atom[i] for i in specs])
+        return _fill_basis_empty(np.array(species), atom)
 
 
 add_sile("ORB_INDX", orbindxSileSiesta, gzip=True)

src/sisl/io/siesta/tests/test_orb_indx.py

@@ -8,6 +8,7 @@
 import numpy as np
 import pytest
 
+from sisl import Atom, Atoms
 from sisl.io.siesta.orb_indx import *
 
 pytestmark = [pytest.mark.io, pytest.mark.siesta]
@@ -22,6 +23,7 @@ def test_si_pdos_kgrid_orb_indx(sisl_files):
 
     assert len(atoms) == 2
     assert atoms.nspecies == 1
+    assert atoms.atom[0].Z == 14
     assert len(atoms[0]) == 13
     assert len(atoms[1]) == 13
 
@@ -34,5 +36,80 @@ def test_sih_orb_indx(sisl_files):
 
     assert len(atoms) == 65
     assert atoms.nspecies == 2
+    Z = np.zeros(len(atoms))
+    Z[:] = 14
+    Z[-1] = 1
+    assert np.allclose(atoms.Z, Z)
+    # number of orbitals on each atom
     assert len(atoms[0]) == 4
     assert len(atoms[-1]) == 1
+
+
+def test_orb_indx_order(sisl_tmp):
+    f = sisl_tmp("test.ORD_INDX", _dir)
+
+    with open(f, "w") as fh:
+        fh.write(
+            """\
+     29  29 = orbitals in unit cell and supercell. See end of file.
+
+    io    ia is   spec iao  n  l  m  z  p          sym      rc    isc     iuo
+     1     1  2 Fe_SOC   1  4  0  0  1  F            s   7.329  0  0  0     1
+     2     1  2 Fe_SOC   2  4  0  0  2  F            s   6.153  0  0  0     2
+     3     1  2 Fe_SOC   3  3  2 -2  1  F          dxy   4.336  0  0  0     3
+     4     1  2 Fe_SOC   4  3  2 -1  1  F          dyz   4.336  0  0  0     4
+     5     1  2 Fe_SOC   5  3  2  0  1  F          dz2   4.336  0  0  0     5
+     6     1  2 Fe_SOC   6  3  2  1  1  F          dxz   4.336  0  0  0     6
+     7     1  2 Fe_SOC   7  3  2  2  1  F       dx2-y2   4.336  0  0  0     7
+     8     1  2 Fe_SOC   8  3  2 -2  2  F          dxy   2.207  0  0  0     8
+     9     1  2 Fe_SOC   9  3  2 -1  2  F          dyz   2.207  0  0  0     9
+    10     1  2 Fe_SOC  10  3  2  0  2  F          dz2   2.207  0  0  0    10
+    11     1  2 Fe_SOC  11  3  2  1  2  F          dxz   2.207  0  0  0    11
+    12     1  2 Fe_SOC  12  3  2  2  2  F       dx2-y2   2.207  0  0  0    12
+    13     1  2 Fe_SOC  13  4  1 -1  1  T          Ppy   7.329  0  0  0    13
+    14     1  2 Fe_SOC  14  4  1  0  1  T          Ppz   7.329  0  0  0    14
+    15     1  2 Fe_SOC  15  4  1  1  1  T          Ppx   7.329  0  0  0    15
+    16     2  1     Pt   1  6  0  0  1  F            s   7.158  0  0  0    16
+    17     2  1     Pt   2  6  0  0  2  F            s   6.009  0  0  0    17
+    18     2  1     Pt   3  5  2 -2  1  F          dxy   5.044  0  0  0    18
+    19     2  1     Pt   4  5  2 -1  1  F          dyz   5.044  0  0  0    19
+    20     2  1     Pt   5  5  2  0  1  F          dz2   5.044  0  0  0    20
+    21     2  1     Pt   6  5  2  1  1  F          dxz   5.044  0  0  0    21
+    22     2  1     Pt   7  5  2  2  1  F       dx2-y2   5.044  0  0  0    22
+    23     2  1     Pt   8  5  2 -2  2  F          dxy   3.022  0  0  0    23
+    24     2  1     Pt   9  5  2 -1  2  F          dyz   3.022  0  0  0    24
+    25     2  1     Pt  10  5  2  0  2  F          dz2   3.022  0  0  0    25
+    26     2  1     Pt  11  5  2  1  2  F          dxz   3.022  0  0  0    26
+    27     2  1     Pt  12  5  2  2  2  F       dx2-y2   3.022  0  0  0    27
+    28     2  1     Pt  13  6  1 -1  1  T          Ppy   7.158  0  0  0    28
+    29     2  1     Pt  14  6  1  0  1  T          Ppz   7.158  0  0  0    29
+"""
+        )
+
+    atoms = orbindxSileSiesta(f).read_basis()
+
+    assert len(atoms) == 2
+    assert atoms.nspecies == 2
+    Z = [1001, 78]
+    assert np.allclose(atoms.Z, Z)
+    # number of orbitals on each atom
+    assert len(atoms[0]) == 15
+    assert len(atoms[1]) == 14
+    assert atoms.atom[0].tag == "Pt"
+    assert len(atoms.atom[0]) == 14
+    assert atoms.atom[1].tag == "Fe_SOC"
+    assert len(atoms.atom[1]) == 15
+
+    atoms = orbindxSileSiesta(f).read_basis(atoms=Atoms([Atom("Pt"), Atom("Fe")]))
+
+    assert len(atoms) == 2
+    assert atoms.nspecies == 2
+    Z = [26, 78]
+    assert np.allclose(atoms.Z, Z)
+    # number of orbitals on each atom
+    assert len(atoms[0]) == 15
+    assert len(atoms[1]) == 14
+    assert atoms.atom[0].tag == "Pt"
+    assert len(atoms.atom[0]) == 14
+    assert atoms.atom[1].tag == "Fe_SOC"
+    assert len(atoms.atom[1]) == 15