ensured to keep dimensions of sc/uc queries

Now the axyz, [us]c2[su]c routines keeps
the dimensions to ensure one can use
them in more versatile ways.

This of course won't work with unique=True
which ravels data.

Also deprecated the sc2uc and uc2sc
since we want to be explicit. Prefer the asc2uc
and auc2sc.

Added more tests to cover these things.

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

CHANGELOG.md

@@ -51,6 +51,7 @@ we hit release version 1.0.0.
 - A new `AtomicMatrixPlot` to plot sparse matrices, #668
 
 ### Fixed
+- `Geometry.[ao][us]c2[su]c` methods now retains the input shapes (unless `unique=True`)
 - lots of `Lattice` methods did not consistently copy over BC
 - `BrillouinZone.volume` fixed to actually return BZ volume
   use `Lattice.volume` for getting the lattice volume.

src/sisl/_core/_ufuncs_geometry.py

@@ -1042,7 +1042,7 @@ def sub(geometry: Geometry, atoms: AtomsIndex) -> Geometry:
     Lattice.fit : update the supercell according to a reference supercell
     Geometry.remove : the negative of this routine, i.e. remove a subset of atoms
     """
-    atoms = geometry.sc2uc(atoms)
+    atoms = geometry.asc2uc(atoms)
     return geometry.__class__(
         geometry.xyz[atoms, :].copy(),
         atoms=geometry.atoms.sub(atoms),
@@ -1067,7 +1067,7 @@ def remove(geometry: Geometry, atoms: AtomsIndex) -> Geometry:
     --------
     Geometry.sub : the negative of this routine, i.e. retain a subset of atoms
     """
-    atoms = geometry.sc2uc(atoms)
+    atoms = geometry.asc2uc(atoms)
     if atoms.size == 0:
         return geometry.copy()
     atoms = np.delete(_a.arangei(geometry.na), atoms)
