emu-mps 2.0.0__py3-none-any.whl → 2.0.2__py3-none-any.whl

emu_mps/mpo.py

@@ -83,7 +83,7 @@ class MPO(Operator[complex, torch.Tensor, MPS]):
             other.factors,
             config=other.config,
         )
-        return MPS(factors, orthogonality_center=0)
+        return MPS(factors, orthogonality_center=0, eigenstates=other.eigenstates)
 
     def __add__(self, other: MPO) -> MPO:
         """

emu_mps/mps.py

@@ -38,6 +38,7 @@ class MPS(State[complex, torch.Tensor]):
         orthogonality_center: Optional[int] = None,
         config: Optional[MPSConfig] = None,
         num_gpus_to_use: Optional[int] = DEVICE_COUNT,
+        eigenstates: Sequence[Eigenstate] = ("r", "g"),
     ):
         """
         This constructor creates a MPS directly from a list of tensors. It is for internal use only.
@@ -56,7 +57,7 @@ class MPS(State[complex, torch.Tensor]):
             num_gpus_to_use: distribute the factors over this many GPUs
                 0=all factors to cpu, None=keep the existing device assignment.
         """
-        self._eigenstates = ["0", "1"]
+        super().__init__(eigenstates=eigenstates)
         self.config = config if config is not None else MPSConfig()
         assert all(
             factors[i - 1].shape[2] == factors[i].shape[0] for i in range(1, len(factors))
@@ -88,6 +89,7 @@ class MPS(State[complex, torch.Tensor]):
         num_sites: int,
         config: Optional[MPSConfig] = None,
         num_gpus_to_use: int = DEVICE_COUNT,
+        eigenstates: Sequence[Eigenstate] = ["0", "1"],
     ) -> MPS:
         """
         Returns a MPS in ground state |000..0>.
@@ -111,6 +113,7 @@ class MPS(State[complex, torch.Tensor]):
             config=config,
             num_gpus_to_use=num_gpus_to_use,
             orthogonality_center=0,  # Arbitrary: every qubit is an orthogonality center.
+            eigenstates=eigenstates,
         )
 
     def __repr__(self) -> str:
@@ -307,6 +310,25 @@ class MPS(State[complex, torch.Tensor]):
         """
         return torch.abs(self.inner(other)) ** 2  # type: ignore[no-any-return]
 
+    def entanglement_entropy(self, mps_site: int) -> torch.Tensor:
+        """
+        Returns
+        the Von Neumann entanglement entropy of the state `mps` at the bond between sites b and b+1
+        S = -Σᵢsᵢ² log(sᵢ²)),
+        where sᵢ are the singular values at the chosen bond.
+        """
+        self.orthogonalize(mps_site)
+
+        # perform svd on reshaped matrix at site b
+        matrix = self.factors[mps_site].flatten(end_dim=1)
+        s = torch.linalg.svdvals(matrix)
+
+        s_e = torch.Tensor(torch.special.entr(s**2))
+        s_e = torch.sum(s_e)
+
+        self.orthogonalize(0)
+        return s_e.cpu()
+
     def get_memory_footprint(self) -> float:
         """
         Returns the number of MBs of memory occupied to store the state
@@ -331,12 +353,16 @@ class MPS(State[complex, torch.Tensor]):
             the summed state
         """
         assert isinstance(other, MPS), "Other state also needs to be an MPS"
+        assert (
+            self.eigenstates == other.eigenstates
+        ), f"`Other` state has basis {other.eigenstates} != {self.eigenstates}"
         new_tt = add_factors(self.factors, other.factors)
         result = MPS(
             new_tt,
             config=self.config,
             num_gpus_to_use=None,
             orthogonality_center=None,  # Orthogonality is lost.
+            eigenstates=self.eigenstates,
         )
         result.truncate()
         return result
@@ -362,6 +388,7 @@ class MPS(State[complex, torch.Tensor]):
             config=self.config,
             num_gpus_to_use=None,
             orthogonality_center=self.orthogonality_center,
+            eigenstates=self.eigenstates,
         )
 
     def __imul__(self, scalar: complex) -> MPS:
@@ -371,7 +398,7 @@ class MPS(State[complex, torch.Tensor]):
     def _from_state_amplitudes(
         cls,
         *,
-        eigenstates: Sequence[str],
+        eigenstates: Sequence[Eigenstate],
         amplitudes: Mapping[str, complex],
     ) -> tuple[MPS, Mapping[str, complex]]:
         """
