57 wigner distribution

pyproject.toml

@@ -18,6 +18,7 @@ dependencies = [
     "pyyaml",
     "pint",
     "sympy",
+    "qutip",
     "h5py"
 ]
 classifiers = [

src/ionsim/basis.py

@@ -200,6 +200,9 @@ def atomic_structure_DOFs(self):
         """ Returns list of atomic structure degrees of freedom or empty list if none. """
         return [DOF for DOF in self.degrees_of_freedom if isinstance(DOF, AtomicStructure)]
 
+    @property
+    def motional_modes(self):
+        return [dof for dof in self.degrees_of_freedom if dof not in self.atomic_structure_DOFs]
 
 @dataclass(frozen=True, eq=False)
 class ZPauliBasis(StandardBasis):
@@ -229,6 +232,11 @@ def vectors(self):
         pairs = list(itertools.product(*[[plus, minus] for dof in self.degrees_of_freedom]))
         return [np.kron(*pair) for pair in pairs]
 
+    @property
+    def spin_DOFs(self):
+        """ Returns list of spin degrees of freedom """ 
+        return self.degrees_of_freedom 
+
 @dataclass(frozen=True, eq=False)
 class YPauliBasis(Basis):
     """A basis in which the basis vectors correspond to the (plus/minus) eigenstates of the x-Pauli spin matrix."""
@@ -270,3 +278,4 @@ def vectors(self):
             ]
         ))
         return [ft.reduce(np.kron, group) for group in groups]
+

src/ionsim/named_operators.py

@@ -58,6 +58,14 @@ def raising(cls, fock_dimension: int):
     def number(cls, fock_dimension: int):
         return np.diag(np.arange(fock_dimension))
 
+    @classmethod # or "x" or "X"?
+    def position(cls, fock_dimension: int):
+        return np.sqrt(0.5)*(cls.raising(fock_dimension) + cls.lowering(fock_dimension))
+
+    @classmethod # or "p" or "P"?
+    def momentum(cls, fock_dimension: int):
+        return 1j*np.sqrt(0.5)*(cls.raising(fock_dimension) - cls.lowering(fock_dimension))
+
     @staticmethod
     def debye_waller_lowering(fock_dimension: int, lamb_dicke_parameter: float):
         """The lowering operator for a harmonic oscillator."""

src/ionsim/state.py

@@ -8,14 +8,16 @@
 #***************************************************************************************************
 
 from ionsim.basis import Basis, StandardBasis
+from ionsim.operator import Operator, CouplingOperator
 from ionsim.degree_of_freedom import DegreeOfFreedom
 from ionsim.custom_types import Vector, Matrix
 from ionsim.ionsim_error import IonSimError
 from ionsim.hamiltonian import Hamiltonian
 from ionsim.dissipator import Dissipator, Lindbladian
+from ionsim.named_operators import Fock 
+from ionsim.collective_motional_energy_level import CollectiveMotionalEnergyLevel 
 
 import numpy as np
-# from typing import Any
 from dataclasses import dataclass
 from numpy.linalg import multi_dot
 
@@ -32,6 +34,19 @@ class State:
     def supervector(self) -> Vector:
         return self.basis.compute_supervector_from_density_matrix(self.density_matrix)
 
+    @property
+    def motional_state(self):
+        """ Motional portion of the state, if it exists. Returns None if no motional DOFs """ 
+        if not self.basis.motional_modes: 
+            return None 
+
+        # Trace out atomic structure DOFs to obtain a purely motional state  
+        state = self
+        for dof in self.basis.atomic_structure_DOFs:
+            state = state.trace_out_degree_of_freedom(dof)
+
+        return state 
+
     @classmethod
     def from_coefficients(cls, basis: Basis, coefficients: list[float]): 
         """Build a state from a list of basis-state coefficients."""
@@ -164,19 +179,134 @@ def trace_out_degree_of_freedom(self, degree_of_freedom: DegreeOfFreedom):
         basis = StandardBasis([dof for dof in self.basis.degrees_of_freedom if dof is not degree_of_freedom])
         return State.from_density_matrix(basis, density_matrix)
 
