feat: channels can be extended with qubit-specific T1 and T2 values/lists

src/notebooks/noisy_simulation/02_qutrit_gad.ipynb

@@ -2,7 +2,7 @@
   "cells": [
     {
       "cell_type": "code",
-      "execution_count": 18,
+      "execution_count": 1,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -23,7 +23,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 19,
+      "execution_count": 2,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -48,7 +48,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 20,
+      "execution_count": 3,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -75,7 +75,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 21,
+      "execution_count": 4,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -107,21 +107,21 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 22,
+      "execution_count": 5,
       "metadata": {},
       "outputs": [
         {
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            " 12%|\u2588\u258e        | 5/40 [00:00<00:00, 43.82it/s]"
+            " 12%|\u2588\u258e        | 5/40 [00:00<00:00, 46.90it/s]"
           ]
         },
         {
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 40/40 [00:02<00:00, 13.54it/s]\n"
+            "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 40/40 [00:03<00:00, 13.05it/s]\n"
           ]
         },
         {
@@ -148,7 +148,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 23,
+      "execution_count": 6,
       "metadata": {},
       "outputs": [
         {
@@ -162,7 +162,7 @@
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 40/40 [00:02<00:00, 13.96it/s]\n"
+            "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 40/40 [00:03<00:00, 12.68it/s]\n"
           ]
         },
         {

src/quantum_logical/channel.py

@@ -1,5 +1,6 @@
 """Error channels for quantum systems."""
 from abc import ABC, abstractmethod
+from typing import List, Union
 
 import numpy as np
 from qutip import Qobj, qeye, tensor
@@ -43,12 +44,25 @@ def _init_kraus_operators(self):
 class Channel(CPTPMap):
     """Base class for quantum error channels."""
 
-    def __init__(self, dims):
+    def __init__(self, num_qubits, hilbert_space_dim, **kwargs):
         """Initialize the channel with specified dimensions."""
-        super().__init__(dims)
+        super().__init__(dims=hilbert_space_dim**num_qubits)
+        self.num_qubits = num_qubits
+        self.hilbert_space_dim = hilbert_space_dim
+        self.params = kwargs
         self._trotter_dt = None
         self._E = None
 
+        # Each param in params should either be a single value or an iterable
+        # If iterable, then it should have length equal to num_qubits
+        for key, value in kwargs.items():
+            if isinstance(value, (int, float)):
+                self.params[key] = [value] * num_qubits
+            elif len(value) != num_qubits:
+                raise ValueError(
+                    f"Number of parameters for {key} must equal number of qubits."
+                )
+
     @property
     def E(self):
         """Return the Kraus operators for the channel."""
@@ -64,16 +78,18 @@ def set_trotter_dt(self, trotter_dt):
 
     def _init_kraus_operators(self):
         """Initialize and extend Kraus operators to multiple qubits."""
-        single_qubit_operators = self._create_single_qubit_operators()
-        self._E = self._extend_kraus_operators_to_multiple_qubits(
-            single_qubit_operators
-        )
+        qubit_operators = []
+
+        # Iterate over all qubits
+        for qubit in range(self.num_qubits):
+            qubit_params = {key: value[qubit] for key, value in self.params.items()}
+            qubit_operators.append(self._create_single_qubit_operators(**qubit_params))
+        self._E = self._extend_kraus_operators_to_multiple(qubit_operators)
         self._E = [np.array(E, dtype=complex) for E in self._E]
         self._verify_completeness()
         return self._E
 
-    # TODO: Make this so we could extend with varying T1 and T2 between qubits
-    def _extend_kraus_operators_to_multiple_qubits(self, single_qubit_operators):
+    def _extend_kraus_operators_to_multiple(self, qubit_operators):
         """Extend single-qubit Kraus operators to multiple qubits.
 
         This method assumes that errors occur independently on each
@@ -85,18 +101,34 @@ def _extend_kraus_operators_to_multiple_qubits(self, single_qubit_operators):
         errors on multiple qubits (higher-order errors) are not
         explicitly modeled, which is a common assumption in many quantum
         error correction scenarios.
+
+        Args:
+            qubit_operators (list): List of single-qubit Kraus operators, where each
+            sublist contains the Kraus operators for a specific qubit.
+
+        Returns:
+            list: List of extended Kraus operators for the full system.
         """
         if self.num_qubits == 1:
-            return single_qubit_operators
+            return qubit_operators[0]
 
         identity = qeye(self.hilbert_space_dim)
         extended_operators = []
+
         for qubit in range(self.num_qubits):
+            # Get the Kraus operators for the current qubit
+            single_qubit_operators = qubit_operators[qubit]
+
             for op in single_qubit_operators:
-                operators = [
-                    Qobj(op) if i == qubit else identity for i in range(self.num_qubits)
-                ]
-                extended_operators.append(tensor(*operators))
+                # Create a list of identity operators for all qubits
+                operators_on_all_qubits = [identity for _ in range(self.num_qubits)]
+
+                # Replace the identity operator with the actual Kraus operator for the current qubit
+                operators_on_all_qubits[qubit] = Qobj(op)
+
+                # Tensor the operators together
+                extended_operator = tensor(*operators_on_all_qubits)
+                extended_operators.append(extended_operator)
 
         # Apply normalization
         normalization_factor = np.sqrt(self.num_qubits)
@@ -105,7 +137,7 @@ def _extend_kraus_operators_to_multiple_qubits(self, single_qubit_operators):
         return extended_operators
 
     @abstractmethod
-    def _create_single_qubit_operators(self):
+    def _create_single_qubit_operators(self, **kwargs):
         """Create and return single-qubit Kraus operators.
 
         This method should be implemented by each subclass.
@@ -116,17 +148,25 @@ def _create_single_qubit_operators(self):
 class AmplitudeDamping(Channel):
     """Amplitude damping channel for qubits."""
 
-    def __init__(self, T1, num_qubits=1, hilbert_space_dim=2):
-        """Initialize with a given T1 relaxation time and number of qubits."""
-        self.T1 = T1
-        self.num_qubits = num_qubits
-        self.hilbert_space_dim = hilbert_space_dim
-        super().__init__(dims=hilbert_space_dim**num_qubits)
+    def __init__(
+        self,
+        T1: Union[float, List[float]],
+        num_qubits: int = 1,
+        hilbert_space_dim: int = 2,
+    ):
+        """Initialize the amplitude damping channel.
+
+        Args:
+            T1 (float or list of floats): The T1 energy relaxation rate for the qubits.
+            num_qubits (int): The number of qubits on which the channel acts.
+            hilbert_space_dim (int): The dimension of the Hilbert space of each qubit.
+        """
+        super().__init__(num_qubits, hilbert_space_dim, T1=T1)
 
-    def _create_single_qubit_operators(self):
+    def _create_single_qubit_operators(self, T1):
         """Create single-qubit Kraus operators for amplitude damping."""
         if self.hilbert_space_dim == 2:  # standard qubit case
-            _gamma = 1 - np.exp(-self._trotter_dt / self.T1)
+            _gamma = 1 - np.exp(-self._trotter_dt / T1)
             E0_single = np.array([[1, 0], [0, np.sqrt(1 - _gamma)]])
             E1_single = np.array([[0, np.sqrt(_gamma)], [0, 0]])
             return [E0_single, E1_single]
@@ -137,7 +177,7 @@ def _create_single_qubit_operators(self):
             # assuming f->e and e->g are the same and f->g is 0
 
             # Simplified transition rates
-            _gamma_fe = _gamma_eg = 1 - np.exp(-self._trotter_dt / self.T1)
+            _gamma_fe = _gamma_eg = 1 - np.exp(-self._trotter_dt / T1)
             _gamma_fg = 0  # Neglected direct transition from f to g
 
             # Simplified Kraus operators
@@ -160,23 +200,42 @@ def _create_single_qubit_operators(self):
 class PhaseDamping(Channel):
     """Phase damping channel for qubits."""
 
-    def __init__(self, T1, T2, num_qubits=1, hilbert_space_dim=2):
-        """Initialize with a given T_phi and number of qubits."""
-        self.Tphi = 1 / (1 / T2 - 1 / (2 * T1))
-        self.num_qubits = num_qubits
-        self.hilbert_space_dim = hilbert_space_dim
-        super().__init__(dims=hilbert_space_dim**num_qubits)
+    def __init__(
+        self,
+        T1: Union[float, List[float]],
+        T2: Union[float, List[float]],
+        num_qubits: int = 1,
+        hilbert_space_dim: int = 2,
+    ):
+        """Initialize the phase damping channel.
+
+        Args:
+            T1 (float or list of floats): The T1 energy relaxation rate for the qubits.
+            T2 (float or list of floats): The T2 dephrasing rate for the qubits.
+            num_qubits (int): The number of qubits on which the channel acts.
+            hilbert_space_dim (int): The dimension of the Hilbert space of each qubit.
+        """
+        Tphi = self._Tphi(T1, T2)
+        super().__init__(num_qubits, hilbert_space_dim, Tphi=Tphi)
+
+    @staticmethod
+    def _Tphi(T1, T2):
+        """Calculate the dephasing time T_phi from T1 and T2."""
+        if isinstance(T1, list) and isinstance(T2, list):
+            return [1 / (1 / t2 - 1 / (2 * t1)) for t1, t2 in zip(T1, T2)]
+        else:
+            return 1 / (1 / T2 - 1 / (2 * T1))
 
-    def _create_single_qubit_operators(self):
+    def _create_single_qubit_operators(self, Tphi):
         """Create single-qubit Kraus operators for phase damping."""
         if self.hilbert_space_dim == 2:  # standard qubit case
-            _gamma = 1 - np.exp(-self._trotter_dt / self.Tphi)
+            _gamma = 1 - np.exp(-self._trotter_dt / Tphi)
             E0_single = np.array([[1, 0], [0, np.sqrt(1 - _gamma)]])
             E1_single = np.array([[0, 0], [0, np.sqrt(_gamma)]])
             return [E0_single, E1_single]
 
         elif self.hilbert_space_dim == 3:  # qutrit case
-            _gamma = 1 - np.exp(-self._trotter_dt / self.Tphi)
+            _gamma = 1 - np.exp(-self._trotter_dt / Tphi)
             E0_single = np.array(
                 [[1, 0, 0], [0, np.sqrt(1 - _gamma), 0], [0, 0, np.sqrt(1 - _gamma)]]
             )

src/quantum_logical/trotter.py

@@ -30,9 +30,6 @@ def _compose(self, operator):
         if self.continuous_operators:
             if operator.dims != self.continuous_operators[0].dims:
                 raise ValueError("Dimension mismatch among operators in TrotterGroup.")
-
-        if isinstance(operator, TrotterGroup):
-            self.continuous_operators.extend(operator.operators)
         elif isinstance(operator, Channel):
             if operator._trotter_dt != self.trotter_dt:
                 operator.set_trotter_dt(self.trotter_dt)