@@ -385,8 +412,6 @@ class MPS(State[complex, torch.Tensor]):
         Returns:
             The resulting MPS representation of the state.s
         """
-
-        nqubits = len(next(iter(amplitudes.keys())))
         basis = set(eigenstates)
         if basis == {"r", "g"}:
             one = "r"
@@ -395,17 +420,20 @@ class MPS(State[complex, torch.Tensor]):
         else:
             raise ValueError("Unsupported basis provided")
 
+        nqubits = cls._validate_amplitudes(amplitudes, eigenstates)
+
         basis_0 = torch.tensor([[[1.0], [0.0]]], dtype=torch.complex128)  # ground state
         basis_1 = torch.tensor([[[0.0], [1.0]]], dtype=torch.complex128)  # excited state
 
         accum_mps = MPS(
             [torch.zeros((1, 2, 1), dtype=torch.complex128)] * nqubits,
             orthogonality_center=0,
+            eigenstates=eigenstates,
         )
 
         for state, amplitude in amplitudes.items():
             factors = [basis_1 if ch == one else basis_0 for ch in state]
-            accum_mps += amplitude * MPS(factors)
+            accum_mps += amplitude * MPS(factors, eigenstates=eigenstates)
         norm = accum_mps.norm()
         if not math.isclose(1.0, norm, rel_tol=1e-5, abs_tol=0.0):
             print("\nThe state is not normalized, normalizing it for you.")

emu_mps/mps_backend_impl.py

@@ -151,14 +151,14 @@ class MPSBackendImpl:
     def __getstate__(self) -> dict:
         for obs in self.config.observables:
             obs.apply = MethodType(type(obs).apply, obs)  # type: ignore[method-assign]
-        d = self.__dict__
+        d = self.__dict__.copy()
         # mypy thinks the method below is an attribute, because of the __getattr__ override
         d["results"] = self.results._to_abstract_repr()  # type: ignore[operator]
         return d
 
     def __setstate__(self, d: dict) -> None:
-        d["results"] = Results._from_abstract_repr(d["results"])  # type: ignore [attr-defined]
         self.__dict__ = d
+        self.results = Results._from_abstract_repr(d["results"])  # type: ignore [attr-defined]
         self.config.monkeypatch_observables()
 
     @staticmethod
@@ -208,6 +208,7 @@ class MPSBackendImpl:
             [f.clone().detach() for f in initial_state.factors],
             config=self.config,
             num_gpus_to_use=self.config.num_gpus_to_use,
+            eigenstates=initial_state.eigenstates,
         )
         initial_state.truncate()
         initial_state *= 1 / initial_state.norm()
@@ -444,7 +445,6 @@ class MPSBackendImpl:
         basename = self.autosave_file
         with open(basename.with_suffix(".new"), "wb") as file_handle:
             pickle.dump(self, file_handle)
-
         if basename.is_file():
             os.rename(basename, basename.with_suffix(".bak"))
 
@@ -497,13 +497,15 @@ class MPSBackendImpl:
                     )
                     full_state = MPS(
                         extended_mps_factors(
-                            normalized_state.factors, self.well_prepared_qubits_filter
+                            normalized_state.factors,
+                            self.well_prepared_qubits_filter,
                         ),
                         num_gpus_to_use=None,  # Keep the already assigned devices.
                         orthogonality_center=get_extended_site_index(
                             self.well_prepared_qubits_filter,
                             normalized_state.orthogonality_center,
                         ),
+                        eigenstates=normalized_state.eigenstates,
                     )
 
                 callback(self.config, fractional_time, full_state, full_mpo, self.results)
@@ -536,7 +538,7 @@ class NoisyMPSBackendImpl(MPSBackendImpl):
 
     def set_jump_threshold(self, bound: float) -> None:
         self.jump_threshold = random.uniform(0.0, bound)
-        self.norm_gap_before_jump = self.state.norm() ** 2 - self.jump_threshold
+        self.norm_gap_before_jump = self.state.norm().item() ** 2 - self.jump_threshold
 
     def init(self) -> None:
         self.init_lindblad_noise()
@@ -547,7 +549,7 @@ class NoisyMPSBackendImpl(MPSBackendImpl):
         previous_time = self.current_time
         self.current_time = self.target_time
         previous_norm_gap_before_jump = self.norm_gap_before_jump
-        self.norm_gap_before_jump = self.state.norm() ** 2 - self.jump_threshold
+        self.norm_gap_before_jump = self.state.norm().item() ** 2 - self.jump_threshold
 
         if self.root_finder is None:
             # No quantum jump location finding in progress
@@ -567,7 +569,7 @@ class NoisyMPSBackendImpl(MPSBackendImpl):
 
             return
 
-        self.norm_gap_before_jump = self.state.norm() ** 2 - self.jump_threshold
+        self.norm_gap_before_jump = self.state.norm().item() ** 2 - self.jump_threshold
         self.root_finder.provide_ordinate(self.current_time, self.norm_gap_before_jump)
 
         if self.root_finder.is_converged(tolerance=1):
@@ -593,7 +595,7 @@ class NoisyMPSBackendImpl(MPSBackendImpl):
         self.state *= 1 / self.state.norm()
         self.init_baths()
 
-        norm_after_normalizing = self.state.norm()
+        norm_after_normalizing = self.state.norm().item()
         assert math.isclose(norm_after_normalizing, 1, abs_tol=1e-10)
         self.set_jump_threshold(norm_after_normalizing**2)
 

emu_mps/observables.py

@@ -0,0 +1,40 @@
+from pulser.backend.state import State
+from pulser.backend.observable import Observable
+from emu_mps.mps import MPS
+from typing import Sequence, Any
+import torch
+
+
+class EntanglementEntropy(Observable):
+    """Entanglement Entropy subclass used only in emu_mps"""
+
+    def __init__(
+        self,
+        mps_site: int,
+        *,
+        evaluation_times: Sequence[float] | None = None,
+        tag_suffix: str | None = None,
+    ):
+        super().__init__(evaluation_times=evaluation_times, tag_suffix=tag_suffix)
+        self.mps_site = mps_site
+
+    @property
+    def _base_tag(self) -> str:
+        return "entanglement_entropy"
+
+    def _to_abstract_repr(self) -> dict[str, Any]:
+        repr = super()._to_abstract_repr()
+        repr["mps_site"] = self.mps_site
+        return repr
+
+    def apply(self, *, state: State, **kwargs: Any) -> torch.Tensor:
+        if not isinstance(state, MPS):
+            raise NotImplementedError(
+                "Entanglement entropy observable is only available for emu_mps emulator."
+            )
+        if not (0 <= self.mps_site <= len(state.factors) - 2):
+            raise ValueError(
+                f"Invalid bond index {self.mps_site}. "
+                f"Expected value in range 0 <= bond_index <= {len(state.factors)-2}."
+            )
+        return state.entanglement_entropy(self.mps_site)

emu_mps/optimatrix/__init__.py

@@ -1,9 +1,9 @@
 from .optimiser import minimize_bandwidth
-from .permutations import permute_list, permute_matrix, invert_permutation
+from .permutations import permute_tensor, inv_permutation, permute_string
 
 __all__ = [
     "minimize_bandwidth",
-    "permute_list",
-    "permute_matrix",
-    "invert_permutation",
+    "permute_string",
+    "permute_tensor",
+    "inv_permutation",
 ]

emu_mps/optimatrix/optimiser.py

@@ -1,21 +1,18 @@
+import itertools
+
+import torch
 from scipy.sparse import csr_matrix
 from scipy.sparse.csgraph import reverse_cuthill_mckee
-import numpy as np
-from emu_mps.optimatrix.permutations import permute_matrix, permute_list
-import itertools
 
+from emu_mps.optimatrix.permutations import permute_tensor
 
-def is_symmetric(mat: np.ndarray) -> bool:
-    if mat.shape[0] != mat.shape[1]:
-        return False
-    if not np.allclose(mat, mat.T, atol=1e-8):
-        return False
 
-    return True
+def is_symmetric(matrix: torch.Tensor, tol: float = 1e-8) -> bool:
+    return torch.allclose(matrix, matrix.T, atol=tol)
 
 
-def matrix_bandwidth(mat: np.ndarray) -> float:
-    """matrix_bandwidth(matrix: np.ndarray) -> float
+def matrix_bandwidth(mat: torch.Tensor) -> float:
+    """matrix_bandwidth(matrix: torch.tensor) -> torch.Tensor
 
     Computes bandwidth as max weighted distance between columns of
     a square matrix as `max (abs(matrix[i, j] * (j - i))`.
@@ -45,19 +42,27 @@ def matrix_bandwidth(mat: np.ndarray) -> float:
 
     Example:
     -------
-    >>> matrix = np.array([
-    ...    [  1, -17, 2.4],
-    ...    [  9,   1, -10],
-    ...    [-15,  20,   1],])
-    >>> matrix_bandwidth(matrix) # 30.0 because abs(-15 * (2-0) == 30)
+    >>> matrix = torch.tensor([
+    ...     [1.0, -17.0, 2.4],
+    ...     [9.0, 1.0, -10.0],
+    ...     [-15.0, 20.0, 1.0]
+    ... ])
+    >>> matrix_bandwidth(matrix)  # because abs(-15 * (0 - 2)) = 30.0
     30.0
     """
 
