emu-mps 1.2.5__py3-none-any.whl → 1.2.7__py3-none-any.whl

emu_mps/__init__.py

@@ -35,4 +35,4 @@ __all__ = [
     "SecondMomentOfEnergy",
 ]
 
-__version__ = "1.2.5"
+__version__ = "1.2.7"

emu_mps/mpo.py

@@ -6,10 +6,9 @@ import torch
 
 from emu_mps.algebra import add_factors, scale_factors, zip_right
 from emu_base.base_classes.operator import FullOp, QuditOp
-from emu_base import Operator, State
+from emu_base import Operator, State, DEVICE_COUNT
 from emu_mps.mps import MPS
 from emu_mps.utils import new_left_bath, assign_devices
-from emu_mps.constants import DEVICE_COUNT
 
 
 def _validate_operator_targets(operations: FullOp, nqubits: int) -> None:

emu_mps/mps.py

@@ -6,7 +6,7 @@ from typing import Any, List, Optional, Iterable
 
 import torch
 
-from emu_base import State
+from emu_base import State, DEVICE_COUNT
 from emu_mps import MPSConfig
 from emu_mps.algebra import add_factors, scale_factors
 from emu_mps.utils import (
@@ -16,7 +16,6 @@ from emu_mps.utils import (
     tensor_trace,
     n_operator,
 )
-from emu_mps.constants import DEVICE_COUNT
 
 
 class MPS(State):
@@ -190,11 +189,58 @@ class MPS(State):
         """
         self.orthogonalize(0)
 
-        num_qubits = len(self.factors)
-        rnd_matrix = torch.rand(num_shots, num_qubits)
-        bitstrings = Counter(
-            self._sample_implementation(rnd_matrix[x, :]) for x in range(num_shots)
-        )
+        rnd_matrix = torch.rand(num_shots, self.num_sites).to(self.factors[0].device)
+
+        bitstrings: Counter[str] = Counter()
+
+        # Shots are performed in batches.
+        # Larger max_batch_size is faster but uses more memory.
+        max_batch_size = 32
+
+        shots_done = 0
+        while shots_done < num_shots:
+            batch_size = min(max_batch_size, num_shots - shots_done)
+            batched_accumulator = torch.ones(
+                batch_size, 1, dtype=torch.complex128, device=self.factors[0].device
+            )
+
+            batch_outcomes = torch.empty(batch_size, self.num_sites, dtype=torch.bool)
+
+            for qubit, factor in enumerate(self.factors):
+                batched_accumulator = torch.tensordot(
+                    batched_accumulator.to(factor.device), factor, dims=1
+                )
+
+                # Probability of measuring qubit == 0 for each shot in the batch
+                probas = (
+                    torch.linalg.vector_norm(batched_accumulator[:, 0, :], dim=1) ** 2
+                )
+
+                outcomes = (
+                    rnd_matrix[shots_done : shots_done + batch_size, qubit].to(
+                        factor.device
+                    )
+                    > probas
+                )
+                batch_outcomes[:, qubit] = outcomes
+
+                # Batch collapse qubit
+                tmp = torch.stack((~outcomes, outcomes), dim=1).to(dtype=torch.complex128)
+
+                batched_accumulator = (
+                    torch.tensordot(batched_accumulator, tmp, dims=([1], [1]))
+                    .diagonal(dim1=0, dim2=2)
+                    .transpose(1, 0)
+                )
+                batched_accumulator /= torch.sqrt(
+                    (~outcomes) * probas + outcomes * (1 - probas)
+                ).unsqueeze(1)
+
+            shots_done += batch_size
+
+            for outcome in batch_outcomes:
+                bitstrings.update(["".join("0" if x == 0 else "1" for x in outcome)])
+
         if p_false_neg > 0 or p_false_pos > 0:
             bitstrings = apply_measurement_errors(
                 bitstrings,
@@ -215,42 +261,6 @@ class MPS(State):
             torch.linalg.norm(self.factors[orthogonality_center].to("cpu")).item()
         )
 
-    def _sample_implementation(self, rnd_vector: torch.Tensor) -> str:
-        """
-        Samples this MPS once, returning the resulting bitstring.
-        """
-        assert rnd_vector.shape == (self.num_sites,)
-        assert self.orthogonality_center == 0
-
-        num_qubits = len(self.factors)
-
-        bitstring = ""
-        acc_mps_j: torch.Tensor = self.factors[0]
-
-        for qubit in range(num_qubits):
-            # comp_basis is a projector: 0 is for ket |0> and 1 for ket |1>
-            comp_basis = 0  # check if the qubit is in |0>
-            # Measure the qubit j by applying the projector onto nth comp basis state
-            tensorj_projected_n = acc_mps_j[:, comp_basis, :]
-            probability_n = (tensorj_projected_n.norm() ** 2).item()
-
-            if rnd_vector[qubit] > probability_n:
-                # the qubit is in |1>
-                comp_basis = 1
-                tensorj_projected_n = acc_mps_j[:, comp_basis, :]
-                probability_n = 1 - probability_n
-
-            bitstring += str(comp_basis)
-            if qubit < num_qubits - 1:
-                acc_mps_j = torch.tensordot(
-                    tensorj_projected_n.to(device=self.factors[qubit + 1].device),
-                    self.factors[qubit + 1],
-                    dims=1,
-                )
-                acc_mps_j /= math.sqrt(probability_n)
-
-        return bitstring
-
     def inner(self, other: State) -> float | complex:
         """
         Compute the inner product between this state and other.

emu_mps/mps_backend_impl.py

@@ -11,9 +11,8 @@ import torch
 import time
 from pulser import Sequence
 
-from emu_base import Results, State, PulserData
+from emu_base import Results, State, PulserData, DEVICE_COUNT
 from emu_base.math.brents_root_finding import BrentsRootFinder
-from emu_mps.constants import DEVICE_COUNT
 from emu_mps.hamiltonian import make_H, update_H
 from emu_mps.mpo import MPO
 from emu_mps.mps import MPS
@@ -55,7 +54,8 @@ class MPSBackendImpl:
 
     def __init__(self, mps_config: MPSConfig, pulser_data: PulserData):
         self.config = mps_config
-        self.target_time = float(self.config.dt)
+        self.target_times = pulser_data.target_times
+        self.target_time = self.target_times[1]
         self.pulser_data = pulser_data
         self.qubit_count = pulser_data.qubit_count
         assert self.qubit_count >= 2
@@ -95,18 +95,16 @@ class MPSBackendImpl:
         return pathlib.Path(os.getcwd()) / (autosave_prefix + str(uuid.uuid1()) + ".dat")
 
     def init_dark_qubits(self) -> None:
-        has_state_preparation_error: bool = (
+        # has_state_preparation_error
+        if (
             self.config.noise_model is not None
             and self.config.noise_model.state_prep_error > 0.0
-        )
-
-        self.well_prepared_qubits_filter = (
-            pick_well_prepared_qubits(
+        ):
+            self.well_prepared_qubits_filter = pick_well_prepared_qubits(
                 self.config.noise_model.state_prep_error, self.qubit_count
             )
-            if has_state_preparation_error
-            else None
-        )
+        else:
+            self.well_prepared_qubits_filter = None
 
         if self.well_prepared_qubits_filter is not None:
             self.qubit_count = sum(1 for x in self.well_prepared_qubits_filter if x)
@@ -153,9 +151,11 @@ class MPSBackendImpl:
         too many factors are put in the Hamiltonian
         """
         self.hamiltonian = make_H(
-            interaction_matrix=self.masked_interaction_matrix
-            if self.is_masked
-            else self.full_interaction_matrix,
+            interaction_matrix=(
+                self.masked_interaction_matrix
+                if self.is_masked
+                else self.full_interaction_matrix
+            ),
             hamiltonian_type=self.hamiltonian_type,
             num_gpus_to_use=self.config.num_gpus_to_use,
         )
@@ -177,6 +177,12 @@ class MPSBackendImpl:
         self.right_baths = right_baths(self.state, self.hamiltonian, final_qubit=2)
         assert len(self.right_baths) == self.qubit_count - 1
 
+    def get_current_right_bath(self) -> torch.Tensor:
+        return self.right_baths[-1]
+
+    def get_current_left_bath(self) -> torch.Tensor:
+        return self.left_baths[-1]
+
     def init(self) -> None:
         self.init_dark_qubits()
         self.init_initial_state(self.config.initial_state)
@@ -197,7 +203,7 @@ class MPSBackendImpl:
         """
         assert 1 <= len(indices) <= 2
 
-        baths = (self.left_baths[-1], self.right_baths[-1])
+        baths = (self.get_current_left_bath(), self.get_current_right_bath())
 
         if len(indices) == 1:
             assert orth_center_right is None
@@ -269,10 +275,10 @@ class MPSBackendImpl:
             )
             self.left_baths.append(
                 new_left_bath(
-                    self.left_baths[-1],
+                    self.get_current_left_bath(),
                     self.state.factors[self.tdvp_index],
                     self.hamiltonian.factors[self.tdvp_index],
-                )
+                ).to(self.state.factors[self.tdvp_index + 1].device)
             )
             self._evolve(self.tdvp_index + 1, dt=-delta_time / 2)
             self.right_baths.pop()
