diff --git a/pyproject.toml b/pyproject.toml index ccb02c8..00f9fa5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,7 @@ dependencies = [ "pyyaml", "pint", "sympy", + "qutip", "h5py" ] classifiers = [ diff --git a/src/ionsim/basis.py b/src/ionsim/basis.py index d3e98b4..4363198 100644 --- a/src/ionsim/basis.py +++ b/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] + diff --git a/src/ionsim/named_operators.py b/src/ionsim/named_operators.py index 1410c92..62afecb 100644 --- a/src/ionsim/named_operators.py +++ b/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.""" diff --git a/src/ionsim/state.py b/src/ionsim/state.py index 0b80cc8..988865a 100644 --- a/src/ionsim/state.py +++ b/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 . + + - 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 . + + - 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 and

for each motional mode in the state. + - fails if there are no motional DOFs + - returns lists of and

; each list element corresponds to 1 motional mode + - list order matches DOF order in the state's basis + - option to return and , 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): diff --git a/tests/test_dissipators.py b/tests/test_dissipators.py index b2bffbc..edddad0 100644 --- a/tests/test_dissipators.py +++ b/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) diff --git a/tests/test_state.py b/tests/test_state.py index bc433c4..e4c37cb 100644 --- a/tests/test_state.py +++ b/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()