-    bandwidth = max(abs(el * (index[0] - index[1])) for index, el in np.ndenumerate(mat))
-    return float(bandwidth)
+    n = mat.shape[0]
 
+    i_arr = torch.arange(n).view(-1, 1)  # shape (n, 1)
+    j_arr = torch.arange(n).view(1, -1)  # shape (1, n)
 
-def minimize_bandwidth_above_threshold(mat: np.ndarray, threshold: float) -> np.ndarray:
+    weighted = torch.abs(mat * (j_arr - i_arr))
+    return torch.max(weighted).to(mat.dtype).item()
+
+
+def minimize_bandwidth_above_threshold(
+    mat: torch.Tensor, threshold: float
+) -> torch.Tensor:
     """
     minimize_bandwidth_above_threshold(matrix, trunc) -> permutation_lists
 
@@ -78,24 +83,25 @@ def minimize_bandwidth_above_threshold(mat: np.ndarray, threshold: float) -> np.
 
     Example:
     -------
-    >>> matrix = np.array([
-    ...    [1, 2, 3],
-    ...    [2, 5, 6],
-    ...    [3, 6, 9]])
+    >>> matrix = torch.tensor([
+    ...     [1, 2, 3],
+    ...     [2, 5, 6],
+    ...     [3, 6, 9]
+    ... ], dtype=torch.float32)
     >>> threshold = 3
     >>> minimize_bandwidth_above_threshold(matrix, threshold)
-    array([1, 2, 0], dtype=int32)
+    tensor([1, 2, 0], dtype=torch.int32)
     """
 
