diff --git a/src/ionsim/atomic_internal_energy_level.py b/src/ionsim/atomic_internal_energy_level.py index 944b35d..f475718 100644 --- a/src/ionsim/atomic_internal_energy_level.py +++ b/src/ionsim/atomic_internal_energy_level.py @@ -8,7 +8,7 @@ #*************************************************************************************************** from abc import ABC, abstractmethod -from dataclasses import dataclass +from dataclasses import dataclass, field from fractions import Fraction from sympy.physics.wigner import wigner_3j, wigner_6j import sympy @@ -25,6 +25,7 @@ class AtomicInternalEnergyLevel(EnergyLevel): term_symbol: str fine_energy: float hyperfine_A: float + alias: str | None = field(default=None, kw_only=True) @property @abstractmethod diff --git a/src/ionsim/collective_motional_energy_level.py b/src/ionsim/collective_motional_energy_level.py index cd4f4de..843b45c 100644 --- a/src/ionsim/collective_motional_energy_level.py +++ b/src/ionsim/collective_motional_energy_level.py @@ -20,6 +20,7 @@ class CollectiveMotionalEnergyLevel(EnergyLevel): """An energy level of a normal mode of motion of an ion chain.""" mode_frequency: float fock_number: float # TODO: someday we'll need to make a motional basis of coherent states + alias: str | None=None @property def energy(self): @@ -30,4 +31,4 @@ def energy(self): def name(self): """A unique name for the collective motional energy level.""" return str(self.fock_number) - \ No newline at end of file + diff --git a/src/ionsim/degree_of_freedom.py b/src/ionsim/degree_of_freedom.py index 3e1a184..afb978d 100644 --- a/src/ionsim/degree_of_freedom.py +++ b/src/ionsim/degree_of_freedom.py @@ -15,7 +15,7 @@ import importlib.resources from pathlib import Path -from dataclasses import dataclass +from dataclasses import dataclass, replace from abc import ABC import yaml from fractions import Fraction @@ -36,9 +36,9 @@ class AtomicStructure(DegreeOfFreedom): energy_levels: list[AtomicInternalEnergyLevel] @classmethod - def from_species(cls, species: str, term_symbols: list[str] | None = None, level_names: list[str] | None = None, - name: str | None = None, magnetic_field: float=0., **kwargs): - """Build the atomic spin degree of freedom for a particular species of atom.""" + def from_species(cls, species: str, term_symbols: list[str] | None = None, level_names: list[str] | None = None, + level_aliases: list[str] | None=None, name: str | None = None, magnetic_field: float=0., **kwargs): + """Build the atomic structure degree of freedom for a particular species of atom.""" config_data = cls.get_config_data(species) nuclear_spin = config_data['nuclear_spin'] levels_data = config_data['levels'] @@ -50,6 +50,9 @@ def from_species(cls, species: str, term_symbols: list[str] | None = None, level if term_symbols is not None: levels_data = cls.select_some_data(term_symbols, levels_data) + if level_aliases and len(level_aliases) != len(level_names): + raise IonSimError(f'User should specify a level alias for each level in the atomic structure. Expected {len(level_names)} but have {level_aliases}') + levels = [] for level_data in levels_data: @@ -102,6 +105,11 @@ def from_species(cls, species: str, term_symbols: list[str] | None = None, level # Create the level level = FineLevel(**fine_data, mj=mj, external_energy_shift=zeeman_shift_energy*np.pi*2.) if level_names is None or level.name in level_names: + if level_aliases: + level_name_index = level_names.index(level.name) + level_alias = level_aliases[level_name_index] + # Overwrite the level to include its alias + level = replace(level, alias = level_alias) levels.append(level) else: for f in np.arange(np.abs(j - nuclear_spin), j + nuclear_spin + 1): @@ -115,6 +123,11 @@ def from_species(cls, species: str, term_symbols: list[str] | None = None, level # Create the level level = HyperfineLevel(**fine_data, i=nuclear_spin, f=f, mf=mf, external_energy_shift = zeeman_shift_energy * np.pi * 2.) if level_names is None or level.name in level_names: + if level_aliases: + level_name_index = level_names.index(level.name) + level_alias = level_aliases[level_name_index] + # Overwrite the level to include its alias + level = replace(level, alias = level_alias) levels.append(level) return cls(levels, name) @@ -211,7 +224,10 @@ class MotionalMode(DegreeOfFreedom): energy_levels: list[CollectiveMotionalEnergyLevel] @classmethod - def from_frequency(cls, frequency: float, fock_dimension: int, name: str | None = None): + def from_frequency(cls, frequency: float, fock_dimension: int, name: str | None = None, level_aliases: list[str] | None=None): """Build a motional normal-mode degree of freedom for an ion chain.""" - levels = [CollectiveMotionalEnergyLevel(frequency, fock_number) for fock_number in range(fock_dimension)] + if level_aliases: + levels = [CollectiveMotionalEnergyLevel(frequency, fock_number, alias = level_aliases[fock_number]) for fock_number in range(fock_dimension)] + else: + levels = [CollectiveMotionalEnergyLevel(frequency, fock_number) for fock_number in range(fock_dimension)] return cls(levels, name) diff --git a/src/ionsim/energy_level.py b/src/ionsim/energy_level.py index 6e6d673..32c1c52 100644 --- a/src/ionsim/energy_level.py +++ b/src/ionsim/energy_level.py @@ -40,6 +40,10 @@ def name(self): """A unique name for the state.""" return ' : '.join([component.name for component in self.components]) + def alias(self, join_by: str=''): + """An alias for the state, allowing user specification on how to join state aliases.""" + return join_by.join([component.alias for component in self.components]) + @dataclass(frozen=True, eq=False) class InternalEnergyLevel(EnergyLevel): """ A simple energy level that is agnostic to physical qubit (e.g. atom or ion) details. """ @@ -53,3 +57,4 @@ def energy(self): def name(self): """ The name of the level""" return self.name + diff --git a/tests/test_basis.py b/tests/test_basis.py index 3431e1b..32369ae 100644 --- a/tests/test_basis.py +++ b/tests/test_basis.py @@ -18,10 +18,10 @@ class TestBasis(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.mode_0 = MotionalMode.from_frequency(frequency=3e6*2*np.pi, fock_dimension=3) - self.mode_1 = MotionalMode.from_frequency(frequency=4e6*2*np.pi, fock_dimension=2) + self.spin_a = AtomicStructure.from_species(species='171Yb+', term_symbols=['S1/2'], level_names=['S1/2,0,0', 'S1/2,1,0'], level_aliases=['0', '1']) + self.spin_b = AtomicStructure.from_species(species='171Yb+', term_symbols=['S1/2'], level_names=['S1/2,0,0', 'S1/2,1,0'], level_aliases=['0', '1']) + self.mode_0 = MotionalMode.from_frequency(frequency=3e6*2*np.pi, fock_dimension=3, level_aliases = ['Mode 0, n = ' + str(n) for n in range(3)]) + self.mode_1 = MotionalMode.from_frequency(frequency=4e6*2*np.pi, fock_dimension=2, level_aliases = ['Mode 1, n = ' + str(n) for n in range(2)]) def test_spin_basis_states(self): """Test the states of the spin basis.""" @@ -35,6 +35,12 @@ def test_spin_basis_states(self): actual_states = [state.name for state in spin_basis.states] self.assertEqual(actual_states, expected_states) + expected_aliases = ['00', '01', '10', '11'] + actual_aliases = [state.alias() for state in spin_basis.states] + + # Test aliases for spin basis states + self.assertEqual(expected_aliases, actual_aliases) + def test_motional_basis_states(self): """Test the states of the motional basis.""" motional_basis = StandardBasis([self.mode_0, self.mode_1]) @@ -49,6 +55,9 @@ def test_motional_basis_states(self): actual_states = [state.name for state in motional_basis.states] self.assertEqual(actual_states, expected_states) + # Test alias for motional basis state + self.assertEqual(motional_basis.states[4].alias('; '), 'Mode 0, n = 2; Mode 1, n = 0') + def test_full_basis_states(self): """Test the states of the full basis.""" full_basis = StandardBasis([self.spin_a, self.spin_b, self.mode_0, self.mode_1]) @@ -80,6 +89,8 @@ def test_full_basis_states(self): ] actual_states = [state.name for state in full_basis.states] self.assertEqual(actual_states, expected_states) + # Test alias for full basis state + self.assertEqual(full_basis.states[3].alias('; '), '0; 0; Mode 0, n = 1; Mode 1, n = 1') def test_z_pauli_basis_vectors(self): """Test the vectors of the Z Pauli basis."""