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

emu_mps/__init__.py

@@ -37,4 +37,4 @@ __all__ = [
     "EntanglementEntropy",
 ]
 
-__version__ = "2.0.2"
+__version__ = "2.0.3"

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,6 +221,18 @@ 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],
@@ -510,6 +540,43 @@ class MPSBackendImpl:
 
                 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):
     """

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_tensor, inv_permutation, permute_string
+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_tuple",
 ]

emu_mps/optimatrix/permutations.py

@@ -1,6 +1,54 @@
 import torch
 
 
+def eye_permutation(n: int) -> torch.Tensor:
+    """
+    Returns toch.tensor([0, 1, 2, .., n-1])
+    """
+    return torch.arange(n)
+
+
+def permute_list(input_list: list, perm: torch.Tensor) -> list:
+    """
+    Permutes the input list according to the given permutation.
+    Parameters
+    -------
+    input_list :
+        A list to permute.
+    permutation :
+        A list of indices representing the new order.
+    Returns
+    -------
+        The permuted list.
+    Example
+    -------
+    >>> permute_list(['a','b','c'], torch.tensor([2, 0, 1]))
+    ['c', 'a', 'b']
+    """
+    return [input_list[i] for i in perm.tolist()]
+
+
+def permute_tuple(input_tuple: tuple, perm: torch.Tensor) -> tuple:
+    """
+    Permutes the input tuple according to the given permutation.
+    Parameters
+    -------
+    input_tuple :
+        A tuple to permute.
+    permutation :
+        A tuple of indices representing the new order.
+    Returns
+    -------
+        The permuted tuple.
+    Example
+    -------
+    >>> permute_tuple(('a','b','c'), torch.tensor([2, 0, 1]))
+    ('c', 'a', 'b')
+    """
+    lst_elem = list(input_tuple)
+    return tuple(permute_list(lst_elem, perm))
+
+
 def permute_string(input_str: str, perm: torch.Tensor) -> str:
     """
     Permutes the input string according to the given permutation.