-    matrix_truncated = mat.copy()
-    matrix_truncated[mat < threshold] = 0
-    rcm_permutation = reverse_cuthill_mckee(
-        csr_matrix(matrix_truncated), symmetric_mode=True
-    )
-    return np.array(rcm_permutation)
+    m_trunc = mat.clone()
+    m_trunc[mat < threshold] = 0.0
 
+    matrix_np = csr_matrix(m_trunc.numpy())  # SciPy's RCM compatibility
+    rcm_perm = reverse_cuthill_mckee(matrix_np, symmetric_mode=True)
+    return torch.from_numpy(rcm_perm.copy())  # translation requires copy
 
-def minimize_bandwidth_global(mat: np.ndarray) -> list[int]:
+
+def minimize_bandwidth_global(mat: torch.Tensor) -> torch.Tensor:
     """
     minimize_bandwidth_global(matrix) -> list
 
@@ -111,74 +117,78 @@ def minimize_bandwidth_global(mat: np.ndarray) -> list[int]:
     -------
         permutation order that minimizes matrix bandwidth
 
-    Example:
+    Example
     -------
-    >>> matrix = np.array([
-    ...    [1, 2, 3],
-    ...    [2, 5, 6],
-    ...    [3, 6, 9]])
+    >>> matrix = torch.tensor([
+    ...     [1, 2, 3],
+    ...     [2, 5, 6],
+    ...     [3, 6, 9]
+    ... ], dtype=torch.float32)
     >>> minimize_bandwidth_global(matrix)
-    [2, 1, 0]
+    tensor([2, 1, 0], dtype=torch.int32)
     """
