added astype for sparse spin matrices

This means that when transforming spin
matrices it, converts the spin data
to the resulting real + imag parts, and vice
versa.

This should make things much simpler since codes
can re-use these transitions without a siesta
specific code.

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

CHANGELOG.md

@@ -9,6 +9,9 @@ we hit release version 1.0.0.
 ## [0.15.3] - YYYY-MM-DD
 
 ### Added
+- `astype` for sparse matrices, enables one to change data-types, #865
+  This should be preferred over transform which can't do real->complex
+  of spin matrices.
 - added ADOS extraction of PDOS data in `sisl.viz`
 - enabled submodule access without imports:
 

src/sisl/_core/sparse.py

@@ -1609,6 +1609,28 @@ def transform(self, matrix, dtype=None):
 
         return new
 
+    def astype(self, dtype, copy: bool = True) -> Self:
+        """Convert the stored data-type to something else
+
+        Parameters
+        ----------
+        dtype :
+            the new dtype for the sparse matrix
+        copy :
+            copy when needed, or do not copy when not needed.
+        """
+        old_dtype = np.dtype(self.dtype)
+        new_dtype = np.dtype(dtype)
+
+        if old_dtype == new_dtype:
+            if copy:
+                return self.copy()
+            return self
+
+        new = self.copy()
+        new._D = new._D.astype(dtype, copy=copy)
+        return new
+
     @classmethod
     def fromsp(cls, *sps, dtype=None):
         """Combine multiple single-dimension sparse matrices into one SparseCSR matrix

src/sisl/io/siesta/_help.py

@@ -98,129 +98,7 @@ def _csr_from(col_from, csr):
     csr.translate_columns(col_from, col_to)
 
 
-def _mat2dtype(M, dtype: np.dtype) -> None:
-    """Change the internal CSR matrix in `M` to a follow `dtype`"""
-
-    if M.dtype == dtype:
-        return M
-
-    spin = M.spin
-    csr = M._csr
-    shape = csr._D.shape
-
-    # Change details
-    old_dtype = np.dtype(M.dtype)
-    new_dtype = np.dtype(dtype)
-
-    def toc(D, re, im):
-        return (D[..., re] + 1j * D[..., im]).astype(dtype, copy=False)
-
-    if old_dtype.kind in ("f", "i"):
-        if new_dtype.kind in ("f", "i"):
-            # this is just simple casting
-            csr._D = csr._D.astype(dtype)
-        elif new_dtype.kind == "c":
-            # we need to *collect it
-            if spin.is_diagonal:
-                # this is just simple casting,
-                # each diagonal component has its own index
-                csr._D = csr._D.astype(dtype)
-            elif spin.is_noncolinear:
-                D = np.empty(shape[:-1] + (shape[-1] - 1,), dtype=dtype)
-                # These should be real only anyways!
-                D[..., [0, 1]] = csr._D[..., [0, 1]].real.astype(dtype)
-                D[..., 2] = toc(csr._D, 2, 3)
-                if D.shape[-1] > 4:
-                    D[..., 3:] = csr._D[..., 4:].astype(dtype)
-                csr._D = D
-            elif spin.is_spinorbit:
-                D = np.empty(shape[:-1] + (shape[-1] - 4,), dtype=dtype)
-                D[..., 0] = toc(csr._D, 0, 4)
-                D[..., 1] = toc(csr._D, 1, 5)
-                D[..., 2] = toc(csr._D, 2, 3)
-                D[..., 3] = toc(csr._D, 6, 7)
-                if D.shape[-1] > 4:
-                    D[..., 4:] = csr._D[..., 8:].astype(dtype)
-                csr._D = D
-            elif spin.is_nambu:
-                D = np.empty(shape[:-1] + (shape[-1] - 8,), dtype=dtype)
-                D[..., 0] = toc(csr._D, 0, 4)
-                D[..., 1] = toc(csr._D, 1, 5)
-                D[..., 2] = toc(csr._D, 2, 3)
-                D[..., 3] = toc(csr._D, 6, 7)
-                D[..., 4] = toc(csr._D, 8, 9)  # S
-                D[..., 5] = toc(csr._D, 10, 11)  # Tuu
-                D[..., 6] = toc(csr._D, 12, 13)  # Tdd
-                D[..., 7] = toc(csr._D, 14, 15)  # T0
-                if D.shape[-1] > 8:
-                    D[..., 8:] = csr._D[..., 16:].astype(dtype)
-                csr._D = D
-            else:
-                raise NotImplementedError
-        else:
-            raise NotImplementedError
-
-    elif old_dtype.kind == "c":
-        if new_dtype.kind == "c":
-            # this is just simple casting
-            csr._D = csr._D.astype(dtype)
-        elif new_dtype.kind in ("f", "i"):
-            # we need to *collect it
-            if spin.is_diagonal:
-                # this is just simple casting,
-                # each diagonal component has its own index
-                csr._D = csr._D.astype(dtype)
-            elif spin.is_noncolinear:
-                D = np.empty(shape[:-1] + (shape[-1] + 1,), dtype=dtype)
-                # These should be real only anyways!
-                D[..., [0, 1]] = csr._D[..., [0, 1]].real.astype(dtype)
-                D[..., 2] = csr._D[..., 2].real.astype(dtype)
-                D[..., 3] = csr._D[..., 2].imag.astype(dtype)
-                if D.shape[-1] > 4:
-                    D[..., 4:] = csr._D[..., 3:].real.astype(dtype)
-                csr._D = D
-            elif spin.is_spinorbit:
-                D = np.empty(shape[:-1] + (shape[-1] + 4,), dtype=dtype)
-                D[..., 0] = csr._D[..., 0].real.astype(dtype)
-                D[..., 1] = csr._D[..., 1].real.astype(dtype)
-                D[..., 2] = csr._D[..., 2].real.astype(dtype)
-                D[..., 3] = csr._D[..., 2].imag.astype(dtype)
-                D[..., 4] = csr._D[..., 0].imag.astype(dtype)
-                D[..., 5] = csr._D[..., 1].imag.astype(dtype)
-                D[..., 6] = csr._D[..., 3].real.astype(dtype)
-                D[..., 7] = csr._D[..., 3].imag.astype(dtype)
-                if D.shape[-1] > 8:
-                    D[..., 8:] = csr._D[..., 4:].real.astype(dtype)
-                csr._D = D
-            elif spin.is_nambu:
-                D = np.empty(shape[:-1] + (shape[-1] + 8,), dtype=dtype)
-                D[..., 0] = csr._D[..., 0].real.astype(dtype)
-                D[..., 1] = csr._D[..., 1].real.astype(dtype)
-                D[..., 2] = csr._D[..., 2].real.astype(dtype)
-                D[..., 3] = csr._D[..., 2].imag.astype(dtype)
-                D[..., 4] = csr._D[..., 0].imag.astype(dtype)
-                D[..., 5] = csr._D[..., 1].imag.astype(dtype)
-                D[..., 6] = csr._D[..., 3].real.astype(dtype)
-                D[..., 7] = csr._D[..., 3].imag.astype(dtype)
-                D[..., 8] = csr._D[..., 4].real.astype(dtype)  # S
-                D[..., 9] = csr._D[..., 4].imag.astype(dtype)
-                D[..., 10] = csr._D[..., 5].real.astype(dtype)  # Tuu
-                D[..., 11] = csr._D[..., 5].imag.astype(dtype)
-                D[..., 12] = csr._D[..., 6].real.astype(dtype)  # Tdd
-                D[..., 13] = csr._D[..., 6].imag.astype(dtype)
-                D[..., 14] = csr._D[..., 7].real.astype(dtype)  # T0
-                D[..., 15] = csr._D[..., 7].imag.astype(dtype)
-                if D.shape[-1] > 16:
-                    D[..., 16:] = csr._D[..., 8:].real.astype(dtype)
-                csr._D = D
-            else:
-                raise NotImplementedError
-        else:
-            raise NotImplementedError
-    M._reset()
-
-
-def _mat_siesta2sisl(M, dtype: Optional[np.dtype] = None) -> None:
+def _mat_siesta2sisl(M) -> None:
     """Conversion of Siesta spin matrices to sisl spin matrices
 
     The matrices from Siesta are given in a format adheering to the following
