emu-base 2.2.1__py3-none-any.whl → 2.4.0__py3-none-any.whl

emu_base/__init__.py

@@ -4,20 +4,21 @@ from .math.brents_root_finding import find_root_brents
 from .math.krylov_exp import krylov_exp, DEFAULT_MAX_KRYLOV_DIM
 from .jump_lindblad_operators import compute_noise_from_lindbladians
 from .math.matmul import matmul_2x2_with_batched
-from .aggregators import AggregationType, aggregate
+from .utils import get_max_rss, apply_measurement_errors, unix_like
 
 __all__ = [
     "__version__",
+    "get_max_rss",
     "compute_noise_from_lindbladians",
     "matmul_2x2_with_batched",
-    "AggregationType",
-    "aggregate",
     "PulserData",
     "find_root_brents",
     "krylov_exp",
     "HamiltonianType",
     "DEFAULT_MAX_KRYLOV_DIM",
     "DEVICE_COUNT",
+    "apply_measurement_errors",
+    "unix_like",
 ]
 
-__version__ = "2.2.1"
+__version__ = "2.4.0"

emu_base/math/krylov_energy_min.py

@@ -1,7 +1,7 @@
 import torch
 from typing import Callable, Tuple
 
-DEFAULT_MAX_KRYLOV_DIM: int = 100
+DEFAULT_MAX_KRYLOV_DIM: int = 200
 
 
 def _lowest_eigen_pair(

emu_base/pulser_adapter.py

@@ -1,242 +1,90 @@
-from typing import Tuple, Sequence, Any
+from typing import Sequence
 from enum import Enum
 import torch
 import math
 import pulser
+from pulser.sampler import SequenceSamples
 from pulser.noise_model import NoiseModel
-from pulser.register.base_register import BaseRegister, QubitId
+from pulser.register.base_register import QubitId
 from pulser.backend.config import EmulationConfig
+from pulser._hamiltonian_data import HamiltonianData
 from emu_base.jump_lindblad_operators import get_lindblad_operators
 
-KB_PER_RUBIDIUM_MASS = 95.17241379310344  # J/K/kg
-KEFF = 8.7  # µm^-1, conversion from atom velocity to detuning
-
 
 class HamiltonianType(Enum):
     Rydberg = 1
     XY = 2
 
 
-SUPPORTED_NOISES: dict = {
-    HamiltonianType.Rydberg: {
-        "amplitude",
-        "dephasing",
-        "relaxation",
-        "depolarizing",
-        "doppler",
-        "eff_noise",
-        "SPAM",
-        # "leakage",
-    },
-    HamiltonianType.XY: {
-        "dephasing",
-        "depolarizing",
-        "eff_noise",
-        "SPAM",
-    },  # , "leakage"},
-}
+_NON_LINDBLADIAN_NOISE = {"SPAM", "doppler", "amplitude", "detuning", "register"}
 
 
-def _get_qubit_positions(
-    register: BaseRegister,
+def _get_all_lindblad_noise_operators(
+    noise_model: NoiseModel | None,
 ) -> list[torch.Tensor]:
-    """Conversion from pulser Register to emu-mps register (torch type).
-    Each element will be given as [Rx,Ry,Rz]"""
+    if noise_model is None:
+        return []
 
-    positions = [
-        position.as_tensor().to(dtype=torch.float64)
-        for position in register.qubits.values()
+    return [
+        op
+        for noise_type in noise_model.noise_types
+        if noise_type not in _NON_LINDBLADIAN_NOISE
+        for op in get_lindblad_operators(noise_type=noise_type, noise_model=noise_model)
     ]
-    if len(positions[0]) == 2:
-        return [torch.cat((position, torch.zeros(1))) for position in positions]
-    return positions
-
 
-def _rydberg_interaction(sequence: pulser.Sequence) -> torch.Tensor:
-    """
-    Returns the Rydberg interaction matrix from the qubit positions.
-        Uᵢⱼ=C₆/|rᵢ-rⱼ|⁶
 
-    see Pulser
-    [documentation](https://pulser.readthedocs.io/en/stable/conventions.html#interaction-hamiltonian).
-    """
+def _get_target_times(
+    sequence: pulser.Sequence, config: EmulationConfig, dt: int
+) -> list[int]:
+    sequence_duration = sequence.get_duration(include_fall_time=config.with_modulation)
 
-    nqubits = len(sequence.register.qubit_ids)
-    c6 = sequence.device.interaction_coeff
-    positions = _get_qubit_positions(sequence.register)
+    observable_times = set(range(0, sequence_duration, dt))
+    observable_times.add(sequence_duration)
+    for obs in config.observables:
+        times: Sequence[float]
+        if obs.evaluation_times is not None:
+            times = obs.evaluation_times
+        elif config.default_evaluation_times != "Full":
+            times = config.default_evaluation_times.tolist()  # type: ignore[union-attr,assignment]
+        observable_times |= set([round(time * sequence_duration) for time in times])
 
-    interaction_matrix = torch.zeros(nqubits, nqubits, dtype=torch.float64)
-    for i in range(nqubits):
-        for j in range(i + 1, nqubits):
-            rij = torch.dist(positions[i], positions[j])
-            interaction_matrix[[i, j], [j, i]] = c6 / rij**6
-    return interaction_matrix
-
-
-def _xy_interaction(sequence: pulser.Sequence) -> torch.Tensor:
-    """
-    Returns the XY interaction matrix from the qubit positions.
-        Uᵢⱼ=C₃(1−3cos(𝜃ᵢⱼ)²)/|rᵢ-rⱼ|³
-    with
-        cos(𝜃ᵢⱼ) = (rᵢ-rⱼ)·m/|m||rᵢ-rⱼ|
-
-    see Pulser
-    [documentation](https://pulser.readthedocs.io/en/stable/conventions.html#interaction-hamiltonian).
-    """
-
-    nqubits = len(sequence.register.qubit_ids)
-    c3 = sequence.device.interaction_coeff_xy
-    mag_field = torch.tensor(sequence.magnetic_field, dtype=torch.float64)
-    mag_field /= mag_field.norm()
-    positions = _get_qubit_positions(sequence.register)
-
-    interaction_matrix = torch.zeros(nqubits, nqubits, dtype=torch.float64)
-    for i in range(nqubits):
-        for j in range(i + 1, nqubits):
-            rij = torch.dist(positions[i], positions[j])
-            cos_ij = torch.dot(positions[i] - positions[j], mag_field) / rij
-            interaction_matrix[[i, j], [j, i]] = c3 * (1 - 3 * cos_ij**2) / rij**3
-    return interaction_matrix
-
-
-def _get_amp_factors(
-    samples: pulser.sampler.SequenceSamples,
-    amp_sigma: float,
-    laser_waist: float | None,
-    qubit_positions: list[torch.Tensor],
-    q_ids: tuple[str, ...],
-) -> dict[int, torch.Tensor]:
-    def perp_dist(pos: torch.Tensor, axis: torch.Tensor) -> Any:
-        return torch.linalg.vector_norm(pos - torch.vdot(pos, axis) * axis)
-
-    times_to_amp_factors: dict[int, torch.Tensor] = {}
-    for ch, ch_samples in samples.channel_samples.items():
-        ch_obj = samples._ch_objs[ch]
-        prop_dir = torch.tensor(
-            ch_obj.propagation_dir or [0.0, 1.0, 0.0], dtype=torch.float64
-        )
-        prop_dir /= prop_dir.norm()
-
-        # each channel has a noise on its laser amplitude
-        # we assume each channel has the same noise amplitude currently
-        # the hardware currently has only a global channel anyway
-        sigma_factor = (
-            1.0
-            if amp_sigma == 0.0
-            else torch.max(torch.tensor(0), torch.normal(1.0, amp_sigma, (1,))).item()
-        )
-        for slot in ch_samples.slots:
-            factors = (
-                torch.tensor(
-                    [
-                        math.exp(-((perp_dist(x, prop_dir) / laser_waist) ** 2))
-                        for x in qubit_positions
-                    ],
-                    dtype=torch.float64,
-                )  # the lasers have a gaussian profile perpendicular to the propagation direction
-                if laser_waist and ch_obj.addressing == "Global"
-                else torch.ones(
-                    len(q_ids), dtype=torch.float64
-                )  # but for a local channel, this does not matter
-            )
-
-            # add the amplitude noise for the targeted qubits
-            factors[[x in slot.targets for x in q_ids]] *= sigma_factor
-
-            for i in range(slot.ti, slot.tf):
-                if i in times_to_amp_factors:  # multiple local channels at the same time
-                    # pulser enforces that no two lasers target the same qubit simultaneously
-                    # so only a single factor will be != 1.0 for each qubit
-                    times_to_amp_factors[i] = factors * times_to_amp_factors[i]
-                else:
-                    times_to_amp_factors[i] = factors
-    return times_to_amp_factors
-
-
-def _get_delta_offset(nqubits: int, temperature: float) -> torch.Tensor:
-    """
-    The delta values are shifted due to atomic velocities.
-    The atomic velocities follow the Maxwell distribution
-    https://en.wikipedia.org/wiki/Maxwell%E2%80%93Boltzmann_distribution
-    and then a given residual velocity is converted to a delta offset per
-    https://en.wikipedia.org/wiki/Doppler_broadening
-    """
-    if temperature == 0.0:
-        return torch.zeros(nqubits, dtype=torch.float64)
-    t = temperature * 1e-6  # microKelvin -> Kelvin
-    sigma = KEFF * math.sqrt(KB_PER_RUBIDIUM_MASS * t)
-    return torch.normal(0.0, sigma, (nqubits,))
+    target_times: list[int] = list(observable_times)
+    target_times.sort()
+    return target_times
 
 
 def _extract_omega_delta_phi(
-    *,
-    sequence: pulser.Sequence,
+    noisy_samples: SequenceSamples,
+    qubit_ids: tuple[str, ...],
     target_times: list[int],
-    with_modulation: bool,
-    laser_waist: float | None,
-    amp_sigma: float,
-    temperature: float,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    Samples the Pulser sequence and returns a tuple of tensors (omega, delta, phi)
-    containing:
-    - omega[i, q] = amplitude at time i * dt for qubit q
-    - delta[i, q] = detuning at time i * dt for qubit q
-    - phi[i, q] = phase at time i * dt for qubit q
-
-    if laser_waist is w_0 != None, the omega values coming from the global pulse channel
-    will me modulated as $\\Omega_i=\\Omega_i e^{-r_i^2/w_0^2}$
-    """
-
-    if with_modulation and sequence._slm_mask_targets:
-        raise NotImplementedError(
-            "Simulation of sequences combining an SLM mask and output "
-            "modulation is not supported."
-        )
-
-    q_ids = sequence.register.qubit_ids
-
-    samples = pulser.sampler.sample(
-        sequence,
-        modulation=with_modulation,
-        extended_duration=sequence.get_duration(include_fall_time=with_modulation),
-    )
-    sequence_dict = samples.to_nested_dict(all_local=True, samples_type="tensor")["Local"]
-
-    if "ground-rydberg" in sequence_dict and len(sequence_dict) == 1:
-        locals_a_d_p = sequence_dict["ground-rydberg"]
-    elif "XY" in sequence_dict and len(sequence_dict) == 1:
-        locals_a_d_p = sequence_dict["XY"]
-    else:
-        raise ValueError("Only `ground-rydberg` and `mw_global` channels are supported.")
-
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    sequence_dict = noisy_samples.to_nested_dict(all_local=True, samples_type="tensor")[
+        "Local"
+    ]
     nsamples = len(target_times) - 1
     omega = torch.zeros(
         nsamples,
-        len(q_ids),
+        len(qubit_ids),
         dtype=torch.complex128,
     )
-
     delta = torch.zeros(
         nsamples,
-        len(q_ids),
+        len(qubit_ids),
         dtype=torch.complex128,
     )
     phi = torch.zeros(
         nsamples,
-        len(q_ids),
+        len(qubit_ids),
         dtype=torch.complex128,
     )
-    qubit_positions = _get_qubit_positions(sequence.register)
-
-    times_to_amp_factors = _get_amp_factors(
-        samples, amp_sigma, laser_waist, qubit_positions, q_ids
-    )
-
-    omega_1 = torch.zeros_like(omega[0])
-    omega_2 = torch.zeros_like(omega[0])
-    max_duration = sequence.get_duration(include_fall_time=with_modulation)
+    max_duration = noisy_samples.max_duration
 
+    if "ground-rydberg" in sequence_dict and len(sequence_dict) == 1:
+        locals_a_d_p = sequence_dict["ground-rydberg"]
+    elif "XY" in sequence_dict and len(sequence_dict) == 1:
+        locals_a_d_p = sequence_dict["XY"]
+    else:
+        raise ValueError("Only `ground-rydberg` and `mw_global` channels are supported.")
     for i in range(nsamples):
         t = (target_times[i] + target_times[i + 1]) / 2
         # The sampled values correspond to the start of each interval
@@ -246,23 +94,20 @@ def _extract_omega_delta_phi(
             # Note that for dt even, t1=t2
             t1 = math.floor(t)
             t2 = math.ceil(t)
-            for q_pos, q_id in enumerate(q_ids):
-                omega_1[q_pos] = locals_a_d_p[q_id]["amp"][t1]
-                omega_2[q_pos] = locals_a_d_p[q_id]["amp"][t2]
+            for q_pos, q_id in enumerate(qubit_ids):
+                omega[i, q_pos] = (
+                    locals_a_d_p[q_id]["amp"][t1] + locals_a_d_p[q_id]["amp"][t2]
+                ) / 2.0
                 delta[i, q_pos] = (
                     locals_a_d_p[q_id]["det"][t1] + locals_a_d_p[q_id]["det"][t2]
                 ) / 2.0
                 phi[i, q_pos] = (
                     locals_a_d_p[q_id]["phase"][t1] + locals_a_d_p[q_id]["phase"][t2]
                 ) / 2.0
-            # omegas at different times need to have the laser waist applied independently
-            omega_1 *= times_to_amp_factors.get(t1, 1.0)
-            omega_2 *= times_to_amp_factors.get(t2, 1.0)
-            omega[i] = 0.5 * (omega_1 + omega_2)
         else:
             # We're in the final step and dt=1, approximate this using linear extrapolation
             # we can reuse omega_1 and omega_2 from before
-            for q_pos, q_id in enumerate(q_ids):
+            for q_pos, q_id in enumerate(qubit_ids):
                 delta[i, q_pos] = (
                     3.0 * locals_a_d_p[q_id]["det"][t2] - locals_a_d_p[q_id]["det"][t1]
                 ) / 2.0
@@ -270,30 +115,14 @@ def _extract_omega_delta_phi(
                     3.0 * locals_a_d_p[q_id]["phase"][t2]
                     - locals_a_d_p[q_id]["phase"][t1]
                 ) / 2.0
-            omega[i] = torch.clamp(0.5 * (3 * omega_2 - omega_1).real, min=0.0)
-
-    doppler_offset = _get_delta_offset(len(q_ids), temperature)
-    delta += doppler_offset
+                omega[i, q_pos] = max(
+                    (3.0 * locals_a_d_p[q_id]["amp"][t2] - locals_a_d_p[q_id]["amp"][t1])
+                    / 2.0,
+                    0.0,
+                )
     return omega, delta, phi
 
 
-_NON_LINDBLADIAN_NOISE = {"SPAM", "doppler", "amplitude"}
-
-
-def _get_all_lindblad_noise_operators(
-    noise_model: NoiseModel | None,
-) -> list[torch.Tensor]:
-    if noise_model is None:
-        return []
-
-    return [
-        op
-        for noise_type in noise_model.noise_types
-        if noise_type not in _NON_LINDBLADIAN_NOISE
-        for op in get_lindblad_operators(noise_type=noise_type, noise_model=noise_model)
-    ]
-
-
 class PulserData:
     slm_end_time: float
     full_interaction_matrix: torch.Tensor
@@ -308,51 +137,24 @@ class PulserData:
     def __init__(self, *, sequence: pulser.Sequence, config: EmulationConfig, dt: int):
         self.qubit_ids = sequence.register.qubit_ids
         self.qubit_count = len(self.qubit_ids)
-        sequence_duration = sequence.get_duration()
-        # the end value is exclusive, so add +1
-        observable_times = set(torch.arange(0, sequence.get_duration() + 1, dt).tolist())
-        observable_times.add(sequence.get_duration())
-        for obs in config.observables:
-            times: Sequence[float]
-            if obs.evaluation_times is not None:
-                times = obs.evaluation_times
-            elif config.default_evaluation_times != "Full":
-                times = (
-                    config.default_evaluation_times.tolist()  # type: ignore[union-attr,assignment]
-                )
-            observable_times |= set([round(time * sequence_duration) for time in times])
-
-        self.target_times: list[int] = list(observable_times)
-        self.target_times.sort()
+        self.target_times = _get_target_times(sequence=sequence, config=config, dt=dt)
+        self.hamiltonian = HamiltonianData.from_sequence(
+            sequence,
+            with_modulation=config.with_modulation,
+            noise_model=config.noise_model,
+        )
 
-        laser_waist = config.noise_model.laser_waist
-        amp_sigma = config.noise_model.amp_sigma
-        temperature = config.noise_model.temperature
         self.omega, self.delta, self.phi = _extract_omega_delta_phi(
-            sequence=sequence,
-            target_times=self.target_times,
-            with_modulation=config.with_modulation,
-            laser_waist=laser_waist,
-            amp_sigma=amp_sigma,
-            temperature=temperature,
+            self.hamiltonian.noisy_samples, self.qubit_ids, self.target_times
         )
 
-        addressed_basis = sequence.get_addressed_bases()[0]
-        if addressed_basis == "ground-rydberg":  # for local and global
+        int_type = self.hamiltonian.interaction_type
+        if int_type == "ising":  # for local and global
             self.hamiltonian_type = HamiltonianType.Rydberg
-        elif addressed_basis == "XY":
+        elif int_type == "XY":
             self.hamiltonian_type = HamiltonianType.XY
         else:
-            raise ValueError(f"Unsupported basis: {addressed_basis}")
-
-        not_supported = (
-            set(config.noise_model.noise_types) - SUPPORTED_NOISES[self.hamiltonian_type]
-        )
-        if not_supported:
-            raise NotImplementedError(
-                f"Interaction mode '{self.hamiltonian_type}' does not support "
-                f"simulation of noise types: {', '.join(not_supported)}."
-            )
+            raise ValueError(f"Unsupported basis: {int_type}")
 
         self.lindblad_ops = _get_all_lindblad_noise_operators(config.noise_model)
         self.has_lindblad_noise: bool = self.lindblad_ops != []
@@ -364,10 +166,11 @@ class PulserData:
             )
 
             self.full_interaction_matrix = config.interaction_matrix.as_tensor()
-        elif self.hamiltonian_type == HamiltonianType.Rydberg:
-            self.full_interaction_matrix = _rydberg_interaction(sequence)
-        elif self.hamiltonian_type == HamiltonianType.XY:
-            self.full_interaction_matrix = _xy_interaction(sequence)
+        else:
+            self.full_interaction_matrix = (
+                self.hamiltonian.noisy_interaction_matrix.as_tensor()
+            )
+
         self.full_interaction_matrix[
             torch.abs(self.full_interaction_matrix) < config.interaction_cutoff
         ] = 0.0

emu_base/utils.py

@@ -1,4 +1,11 @@
+from collections import Counter
+import random
 import torch
+import os
+
+unix_like = os.name != "nt"
+if unix_like:
+    from resource import RUSAGE_SELF, getrusage
 
 
 def deallocate_tensor(t: torch.Tensor) -> None:
@@ -33,3 +40,53 @@ def deallocate_tensor(t: torch.Tensor) -> None:
     if t._base is not None:
         t._base.resize_(0)
         t._base.set_(source=replacement_storage)
+
+
+def get_max_rss(gpu: bool) -> float:
+    if gpu:
+        max_mem_per_device = (
+            torch.cuda.max_memory_allocated(device) * 1e-6
+            for device in range(torch.cuda.device_count())
+        )
+        max_mem = max(max_mem_per_device)
+    elif unix_like:
+        max_mem = getrusage(RUSAGE_SELF).ru_maxrss * 1e-3
+    else:
+        return 0.0
+    return max_mem
+
+
+def readout_with_error(c: str, *, p_false_pos: float, p_false_neg: float) -> str:
+    # p_false_pos = false positive, p_false_neg = false negative
+    r = random.random()
+    if c == "0" and r < p_false_pos:
+        return "1"
+
+    if c == "1" and r < p_false_neg:
+        return "0"
+
+    return c
+
+
+def apply_measurement_errors(
+    bitstrings: Counter[str], *, p_false_pos: float, p_false_neg: float
+) -> Counter[str]:
+    """
+    Given a bag of sampled bitstrings, returns another bag of bitstrings
+    sampled with readout/measurement errors.
+
+        p_false_pos: probability of false positive
+        p_false_neg: probability of false negative
+    """
+
+    result: Counter[str] = Counter()
+    for bitstring, count in bitstrings.items():
+        for _ in range(count):
+            bitstring_with_error = "".join(
+                readout_with_error(c, p_false_pos=p_false_pos, p_false_neg=p_false_neg)
+                for c in bitstring
+            )
+
+            result[bitstring_with_error] += 1
+
+    return result

{emu_base-2.2.1.dist-info → emu_base-2.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: emu-base
-Version: 2.2.1
+Version: 2.4.0
 Summary: Pasqal base classes for emulators
 Project-URL: Documentation, https://pasqal-io.github.io/emulators/
 Project-URL: Repository, https://github.com/pasqal-io/emulators
@@ -25,8 +25,8 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
 Requires-Python: >=3.10
-Requires-Dist: pulser-core==1.5.*
-Requires-Dist: torch==2.7.0
+Requires-Dist: pulser-core==1.6.*
+Requires-Dist: torch==2.8.0
 Description-Content-Type: text/markdown
 
 <div align="center">

{emu_base-2.2.1.dist-info → emu_base-2.4.0.dist-info}/RECORD

@@ -1,15 +1,14 @@
-emu_base/__init__.py,sha256=lXgYx3VbA0qNMwDZyV4HpFcr6aUDxhqWjy8Um73JXrs,681
-emu_base/aggregators.py,sha256=bB-rldoDAErxQMpL715K5lpiabGOpkCY0GyxW7mfHuc,5000
+emu_base/__init__.py,sha256=eMd_qHkRgYxhb7FHPekZkxCr7JENqGH1EtqRhP98KzA,725
 emu_base/constants.py,sha256=41LYkKLUCz-oxPbd-j7nUDZuhIbUrnez6prT0uR0jcE,56
 emu_base/jump_lindblad_operators.py,sha256=Y30f8emVFS4Dazljc_Rh4lX9qU4QQY_AxPNahnzcsfY,2101
-emu_base/pulser_adapter.py,sha256=WJX8y6iBYg8FX--q0Y10haz-35TbmUGoHbMDQEDuZIY,14534
-emu_base/utils.py,sha256=-nIoJuu1pOxkPc2tiJTLjQ0ONsPR43CCB6vOJqdWaUc,1425
+emu_base/pulser_adapter.py,sha256=ss3EIgIq_hZJpcbwT5FBhQ2dxXqyncQJ8fQrDsAZ2qM,6965
+emu_base/utils.py,sha256=OoJ0GjL1K8m9Jq-BeBfmbWuFp9bvwXp8UcI3BMW94Js,2974
 emu_base/math/__init__.py,sha256=6BbIytYV5uC-e5jLMtIErkcUl_PvfSNnhmVFY9Il8uQ,97
 emu_base/math/brents_root_finding.py,sha256=AVx6L1Il6rpPJWrLJ7cn6oNmJyZOPRgEaaZaubC9lsU,3711
 emu_base/math/double_krylov.py,sha256=X16dyCbyzdP7fFK-hmKS03Q-DJtC6TZ8sJrGTJ6akIc,3708
-emu_base/math/krylov_energy_min.py,sha256=hm_B5qtBXHY1hl-r_LgDUKNDsdqVCDBHprQB3D-UFR8,4009
+emu_base/math/krylov_energy_min.py,sha256=iR4hmE0eXptbAg3opikd5d4Zv7dhnDrawH-n_4KG-cc,4009
 emu_base/math/krylov_exp.py,sha256=mGFddVQ8mEbwypbZtnlRPFpi4Nf8JZT6OKLHloIwCDQ,3934
 emu_base/math/matmul.py,sha256=lEAnV0b5z_f1xEA-9p-WXxA8bM3QbShiHdXQ3ZkZFcQ,877
-emu_base-2.2.1.dist-info/METADATA,sha256=Rs1xk1Xu4iB3-sWXgq0ASiM2bULwk0DVEfC-AclHIGs,3604
-emu_base-2.2.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-emu_base-2.2.1.dist-info/RECORD,,
+emu_base-2.4.0.dist-info/METADATA,sha256=q4KRKKVsXCzNSsDjSNmeMAZDsNmxufLy9kSSsxHPdTc,3604
+emu_base-2.4.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+emu_base-2.4.0.dist-info/RECORD,,

emu_base/aggregators.py

@@ -1,158 +0,0 @@
-import statistics
-import torch
-from typing import Any, Callable
-import collections
-from enum import Enum, auto
-from pulser.backend import (
-    Results,
-)
-import logging
-
-
-_NUMERIC_TYPES = {int, float, complex}
-
-
-class AggregationType(Enum):
-    """
-    Defines how to combine multiple values from different simulation results.
-    """
-
-    MEAN = auto()  # statistics.fmean or list/matrix-wise equivalent
-    BAG_UNION = auto()  # Counter.__add__
-
-
-def mean_aggregator(
-    values: list[Any],
-) -> (
-    complex
-    | float
-    | list[complex]
-    | list[float]
-    | list[list[complex]]
-    | list[list[float]]
-    | torch.Tensor
-):  # FIXME: support tuples?
-    if values == []:
-        raise ValueError("Cannot average 0 samples")
-
-    element_type = type(values[0])
-
-    if element_type in _NUMERIC_TYPES:
-        return statistics.fmean(values)
-
-    if element_type == torch.Tensor:
-        acc = torch.zeros_like(values[0])
-        for ten in values:
-            acc += ten
-        return acc / len(values)
-
-    if element_type != list:
-        raise NotImplementedError("Cannot average this type of data")
-
-    if values[0] == []:
-        raise ValueError("Cannot average list of empty lists")
-
-    sub_element_type = type(values[0][0])
-
-    if sub_element_type in _NUMERIC_TYPES:
-        dim = len(values[0])
-        return [statistics.fmean(value[i] for value in values) for i in range(dim)]
-
-    if sub_element_type != list:  # FIXME: ABC.Iterable? Collection? subclass?
-        raise ValueError(f"Cannot average list of lists of {sub_element_type}")
-
-    if values[0][0] == []:
-        raise ValueError("Cannot average list of matrices with no columns")
-
-    if (sub_sub_element_type := type(values[0][0][0])) not in _NUMERIC_TYPES:
-        raise ValueError(f"Cannot average list of matrices of {sub_sub_element_type}")
-
-    dim1 = len(values[0])
-    dim2 = len(values[0][0])
-    return [
-        [statistics.fmean(value[i][j] for value in values) for j in range(dim2)]
-        for i in range(dim1)
-    ]
-
-
-def bag_union_aggregator(values: list[collections.Counter]) -> collections.Counter:
-    return sum(values, start=collections.Counter())
-
-
-aggregation_types_definitions: dict[AggregationType, Callable] = {
-    AggregationType.MEAN: mean_aggregator,
-    AggregationType.BAG_UNION: bag_union_aggregator,
-}
-
-
-def _get_aggregation_type(tag: str) -> AggregationType | None:
-    if tag.startswith("bitstrings"):
-        return AggregationType.BAG_UNION
-    if tag.startswith("expectation"):
-        return AggregationType.MEAN
-    if tag.startswith("fidelity"):
-        return AggregationType.MEAN
-    if tag.startswith("correlation_matrix"):
-        return AggregationType.MEAN
-    if tag.startswith("occupation"):
-        return AggregationType.MEAN
-    if tag.startswith("energy"):
-        return AggregationType.MEAN
-    if tag.startswith("energy_second_moment"):
-        return AggregationType.MEAN
-    else:
-        return None
-
-
-def aggregate(
-    results_to_aggregate: list[Results],
-    **aggregator_functions: Callable[[Any], Any],
-) -> Results:
-    if len(results_to_aggregate) == 0:
-        raise ValueError("no results to aggregate")
-    if len(results_to_aggregate) == 1:
-        return results_to_aggregate[0]
-    stored_callbacks = set(results_to_aggregate[0].get_result_tags())
-    if not all(
-        set(results.get_result_tags()) == stored_callbacks
-        for results in results_to_aggregate
-    ):
-        raise ValueError(
-            "Monte-Carlo results seem to provide from incompatible simulations: "
-            "they do not all contain the same observables"
-        )
-    aggregated = Results(
-        atom_order=results_to_aggregate[0].atom_order,
-        total_duration=results_to_aggregate[0].total_duration,
-    )
-    for tag in stored_callbacks:
-        aggregation_type = aggregator_functions.get(
-            tag,
-            _get_aggregation_type(tag),
-        )
-        if aggregation_type is None:
-            logging.getLogger("global_logger").warning(f"Skipping aggregation of `{tag}`")
-            continue
-        aggregation_function: Any = (
-            aggregation_type
-            if callable(aggregation_type)
-            else aggregation_types_definitions[aggregation_type]
-        )
-        evaluation_times = results_to_aggregate[0].get_result_times(tag)
-        if not all(
-            results.get_result_times(tag) == evaluation_times
-            for results in results_to_aggregate
-        ):
-            raise ValueError(
-                "Monte-Carlo results seem to provide from incompatible simulations: "
-                "the callbacks are not stored at the same times"
-            )
-
-        uuid = results_to_aggregate[0]._find_uuid(tag)
-        for t in results_to_aggregate[0].get_result_times(tag):
-            v = aggregation_function(
-                [result.get_result(tag, t) for result in results_to_aggregate]
-            )
-            aggregated._store_raw(uuid=uuid, tag=tag, time=t, value=v)
-
-    return aggregated