-    mat_amplitude = np.max(np.abs(mat))
-    # Search from 1.0 to 0.1 doesn't change result
+    mat_amplitude = torch.max(torch.abs(mat))
+
     permutations = (
-        minimize_bandwidth_above_threshold(mat, trunc * mat_amplitude)
-        for trunc in np.arange(start=0.1, stop=1.0, step=0.01)
+        minimize_bandwidth_above_threshold(mat, trunc.item() * mat_amplitude)
+        for trunc in torch.arange(0.1, 1.0, 0.01)
     )
 
     opt_permutation = min(
-        permutations, key=lambda perm: matrix_bandwidth(permute_matrix(mat, list(perm)))
+        permutations, key=lambda perm: matrix_bandwidth(permute_tensor(mat, perm))
     )
-    return list(opt_permutation)  # opt_permutation is np.ndarray
+
+    return opt_permutation
 
 
 def minimize_bandwidth_impl(
-    matrix: np.ndarray, initial_perm: list[int]
-) -> tuple[list[int], float]:
+    matrix: torch.Tensor, initial_perm: torch.Tensor
+) -> tuple[torch.Tensor, float]:
     """
-    minimize_bandwidth_impl(matrix, initial_perm) -> list
+    minimize_bandwidth_impl(matrix, initial_perm) -> (optimal_perm, bandwidth)
 
     Applies initial_perm to a matrix and
-    finds the permutation list for a symmetric matrix that iteratively minimizes matrix bandwidth.
+    finds the permutation list for a symmetric matrix
+    that iteratively minimizes matrix bandwidth.
 
     Parameters
     -------
     matrix :
         symmetric square matrix
-    initial_perm: list of integers
+    initial_perm: torch list of integers
 
 
     Returns
     -------
-        permutation order that minimizes matrix bandwidth
+        optimal permutation and optimal matrix bandwidth
 
     Example:
     -------
     Periodic 1D chain
-    >>> matrix = np.array([
+    >>> matrix = torch.tensor([
     ...    [0, 1, 0, 0, 1],
     ...    [1, 0, 1, 0, 0],
     ...    [0, 1, 0, 1, 0],
     ...    [0, 0, 1, 0, 1],
-    ...    [1, 0, 0, 1, 0]])
-    >>> id_perm = list(range(matrix.shape[0]))
+    ...    [1, 0, 0, 1, 0]], dtype=torch.float32)
+    >>> id_perm = torch.arange(matrix.shape[0])
     >>> minimize_bandwidth_impl(matrix, id_perm) # [3, 2, 4, 1, 0] does zig-zag
-    ([3, 2, 4, 1, 0], 2.0)
+    (tensor([3, 2, 4, 1, 0]), 2.0)
 
     Simple 1D chain. Cannot be optimised further
-    >>> matrix = np.array([
+    >>> matrix = torch.tensor([
     ...    [0, 1, 0, 0, 0],
     ...    [1, 0, 1, 0, 0],
     ...    [0, 1, 0, 1, 0],
     ...    [0, 0, 1, 0, 1],
-    ...    [0, 0, 0, 1, 0]])
-    >>> id_perm = list(range(matrix.shape[0]))
+    ...    [0, 0, 0, 1, 0]], dtype=torch.float32)
+    >>> id_perm = torch.arange(matrix.shape[0])
     >>> minimize_bandwidth_impl(matrix, id_perm)
-    ([0, 1, 2, 3, 4], 1.0)
+    (tensor([0, 1, 2, 3, 4]), 1.0)
     """
