feat: #993 added ipr_tensor

changes/993.feat.rst

@@ -0,0 +1,6 @@
+implemented a tensor-IPR method
+
+It retains the complex part as it is not
+guaranteed to be a real quantity. Note, depending
+on the size of the states, this can yield a really
+big matrix.

src/sisl/physics/state.py

@@ -572,6 +572,46 @@ def ipr(self, q: int = 2):
         # abs2 is already having the exponent 2
         return (state_abs2**q).sum(-1) / state_abs2.sum(-1) ** q
 
+    def ipr_tensor(self, tensor: str = "abcd") -> np.ndarray:
+        r"""Calculate the full tensor IPR quantity.
+
+        This can be written as:
+
+        .. math::
+
+           I^T_{\alpha\beta\gamma\delta} \propto
+           \sum_i \psi_{\alpha,i}^* \psi_{\beta,i}^* \psi_{\gamma,i} \psi_{\delta,i}
+
+        Parameters
+        ----------
+        tensor:
+            the tensor product.
+
+            If you want to extract the 2-index tensor product, use ``tensor=abab``
+            (or ``tensor=ab``).
+        """
+        state_c = self.state.conj()
+        state = self.state
+
+        if "i" in tensor:
+            raise ValueError(
+                f"{self.__class__.__name__}.ipr_tensor only accepts tensor descriptors "
+                "different from 'i'."
+            )
+
+        # Normalize the tensor description
+        ntensor = len(tensor)
+        if ntensor == 2:
+            tensor = tensor * 2
+            ntensor = len(tensor)
+        assert ntensor % 2 == 0, "Tensor must be divisible by 2."
+        args = [state_c] * (ntensor // 2)
+        args += [state] * (ntensor // 2)
+        final_path = "".join(sorted(set(tensor), key=lambda x: tensor.index(x)))
+        # Now sort according to old index
+        path = "i,".join(tensor) + "i->" + final_path
+        return np.einsum(path, *args, optimize=True)
+
     def normalize(self):
         r"""Return a normalized state where each state has :math:`|\psi|^2=1`
 

src/sisl/physics/tests/test_state.py

@@ -300,6 +300,29 @@ def test_state_ipr():
     assert ipr.shape == (10,)
 
 
+def test_state_ipr_tensor():
+    state = State(ortho_matrix(10))
+    ipr = state.ipr_tensor()
+    assert ipr.shape == (10, 10, 10, 10)
+
+
+def test_state_ipr_tensor_2():
+    N = 10
+    state = State(ortho_matrix(N))
+    ipr2 = state.ipr_tensor("ab")
+    assert ipr2.shape == (N, N)
+    ipr4 = state.ipr_tensor("abab")
+    assert np.allclose(ipr2, ipr4)
+    del ipr4
+    ipr = state.ipr_tensor()
+
+    # Extract the diagonals
+    idx = np.arange(N * N)
+    idx1, idx2 = np.unravel_index(idx, shape=(N, N))
+    ipr = np.reshape(ipr[idx1, idx2, idx1, idx2], (N, N))
+    assert np.allclose(ipr2, ipr)
+
+
 def test_state_align_norm():
     state = ortho_matrix(10)
     state1 = State(state)