From c8ab54c2d26769050d400d2d74ea7be7c01d6ede Mon Sep 17 00:00:00 2001 From: Nick Papior Date: Tue, 21 May 2024 11:03:52 +0200 Subject: [PATCH] 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 --- CHANGELOG.md | 1 + src/sisl/_core/_ufuncs_geometry.py | 6 +- src/sisl/_core/_ufuncs_sparse_geometry.py | 6 +- src/sisl/_core/geometry.py | 51 +++++++----- src/sisl/_core/sparse_geometry.py | 15 ++-- src/sisl/_core/tests/test_geometry.py | 83 +++++++++++++------ .../_category/tests/test_geom_category.py | 2 +- src/sisl/viz/data_sources/bond_data.py | 4 +- 8 files changed, 101 insertions(+), 67 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a8324a13..4926c5235 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/src/sisl/_core/_ufuncs_geometry.py b/src/sisl/_core/_ufuncs_geometry.py index 043dd2615..242da71dc 100644 --- a/src/sisl/_core/_ufuncs_geometry.py +++ b/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])) diff --git a/src/sisl/_core/_ufuncs_sparse_geometry.py b/src/sisl/_core/_ufuncs_sparse_geometry.py index 8c4f9b83c..056cbc554 100644 --- a/src/sisl/_core/_ufuncs_sparse_geometry.py +++ b/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) diff --git a/src/sisl/_core/geometry.py b/src/sisl/_core/geometry.py index 6a7bd3ed5..24d79ec59 100644 --- a/src/sisl/_core/geometry.py +++ b/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): diff --git a/src/sisl/_core/sparse_geometry.py b/src/sisl/_core/sparse_geometry.py index 6d2a76cd9..463355fca 100644 --- a/src/sisl/_core/sparse_geometry.py +++ b/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) diff --git a/src/sisl/_core/tests/test_geometry.py b/src/sisl/_core/tests/test_geometry.py index 26fa01512..77c2c2a9d 100644 --- a/src/sisl/_core/tests/test_geometry.py +++ b/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 diff --git a/src/sisl/geom/_category/tests/test_geom_category.py b/src/sisl/geom/_category/tests/test_geom_category.py index 7f77b9d46..bfd92c2a5 100644 --- a/src/sisl/geom/_category/tests/test_geom_category.py +++ b/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"): diff --git a/src/sisl/viz/data_sources/bond_data.py b/src/sisl/viz/data_sources/bond_data.py index ce75ee5d0..0009c9e70 100644 --- a/src/sisl/viz/data_sources/bond_data.py +++ b/src/sisl/viz/data_sources/bond_data.py @@ -54,7 +54,7 @@ def bond_data_from_atom( fold_to_uc: bool = False, ): if fold_to_uc: - bonds = geometry.sc2uc(bonds) + bonds = geometry.asc2uc(bonds) return atom_data[bonds[:, 0]] @@ -63,7 +63,7 @@ def bond_data_from_matrix( matrix, geometry: sisl.Geometry, bonds: np.ndarray, fold_to_uc: bool = False ): if fold_to_uc: - bonds = geometry.sc2uc(bonds) + bonds = geometry.asc2uc(bonds) return matrix[bonds[:, 0], bonds[:, 1]]