@@ -1154,7 +1154,7 @@ def rotate(
         if what is None:
             what = "xyz"
         # Only rotate the unique values
-        atoms = geometry.sc2uc(atoms, unique=True)
+        atoms = geometry.asc2uc(atoms, unique=True)
 
     if isinstance(v, Integral):
         v = direction(v, abc=geometry.cell, xyz=np.diag([1, 1, 1]))

src/sisl/_core/_ufuncs_sparse_geometry.py

@@ -462,7 +462,7 @@ def sub(SA: SparseAtom, atoms: AtomsIndex) -> SparseAtom:
     Geometry.sub : equivalent to the resulting `Geometry` from this routine
     SparseAtom.remove : the negative of `sub`, i.e. remove a subset of atoms
     """
-    atoms = SA.sc2uc(atoms)
+    atoms = SA.asc2uc(atoms)
     geom = SA.geometry.sub(atoms)
 
     idx = np.tile(atoms, SA.n_s)
@@ -503,7 +503,7 @@ def sub(SO: SparseOrbital, atoms: AtomsIndex) -> SparseOrbital:
     Geometry.sub : equivalent to the resulting `Geometry` from this routine
     SparseOrbital.remove : the negative of `sub`, i.e. remove a subset of atoms
     """
-    atoms = SO.sc2uc(atoms)
+    atoms = SO.asc2uc(atoms)
     orbs = SO.a2o(atoms, all=True)
     geom = SO.geometry.sub(atoms)
 
@@ -538,6 +538,6 @@ def remove(S: _SparseGeometry, atoms: AtomsIndex) -> _SparseGeometry:
     Geometry.sub : the negative of `Geometry.remove`
     sub : the opposite of `remove`, i.e. retain a subset of atoms
     """
-    atoms = S.sc2uc(atoms)
+    atoms = S.asc2uc(atoms)
     atoms = np.delete(_a.arangei(S.na), atoms)
     return S.sub(atoms)

src/sisl/_core/geometry.py

@@ -47,7 +47,7 @@
 from sisl._internal import set_module
 from sisl._math_small import cross3, is_ascending
 from sisl._namedindex import NamedIndex
-from sisl.messages import SislError, deprecate_argument, info, warn
+from sisl.messages import SislError, deprecate_argument, deprecation, info, warn
 from sisl.shape import Cube, Shape, Sphere
 from sisl.typing import (
     ArrayLike,
@@ -1004,7 +1004,8 @@ def iter_block_rand(
             # Get unit-cell atoms, we are drawing a circle, and this
             # circle only encompasses those already in the unit-cell.
             all_idx[1] = np.union1d(
-                self.sc2uc(all_idx[0], unique=True), self.sc2uc(all_idx[1], unique=True)
+                self.asc2uc(all_idx[0], unique=True),
+                self.asc2uc(all_idx[1], unique=True),
             )
             # If we translated stuff into the unit-cell, we could end up in situations
             # where the supercell atom is in the circle, but not the UC-equivalent
@@ -1134,7 +1135,8 @@ def iter_block_shape(
             # Get unit-cell atoms, we are drawing a circle, and this
             # circle only encompasses those already in the unit-cell.
             all_idx[1] = np.union1d(
-                self.sc2uc(all_idx[0], unique=True), self.sc2uc(all_idx[1], unique=True)
+                self.asc2uc(all_idx[0], unique=True),
+                self.asc2uc(all_idx[1], unique=True),
             )
             # If we translated stuff into the unit-cell, we could end up in situations
             # where the supercell atom is in the circle, but not the UC-equivalent
@@ -2105,7 +2107,7 @@ def axyz(self, atoms: AtomsIndex = None, isc=None) -> ndarray:
             # get offsets from atomic indices (note that this will be per atom)
             isc = self.a2isc(atoms)
             offset = self.lattice.offset(isc)
-            return self.xyz[self.sc2uc(atoms)] + offset
+            return self.xyz[self.asc2uc(atoms)] + offset
 
         # Neither of atoms, or isc are `None`, we add the offset to all coordinates
         return self.axyz(atoms) + self.lattice.offset(isc)
@@ -2512,7 +2514,7 @@ def bond_correct(
             )
             i = np.argmin(d[1])
             # Convert to unitcell atom (and get the one atom)
-            atoms = self.sc2uc(atoms[1][i])
+            atoms = self.asc2uc(atoms[1][i])
             c = c[1][i]
             d = d[1][i]
 
@@ -2995,13 +2997,13 @@ def o2a(self, orbitals: OrbitalsIndex, unique: bool = False) -> ndarray:
                 + (orbitals // self.no) * self.na
             )
 
-        isc, orbitals = np.divmod(_a.asarrayi(orbitals.ravel()), self.no)
-        a = list_index_le(orbitals, self.lasto)
+        isc, orbitals = np.divmod(_a.asarrayi(orbitals), self.no)
+        a = list_index_le(orbitals.ravel(), self.lasto).reshape(orbitals.shape)
         if unique:
             return np.unique(a + isc * self.na)
         return a + isc * self.na
 
-    def uc2sc(self, atoms: AtomsIndex, unique: bool = False) -> ndarray:
+    def auc2sc(self, atoms: AtomsIndex, unique: bool = False) -> ndarray:
         """Returns atom from unit-cell indices to supercell indices, possibly removing dublicates
 
         Parameters
@@ -3012,16 +3014,20 @@ def uc2sc(self, atoms: AtomsIndex, unique: bool = False) -> ndarray:
            If True the returned indices are unique and sorted.
         """
         atoms = self._sanitize_atoms(atoms) % self.na
-        atoms = (
-            atoms.reshape(1, -1) + _a.arangei(self.n_s).reshape(-1, 1) * self.na
-        ).ravel()
+        atoms = (atoms[..., None] + _a.arangei(self.n_s) * self.na).reshape(
+            *atoms.shape[:-1], -1
+        )
         if unique:
             return np.unique(atoms)
         return atoms
 
-    auc2sc = uc2sc
+    uc2sc = deprecation(
+        "uc2sc is deprecated, update the code to use the explicit form auc2sc",
+        "0.15.0",
+        "0.16.0",
+    )(auc2sc)
 
-    def sc2uc(self, atoms: AtomsIndex, unique: bool = False) -> ndarray:
+    def asc2uc(self, atoms: AtomsIndex, unique: bool = False) -> ndarray:
         """Returns atoms from supercell indices to unit-cell indices, possibly removing dublicates
 
         Parameters
@@ -3036,7 +3042,11 @@ def sc2uc(self, atoms: AtomsIndex, unique: bool = False) -> ndarray:
             return np.unique(atoms)
         return atoms
 
-    asc2uc = sc2uc
+    sc2uc = deprecation(
+        "sc2uc is deprecated, update the code to use the explicit form asc2uc",
+        "0.15.0",
+        "0.16.0",
+    )(asc2uc)
 
     def osc2uc(self, orbitals: OrbitalsIndex, unique: bool = False) -> ndarray:
         """Orbitals from supercell indices to unit-cell indices, possibly removing dublicates
@@ -3064,10 +3074,9 @@ def ouc2sc(self, orbitals: OrbitalsIndex, unique: bool = False) -> ndarray:
            If True the returned indices are unique and sorted.
         """
         orbitals = self._sanitize_orbs(orbitals) % self.no
-        orbitals = (
-            orbitals.reshape(1, *orbitals.shape)
-            + _a.arangei(self.n_s).reshape(-1, *([1] * orbitals.ndim)) * self.no
-        ).ravel()
+        orbitals = (orbitals[..., None] + _a.arangei(self.n_s) * self.no).reshape(
+            *orbitals.shape[:-1], -1
+        )
         if unique:
             return np.unique(orbitals)
         return orbitals
@@ -3086,8 +3095,6 @@ def a2isc(self, atoms: AtomsIndex) -> ndarray:
             atom indices to extract the supercell locations of
         """
         atoms = self._sanitize_atoms(atoms) // self.na
-        if atoms.ndim > 1:
-            atoms = atoms.ravel()
         return self.lattice.sc_off[atoms, :]
 
     # This function is a bit weird, it returns a real array,
@@ -3110,8 +3117,6 @@ def o2isc(self, orbitals: OrbitalsIndex) -> ndarray:
         Returns a vector of 3 numbers with integers.
         """
         orbitals = self._sanitize_orbs(orbitals) // self.no
-        if orbitals.ndim > 1:
-            orbitals = orbitals.ravel()
         return self.lattice.sc_off[orbitals, :]
 
     def o2sc(self, orbitals: OrbitalsIndex) -> ndarray:
@@ -3524,7 +3529,7 @@ def within_inf(
 
         # Convert indices to unit-cell indices and also return coordinates and
         # infinite supercell indices
-        return self.sc2uc(idx), xyz, isc
+        return self.asc2uc(idx), xyz, isc
 
     # Create pickling routines
     def __getstate__(self):

src/sisl/_core/sparse_geometry.py

@@ -1155,7 +1155,7 @@ def __getitem__(self, key):
                 # We guess it is the supercell index
                 off = self.geometry.sc_index(key[-1]) * self.na
                 key = [el for el in key[:-1]]
-                key[1] = self.geometry.sc2uc(key[1]) + off
+                key[1] = self.geometry.asc2uc(key[1]) + off
         if dd >= 0:
             key = tuple(key) + (dd,)
             self._def_dim = -1
@@ -1174,7 +1174,7 @@ def __setitem__(self, key, val):
                 # We guess it is the supercell index
                 off = self.geometry.sc_index(key[-1]) * self.na
                 key = [el for el in key[:-1]]
-                key[1] = self.geometry.sc2uc(key[1]) + off
+                key[1] = self.geometry.asc2uc(key[1]) + off
         key = tuple(
             self.geometry._sanitize_atoms(k) if i < 2 else k for i, k in enumerate(key)
         )
@@ -2186,7 +2186,7 @@ def _check_edges_and_coordinates(spgeom, atoms, isc, err_help):
             edges_sc = geom.o2a(
                 spgeom.edges(orbitals=_a.arangei(geom.no), exclude=exclude), True
             )
-            edges_uc = geom.sc2uc(edges_sc, True)
+            edges_uc = geom.asc2uc(edges_sc, True)
             edges_valid = np.isin(edges_uc, atoms, assume_unique=True)
             if not np.all(edges_valid):
                 edges_uc = edges_sc % geom.na
@@ -2517,7 +2517,7 @@ def get_reduced_system(sp, atoms):
         def create_geometry(geom, atoms):
             """Create the supercell geometry with coordinates as given"""
             xyz = geom.axyz(atoms)
-            uc_atoms = geom.sc2uc(atoms)
+            uc_atoms = geom.asc2uc(atoms)
             return Geometry(xyz, atoms=geom.atoms[uc_atoms])
 
         # We know that the *IN* connections are in the primary unit-cell
@@ -2761,11 +2761,10 @@ def a2o(geom, atoms, sc=True):
         # 3: couplings from *inside* to *inside* (no scale)
         # 4: couplings from *inside* to *outside* (scaled)
         convert = [[], []]
-        conc = np.concatenate
 
         def assert_unique(old, new):
-            old = conc(old)
-            new = conc(new)
+            old = concatenate(old)
+            new = concatenate(new)
             assert len(unique(old)) == len(old)
             assert len(unique(new)) == len(new)
             return old, new
@@ -2891,7 +2890,7 @@ def toSparseAtom(self, dim: int = None, dtype=None):
             ptr[ia + 1] = ptr[ia] + len(acol)
 
         # Now we can create the sparse atomic
-        col = np.concatenate(col, axis=0).astype(int32, copy=False)
+        col = concatenate(col, axis=0).astype(int32, copy=False)
         spAtom = SparseAtom(geom, dim=dim, dtype=dtype, nnzpr=0)
         spAtom._csr.ptr[:] = ptr[:]
         spAtom._csr.ncol[:] = diff(ptr)

src/sisl/_core/tests/test_geometry.py

@@ -260,14 +260,18 @@ def test_fxyz(self, setup):
         assert np.allclose(np.dot(fxyz, setup.g.cell), setup.g.xyz)
 
     def test_axyz(self, setup):
-        assert np.allclose(setup.g[:], setup.g.xyz[:])
-        assert np.allclose(setup.g[0], setup.g.xyz[0, :])
-        assert np.allclose(setup.g[2], setup.g.axyz(2))
-        isc = setup.g.a2isc(2)
-        off = setup.g.lattice.offset(isc)
-        assert np.allclose(setup.g.xyz[0] + off, setup.g.axyz(2))
-        assert np.allclose(setup.g.xyz[0] + off, setup.g.axyz(0, isc))
-        assert np.allclose(setup.g.xyz[0] + off, setup.g.axyz(isc=isc)[0])
+        g = setup.g
+        assert np.allclose(g[:], g.xyz[:])
+        assert np.allclose(g[0], g.xyz[0, :])
+        assert np.allclose(g[2], g.axyz(2))
+        isc = g.a2isc(2)
+        off = g.lattice.offset(isc)
+        assert np.allclose(g.xyz[0] + off, g.axyz(2))
+        assert np.allclose(g.xyz[0] + off, g.axyz(0, isc))
+        assert np.allclose(g.xyz[0] + off, g.axyz(isc=isc)[0])
+
+        # Also check multidimensional things
+        assert g.axyz([[0], [1]]).shape == (2, 1, 3)
 
     def test_atranspose_indices(self, setup):
         g = setup.g
@@ -296,7 +300,7 @@ def test_otranspose_indices(self, setup):
     def test_auc2sc(self, setup):
         g = setup.g
         # All supercell indices
-        asc = g.uc2sc(0)
+        asc = g.auc2sc(0)
         assert asc.size == g.n_s
         assert (asc % g.na == 0).sum() == g.n_s
 
@@ -738,11 +742,14 @@ def test_a2o(self, setup):
         setup.g.reorder()
 
     def test_o2a(self, setup):
+        g = setup.g
         # There are 2 orbitals per C atom
-        assert setup.g.o2a(2) == 1
-        assert setup.g.o2a(3) == 1
-        assert setup.g.o2a(4) == 2
-        assert np.all(setup.g.o2a([0, 2, 4]) == [0, 1, 2])
+        assert g.o2a(2) == 1
+        assert g.o2a(3) == 1
+        assert g.o2a(4) == 2
+        assert np.all(g.o2a([0, 2, 4]) == [0, 1, 2])
+        assert np.all(g.o2a([[0], [2], [4]]) == [[0], [1], [2]])
+        assert np.all(g.o2a([[0], [2], [4]], unique=True) == [0, 1, 2])
 
     def test_angle(self, setup):
         # There are 2 orbitals per C atom
@@ -755,30 +762,52 @@ def test_angle(self, setup):
 
     def test_2uc(self, setup):
         # functions for any-thing to UC
-        assert setup.g.sc2uc(2) == 0
-        assert np.all(setup.g.sc2uc([2, 3]) == [0, 1])
-        assert setup.g.asc2uc(2) == 0
-        assert np.all(setup.g.asc2uc([2, 3]) == [0, 1])
-        assert setup.g.osc2uc(4) == 0
-        assert setup.g.osc2uc(5) == 1
-        assert np.all(setup.g.osc2uc([4, 5]) == [0, 1])
+        g = setup.g
+        assert g.asc2uc(2) == 0
+        assert np.all(g.asc2uc([2, 3]) == [0, 1])
+        assert g.asc2uc(2) == 0
+        assert np.all(g.asc2uc([2, 3]) == [0, 1])
+        assert g.osc2uc(4) == 0
+        assert g.osc2uc(5) == 1
+        assert np.all(g.osc2uc([4, 5]) == [0, 1])
+
+    def test_2uc_many_axes(self, setup):
+        # 2 orbitals per atom
+        g = setup.g
+        # functions for any-thing to SC
+        idx = [[1], [2]]
+        assert np.all(g.asc2uc(idx) == [[1], [0]])
+        idx = [[2], [4]]
+        assert np.all(g.osc2uc(idx) == [[2], [0]])
 
     def test_2sc(self, setup):
         # functions for any-thing to SC
-        c = setup.g.cell
+        g = setup.g
+        c = g.cell
 
         # check indices
-        assert np.all(setup.g.a2isc([1, 2]) == [[0, 0, 0], [-1, -1, 0]])
-        assert np.all(setup.g.a2isc(2) == [-1, -1, 0])
-        assert np.allclose(setup.g.a2sc(2), -c[0, :] - c[1, :])
-        assert np.all(setup.g.o2isc([1, 5]) == [[0, 0, 0], [-1, -1, 0]])
-        assert np.all(setup.g.o2isc(5) == [-1, -1, 0])
-        assert np.allclose(setup.g.o2sc(5), -c[0, :] - c[1, :])
+        assert np.all(g.a2isc([1, 2]) == [[0, 0, 0], [-1, -1, 0]])
+        assert np.all(g.a2isc(2) == [-1, -1, 0])
+        assert np.allclose(g.a2sc(2), -c[0, :] - c[1, :])
+        assert np.all(g.o2isc([1, 5]) == [[0, 0, 0], [-1, -1, 0]])
+        assert np.all(g.o2isc(5) == [-1, -1, 0])
+        assert np.allclose(g.o2sc(5), -c[0, :] - c[1, :])
 
         # Check off-sets
         assert np.allclose(setup.g.a2sc([1, 2]), [[0.0, 0.0, 0.0], -c[0, :] - c[1, :]])
         assert np.allclose(setup.g.o2sc([1, 5]), [[0.0, 0.0, 0.0], -c[0, :] - c[1, :]])
 
+    def test_2sc_many_axes(self, setup):
+        # 2 orbitals per atom
+        g = setup.g
+        # functions for any-thing to SC
+        idx = [[1], [2]]
+        assert np.all(g.a2isc(idx) == [[[0, 0, 0]], [[-1, -1, 0]]])
+        assert g.auc2sc(idx).shape == (2, g.n_s)
+        idx = [[2], [4]]
+        assert np.all(g.o2isc(idx) == [[[0, 0, 0]], [[-1, -1, 0]]])
+        assert g.ouc2sc(idx).shape == (2, g.n_s)
+
     def test_reverse(self, setup):
         rev = setup.g.reverse()
         assert len(rev) == 2

src/sisl/geom/_category/tests/test_geom_category.py

@@ -258,7 +258,7 @@ def test_geom_category_xyz_meta():
     We check that the metaclass defined for individual direction categories works.
     """
     hBN_gr = bilayer(1.42, Atom[5, 7], Atom[6]) * (4, 5, 1)
-    sc2uc = hBN_gr.sc2uc
+    sc2uc = hBN_gr.asc2uc
 
     # Check that all classes work
     for key in ("x", "y", "z", "f_x", "f_y", "f_z", "a_x", "a_y", "a_z"):