-    def compute_coherent_displacements(self, spin_dofs: list[DegreeOfFreedom], motional_dof: DegreeOfFreedom):
-        """Compute the coherent displacement (expectation value of the lowering operator) for each spin state."""
-        assert(len(self.basis.degrees_of_freedom) == len(spin_dofs) + 1) # TODO: trace out other degrees of freedom
-        spin_basis = StandardBasis(spin_dofs)
-        lowering = lowering_motion(len(motional_dof.energy_levels))
+    def compute_coherent_displacements(self, atomic_structure_dofs: list[DegreeOfFreedom], motional_dof: DegreeOfFreedom):
+        """Compute the coherent displacement (expectation value of the lowering operator) for each angular momentum state."""
+        assert(len(self.basis.degrees_of_freedom) == len(atomic_structure_dofs) + 1) # TODO: trace out other degrees of freedom
+        structure_basis = StandardBasis(atomic_structure_dofs)
+        lowering = Fock.lowering(len(motional_dof.energy_levels))
         diplacements = []
-        for vector in spin_basis.vectors:
-            spin_proj = spin_basis.compute_projector_matrix(vector)
-            displacement = np.trace(np.kron(spin_proj, lowering).dot(self.density_matrix))
+        for vector in structure_basis.vectors:
+            structure_proj = structure_basis.compute_projector_matrix(vector)
+            displacement = np.trace(np.kron(structure_proj, lowering).dot(self.density_matrix))
             diplacements.append(displacement)
         return diplacements
 
-    # def transform_to_spin_eigenbasis(self):
+    ### Added methods by ECM in this branch:### 
+    def compute_matrix_observable_expectation(self, observable_operator: Matrix) -> Matrix:
+        """ Compute the expectation value of an operator observable, represented by <O>. 
+
+            - Computes via Tr[O rho], where rho is the density matrix. 
+            - O and rho must be in compatable bases.
+
+        """
+        # Check that observable and density matrix are compatible 
+        basis_size = len(self.basis.vectors)
+        if observable_operator.shape != (basis_size, basis_size):
+            raise IonSimError("Observable operator must correspond with a matrix of shape: {(basis_size, basis_size)}.")
+
+        return np.trace(observable_operator.dot(self.density_matrix))
+
+    def compute_operator_observable_expectation(self, observable_operator: Operator) -> Matrix:
+        """ Compute the expectation value of an operator observable, represented by <O>. 
+
+            - Computes via Tr[O rho], where rho is the density matrix. 
+            - O and rho must be in compatable bases.
+
+        """
+        # Check that observable and density matrix are compatible 
+        if isinstance(observable_operator, CouplingOperator) and observable_operator.oscillation_rate != 0.:
+            raise IonSimeError(f"Observable operator must be a static operator, but oscillation rate is {observable_operator.oscillation_rate}.")
+
+        return self.compute_matrix_observable_expectation(observable_operator.static_matrix) 
+
+    def compute_quadratures(self, enable_squared_quadratures: bool=False):
+        """ Returns quadrature expectation values <x> and <p> for each motional mode in the state. 
+            - fails if there are no motional DOFs 
+            - returns lists of <x> and <p>; each list element corresponds to 1 motional mode 
+            - list order matches DOF order in the state's basis 
+            - option to return <x^2> and <p^2>, the expectation of the squared quadrature operators.
+        """ 
+        x = []        
+        p = []        
+        if enable_squared_quadratures:
+            x2 = []        
+            p2 = []        
+
 