-    if initial_perm != list(range(matrix.shape[0])):
-        matrix = permute_matrix(matrix, initial_perm)
+    L = matrix.shape[0]
+    if not torch.equal(initial_perm, torch.arange(L)):
+        matrix = permute_tensor(matrix, initial_perm)
     bandwidth = matrix_bandwidth(matrix)
     acc_permutation = initial_perm
 
@@ -191,28 +201,28 @@ def minimize_bandwidth_impl(
             )
 
         optimal_perm = minimize_bandwidth_global(matrix)
-        test_mat = permute_matrix(matrix, optimal_perm)
+        test_mat = permute_tensor(matrix, optimal_perm)
         new_bandwidth = matrix_bandwidth(test_mat)
 
         if bandwidth <= new_bandwidth:
             break
 
         matrix = test_mat
-        acc_permutation = permute_list(acc_permutation, optimal_perm)
+        acc_permutation = permute_tensor(acc_permutation, optimal_perm)
         bandwidth = new_bandwidth
 
     return acc_permutation, bandwidth
 
 
-def minimize_bandwidth(input_matrix: np.ndarray, samples: int = 100) -> list[int]:
+def minimize_bandwidth(input_matrix: torch.Tensor, samples: int = 100) -> torch.Tensor:
     assert is_symmetric(input_matrix), "Input matrix is not symmetric"
-    input_mat = abs(input_matrix)
+    input_mat = torch.abs(input_matrix)
     # We are interested in strength of the interaction, not sign
 
     L = input_mat.shape[0]
