emu-mps 1.2.1__py3-none-any.whl

emu_mps/tdvp.py

@@ -0,0 +1,209 @@
+import torch
+
+from emu_base import krylov_exp
+from emu_mps import MPS, MPO
+from emu_mps.utils import split_tensor
+
+
+def new_right_bath(
+    bath: torch.Tensor, state: torch.Tensor, op: torch.Tensor
+) -> torch.Tensor:
+    bath = torch.tensordot(state, bath, ([2], [2]))
+    bath = torch.tensordot(op.to(bath.device), bath, ([2, 3], [1, 3]))
+    return torch.tensordot(state.conj(), bath, ([1, 2], [1, 3]))
+
+
+"""
+function to compute the right baths. The three indices in the bath are as follows:
+(bond of state conj, bond of operator, bond of state)
+The baths have shape
+-xx
+-xx
+-xx
+with the index ordering (top, middle, bottom)
+bath tensors are put on the device of the factor to the left
+"""
+
+
+def right_baths(state: MPS, op: MPO, final_qubit: int) -> list[torch.Tensor]:
+    state_factor = state.factors[-1]
+    bath = torch.ones(1, 1, 1, device=state_factor.device, dtype=state_factor.dtype)
+    baths = [bath]
+    for i in range(len(state.factors) - 1, final_qubit - 1, -1):
+        bath = new_right_bath(bath, state.factors[i], op.factors[i])
+        bath = bath.to(state.factors[i - 1].device)
+        baths.append(bath)
+    return baths
+
+
+"""
+Computes H(psi) where
+    x-    -x
+    x  ||  x             ||
+H = x- xx -x  and psi = -xx-
+    x  ||  x
+    x-    -x
+
+Expects the two qubit factors of the MPS precontracted,
+with one 'fat' physical index of dim 4 and index ordering
+(left bond, physical index, right bond):
+         ||
+      -xxxxxx-
+The Hamiltonian should have an index ordering of
+(left bond, out, in, right bond).
+The baths must have shape (top, middle, bottom).
+All tensors must be on the same device
+"""
+
+
+def apply_effective_Hamiltonian(
+    state: torch.Tensor,
+    ham: torch.Tensor,
+    left_bath: torch.Tensor,
+    right_bath: torch.Tensor,
+) -> torch.Tensor:
+    assert left_bath.ndim == 3 and left_bath.shape[0] == left_bath.shape[2]
+    assert right_bath.ndim == 3 and right_bath.shape[0] == right_bath.shape[2]
+    assert left_bath.shape[2] == state.shape[0] and right_bath.shape[2] == state.shape[2]
+    assert left_bath.shape[1] == ham.shape[0] and right_bath.shape[1] == ham.shape[3]
+
+    # the optimal contraction order depends on the details
+    # 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]))
+    return state
+
+
+_TIME_CONVERSION_COEFF = 0.001  # Omega and delta are given in rad/ms, dt in ns
+
+
+class EvolveConfig:
+    def __init__(
+        self,
+        *,
+        exp_tolerance: float,
+        norm_tolerance: float,
+        max_krylov_dim: int,
+        is_hermitian: bool,
+        max_error: float,
+        max_rank: int
+    ) -> None:
+        self.exp_tolerance = exp_tolerance
+        self.norm_tolerance = norm_tolerance
+        self.max_krylov_dim = max_krylov_dim
+        self.is_hermitian = is_hermitian
+        self.max_error = (
+            max_error  # FIXME: max_error and max_rank are irrelevant for evolve_single
+        )
+        self.max_rank = max_rank
+
+
+def evolve_pair(
+    *,
+    state_factors: list[torch.Tensor],
+    baths: tuple[torch.Tensor, torch.Tensor],
+    ham_factors: list[torch.Tensor],
+    dt: float,
+    orth_center_right: bool,
+    config: EvolveConfig
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Time evolution of a pair of tensors of a tensor train using baths and truncated SVD.
+    Returned state tensors are kept on their respective devices.
+    """
+    assert len(state_factors) == 2
+    assert len(baths) == 2
+    assert len(ham_factors) == 2
+
+    left_state_factor, right_state_factor = state_factors
+    left_bath, right_bath = baths
+    left_ham_factor, right_ham_factor = ham_factors
+
+    left_device = left_state_factor.device
+    right_device = right_state_factor.device
+
+    left_bond_dim = left_state_factor.shape[0]
+    right_bond_dim = right_state_factor.shape[-1]
+
+    # Computation is done on left_device (arbitrary)
+
+    combined_state_factors = torch.tensordot(
+        left_state_factor, right_state_factor.to(left_device), dims=1
+    ).reshape(left_bond_dim, 4, right_bond_dim)
+
+    left_ham_factor = left_ham_factor.to(left_device)
+    right_ham_factor = right_ham_factor.to(left_device)
+
+    combined_hamiltonian_factors = (
+        torch.tensordot(left_ham_factor, right_ham_factor, dims=1)
+        .transpose(2, 3)
+        .reshape(left_ham_factor.shape[0], 4, 4, -1)
+    )
+
+    op = (
+        lambda x: -_TIME_CONVERSION_COEFF
+        * 1j
+        * dt
+        * apply_effective_Hamiltonian(
+            x, combined_hamiltonian_factors, left_bath, right_bath
+        )
+    )
+
+    evol = krylov_exp(
+        op,
+        combined_state_factors,
+        exp_tolerance=config.exp_tolerance,
+        norm_tolerance=config.norm_tolerance,
+        max_krylov_dim=config.max_krylov_dim,
+        is_hermitian=config.is_hermitian,
+    ).reshape(left_bond_dim * 2, 2 * right_bond_dim)
+
+    l, r = split_tensor(
+        evol,
+        max_error=config.max_error,
+        max_rank=config.max_rank,
+        orth_center_right=orth_center_right,
+    )
+
+    return l.reshape(left_bond_dim, 2, -1), r.reshape(-1, 2, right_bond_dim).to(
+        right_device
+    )
+
+
+def evolve_single(
+    *,
+    state_factor: torch.Tensor,
+    baths: tuple[torch.Tensor, torch.Tensor],
+    ham_factor: torch.Tensor,
+    dt: float,
+    config: EvolveConfig
+) -> torch.Tensor:
+    """
+    Time evolution of a single tensor of a tensor train using baths.
+    """
+    assert len(baths) == 2
+
+    left_bath, right_bath = baths
+
+    op = (
+        lambda x: -_TIME_CONVERSION_COEFF
+        * 1j
+        * dt
+        * apply_effective_Hamiltonian(
+            x,
+            ham_factor,
+            left_bath,
+            right_bath,
+        )
+    )
+
+    return krylov_exp(
+        op,
+        state_factor,
+        exp_tolerance=config.exp_tolerance,
+        norm_tolerance=config.norm_tolerance,
+        max_krylov_dim=config.max_krylov_dim,
+        is_hermitian=config.is_hermitian,
+    )

emu_mps/utils.py

@@ -0,0 +1,258 @@
+from typing import List, Optional
+import torch
+import random
+from collections import Counter
+
+DEVICE_COUNT = torch.cuda.device_count()
+
+
+def new_left_bath(
+    bath: torch.Tensor, state: torch.Tensor, op: torch.Tensor
+) -> torch.Tensor:
+    # this order is more efficient than contracting the op first in general
+    bath = torch.tensordot(bath, state.conj(), ([0], [0]))
+    bath = torch.tensordot(bath, op.to(bath.device), ([0, 2], [0, 1]))
+    return torch.tensordot(bath, state, ([0, 2], [0, 1]))
+
+
+def _determine_cutoff_index(d: torch.Tensor, max_error: float) -> int:
+    assert max_error > 0
+    squared_max_error = max_error * max_error
+    acc = 0
+    for i in range(d.shape[0]):
+        acc += d[i]
+        if acc > squared_max_error:
+            return i
+    return 0  # type: ignore[no-any-return]
+
+
+def split_tensor(
+    m: torch.Tensor,
+    max_error: float = 1e-5,
+    max_rank: int = 1024,
+    orth_center_right: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Computes a low-rank approximation split of m using the Eckart-Young-Mirsky theorem.
+    """
+    assert m.ndim == 2
+
+    if orth_center_right:
+        d, q = torch.linalg.eigh(m @ m.T.conj())
+        max_bond = max(
+            _determine_cutoff_index(d, max_error),
+            d.shape[0] - max_rank,
+        )
+        left = q[:, max_bond:]
+        right = q.T.conj() @ m
+        right = right[max_bond:, :]
+    else:
+        d, q = torch.linalg.eigh(m.T.conj() @ m)
+        max_bond = max(
+            _determine_cutoff_index(d, max_error),
+            d.shape[0] - max_rank,
+        )
+        right = q.T.conj()[max_bond:, :]
+        left = m @ q
+        left = left[:, max_bond:]
+
+    return left, right
+
+
+def truncate_impl(
+    factors: list[torch.tensor],
+    max_error: float = 1e-5,
+    max_rank: int = 1024,
+) -> None:
+    """
+    Eigenvalues-based truncation of a matrix product.
+    An in-place operation.
+
+    Note:
+        Sweeps from right to left.
+        Requires the matrix product to be orthogonalized on the last element.
+        At each step moves the orthogonal center to the left while truncating.
+    """
+    for i in range(len(factors) - 1, 0, -1):
+        factor_shape = factors[i].shape
+
+        l, r = split_tensor(
+            factors[i].reshape(factor_shape[0], -1),
+            max_error=max_error,
+            max_rank=max_rank,
+            orth_center_right=False,
+        )
+
+        factors[i] = r.reshape(-1, *factor_shape[1:])
+        factors[i - 1] = torch.tensordot(
+            factors[i - 1], l.to(factors[i - 1].device), dims=1
+        )
+
+
+def assign_devices(tensors: List[torch.Tensor], num_gpus_to_use: int) -> None:
+    """
+    Evenly distributes each tensor in the list to a device.
+    If num_gpus_to_use is 0, then all tensors go to CPU.
+    """
+    num_gpus_to_use = min(len(tensors), num_gpus_to_use)
+
+    if num_gpus_to_use <= 0:
+        for i in range(len(tensors)):
+            tensors[i] = tensors[i].to("cpu")
+        return
+
+    tensors_per_device = len(tensors) // num_gpus_to_use
+
+    if len(tensors) % num_gpus_to_use != 0:
+        tensors_per_device += 1
+
+    for i in range(len(tensors)):
+        tensors[i] = tensors[i].to(f"cuda:{i // tensors_per_device}")
+
+
+def extended_mps_factors(
+    mps_factors: list[torch.Tensor], where: list[bool]
+) -> list[torch.Tensor]:
+    """
+    Given a valid list of MPS factors, accounting for qubits marked as `True` in `where`,
+    fills the `False` positions with new qubits in the |0> state.
+    """
+    assert len(mps_factors) == sum(1 for b in where if b)
+
+    bond_dimension = 1
+    factor_index = 0
+    result = []
+    for is_factor in where:
+        assert 0 <= factor_index <= len(mps_factors)
+
+        if is_factor:
+            result.append(mps_factors[factor_index])
+            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[:, 0, :] = torch.eye(bond_dimension, 1)
+            bond_dimension = 1
+            result.append(factor)
+        else:
+            factor = torch.zeros(
+                bond_dimension, 2, bond_dimension, dtype=torch.complex128
+            )
+            factor[:, 0, :] = torch.eye(bond_dimension, bond_dimension)
+            result.append(factor)
+    return result
+
+
+def extended_mpo_factors(
+    mpo_factors: list[torch.Tensor], where: list[bool]
+) -> list[torch.Tensor]:
+    """
+    Given a valid list of MPO factors, accounting for qubits marked as `True` in `where`,
+    fills the `False` positions with new MPO identity factors.
+    """
+    assert len(mpo_factors) == sum(1 for b in where if b)
+
+    bond_dimension = 1
+    factor_index = 0
+    result = []
+    for is_factor in where:
+        assert 0 <= factor_index <= len(mpo_factors)
+
+        if is_factor:
+            result.append(mpo_factors[factor_index])
+            bond_dimension = mpo_factors[factor_index].shape[3]
+            factor_index += 1
+        elif factor_index == len(mpo_factors):
+            factor = torch.zeros(bond_dimension, 2, 2, 1, dtype=torch.complex128)
+            factor[:, 0, 0, :] = torch.eye(bond_dimension, 1)
+            factor[:, 1, 1, :] = torch.eye(bond_dimension, 1)
+            bond_dimension = 1
+            result.append(factor)
+        else:
+            factor = torch.zeros(
+                bond_dimension, 2, 2, bond_dimension, dtype=torch.complex128
+            )
+            factor[:, 0, 0, :] = torch.eye(bond_dimension, bond_dimension)
+            factor[:, 1, 1, :] = torch.eye(bond_dimension, bond_dimension)
+            result.append(factor)
+    return result
+
+
+def get_extended_site_index(
+    where: list[bool], desired_index: Optional[int]
+) -> Optional[int]:
+    """
+    Returns the index in `where` that has `desired_index` preceding True elements.
+
+    This function is used to find the index of the orthogonality center in an MPS obtained
+    with `extended_mps_factors` in the presence of dark qubits:
+        `where` is the mask specifying whether qubits are well-prepared.
+        `desired_index` is the index of the orthogonality center of the MPS without dark qubits.
+        The return value is then the index of the orthogonality center
+        in the full MPS with added dark qubits.
+    """
+
+    if desired_index is None:
+        return None
+
+    index = -1
+    for extended_index, boolean_value in enumerate(where):
+        if boolean_value:
+            index += 1
+            if index == desired_index:
+                return extended_index
+
+    raise ValueError(f"Index {desired_index} does not exist")
+
+
+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
+
+
+n_operator = torch.tensor(
+    [
+        [0, 0],
+        [0, 1],
+    ],
+    dtype=torch.complex128,
+)
+
+
+def tensor_trace(tensor: torch.Tensor, dim1: int, dim2: int) -> torch.Tensor:
+    """
+    Contract two legs of a single tensor.
+    """
+    assert tensor.shape[dim1] == tensor.shape[dim2], "dimensions should match"
+    return tensor.diagonal(offset=0, dim1=dim1, dim2=dim2).sum(-1)

emu_mps-1.2.1.dist-info/METADATA

@@ -0,0 +1,133 @@
+Metadata-Version: 2.4
+Name: emu-mps
+Version: 1.2.1
+Summary: Pasqal MPS based pulse emulator built on PyTorch
+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: Proprietary
+Classifier: License :: Other/Proprietary License
+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.0
+Description-Content-Type: text/markdown
+
+<div align="center">
+  <img src="docs/logos/LogoTaglineSoftGreen.svg">
+
+  # Emu-MPS
+</div>
+
+**EMU-MPS** is a Pulser backend, designed to **EMU**late the dynamics of programmable arrays of neutral atoms, with matrix product states (**MPS**). It allows users to increase the number of qubits and reduce computation time.
+
+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.
+
+
+## Getting started
+
+You can install from source, or download the package from the private pypi registry that pasqal maintains in gitlab.
+For developers, we recommend installing from source, for users we recommend installing from the registry.
+
+**Warning:** installing emu-mps will update pulser-core
+
+We always recommend using a virtual environment.
+
+<details>
+  <summary>Click me to see how it is done</summary>
+
+  #### Create a virtual environment using python
+
+  ```
+  python -m venv .venv
+  ```
+
+  Or
+
+  ```
+  python -m venv /path/to/new/virtual/environment
+  ```
+
+  Replace `/path/to/new/virtual/environment` with your desired directory path.
+
+  Then activate the environment On linux or MacOS
+
+  ```
+  source /path/to/new/virtual/environment/bin/activate
+  ```
+
+  While on Windows it's
+
+  ```
+  C:\> /path/to/new/virtual/environment/Scripts/activate
+  ```
+
+  Remember to replace `/path/to/new/virtual/environment` with the actual path to your virtual environment. Once the environment is activated, you can clone emu_mps and install it using
+
+</details>
+
+### installing from the registry
+
+When pip is configured to know about the pasqal registry, Emu-MPS installs as
+
+```bash
+pip install emu-mps
+```
+When pip is not already configured, the easiest way to do so, is to add a file `~/.config/pip/pip.conf` containing:
+
+```
+[global]
+extra-index-url=https://gitlab.pasqal.com/api/v4/projects/597/packages/pypi/simple
+                possible.other.urls
+```
+
+As this shows, it is also possible to have multiple extra repositories configured. Note that the order is not important.
+
+It is also possible to add the `extra-index-url` to the `pip install` command directly, if you somehow don't want to create a `pip.conf` file.
+
+### installing from source
+git clone this [repository ](https://gitlab.pasqal.com/emulation/rydberg-atoms/emu-ct) or download
+
+
+Then, `cd` into the root folder of the repo and type
+
+```bash
+pip install -e .
+```
+
+<details>
+  <summary>Guidelines for developers </summary>
+  We recommend using an environment, git clone the repository, then inside the `emu_mps` folder
+
+```bash
+pip install -e .
+```
+
+  Also, the installation of pytest, nbmake, pre-commit.
+
+  Do not forget to run the unit test suite by simply running `pytest` command.
+
+  Another way can be using hatch.
+
+  #### virtual environment with `hatch`
+
+  ```bash
+  python -m pip install hatch
+  python -m hatch -v shell
+  ```
+
+  When inside the shell with development dependencies, install first the pre-commit hook:
+  ```
+  pre-commit install
+  ```
+</details>
+
+## Check the tutorial notebooks and example scripts
+
+For more information, you can check the tutorials and examples located in the [examples folder](https://gitlab.pasqal.com/emulation/rydberg-atoms/emu-ct/-/tree/main/examples?ref_type=heads)
+
+## Documentation
+
+Please check the [documentation](./docs/index.md) page for more info about contributing, the API, benchmarks, etc.
+
+
+![Code Coverage](https://img.shields.io/badge/Coverage-95%25-brightgreen.svg)

emu_mps-1.2.1.dist-info/RECORD

@@ -0,0 +1,14 @@
+emu_mps/__init__.py,sha256=0yj0gJbi5R8_-wdu7olvvJrvALFaNchusJwUCmCyt7M,646
+emu_mps/algebra.py,sha256=fMkhie3iy-zEivE6G3EpJoo2o6Gde0rqrOwBUKOBkME,5384
+emu_mps/hamiltonian.py,sha256=Ojb-T2bpb-Ym99IbFsTC2WqQ4C06IbcYK4VusA2pt_0,15799
+emu_mps/mpo.py,sha256=RA9n50li1NJsIWCid-uOFgYV-usX4S7ZgmRdPLht8hs,8611
+emu_mps/mps.py,sha256=3Oi_gT5KdY46Q3l9uGifWyrb77DUQOHxnu_O9ky7MYo,18338
+emu_mps/mps_backend.py,sha256=ZCZMCsJz4-Oi0Ox2CqCOWqQWdE9ZJj6n_Jssd0g_1T0,972
+emu_mps/mps_backend_impl.py,sha256=dEhBVciS7bdZ0AymbgDVyNUfprxuv-GgUoRcmKBkLHY,19512
+emu_mps/mps_config.py,sha256=-ScHCNiiAwg_kMBSGOsXydjnbo_y3AqrvqpG3RzTYmU,2703
+emu_mps/noise.py,sha256=h4X2EFjoC_Ok0gZ8I9wN77RANXaVehTBbjkcbY_GAmY,784
+emu_mps/tdvp.py,sha256=mDStFKgbG3OXsEG-Ja8NJKsTGyLxsjNS6sQXK7B33EY,6088
+emu_mps/utils.py,sha256=05s5TBE2zhsJwugbQFDCFQi7L0pVX0ukI-WIW2zN5TA,8052
+emu_mps-1.2.1.dist-info/METADATA,sha256=SUITsADc02ZBCuTVgYUzWMwuyAx_AfIPN4mkC0Tt9Nw,4203
+emu_mps-1.2.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+emu_mps-1.2.1.dist-info/RECORD,,

emu_mps-1.2.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any