-def lowering_motion(fock_dimension: int):
-    return np.diag([np.sqrt(n+1) for n in range(fock_dimension-1)], k=1)
+        use_partial_trace = False 
+        # Either compute partial trace on the density matrix or elevate observable to full space 
+        # It may be cheaper to compute the expectation value in the full space 
+        if not use_partial_trace :
+            full_basis = self.basis
+
+            # For each mode, build the full-space raising and lowering operators  
+            for i, mode_i in enumerate(self.basis.motional_modes):
+                if not isinstance(mode_i.energy_levels[0], CollectiveMotionalEnergyLevel):
+                    raise IonSimError("Quadrature calculation assumes motional mode is in the Fock number state basis.")
+
+                # Enlarge the x and p operators for the current mode 
+                Fock_dim = len(mode_i.energy_levels) 
+                x_enlarged = full_basis.enlarge_one_dof_matrix(Fock.position(Fock_dim), mode_i) 
+                p_enlarged = full_basis.enlarge_one_dof_matrix(Fock.momentum(Fock_dim), mode_i) 
+
+                # Compute expectation values: 
+                x.append(self.compute_matrix_observable_expectation(x_enlarged))
+                p.append(self.compute_matrix_observable_expectation(p_enlarged))
+                if enable_squared_quadratures:
+                    x2.append(self.compute_matrix_observable_expectation(x_enlarged @ x_enlarged)) 
+                    p2.append(self.compute_matrix_observable_expectation(p_enlarged @ p_enlarged))
+
+        else:
+            for i, mode_i in enumerate(self.basis.motional_modes):
+                if not isinstance(mode_i.energy_levels[0], CollectiveMotionalEnergyLevel):
+                    raise IonSimError("Quadrature calculation assumes motional mode is in the Fock number state basis.")
+
+                mode_state = self.motional_state # reset the state  
+                # Trace out other modes 
+                for j, mode_j in enumerate(self.basis.motional_modes):
+                    if i == j:
+                        continue 
+                    mode_state = mode_state.trace_out_degree_of_freedom(mode_j) 
+                assert mode_state is not None
+                Fock_dim = len(mode_i.energy_levels) 
+                # Compute expectation values: 
+                x.append(mode_state.compute_matrix_observable_expectation(Fock.position(Fock_dim)))
+                p.append(mode_state.compute_matrix_observable_expectation(Fock.momentum(Fock_dim)))
+                if enable_squared_quadratures:
+                    x2.append(mode_state.compute_matrix_observable_expectation(Fock.position(Fock_dim) @ Fock.position(Fock_dim)))
+                    p2.append(mode_state.compute_matrix_observable_expectation(Fock.momentum(Fock_dim) @ Fock.momentum(Fock_dim)))
+
+        if enable_squared_quadratures:
+            return x, p, x2, p2
+        else:
+            return x, p
+
+
+    def compute_wigner_distribution(self, x_grid: Vector, p_grid: Vector): 
+        """ Computes W(x,p) the Wigner distribution for each motional mode in the basis; assumes a Fock basis for each mode.
+            - requires a specification of the x and p grids as 1-dimensional arrays, determines resolution of Wigner distributions. 
+            - returns a list of Wigner distributions -> one for each mode. 
+            - returns an empty list if there are no modes in the basis.  
+            - assumes the motional modes are in the Fock number state basis  
+        """  
+        from qutip import Qobj, wigner
+        wigner_distributions = []
+
+        # For each mode, trace out all other modes  
+        for i, mode_i in enumerate(self.basis.motional_modes):
+            if not isinstance(mode_i.energy_levels[0], CollectiveMotionalEnergyLevel):
+                raise IonSimError("Wigner distribution calculation assumes motional mode is in the Fock number state basis.")
+
+            mode_state = self.motional_state # reset the state  
+            assert mode_state is not None
+
+            for j, mode_j in enumerate(self.basis.motional_modes):
+                if i == j:
+                    continue 
+                mode_state = mode_state.trace_out_degree_of_freedom(mode_j) 
+            N_fock = len(mode_i.energy_levels) 
+            wigner_distributions.append(wigner(Qobj(mode_state.density_matrix, dims=[[N_fock], [N_fock]]), x_grid, p_grid))
+
+        return wigner_distributions 
+
+    # def transform_to_spin_eigenbasis(self):

tests/test_dissipators.py

@@ -142,10 +142,6 @@ def setUp(self):
 
                 spin_identity = np.sum(spin_identities, axis=0)
 
-                # Tensor raising/lowering operators with spin identity for each mode that is heating  
-                lowering_operators = []
-                raising_operators = []
-
                 # Basis is of the form spins x modes, so we tensor in the that order: 
                 N_thermal = case['N thermal']
                 lowering_operator = np.sqrt(decay_rate*(N_thermal+1)) * skron(spin_identity, mode_lowering_matrix) 

tests/test_state.py

@@ -12,7 +12,11 @@
 import numpy as np
 
 from ionsim.state import State
-from ionsim.degree_of_freedom import AtomicStructure
+#<<<<<<< HEAD
+#from ionsim.degree_of_freedom import AtomicSpin, MotionalMode
+#=======
+from ionsim.degree_of_freedom import AtomicStructure, MotionalMode
+#>>>>>>> main
 from ionsim.basis import StandardBasis, XPauliBasis
 from ionsim.named_operators import Pauli
 from ionsim.testing import assert_array_close
@@ -21,13 +25,20 @@ class TestState(unittest.TestCase):
 
     def setUp(self):
         """Set up the necessary objects for testing."""