@@ -266,9 +144,6 @@ def _mat_siesta2sisl(M, dtype: Optional[np.dtype] = None) -> None:
     On top of this it depends on whether the data-type is complex
     or not.
     """
-    if dtype is None:
-        dtype = M.dtype
-
     spin = M.spin
 
     if spin.kind in (spin.NONCOLINEAR, spin.SPINORBIT, spin.NAMBU):
@@ -277,17 +152,9 @@ def _mat_siesta2sisl(M, dtype: Optional[np.dtype] = None) -> None:
         else:
             M._csr._D[:, 2] = M._csr._D[:, 2].conj()
 
-    _mat2dtype(M, dtype)
-
 
-def _mat_sisl2siesta(M, dtype: Optional[np.dtype] = None) -> None:
+def _mat_sisl2siesta(M) -> None:
     """Conversion of sisl to Siesta spin matrices"""
-    if dtype is None:
-        dtype = M.dtype
-
-    # convert to float
-    _mat2dtype(M, dtype)
-
     spin = M.spin
 
     if spin.kind in (spin.NONCOLINEAR, spin.SPINORBIT, spin.NAMBU):

src/sisl/io/siesta/binaries.py

@@ -420,7 +420,8 @@ def read_hamiltonian(self, geometry=None, **kwargs) -> Hamiltonian:
             H._csr._D[:, :spin] = dH[:, :] * _Ry2eV
             H._csr._D[:, spin] = dS[:]
 
-        _mat_siesta2sisl(H, dtype=kwargs.get("dtype"))
+        _mat_siesta2sisl(H)
+        H = H.astype(dtype=kwargs.get("dtype"), copy=False)
 
         # Convert to sisl supercell
         # equivalent as _csr_from_siesta with explicit isc from file
@@ -445,8 +446,7 @@ def write_hamiltonian(self, H, **kwargs):
         # we sort below, so no need to do it here
         # see onlysSileSiesta.read_overlap for .transpose()
         H = H.transpose(spin=False, sort=False)
-        csr = H._csr
-        if csr.nnz == 0:
+        if H._csr.nnz == 0:
             raise SileError(
                 f"{self!r}.write_hamiltonian cannot write "
                 "a zero element sparse matrix!"
@@ -455,9 +455,12 @@ def write_hamiltonian(self, H, **kwargs):
         sort = kwargs.get("sort", True)
 
         # Convert to siesta CSR
-        _csr_to_siesta(H.geometry, csr, diag=True)
-        csr.finalize(sort=sort)
-        _mat_sisl2siesta(H, dtype=np.float64)
+        _csr_to_siesta(H.geometry, H._csr, diag=True)
+        H.finalize(sort=sort)
+
+        H = H.astype(dtype=np.float64, copy=False)
+        _mat_sisl2siesta(H)
+        csr = H._csr
 
         # Extract the data to pass to the fortran routine
         cell = H.geometry.cell
@@ -569,7 +572,8 @@ def read_density_matrix(self, **kwargs) -> DensityMatrix:
         # DM file does not contain overlap matrix... so neglect it for now.
         DM._csr._D[:, spin] = 0.0
 
-        _mat_siesta2sisl(DM, dtype=kwargs.get("dtype"))
+        _mat_siesta2sisl(DM)
+        DM = DM.astype(dtype=kwargs.get("dtype"), copy=False)
 
         # Convert the supercells to sisl supercells
         if nsc[0] != 0 or geom.no_s >= col.max():
@@ -588,20 +592,21 @@ def read_density_matrix(self, **kwargs) -> DensityMatrix:
     def write_density_matrix(self, DM, **kwargs):
         """Writes the density matrix to a siesta.DM file"""
         DM = DM.transpose(spin=False, sort=False)
-        csr = DM._csr
         # This ensures that we don"t have any *empty* elements
-        if csr.nnz == 0:
+        if DM._csr.nnz == 0:
             raise SileError(
                 f"{self!r}.write_density_matrix cannot write "
                 "a zero element sparse matrix!"
             )
 
-        _csr_to_siesta(DM.geometry, csr)
+        _csr_to_siesta(DM.geometry, DM._csr)
         # We do not really need to sort this one, but we do for consistency
         # of the interface.
-        csr.finalize(sort=kwargs.get("sort", True))
+        DM.finalize(sort=kwargs.get("sort", True))
 
-        _mat_sisl2siesta(DM, dtype=np.float64)
+        DM = DM.astype(dtype=np.float64, copy=False)
+        _mat_sisl2siesta(DM)
+        csr = DM._csr
 
         # Get DM
         if DM.orthogonal:
@@ -679,7 +684,8 @@ def read_energy_density_matrix(self, **kwargs) -> EnergyDensityMatrix:
         # EDM file does not contain overlap matrix... so neglect it for now.
         EDM._csr._D[:, spin] = 0.0
 
-        _mat_siesta2sisl(EDM, dtype=kwargs.get("dtype"))
+        _mat_siesta2sisl(EDM)
+        EDM = EDM.astype(dtype=kwargs.get("dtype"), copy=False)
 
         # Convert the supercells to sisl supercells
         if nsc[0] != 0 or geom.no_s >= col.max():
@@ -737,8 +743,11 @@ def write_density_matrices(self, DM, EDM, Ef: float = 0.0, **kwargs):
         _csr_to_siesta(DM.geometry, EDM._csr)
         DM._csr.finalize(sort=sort)
         EDM._csr.finalize(sort=sort)
-        _mat_sisl2siesta(DM, dtype=np.float64)
-        _mat_sisl2siesta(EDM, dtype=np.float64)
+
+        DM = DM.astype(dtype=np.float64, copy=False)
+        _mat_sisl2siesta(DM)
+        EDM = EDM.astype(dtype=np.float64, copy=False)
+        _mat_sisl2siesta(EDM)
 
         # Ensure everything is correct
         if not (
@@ -1367,7 +1376,8 @@ def _r_hamiltonian_v0(self, **kwargs):
         H._csr._D[:, :spin] = dH[:, :] * _Ry2eV
         H._csr._D[:, spin] = dS[:]
 
-        _mat_siesta2sisl(H, dtype=kwargs.get("dtype"))
+        _mat_siesta2sisl(H)
+        H = H.astype(dtype=kwargs.get("dtype"), copy=False)
 
         # Convert the supercells to sisl supercells
         if no_s // no == np.prod(geom.nsc):
@@ -1412,7 +1422,8 @@ def _r_hamiltonian_v1(self, **kwargs):
         H._csr._D[:, :spin] = dH[:, :] * _Ry2eV
         H._csr._D[:, spin] = dS[:]
 
-        _mat_siesta2sisl(H, dtype=kwargs.get("dtype"))
+        _mat_siesta2sisl(H)
+        H = H.astype(dtype=kwargs.get("dtype"), copy=False)
 
         # Convert the supercells to sisl supercells
         _csr_from_sc_off(H.geometry, isc.T, H._csr)