@@ -18,8 +66,7 @@ def permute_string(input_str: str, perm: torch.Tensor) -> str:
     >>> permute_string("abc", torch.tensor([2, 0, 1]))
     'cab'
     """
-    char_list = list(input_str)
-    permuted = [char_list[i] for i in perm.tolist()]
+    permuted = permute_list(list(input_str), perm)
     return "".join(permuted)
 
 

emu_mps/tdvp.py

@@ -73,8 +73,16 @@ def apply_effective_Hamiltonian(
     # this order seems to be pretty balanced, but needs to be
     # revisited when use-cases are more well-known
     state = torch.tensordot(left_bath, state, 1)
-    state = torch.tensordot(state, ham, ([1, 2], [0, 2]))
-    state = torch.tensordot(state, right_bath, ([3, 1], [1, 2]))
+    state = state.permute(0, 3, 1, 2)
+    ham = ham.permute(0, 2, 1, 3)
+    state = state.reshape(state.shape[0], state.shape[1], -1).contiguous()
+    ham = ham.reshape(-1, ham.shape[2], ham.shape[3]).contiguous()
+    state = torch.tensordot(state, ham, 1)
+    state = state.permute(0, 2, 1, 3)
+    state = state.reshape(state.shape[0], state.shape[1], -1).contiguous()
+    right_bath = right_bath.permute(2, 1, 0)
+    right_bath = right_bath.reshape(-1, right_bath.shape[2]).contiguous()
+    state = torch.tensordot(state, right_bath, 1)
     return state
 
 

{emu_mps-2.0.2.dist-info → emu_mps-2.0.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: emu-mps
-Version: 2.0.2
+Version: 2.0.3
 Summary: Pasqal MPS based pulse emulator built on PyTorch
 Project-URL: Documentation, https://pasqal-io.github.io/emulators/
 Project-URL: Repository, https://github.com/pasqal-io/emulators
@@ -25,7 +25,7 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
 Requires-Python: >=3.10
-Requires-Dist: emu-base==2.0.2
+Requires-Dist: emu-base==2.0.3
 Description-Content-Type: text/markdown
 
 <div align="center">

emu_mps-2.0.3.dist-info/RECORD

@@ -0,0 +1,19 @@
+emu_mps/__init__.py,sha256=wGG0u1wHUcq_MuDjygyApshJY_ome-TZMVggIfHcgns,734
+emu_mps/algebra.py,sha256=ngPtTH-j2ZCBWoaJZXlkUyIlug7dY7Q92gzfnRlpPMA,5485
+emu_mps/custom_callback_implementations.py,sha256=CUs0kW3HRaPE7UeFNQOFbeWJMsz4hS2q4rgS57BBp-A,2411
+emu_mps/hamiltonian.py,sha256=gOPxNOBmk6jRPPjevERuCP_scGv0EKYeAJ0uxooihes,15622
+emu_mps/mpo.py,sha256=1ogQ25GZCwMzZ_m449oGHcYyDKrofBCr1eyzzrIPMhQ,8824
+emu_mps/mps.py,sha256=GIiWxctNmHARgf-PgQc6IHKNCe5HYSnbtlXI6Hc-0wI,20085
+emu_mps/mps_backend.py,sha256=bS83qFxvdoK-c12_1WaPw6O7xUc7vdWifZNHUzNP5sM,2091
+emu_mps/mps_backend_impl.py,sha256=8aPSn6ZGFf4DtTodzajX6WJtQHI9F8eMUN1AbrPzbgw,25611
+emu_mps/mps_config.py,sha256=JCK_frUVyMRJjI0NPU-dpCg-uO3JKeVzQMOLDkQ9XVk,8104
+emu_mps/noise.py,sha256=h4X2EFjoC_Ok0gZ8I9wN77RANXaVehTBbjkcbY_GAmY,784
+emu_mps/observables.py,sha256=7GQDH5kyaVNrwckk2f8ZJRV9Ca4jKhWWDsOCqYWsoEk,1349
+emu_mps/tdvp.py,sha256=0qTw9qhg0WbaAyBgeTpULHrNL0ytj80ZUb1P6GKD7Ww,6172
+emu_mps/utils.py,sha256=BqRJYAcXqprtZVJ0V_j954ON2bhTdtZiaTojsYyrWrg,8193
+emu_mps/optimatrix/__init__.py,sha256=fBXQ7-rgDro4hcaBijCGhx3J69W96qcw5_3mWc7tND4,364
+emu_mps/optimatrix/optimiser.py,sha256=k9suYmKLKlaZ7ozFuIqvXHyCBoCtGgkX1mpen9GOdOo,6977
+emu_mps/optimatrix/permutations.py,sha256=9DDMZtrGGZ01b9F3GkzHR3paX4qNtZiPoI7Z_Kia3Lc,3727
+emu_mps-2.0.3.dist-info/METADATA,sha256=ZU1USoiZ12em9lo4GgAWFiGE0oBT83yXitpm2fx8LvA,3505
+emu_mps-2.0.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+emu_mps-2.0.3.dist-info/RECORD,,

emu_mps-2.0.2.dist-info/RECORD

@@ -1,19 +0,0 @@
-emu_mps/__init__.py,sha256=cjDXI6YNKEZ_qT8eIT8z9HUowyx4Mh8fNTFciJpMm9k,734
-emu_mps/algebra.py,sha256=ngPtTH-j2ZCBWoaJZXlkUyIlug7dY7Q92gzfnRlpPMA,5485
-emu_mps/custom_callback_implementations.py,sha256=CUs0kW3HRaPE7UeFNQOFbeWJMsz4hS2q4rgS57BBp-A,2411
-emu_mps/hamiltonian.py,sha256=gOPxNOBmk6jRPPjevERuCP_scGv0EKYeAJ0uxooihes,15622
-emu_mps/mpo.py,sha256=1ogQ25GZCwMzZ_m449oGHcYyDKrofBCr1eyzzrIPMhQ,8824
-emu_mps/mps.py,sha256=GIiWxctNmHARgf-PgQc6IHKNCe5HYSnbtlXI6Hc-0wI,20085
-emu_mps/mps_backend.py,sha256=_3rlg6XeI4fHaDiJRfPL6pDkX9k48hAHKXd8fkvkOFs,2004
-emu_mps/mps_backend_impl.py,sha256=P8GYzM4C7r6Sgf9eSPQuwzQup4nwy3tPzlcQgGFkxpI,23132
-emu_mps/mps_config.py,sha256=89nu5OhNUX31eAeeYvvKnAHegpPVD43jH5Nmp635HVU,6984
-emu_mps/noise.py,sha256=h4X2EFjoC_Ok0gZ8I9wN77RANXaVehTBbjkcbY_GAmY,784
-emu_mps/observables.py,sha256=7GQDH5kyaVNrwckk2f8ZJRV9Ca4jKhWWDsOCqYWsoEk,1349
-emu_mps/tdvp.py,sha256=MMkYtVRNovNZiLO6BO1mmr6I6yYDYUJC3MGapfGlrQw,5756
-emu_mps/utils.py,sha256=BqRJYAcXqprtZVJ0V_j954ON2bhTdtZiaTojsYyrWrg,8193
-emu_mps/optimatrix/__init__.py,sha256=iFftb0AtiT_va153B_rZuY_PWWCo1SHnS8sRKjz8cS0,224
-emu_mps/optimatrix/optimiser.py,sha256=k9suYmKLKlaZ7ozFuIqvXHyCBoCtGgkX1mpen9GOdOo,6977
-emu_mps/optimatrix/permutations.py,sha256=kQKPF1Owe_f63zEhfaDadlxkqotIkAPK4kSzfqELIaM,2597
-emu_mps-2.0.2.dist-info/METADATA,sha256=p6h2IcYP-wiR9yrs6Q_9lzgz5DpDn9YvM8dpLQmgbhI,3505
-emu_mps-2.0.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-emu_mps-2.0.2.dist-info/RECORD,,