-    rnd_permutations: itertools.chain[list[int]] = itertools.chain(
-        [list(range(L))],  # First element is always the identity list
-        (np.random.permutation(L).tolist() for _ in range(samples)),  # type: ignore[misc]
+    rnd_permutations = itertools.chain(
+        [torch.arange(L)],  # identity permutation
+        [torch.randperm(L) for _ in range(samples)],  # list of random permutations
     )
 
     opt_permutations_and_opt_bandwidth = (

emu_mps/optimatrix/permutations.py

@@ -1,36 +1,31 @@
-import numpy as np
+import torch
 
 
-def permute_list(input_list: list, permutation: list[int]) -> list:
+def permute_string(input_str: str, perm: torch.Tensor) -> str:
     """
-    Permutes the input list according to the given permutation.
-
+    Permutes the input string according to the given permutation.
     Parameters
     -------
-    input_list :
-        A list to permute.
+    input_string :
+        A string to permute.
     permutation :
         A list of indices representing the new order.
-
     Returns
     -------
-        The permuted list.
-
+        The permuted string.
     Example
     -------
-    >>> permute_list(['a', 'b', 'c'], [2, 0, 1])
-    ['c', 'a', 'b']
+    >>> permute_string("abc", torch.tensor([2, 0, 1]))
+    'cab'
     """
+    char_list = list(input_str)
+    permuted = [char_list[i] for i in perm.tolist()]
+    return "".join(permuted)
 
-    permuted_list = [None] * len(input_list)
-    for i, p in enumerate(permutation):
-        permuted_list[i] = input_list[p]
-    return permuted_list
 
-
-def invert_permutation(permutation: list[int]) -> list[int]:
+def inv_permutation(permutation: torch.Tensor) -> torch.Tensor:
     """
-    invert_permutation(permutation) -> inv_permutation
+    inv_permutation(permutation) -> inverted_perm
 
     Inverts the input permutation list.
 
@@ -45,47 +40,60 @@ def invert_permutation(permutation: list[int]) -> list[int]:
 
     Example:
     -------
-    >>> invert_permutation([2, 0, 1])
-    [1, 2, 0]
+    >>> inv_permutation(torch.tensor([2, 0, 1]))
+    tensor([1, 2, 0])
     """
-
-    inv_perm = np.empty_like(permutation)
-    inv_perm[permutation] = np.arange(len(permutation))
-    return list(inv_perm)
+    inv_perm = torch.empty_like(permutation)
+    inv_perm[permutation] = torch.arange(len(permutation))
+    return inv_perm
 
 
-def permute_matrix(mat: np.ndarray, permutation: list[int]) -> np.ndarray:
+def permute_tensor(tensor: torch.Tensor, perm: torch.Tensor) -> torch.Tensor:
     """
-    permute_matrix(matrix, permutation_list) -> permuted_matrix
-
-    Simultaneously permutes columns and rows according to a permutation list.
+    Permute a 1D or square 2D torch tensor using the given permutation indices.
+    For 1D tensors, applies the permutation to the elements.
+    For 2D square tensors, applies the same permutation to both rows and columns.
 
     Parameters
-    -------
-    matrix :
-        square matrix nxn
-    permutation :
-        permutation list
+    ----------
+    tensor : torch.Tensor
+        A 1D or 2D square tensor to be permuted.
+    perm : torch.Tensor
+        A 1D tensor of indices specifying the permutation order.
 
     Returns
     -------
-        matrix with permuted columns and rows
-
-    Example:
-    -------
-    >>> matrix = np.array([
-    ...    [1, 2, 3],
-    ...    [4, 5, 6],
-    ...    [7, 8, 9]])
-    >>> permutation = [1, 0, 2]
-    >>> permute_matrix(matrix, permutation)
-    array([[5, 4, 6],
-           [2, 1, 3],
-           [8, 7, 9]])
+    torch.Tensor
+        A new tensor with elements (1D) or rows and columns (2D) permuted according to `perm`.
+
+    Raises
+    ------
+    ValueError
+        If tensor is not 1D or square 2D.
+
+    Examples
+    --------
+    >>> vector = torch.tensor([10, 20, 30])
+    >>> perm = torch.tensor([2, 0, 1])
+    >>> permute_tensor(vector, perm)
+    tensor([30, 10, 20])
+
+    >>> matrix = torch.tensor([
+    ...     [1, 2, 3],
+    ...     [4, 5, 6],
+    ...     [7, 8, 9]])
+    >>> perm = torch.tensor([1, 0, 2])
+    >>> permute_tensor(matrix, perm)
+    tensor([[5, 4, 6],
+            [2, 1, 3],
+            [8, 7, 9]])
     """
-
-    perm = np.array(permutation)
-    return mat[perm, :][:, perm]
+    if tensor.ndim == 1:
+        return tensor[perm]
+    elif tensor.ndim == 2 and tensor.shape[0] == tensor.shape[1]:
+        return tensor[perm][:, perm]
+    else:
+        raise ValueError("Only 1D tensors or square 2D tensors are supported.")
 
 
 if __name__ == "__main__":

emu_mps/tdvp.py

@@ -78,7 +78,7 @@ def apply_effective_Hamiltonian(
     return state
 
 
-_TIME_CONVERSION_COEFF = 0.001  # Omega and delta are given in rad/ms, dt in ns
+_TIME_CONVERSION_COEFF = 0.001  # Omega and delta are given in rad/μs, dt in ns
 
 
 def evolve_pair(