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

emu_mps/__init__.py

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

emu_mps/mps.py

@@ -189,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,
@@ -214,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

@@ -54,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
@@ -94,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)
@@ -152,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,
         )
@@ -176,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)
@@ -196,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
@@ -268,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()
@@ -297,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
@@ -333,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(
@@ -343,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, :],
@@ -479,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
@@ -515,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()
@@ -534,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/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.6.dist-info → emu_mps-1.2.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: emu-mps
-Version: 1.2.6
+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
@@ -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==1.2.6
+Requires-Dist: emu-base==1.2.7
 Description-Content-Type: text/markdown
 
 <div align="center">

{emu_mps-1.2.6.dist-info → emu_mps-1.2.7.dist-info}/RECORD

@@ -1,17 +1,17 @@
-emu_mps/__init__.py,sha256=KphiJUrBt3vjS0e2_tQejeB12X3YwqrEhxwfF1QRXQE,646
+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=CcduX2BC4ArBwwF41w_FQCa6wqmynegQQC9zkK0EmgQ,17826
+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=Rp7WbT0Dto1G4ArqSLEzSHkJAuMIEZfUqUMP9Dyz31M,20838
+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=TH4CcBNczRURXYGPXndWKDs0jWXz_x9ozM961uGiSOw,5711
-emu_mps/utils.py,sha256=n9BcpuIz4Kl6EYlATaK8TKsyF-T7FTwbBo6KSAQYzl8,8066
+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.6.dist-info/METADATA,sha256=iUObGpVmQN3Y6GM_ScEFmFlY80htsWrIjo53ER2Flvw,3505
-emu_mps-1.2.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-emu_mps-1.2.6.dist-info/RECORD,,
+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,,