@@ -298,10 +304,10 @@ class MPSBackendImpl:
             assert self.tdvp_index <= self.qubit_count - 2
             self.right_baths.append(
                 new_right_bath(
-                    self.right_baths[-1],
+                    self.get_current_right_bath(),
                     self.state.factors[self.tdvp_index + 1],
                     self.hamiltonian.factors[self.tdvp_index + 1],
-                )
+                ).to(self.state.factors[self.tdvp_index].device)
             )
             if not self.has_lindblad_noise:
                 # Free memory because it won't be used anymore
@@ -334,7 +340,6 @@ class MPSBackendImpl:
     def timestep_complete(self) -> None:
         self.fill_results()
         self.timestep_index += 1
-        self.target_time = float((self.timestep_index + 1) * self.config.dt)
         if self.is_masked and self.current_time >= self.slm_end_time:
             self.is_masked = False
             self.hamiltonian = make_H(
@@ -344,6 +349,7 @@ class MPSBackendImpl:
             )
 
         if not self.is_finished():
+            self.target_time = self.target_times[self.timestep_index + 1]
             update_H(
                 hamiltonian=self.hamiltonian,
                 omega=self.omega[self.timestep_index, :],
@@ -480,12 +486,15 @@ class NoisyMPSBackendImpl(MPSBackendImpl):
         self.aggregated_lindblad_ops = stacked.conj().transpose(1, 2) @ stacked
 
         self.lindblad_noise = compute_noise_from_lindbladians(self.lindblad_ops)
-        self.jump_threshold = random.random()
+
+    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
 
     def init(self) -> None:
-        super().init()
         self.init_lindblad_noise()
+        super().init()
+        self.set_jump_threshold(1.0)
 
     def tdvp_complete(self) -> None:
         previous_time = self.current_time
@@ -516,7 +525,7 @@ class NoisyMPSBackendImpl(MPSBackendImpl):
 
         if self.root_finder.is_converged(tolerance=1):
             self.do_random_quantum_jump()
-            self.target_time = (self.timestep_index + 1) * self.config.dt
+            self.target_time = self.target_times[self.timestep_index + 1]
             self.root_finder = None
         else:
             self.target_time = self.root_finder.get_next_abscissa()
@@ -535,11 +544,11 @@ class NoisyMPSBackendImpl(MPSBackendImpl):
         self.state.apply(jumped_qubit_index, jump_operator)
         self.state.orthogonalize(0)
         self.state *= 1 / self.state.norm()
+        self.init_baths()
 
         norm_after_normalizing = self.state.norm()
         assert math.isclose(norm_after_normalizing, 1, abs_tol=1e-10)
-        self.jump_threshold = random.uniform(0.0, norm_after_normalizing**2)
-        self.norm_gap_before_jump = norm_after_normalizing**2 - self.jump_threshold
+        self.set_jump_threshold(norm_after_normalizing**2)
 
     def fill_results(self) -> None:
         # Remove the noise from self.hamiltonian for the callbacks.

emu_mps/mps_config.py

@@ -1,7 +1,6 @@
 from typing import Any
 
-from emu_base import BackendConfig, State
-from emu_mps.constants import DEVICE_COUNT
+from emu_base import BackendConfig, State, DEVICE_COUNT
 
 
 class MPSConfig(BackendConfig):

emu_mps/tdvp.py

@@ -117,6 +117,7 @@ def evolve_pair(
 
     left_ham_factor = left_ham_factor.to(left_device)
     right_ham_factor = right_ham_factor.to(left_device)
+    right_bath = right_bath.to(left_device)
 
     combined_hamiltonian_factors = (
         torch.tensordot(left_ham_factor, right_ham_factor, dims=1)

emu_mps/utils.py

@@ -130,13 +130,18 @@ def extended_mps_factors(
             bond_dimension = mps_factors[factor_index].shape[2]
             factor_index += 1
         elif factor_index == len(mps_factors):
-            factor = torch.zeros(bond_dimension, 2, 1, dtype=torch.complex128)
+            factor = torch.zeros(
+                bond_dimension, 2, 1, dtype=torch.complex128
+            )  # FIXME: assign device
             factor[:, 0, :] = torch.eye(bond_dimension, 1)
             bond_dimension = 1
             result.append(factor)
         else:
             factor = torch.zeros(
-                bond_dimension, 2, bond_dimension, dtype=torch.complex128
+                bond_dimension,
+                2,
+                bond_dimension,
+                dtype=torch.complex128,  # FIXME: assign device
             )
             factor[:, 0, :] = torch.eye(bond_dimension, bond_dimension)
             result.append(factor)

{emu_mps-1.2.5.dist-info → emu_mps-1.2.7.dist-info}/METADATA

@@ -1,7 +1,10 @@
 Metadata-Version: 2.4
 Name: emu-mps
-Version: 1.2.5
+Version: 1.2.7
 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
+Project-URL: Issues, https://github.com/pasqal-io/emulators/issues
 Author-email: Anton Quelle <anton.quelle@pasqal.com>, Mauro Mendizabal <mauro.mendizabal-pico@pasqal.com>, Stefano Grava <stefano.grava@pasqal.com>, Pablo Le Henaff <pablo.le-henaff@pasqal.com>
 License: PASQAL OPEN-SOURCE SOFTWARE LICENSE AGREEMENT (MIT-derived)
         
@@ -22,80 +25,13 @@ 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==1.2.5
+Requires-Dist: emu-base==1.2.7
 Description-Content-Type: text/markdown
 
 <div align="center">
   <img src="docs/logos/LogoTaglineSoftGreen.svg">
-
-  # Emu-MPS
 </div>
 
-**Emu-mps** is a backend for the [Pulser low-level Quantum Programming toolkit](https://pulser.readthedocs.io) that lets you run Quantum Algorithms on a simulated device, using GPU acceleration if available. More in depth, emu-mps is designed to **emu**late the dynamics of programmable arrays of neutral atoms, with matrix product states (**mps**). While benchmarking is incomplete as of this writing, early results suggest that this design makes emu-mps faster and more memory-efficient than previous generations of quantum emulators at running simulations with large numbers of qubits.
-
-As of this writing, Emu-MPS is provided for Linux and macOS but will not work under Windows.
-
-## Installation
-
-**Warning:** installing emu-mps will update pulser-core
-
-### Using `hatch`, `uv` or any pyproject-compatible Python manager
-
-To add `emu-mps` to your project, edit your `pyproject.toml` to add the line
-
-```toml
-  "emu-mps"
-```
-
-to the list of `dependencies`.
-
-
-### Using `pip` or `pipx`
-To install the `pipy` package using `pip` or `pipx`
-
-1. Create a `venv` if that's not done yet
-
-```sh
-$ python -m venv venv
-
-```
-
-2. Enter the venv
-
-If you're running Unix:
-
-```sh
-$ . venv/bin/activate
-```
-
-If you're running Windows:
-
-```sh
-C:\> /path/to/new/virtual/environment/Scripts/activate
-```
-
-3. Install the package
-
-```sh
-$ pip install emu-mps
-# or
-$ pipx install emu-mps
-```
-
-
-Join us on [Slack](https://pasqalworkspace.slack.com/archives/C07MUV5K7EU) or by [e-mail](mailto:emulation@pasqal.com) to give us feedback about how you plan to use Emu-MPS or if you require specific feature-upgrades.
-
-## Usage
-
-For the time being, the easiest way to learn how to use this package is to look
-at the [examples](examples/emu_mps_examples) and [notebooks](https://pasqal-io.github.io/emulators/latest/).
-
-See also the [full documentation](https://github.com/pasqal-io/emulators/blob/main/docs/index.md) for
-the API, information about contributing, benchmarks, etc.
-
-
-## Getting in touch
+# Welcome to the Pasqal analog emulators
 
-- [Pasqal Community Portal](https://community.pasqal.com/) (forums, chat, tutorials, examples, code library).
-- [GitHub Repository](https://github.com/pasqal-io/quantum-evolution-kernel) (source code, issue tracker).
-- [Professional Support](https://www.pasqal.com/contact-us/) (if you need tech support, custom licenses, a variant of this library optimized for your workload, your own QPU, remote access to a QPU, ...)
+Welcome, and please see the [GitHub pages](https://pasqal-io.github.io/emulators/) for a landing page to this repo.

emu_mps-1.2.7.dist-info/RECORD

@@ -0,0 +1,17 @@
+emu_mps/__init__.py,sha256=KSTZWIZSKHhjt0yt8-fS23rFFcFVQciNiXBgHS0pnHU,646
+emu_mps/algebra.py,sha256=ngPtTH-j2ZCBWoaJZXlkUyIlug7dY7Q92gzfnRlpPMA,5485
+emu_mps/hamiltonian.py,sha256=LcBs6CKBb643a1e9AAVtQoUfa4L_0dIhLOKecx5OOWs,15864
+emu_mps/mpo.py,sha256=H5vkJvz4AfXfnPbvgWznBWpMUO8LnGL3_NAP3IhxZzQ,8740
+emu_mps/mps.py,sha256=J0I4oQP_F1woEKmnOqnXPOWxx2Y1addxNjosL3yhYAY,18214
+emu_mps/mps_backend.py,sha256=6fVaq-D4xyicYRjGjhqMEieC7---90LpfpbV7ZD7zkQ,2192
+emu_mps/mps_backend_impl.py,sha256=XT2HccHWd6Y1gIAs070pBxjPUPIHBl-hFCuqXJaPS-E,21256
+emu_mps/mps_config.py,sha256=ydKN0OOaWCBcNd9V-4CU5ZZ4w1FRT-bbKyZQD2WCaME,3317
+emu_mps/noise.py,sha256=h4X2EFjoC_Ok0gZ8I9wN77RANXaVehTBbjkcbY_GAmY,784
+emu_mps/tdvp.py,sha256=pIQ2NXA2Mrkp3elhqQbX3pdJVbtKkG3c5r9fFlJo7pI,5755
+emu_mps/utils.py,sha256=BqRJYAcXqprtZVJ0V_j954ON2bhTdtZiaTojsYyrWrg,8193
+emu_mps/optimatrix/__init__.py,sha256=lHWYNeThHp57ZrwTwXd0p8bNvcCv0w_AZ31iCWflBUo,226
+emu_mps/optimatrix/optimiser.py,sha256=cVMdm2r_4OpbthcQuFMrJ9rNR9WEJRga9c_lHrJFkhw,6687
+emu_mps/optimatrix/permutations.py,sha256=JRXGont8B4QgbkV9CzrA0w7uzLgBrmZ1J9aa0G52hPo,1979
+emu_mps-1.2.7.dist-info/METADATA,sha256=-yHfBZrLmNsmc-tA-Yb0KfmxULGgVrcLVUbx4F37oA4,3505
+emu_mps-1.2.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+emu_mps-1.2.7.dist-info/RECORD,,

emu_mps/constants.py

@@ -1,4 +0,0 @@
-import torch
-
-
-DEVICE_COUNT = torch.cuda.device_count()

emu_mps-1.2.5.dist-info/RECORD

@@ -1,18 +0,0 @@
-emu_mps/__init__.py,sha256=EdVngqaa6GtyXmEP2aK3BfTHjYLBJSN3wR0CvltS9NQ,646
-emu_mps/algebra.py,sha256=ngPtTH-j2ZCBWoaJZXlkUyIlug7dY7Q92gzfnRlpPMA,5485
-emu_mps/constants.py,sha256=41LYkKLUCz-oxPbd-j7nUDZuhIbUrnez6prT0uR0jcE,56
-emu_mps/hamiltonian.py,sha256=LcBs6CKBb643a1e9AAVtQoUfa4L_0dIhLOKecx5OOWs,15864
-emu_mps/mpo.py,sha256=7y6q0UIfyX9igQknqtgt6nymuVcgjHlH3-Qv7N7uOZE,8769
-emu_mps/mps.py,sha256=OjG_caqPOioCdOt-bFUkOf2xuNGnKzj0LaMc3EJCHi4,17855
-emu_mps/mps_backend.py,sha256=6fVaq-D4xyicYRjGjhqMEieC7---90LpfpbV7ZD7zkQ,2192
-emu_mps/mps_backend_impl.py,sha256=HKDqUakqSs1THeeEZP5MFZaUyALnaIhrlsZTQ0Qp4qU,20867
-emu_mps/mps_config.py,sha256=MxahrPDaOpfdB6SLG1610iDUOuLR04IaCjKQRk99ICY,3346
-emu_mps/noise.py,sha256=h4X2EFjoC_Ok0gZ8I9wN77RANXaVehTBbjkcbY_GAmY,784
-emu_mps/tdvp.py,sha256=TH4CcBNczRURXYGPXndWKDs0jWXz_x9ozM961uGiSOw,5711
-emu_mps/utils.py,sha256=n9BcpuIz4Kl6EYlATaK8TKsyF-T7FTwbBo6KSAQYzl8,8066
-emu_mps/optimatrix/__init__.py,sha256=lHWYNeThHp57ZrwTwXd0p8bNvcCv0w_AZ31iCWflBUo,226
-emu_mps/optimatrix/optimiser.py,sha256=cVMdm2r_4OpbthcQuFMrJ9rNR9WEJRga9c_lHrJFkhw,6687
-emu_mps/optimatrix/permutations.py,sha256=JRXGont8B4QgbkV9CzrA0w7uzLgBrmZ1J9aa0G52hPo,1979
-emu_mps-1.2.5.dist-info/METADATA,sha256=bOtqQKZZ6ZkK9C6aV4puIipHUm5v1G9-tQrUftkIw0M,5559
-emu_mps-1.2.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-emu_mps-1.2.5.dist-info/RECORD,,