emu-mps 2.0.1__py3-none-any.whl → 2.0.3__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:
@@ -350,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
@@ -381,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:
@@ -390,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]]:
         """
@@ -404,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"
@@ -414,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.py

@@ -55,7 +55,9 @@ class MPSBackend(EmulatorBackend):
         impl = create_impl(self._sequence, self._config)
         impl.init()  # This is separate from the constructor for testing purposes.
 
-        return self._run(impl)
+        results = self._run(impl)
+
+        return impl.permute_results(results, self._config.optimize_qubit_ordering)
 
     @staticmethod
     def _run(impl: MPSBackendImpl) -> Results:

emu_mps/mps_backend_impl.py

@@ -1,29 +1,34 @@
 import math
+import os
 import pathlib
+import pickle
 import random
+import time
+import typing
 import uuid
 
+from collections import Counter
+from enum import Enum, auto
 from resource import RUSAGE_SELF, getrusage
-from typing import Optional, Any
-import typing
-import pickle
-import os
+from types import MethodType
+from typing import Any, Optional
+
 import torch
-import time
 from pulser import Sequence
-from types import MethodType
+from pulser.backend import EmulationConfig, Observable, Results, State
 
-from pulser.backend import State, Observable, EmulationConfig, Results
-from emu_base import PulserData, DEVICE_COUNT
+from emu_base import DEVICE_COUNT, PulserData
 from emu_base.math.brents_root_finding import BrentsRootFinder
+
 from emu_mps.hamiltonian import make_H, update_H
 from emu_mps.mpo import MPO
 from emu_mps.mps import MPS
 from emu_mps.mps_config import MPSConfig
 from emu_mps.noise import compute_noise_from_lindbladians, pick_well_prepared_qubits
+import emu_mps.optimatrix as optimat
 from emu_mps.tdvp import (
-    evolve_single,
     evolve_pair,
+    evolve_single,
     new_right_bath,
     right_baths,
 )
@@ -33,7 +38,6 @@ from emu_mps.utils import (
     get_extended_site_index,
     new_left_bath,
 )
-from enum import Enum, auto
 
 
 class Statistics(Observable):
@@ -118,8 +122,17 @@ class MPSBackendImpl:
         self.timestep_count: int = self.omega.shape[0]
         self.has_lindblad_noise = pulser_data.has_lindblad_noise
         self.lindblad_noise = torch.zeros(2, 2, dtype=torch.complex128)
-        self.full_interaction_matrix = pulser_data.full_interaction_matrix
-        self.masked_interaction_matrix = pulser_data.masked_interaction_matrix
+        self.qubit_permutation = (
+            optimat.minimize_bandwidth(pulser_data.full_interaction_matrix)
+            if self.config.optimize_qubit_ordering
+            else optimat.eye_permutation(self.qubit_count)
+        )
+        self.full_interaction_matrix = optimat.permute_tensor(
+            pulser_data.full_interaction_matrix, self.qubit_permutation
+        )
+        self.masked_interaction_matrix = optimat.permute_tensor(
+            pulser_data.masked_interaction_matrix, self.qubit_permutation
+        )
         self.hamiltonian_type = pulser_data.hamiltonian_type
         self.slm_end_time = pulser_data.slm_end_time
         self.is_masked = self.slm_end_time > 0.0
@@ -128,7 +141,12 @@ class MPSBackendImpl:
         self.swipe_direction = SwipeDirection.LEFT_TO_RIGHT
         self.tdvp_index = 0
         self.timestep_index = 0
-        self.results = Results(atom_order=(), total_duration=self.target_times[-1])
+        self.results = Results(
+            atom_order=optimat.permute_tuple(
+                pulser_data.qubit_ids, self.qubit_permutation
+            ),
+            total_duration=self.target_times[-1],
+        )
         self.statistics = Statistics(
             evaluation_times=[t / self.target_times[-1] for t in self.target_times],
             data=[],
@@ -203,11 +221,24 @@ class MPSBackendImpl:
             )
 
         assert isinstance(initial_state, MPS)
+        if not torch.equal(
+            self.qubit_permutation, optimat.eye_permutation(self.qubit_count)
+        ):
+            # permute the initial state to match with permuted Hamiltonian
+            abstr_repr = initial_state._to_abstract_repr()
+            eigs = abstr_repr["eigenstates"]
+            ampl = {
+                optimat.permute_string(bstr, self.qubit_permutation): amp
+                for bstr, amp in abstr_repr["amplitudes"].items()
+            }
+            initial_state = MPS.from_state_amplitudes(eigenstates=eigs, amplitudes=ampl)
+
         initial_state = MPS(
             # Deep copy of every tensor of the initial state.
             [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()
@@ -496,17 +527,56 @@ 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)
 
+    def permute_results(self, results: Results, permute: bool) -> Results:
+        if permute:
+            inv_perm = optimat.inv_permutation(self.qubit_permutation)
+            permute_bitstrings(results, inv_perm)
+            permute_occupations_and_correlations(results, inv_perm)
+            permute_atom_order(results, inv_perm)
+        return results
+
+
+def permute_bitstrings(results: Results, perm: torch.Tensor) -> None:
+    if "bitstrings" not in results.get_result_tags():
+        return
+    uuid_bs = results._find_uuid("bitstrings")
+
+    results._results[uuid_bs] = [
+        Counter({optimat.permute_string(bstr, perm): c for bstr, c in bs_counter.items()})
+        for bs_counter in results._results[uuid_bs]
+    ]
+
+
+def permute_occupations_and_correlations(results: Results, perm: torch.Tensor) -> None:
+    for corr in ["occupation", "correlation_matrix"]:
+        if corr not in results.get_result_tags():
+            return
+
+        uuid_corr = results._find_uuid(corr)
+        corrs = results._results[uuid_corr]
+        results._results[uuid_corr] = [
+            optimat.permute_tensor(corr, perm) for corr in corrs
+        ]
+
+
+def permute_atom_order(results: Results, perm: torch.Tensor) -> None:
+    at_ord = list(results.atom_order)
+    at_ord = optimat.permute_list(at_ord, perm)
+    results.atom_order = tuple(at_ord)
+
 
 class NoisyMPSBackendImpl(MPSBackendImpl):
     """