-        self.spin_a = AtomicStructure.from_species(species='171Yb+', term_symbols=['S1/2'], level_names=['S1/2,0,0', 'S1/2,1,0'])
-        self.spin_b = AtomicStructure.from_species(species='171Yb+', term_symbols=['S1/2'], level_names=['S1/2,0,0', 'S1/2,1,0'])
-        self.basis = StandardBasis([self.spin_a, self.spin_b])
-        self.basis_x = XPauliBasis([self.spin_a, self.spin_b])
+        self.atomic_structure_a = AtomicStructure.from_species(species='171Yb+', term_symbols=['S1/2'], level_names=['S1/2,0,0', 'S1/2,1,0'])
+        self.atomic_structure_b = AtomicStructure.from_species(species='171Yb+', term_symbols=['S1/2'], level_names=['S1/2,0,0', 'S1/2,1,0'])
+        self.basis = StandardBasis([self.atomic_structure_a, self.atomic_structure_b])
+        self.basis_x = XPauliBasis([self.atomic_structure_a, self.atomic_structure_b])
         self.eig_x = np.array([[1, 0], [0, -1]])
         self.state = State.from_density_matrix(self.basis_x, np.kron(self.eig_x, Pauli.I))
 
+        # Set up atomic structure - motional state with 2 structure and 1 motional mode  
+        self.mode = MotionalMode.from_frequency(frequency=2*np.pi * 3e6, fock_dimension=5)
+        self.atomic_structure_motion_basis = StandardBasis([self.atomic_structure_a, self.atomic_structure_b, self.mode])
+        state_coefficients = np.zeros(len(self.atomic_structure_motion_basis.states))
+        state_coefficients[0] = 1.
+        self.atomic_structure_motion_state = State.from_coefficients(self.atomic_structure_motion_basis, list(state_coefficients))
+
     def test_density_matrix_in_new_basis(self):
         """Test the density matrix in the new basis."""
         rho_p = State.from_state(self.basis, self.state).density_matrix
@@ -36,5 +47,38 @@ def test_density_matrix_in_new_basis(self):
         # Check that the density matrix matches the expected value
         assert_array_close(rho_p, expected_rho_p)
 
+    def test_wigner_function(self):
+        """ Test Wigner function computation from a motional density matrix in the Fock state basis""" 
+        domain_limit = 2.
+        x_grid = np.linspace(-domain_limit, domain_limit, 25)
+        p_grid = np.linspace(-domain_limit, domain_limit, 25)
+        W_distribution = self.atomic_structure_motion_state.compute_wigner_distribution(x_grid, p_grid)[0]        
+
+        # Test just one slice of the array and half of its contents (U(1) symmetric about origin)) 
+        slice_indx = 12
+        W_expected_at_slice = np.array([0.00583005, 0.0110443, 0.0197914, 0.03354962, 0.05379861, 0.08160694, 0.11709966, 0.15894861, 0.20409406, 0.24789999, 0.2848362, 0.30958962, 0.31830989]) 
+        norm = np.trapezoid(np.trapezoid(W_distribution, x_grid, axis=0), p_grid)
+        expected_norm = 0.9902872798196112
+        self.assertAlmostEqual(norm, expected_norm, places=7)
+        expected_W_max = 0.31830988618379075
+        W_max = np.max(W_distribution)
+        self.assertAlmostEqual(W_max, expected_W_max, places=7)
+        assert_array_close(W_distribution[slice_indx, 0:13], W_expected_at_slice, rtol=1e-04, atol=1E-7) 
+
+
+    def test_quadrature_computation(self):
+        """ Yest computation of x, p, and x^2, and p^2 expectations """ 
+        include_squared_quantities = True
+        x, p, x2, p2 = self.atomic_structure_motion_state.compute_quadratures(include_squared_quantities)
+
+        x_expected = 0.
+        p_expected = 0.
+        x2_expected = 0.5
+        p2_expected = 0.5
+        self.assertAlmostEqual(x[0].real, x_expected, places=6) 
+        self.assertAlmostEqual(p[0].real, p_expected, places=6) 
+        self.assertAlmostEqual(x2[0].real, x2_expected, places=6) 
+        self.assertAlmostEqual(p2[0].real, p2_expected, places=6) 
+
 if __name__ == '__main__':
     unittest.main()