@@ -535,7 +605,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()
@@ -546,7 +616,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
@@ -566,7 +636,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):
@@ -592,7 +662,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/mps_config.py

@@ -69,6 +69,7 @@ class MPSConfig(EmulationConfig):
         max_krylov_dim: int = 100,
         extra_krylov_tolerance: float = 1e-3,
         num_gpus_to_use: int = DEVICE_COUNT,
+        optimize_qubit_ordering: bool = False,
         interaction_cutoff: float = 0.0,
         log_level: int = logging.INFO,
         log_file: pathlib.Path | None = None,
@@ -84,6 +85,7 @@ class MPSConfig(EmulationConfig):
             max_krylov_dim=max_krylov_dim,
             extra_krylov_tolerance=extra_krylov_tolerance,
             num_gpus_to_use=num_gpus_to_use,
+            optimize_qubit_ordering=optimize_qubit_ordering,
             interaction_cutoff=interaction_cutoff,
             log_level=log_level,
             log_file=log_file,
@@ -91,6 +93,8 @@ class MPSConfig(EmulationConfig):
             autosave_dt=autosave_dt,
             **kwargs,
         )
+        if self.optimize_qubit_ordering:
+            self.check_permutable_observables()
 
         if "doppler" in self.noise_model.noise_types:
             raise NotImplementedError("Unsupported noise type: doppler")
@@ -136,6 +140,7 @@ class MPSConfig(EmulationConfig):
             "max_krylov_dim",
             "extra_krylov_tolerance",
             "num_gpus_to_use",
+            "optimize_qubit_ordering",
             "interaction_cutoff",
             "log_level",
             "log_file",
@@ -183,3 +188,27 @@ class MPSConfig(EmulationConfig):
                 filemode="w",
                 force=True,
             )
+
+    def check_permutable_observables(self) -> None:
+        allowed_permutable_obs = set(
+            [
+                "bitstrings",
+                "occupation",
+                "correlation_matrix",
+                "statistics",
+                "energy",
+                "energy_variance",
+                "energy_second_moment",
+            ]
+        )
+
+        actual_obs = set([obs._base_tag for obs in self.observables])
+        not_allowed = actual_obs.difference(allowed_permutable_obs)
+        if not_allowed:
+            raise ValueError(
+                f"emu-mp allows only {allowed_permutable_obs} observables with"
+                " `optimize_qubit_ordering = True`."
+                f" you provided unsupported {not_allowed}"
+                " To use other observables, please set"
+                " `optimize_qubit_ordering = False` in `MPSConfig()`."
+            )

emu_mps/optimatrix/__init__.py

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

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 = (