From 764abc56d219b3f7235b97acd71c8ce17da5b369 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 09:54:29 -0600
Subject: [PATCH 01/31]   begin drafting of trapped-ion mode analysis class
 from Wes

---
 src/ionsim/trapped_ion_mode_analysis.py | 695 ++++++++++++++++++++++++
 1 file changed, 695 insertions(+)
 create mode 100644 src/ionsim/trapped_ion_mode_analysis.py

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
new file mode 100644
index 0000000..8a42d45
--- /dev/null
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -0,0 +1,695 @@
+import numpy as np
+from generalized_mode_analysis import GeneralizedModeAnalysis as mode_analyzer
+#plt.style.use('ionsim')
+from time import time as timer
+import scipy.constants as const
+import warnings
+import scipy.optimize as opt    
+
+
+
+def characteristic_length(q, m, wz):
+    k_e = 1 / (4 * np.pi * const.epsilon_0)  # Coulomb constant
+    l0 = ((k_e * q ** 2) / (.5 * m * wz ** 2)) ** (1 / 3)
+    return l0
+
+
+
+def get_norm(en,H):
+    """
+    Get the norm of the eigen vector en w.r.t. the Hamiltonian H.
+    """
+    norm = np.sqrt(en.T.conj() @ H @ en)
+    return norm
+
+
+
+def normalize_eigen_vectors(ens,H,evs=None): 
+    """
+    Rescale the eigen vectors ens w.r.t. the Hamiltonian H.
+    """
+    ens_rescaled = np.zeros_like(ens,dtype=complex)
+    num_coords,num_evs = np.shape(ens)
+    if evs is None:
+        evs = np.ones(num_evs)
+    for i in range(num_evs):
+        en = ens[:,i].reshape(num_coords,1)
+        norm = get_norm(en,H)
+        ens_rescaled[:,i] = en[:,0]/norm
+        ens_rescaled[:,i] *= np.sqrt(evs[i])
+    return ens_rescaled 
+
+
+
+def get_canonical_transformation(H,ens,evs=None):
+    """
+    Given the eigen-solve of the dynamical matrix, get the transform matrix 
+    to the canonical coordinates.
+    X = S X', where X' = (Q,P)^T and X = (q,p)^T
+    """
+    ## TODO: understand the sign in the transformation matrix
+    sign = -1
+    num_coords, num_evs = np.shape(ens)
+    assert num_coords //2 == num_evs    
+    T = np.zeros((num_coords,num_coords),dtype=complex) 
+    ens = normalize_eigen_vectors(ens,H,evs=evs)
+    T = np.sqrt(2)*np.concatenate((np.real(ens), sign*np.imag(ens)), axis=1)
+    # the sign is likely due to the convention of writing the time evolution as exp(-iwt)
+    return T
+
+
+def convert_num_to_array(x):
+    """ Converts a scalar to a vector or returns the vector """ 
+    if not hasattr(x, "__len__"):
+        x = np.ones(N) * x 
+        return x
+    return np.array(x) 
+
+
+class TrappedIonModeAnalysis:
+    def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float, atomic_masses: Vector[float] | float, atomic_numbers: Vector[int] | int): 
+        """ Class for determining properties of trapped-ion phonon modes, requiring the following system parameters:
+
+            - num_ions: the number of ions
+            - omega_x: harmonic trap frequency in the x direction in units of rad/s. 
+            - omega_y: harmonic trap frequency in the y direction in units of rad/s. 
+            - omega_z: harmonic trap frequency in the z direction. This is the "axial" direction used for 1D chains.  
+            - atomic masses: array of atomic masses in amu or a single number => same mass for each ion. 
+            - atomic numbers: array of (Z) atomic numbers (# of protons of an element) or a single number => all ions are the same.  
+        """ 
+        # TODO: Decide how to handle units and how much of pint we should use.  
+        self.num_ions = num_ions
+
+        self.atomic_masses = convert_num_to_array(atomic_masses)
+        self.atomic_numbers = convert_num_to_array(atomic_numbers)
+
+        # Safety checks: 
+        assert len(self.atomic_masses) == num_ions
+        assert len(self.atomic_numbers) == num_ions
+
+        # Proton charge q:
+        #self.q = z * const.e # const.e ==> elementary charge in Coulombs 
+        self.proton_charges = self.atomic_numbers * const.e # Z * e, const.e ==> elementary charge in Coulombs 
+
+        # Trapping frequencies for each ion 
+        self.wx_E = self.calculate_species_trap_frequencies(self.m_E, q_E, omega_x)   
+        self.wy_E = self.calculate_species_trap_frequencies(self.m_E, self.q_E, omega_y)
+        self.wz_E = self.calculate_species_trap_frequencies(self.m_E, self.q_E, omega_z)    
+
+        self.z = ensure_numpy_array(Z,N)
+        self.ionmass_amu = ensure_numpy_array(ionmass_amu,N)
+        self.q_E = self.Z * const.e
+        self.atomic_mass = self.ionmass_amu * const.atomic_mass # amu -> kilograms  
+        self.initial_equilibrium_guess = None
+        self.hasrun = False 
+    
+    def calculate_species_trap_frequencies(self, omega: float) -> Vector[float]: 
+        # assume that the trapping frequency given corresponds to the first ion species 
+        u_r_sqr = np.sqrt(self.atomic_masses / self.proton_charges[0]) * omega
+        omega_E = np.sqrt(self.proton_charges / self.atomic_masses) * u_r_sqr    
+        return omega_E 
+   
+    def dimensionless_parameters(self):
+        # normalize to the first ion species    
+        q0 = self.q_E[0]
+        m0 = self.m_E[0]    
+        w0 = self.wz_E[0]   
+        self.q0 = q0
+        self.m0 = m0
+        self.w0 = w0
+        # ion properties    
+        self.m = self.m_E / m0
+        self.q = self.q_E / q0  
+        # trap frequencies
+        self.wz = self.wz_E / w0 
+        self.wy = self.wy_E / w0 
+        self.wx = self.wx_E / w0 
+        # system parameters
+        self.l0 = characteristic_length(q0, m0, w0)  # characteristic length 
+        self.t0 = 1 / w0  # characteristic time
+        self.v0 = self.l0 / self.t0  # characteristic velocity
+        self.E0 = 0.5 * m0 * self.v0 ** 2  # characteristic energy  
+
+    
+    
+    def trap_is_stable(self):
+        # check that all trap frequencies are positive  
+        return np.all(self.wz_E > 0) and np.all(self.wy_E > 0) and np.all(self.wx_E > 0)
+
+
+
+    def check_for_zero_modes(self):
+        assert np.all(self.evals > 0), "All eigenvalues must be positive"   
+
+    def check_outer_relation(self): 
+        H = self.H_matrix.copy()       
+        D = self.get_symplectic_matrix() @ H
+        Eval, Evec = np.linalg.eig(D)
+        _, en = self.sort_modes(Eval,Evec)
+        en = self.normalize_eigen_vectors(en,H)
+        Outers = np.zeros((6*self.N,6*self.N),dtype=complex)
+        for i in range(6*self.N):
+            norm = get_norm(en[:,i],H)
+            Outers = Outers + np.outer(en[:,i],en[:,i].conj())/norm 
+        I_right = H @ Outers
+        I_left  = Outers @ H
+        eye = np.eye(6*self.N,dtype=complex)
+        np.set_printoptions(precision=2, suppress=True) 
+        try:
+            assert np.allclose(I_left,eye) 
+            assert np.allclose(I_right,eye)
+        except AssertionError:
+            warnings.warn("Outer relation check failed")
+
+
+
+    def has_duplicate_evals(self,evals):
+        evs = evals.copy()    
+        return np.any(np.triu(np.isclose(evs[:, None], evs[None, :], atol=1e-6), k=1))  
+    
+
+
+    def check_diagnolization(self):
+        M = np.linalg.inv(self.T_matrix) @ self.S_matrix    
+        H_diag = M.T @ self.E_matrix @ M    
+        H_diag_check = np.diag(np.tile(self.evals,2)) 
+        np.set_printoptions(precision=2, suppress=True) 
+        try:
+            assert np.allclose(H_diag,H_diag_check)
+        except AssertionError:
+            warnings.warn("Diagnolization check failed")    
+            print("has duplicate evals: ", self.has_duplicate_evals(self.evals))
+
+
+
+    def checks(self):   
+        self.check_outer_relation()
+        self.check_diagnolization()
+
+    def run(self):
+        self.dimensionless_parameters()
+        assert self.trap_is_stable()    
+        
+        self.u = self.calculate_equilibrium_positions()
+        #self.reindex_ions(self.u)
+        self.E_matrix = self.get_E_matrix(self.u)  
+        self.T_matrix = self.get_momentum_transform() 
+        self.H_matrix = self.get_H_matrix(self.T_matrix, self.E_matrix)   
+        self.evals, self.evecs = self.calculate_normal_modes(self.H_matrix)
+        self.evecs_vel = self.get_eigen_vectors_xv_coords(self.T_matrix,self.evecs)    
+        self.check_for_zero_modes() 
+        self.S_matrix = self.get_canonical_transformation() 
+        self.checks() 
+        self.hasrun = True  
+    
+
+
+    def calculate_equilibrium_positions(self):
+
+        if self.initial_equilibrium_guess is None:
+            self.initial_equilibrium_guess = self.get_initial_equilibrium_guess()
+        else:
+            self.u0 = self.initial_equilibrium_guess
+
+        u = self.find_equilibrium_positions(self.u0)
+        self.p0 = self.potential(u)
+        return u   
+
+
+
+    def get_canonical_transformation(self):
+        return get_canonical_transformation(self.H_matrix,self.evecs,evs=self.evals)    
+
+
+
+    def normalize_eigen_vectors(self, evecs, H_matrix,evs=None):
+        return normalize_eigen_vectors(evecs,H_matrix,evs=evs) 
+
+
+
+    def get_eigen_vectors_xv_coords(self,T,ens): 
+        ens_vel = np.zeros_like(ens,dtype=complex)
+        ens_vel[:,:] = np.linalg.inv(T) @ ens
+        return ens_vel 
+
+
+    def calculate_normal_modes(self, H_matrix):
+
+        J = self.get_symplectic_matrix()
+        D_matrix = J @ H_matrix  
+        # u_n(t) = exp(-i w_n t) u_n(0), minus by convention
+        evals, evecs = np.linalg.eig(-D_matrix)   
+        evals, evecs = self.organize_modes(evals, evecs)
+        evecs = self.normalize_eigen_vectors(evecs, H_matrix) 
+
+        return evals, evecs 
+
+
+
+    def get_initial_equilibrium_guess(self):
+        self.u0 = np.zeros(3*self.N)
+        self.u0[:] = (np.random.rand(3*self.N) * 2 - 1) * self.N 
+        return self.u0
+
+
+
+    def reindex_ions(self, u):  
+        # based on the distance from the center of the trap, reindex the ions, smallest index is closest to the center
+        x = u[0:self.N]
+        y = u[self.N:2*self.N]
+        z = u[2*self.N:]
+        r = np.sqrt(x**2 + y**2 + z**2)
+        idx = np.argsort(r)
+        u = np.hstack((x[idx], y[idx], z[idx]))
+        self.u = u 
+        self.m = self.m[idx]
+        self.m_E = self.m_E[idx]
+        self.q = self.q[idx]
+        self.q_E = self.q_E[idx]
+        self.Z = self.Z[idx]
+
+
+
+    def find_equilibrium_positions(self, u0):
+        bfgs_tolerance = 1e-34
+        out = opt.minimize(self.potential, u0, method='BFGS', jac=self.force,
+                                    options={'gtol': bfgs_tolerance, 'disp': False})
+        return out.x
+
+    def potential_trap(self, pos_array):
+        x = pos_array[0:self.N]
+        y = pos_array[self.N:2*self.N]
+        z = pos_array[2*self.N:]
+        V_trap = 0.5 * np.sum((self.m * self.wx ** 2) * x ** 2) + \
+            0.5 * np.sum((self.m * self.wy ** 2) * y ** 2) + \
+                0.5 * np.sum((self.m * self.wz ** 2) * z ** 2)
+        return V_trap
+    
+
+
+    def potential_coulomb(self, pos_array):
+        x = pos_array[0:self.N]
+        y = pos_array[self.N:2*self.N]
+        z = pos_array[2*self.N:]
+
+        dx = x[:, np.newaxis] - x
+        dy = y[:, np.newaxis] - y
+        dz = z[:, np.newaxis] - z
+        rsep = np.sqrt(dx ** 2 + dy ** 2 + dz ** 2).astype(np.float64)
+        qq = (self.q * self.q[:, np.newaxis]).astype(np.float64)    
+
+        with np.errstate(divide='ignore'):
+            V_Coulomb = np.sum( np.where(rsep != 0., qq / rsep, 0) ) / 2 # divide by 2 to avoid double counting
+        V_Coulomb *= .5 
+        return V_Coulomb
+
+    def potential(self, pos_array):
+        return self.potential_trap(pos_array) + self.potential_coulomb(pos_array)   
+
+
+
+    def force_trap(self, pos_array):
+        x = pos_array[0:self.N]
+        y = pos_array[self.N:2*self.N]
+        z = pos_array[2*self.N:]
+
+        Ftrapx = self.m * self.wx**2 * x
+        Ftrapy = self.m * self.wy**2 * y
+        Ftrapz = self.m * self.wz**2 * z
+
+        force_trap = np.hstack((Ftrapx, Ftrapy, Ftrapz))
+        return force_trap
+
+
+
+    def force_coulomb(self, pos_array): 
+        x = pos_array[0:self.N]
+        y = pos_array[self.N:2*self.N]
+        z = pos_array[2*self.N:]
+
+        dx = x[:, np.newaxis] - x
+        dy = y[:, np.newaxis] - y
+        dz = z[:, np.newaxis] - z
+        rsep = np.sqrt(dx ** 2 + dy ** 2 + dz ** 2).astype(np.float64)
+        qq = (self.q * self.q[:, np.newaxis]).astype(np.float64)    
+
+        with np.errstate(divide='ignore', invalid='ignore'):
+            rsep3 = np.where(rsep != 0., rsep ** (-3), 0)
+
+        fx = dx * rsep3 * qq
+        fy = dy * rsep3 * qq
+        fz = dz * rsep3 * qq    
+
+        Fx = -np.sum(fx, axis=1)
+        Fy = -np.sum(fy, axis=1)
+        Fz = -np.sum(fz, axis=1)
+
+        force_coulomb = np.hstack((Fx, Fy, Fz))
+        force_coulomb *= 0.5    
+        return force_coulomb
+
+
+
+    def force(self, pos_array):
+        Force = self.force_coulomb(pos_array) + self.force_trap(pos_array)  
+        return Force
+
+    def force(self, pos_array):
+        Force = self.force_coulomb(pos_array) + self.force_trap(pos_array)  
+        return Force
+
+
+
+    def hessian_trap(self, pos_array):
+        Hxx = np.diag(self.m * (self.wx**2) * np.ones(self.N))
+        Hyy = np.diag(self.m * (self.wy**2) * np.ones(self.N))  
+        Hzz = np.diag(self.m * (self.wz**2) * np.ones(self.N))  
+        zeros = np.zeros((self.N, self.N))  
+        H = np.block([[Hxx, zeros, zeros], [zeros, Hyy, zeros], [zeros, zeros, Hzz]])
+        return H
+
+
+
+    def hessian_coulomb(self, pos_array):
+        x = pos_array[0:self.N]
+        y = pos_array[self.N:2*self.N]
+        z = pos_array[2*self.N:]
+        
+        dx = x[:, np.newaxis] - x
+        dy = y[:, np.newaxis] - y
+        dz = z[:, np.newaxis] - z
+        rsep = np.sqrt(dx ** 2 + dy ** 2 + dz ** 2).astype(np.float64)  
+        qq = (self.q * self.q[:, np.newaxis]).astype(np.float64)    
+
+        with np.errstate(divide='ignore'):
+            rsep5 = np.where(rsep != 0., rsep ** (-5), 0)
+
+        dxsq = dx ** 2
+        dysq = dy ** 2
+        dzsq = dz ** 2
+        rsep2 = rsep ** 2
+
+        # X derivatives, Y derivatives for alpha != beta
+        Hxx = (rsep2 - 3 * dxsq) * rsep5
+        Hyy = (rsep2 - 3 * dysq) * rsep5
+        Hzz = (rsep2 - 3 * dzsq) * rsep5
+
+        # Above, for alpha == beta
+        Hxx[np.diag_indices(self.N)] = -np.sum(Hxx, axis=0)
+        Hyy[np.diag_indices(self.N)] = -np.sum(Hyy, axis=0)
+        Hzz[np.diag_indices(self.N)] = -np.sum(Hzz, axis=0)
+
+        Hxy = -3 * dx * dy * rsep5
+        Hxy[np.diag_indices(self.N)] = 3 * np.sum(dx * dy * rsep5, axis=0)
+        Hxz = -3 * dx * dz * rsep5
+        Hxz[np.diag_indices(self.N)] = 3 * np.sum(dx * dz * rsep5, axis=0)
+        Hyz = -3 * dy * dz * rsep5
+        Hyz[np.diag_indices(self.N)] = 3 * np.sum(dy * dz * rsep5, axis=0)
+        
+        Hxx *= qq
+        Hyy *= qq
+        Hzz *= qq
+        Hxy *= qq
+        Hxz *= qq
+        Hyz *= qq
+
+        H_coulomb = np.block([[Hxx, Hxy, Hxz], [Hxy, Hyy, Hyz], [Hxz, Hyz, Hzz]])
+        H_coulomb /= 2
+        return H_coulomb
+
+
+    def hessian(self, pos_array):
+        H = self.hessian_coulomb(pos_array) + self.hessian_trap(pos_array)  
+        return H
+
+    def get_mass_matrix(self,m):    
+        return np.diag(np.tile(m, 3)) 
+
+    def get_E_matrix(self,u): 
+        PE_matrix = np.zeros((3*self.N, 3*self.N), dtype=np.complex128)
+        KE_matrix = np.zeros((3*self.N, 3*self.N), dtype=np.complex128)
+        E_matrix = np.zeros((6*self.N, 6*self.N), dtype=np.complex128)
+
+        PE_matrix = self.hessian(u)
+        KE_matrix = self.get_mass_matrix(self.m) 
+        zeros = np.zeros((3*self.N, 3*self.N)) 
+        E_matrix = np.block([[PE_matrix, zeros], [zeros, KE_matrix]])
+        return E_matrix
+
+
+
+    def get_H_matrix(self, T_matrix, E_matrix):  
+        T_matrix_inv = np.linalg.inv(T_matrix)  
+        H_matrix = T_matrix_inv.T @ E_matrix @ T_matrix_inv
+        return H_matrix 
+
+
+
+    def get_momentum_transform(self):
+        # assuming no magnetic field
+        mass_matrix = self.get_mass_matrix(self.m)  
+        eye = np.eye(3*self.N)  
+        zeros = np.zeros((3*self.N, 3*self.N))
+        T = np.block([[eye, zeros], [zeros, mass_matrix]])  
+        return T    
+
+
+
+    def get_symplectic_matrix(self):
+        zeros = np.zeros((3*self.N, 3*self.N), dtype=np.complex128)
+        I = np.eye(3*self.N, dtype=np.complex128)
+        J = np.block([[zeros, I], [-I, zeros]])
+        return J    
+
+    def sort_modes(self,evals, evecs):
+       evals = np.imag(evals)
+       sort_dex = np.argsort(evals)
+       evals = evals[sort_dex]
+       evecs = evecs[:,sort_dex]
+       return evals, evecs  
+    
+
+
+    def split_modes(self,evals, evecs): 
+       half = len(evals) // 2
+       evals = evals[half:]
+       evecs = evecs[:,half:]
+       return evals, evecs
+    
+
+
+    def organize_modes(self,evals, evecs):  
+        evals, evecs = self.sort_modes(evals, evecs)    
+        evals, evecs = self.split_modes(evals, evecs)
+        return evals, evecs 
+
+
+ 
+ 
+#classes
+class GeneralizedModeAnalysisWithBranchSortedModes(mode_analyzer):
+ 
+    def _xyz_classify_modes(self, evecs):
+        N_coords, N_modes = np.shape(evecs)
+        N_ions = N_coords // 6
+        classifier = np.zeros(N_modes, dtype=int)
+        for mode_index in range(N_modes):
+            x_score = np.sum(np.abs(evecs[0:N_ions, mode_index])) + np.sum(np.abs(evecs[3*N_ions:4*N_ions, mode_index]))
+            y_score = np.sum(np.abs(evecs[N_ions:2*N_ions, mode_index])) + np.sum(np.abs(evecs[4*N_ions:5*N_ions, mode_index]))
+            z_score = np.sum(np.abs(evecs[2*N_ions:3*N_ions, mode_index])) + np.sum(np.abs(evecs[5*N_ions:6*N_ions, mode_index]))
+            scores = [x_score, y_score, z_score]
+            max_index = np.argmax(scores)
+            classifier[mode_index] = max_index
+        for direction in range(3):
+            count = np.sum(classifier == direction)
+            assert count == N_modes // 3, f"Classification error: direction {direction} has {count} modes, expected {N_modes // 3}"
+        return classifier
+ 
+    def sort_by_branch(self, evals, evecs):
+        classifier = self._xyz_classify_modes(evecs)
+        # within each branch, sort by frequency
+        N_ions = len(evals) // 3
+        sorted_by_branch_evals = np.zeros_like(evals)
+        sorted_by_branch_evecs = np.zeros_like(evecs)
+        for direction in range(3):
+            direction_indices = np.where(classifier == direction)[0]
+            # they are already sorted by frequency from the original mode analysis code, so we can just take them in order
+            sorted_by_branch_evals[direction*N_ions:(direction+1)*N_ions] = evals[direction_indices]
+            sorted_by_branch_evecs[:, direction*N_ions:(direction+1)*N_ions] = evecs[:, direction_indices]
+        return sorted_by_branch_evals, sorted_by_branch_evecs
+ 
+    def organize_modes(self, evals, evecs):        
+        evals, evecs = self.sort_modes(evals, evecs)
+        evals, evecs = self.split_modes(evals, evecs)
+        evals, evecs = self.sort_by_branch(evals, evecs)
+        return evals, evecs
+ 
+ 
+    def reindex_ions_by_z(self, u):  
+        # based on the position along z, lowest i, is 0, up to N - 1
+        x = u[0:self.N]
+        y = u[self.N:2*self.N]
+        z = u[2*self.N:]
+ 
+        idx = np.argsort(z)
+        self.u = np.hstack((x[idx], y[idx], z[idx]))
+ 
+        # reindex the other arrays accordingly
+        self.m = self.m[idx]
+        self.m_E = self.m_E[idx]
+        self.q = self.q[idx]
+        self.q_E = self.q_E[idx]
+        self.Z = self.Z[idx]
+        self.ionmass_amu = self.ionmass_amu[idx]
+        # trapping frequencies
+        self.wz_E = self.wz_E[idx]
+        self.wy_E = self.wy_E[idx]
+        self.wx_E = self.wx_E[idx]
+        # all ions are the same so this is safe
+        self.dimensionless_parameters()
+ 
+ 
+    def run(self):
+        self.dimensionless_parameters()
+        assert self.trap_is_stable()    
+        
+        self.u = self.calculate_equilibrium_positions()
+        self.reindex_ions_by_z(self.u)
+        self.E_matrix = self.get_E_matrix(self.u)  
+        self.T_matrix = self.get_momentum_transform()
+        self.H_matrix = self.get_H_matrix(self.T_matrix, self.E_matrix)   
+        self.evals, self.evecs = self.calculate_normal_modes(self.H_matrix)
+        self.evecs_vel = self.get_eigen_vectors_xv_coords(self.T_matrix,self.evecs)    
+        self.check_for_zero_modes()
+        self.S_matrix = self.get_canonical_transformation()
+        self.checks()
+        self.hasrun = True  
+    
+ 
+#You could redefine the equilibrium finding function to assert that the equilibrium is linear. For example, make a wrapper inside the function for the potential, Jacobian, and Hessian that forces x_i and y_i = 0. 
+ 
+#This is the code for the Lamb-Dicke parameters: (not that overall phases don't matter here, but the relative phase does. We could pin down the phase with some convention like, "each mode's first non-negative LD value is defined positive." I also include some extra helper functions. 
+ 
+
+## helper ? 
+
+def two_central_ion_separation(wz_Hz, l_2_cent):
+    mass_yb_amu = 170.936 # TODO: hardcoded for now
+    four_ion_analysis = GeneralizedModeAnalysisWithBranchSortedModes(N=4, wz=2*np.pi*wz_Hz, wy=2*np.pi*10e6, wx=2*np.pi*11e6, ionmass_amu=mass_yb_amu)
+    four_ion_analysis.run()
+    positions_z = four_ion_analysis.u[2*4:] * four_ion_analysis.l0
+    central_ions = np.argsort(np.abs(positions_z))[:2]
+    dl = np.abs(positions_z[central_ions[0]] - positions_z[central_ions[1]])
+    return dl - l_2_cent
+    
+def find_wz_for_desired_central_ion_separation(l_2_cent, bounds=(0.1e6, 0.5e6)):
+    result = root_scalar(two_central_ion_separation, args=(l_2_cent,), bracket=bounds, method='bisect')
+    if not result.converged:
+        raise ValueError("Root finding did not converge. Try adjusting the bounds or check the function behavior.")
+    return result.root
+    
+ 
+def calc_mode_energies(res, Fock_cutoffs):
+    energies_m1 = np.empty(len(res.states))
+    energies_m2 = np.empty(len(res.states))
+    N_op = qt.num(Fock_cutoffs[0])
+    eye = qt.qeye(Fock_cutoffs[0])
+    spin_eye = qt.qeye(2)
+    E1_op = qt.tensor([spin_eye, N_op, eye])
+    E2_op = qt.tensor([spin_eye, eye, N_op])
+    for k, state in enumerate(res.states):
+        energies_m1[k] = qt.expect(E1_op, state)
+        energies_m2[k] = qt.expect(E2_op, state)
+    return energies_m1, energies_m2
+ 
+def calculate_mode_participation_factors(mode_analysis):
+    evecs = mode_analysis.evecs
+    num_coords, num_modes = np.shape(evecs)
+    num_ions = num_modes // 3
+    mode_participation_factors = np.zeros((3, num_ions, num_modes), dtype = np.complex128)
+    for mode_index in range(num_modes):
+        for pos_coord in range(num_coords//2):
+            direction_index = pos_coord // num_ions
+            ion_index = pos_coord % num_ions
+            factor = np.sqrt(2* mode_analysis.m[ion_index] * mode_analysis.evals[mode_index] )
+            zpm_dimensionful = np.sqrt(const.hbar / (2* mode_analysis.m0 * mode_analysis.w0))
+            mode_participation_factors[direction_index, ion_index, mode_index] = zpm_dimensionful * factor * evecs[pos_coord, mode_index]
+    return mode_participation_factors
+ 
+ 
+ 
+def calc_single_ion_ld_factors(omega, k, mass_amu):
+    z0 = np.sqrt(const.hbar / (2 * mass_amu * const.u * omega))
+    return k * z0 # ignore the phase
+ 
+def check_single_ion_case():
+    # for a single ion, the mode participation factors should just be the LD factors for each mode and direction.
+    # this is a good sanity check to make sure the mode participation factor calculation is correct.
+    wz = 2 * np.pi * .5e6  # axial trap frequency
+    wy = 2 * np.pi * 1.9e6  # radial trap frequency
+    wx = 2 * np.pi * 10e6  # radial trap frequency, something high to avoid any issues with mode ordering
+    mass_yb_amu = 170.936
+    k = 2 * np.pi / 355e-9
+    mode_analysis_one = mode_analyzer(N =1, wz = wz, wy = wy, wx = wx, ionmass_amu= mass_yb_amu)
+    mode_analysis_one.run()
+    mode_participation_factors_one = calculate_mode_participation_factors(mode_analysis_one)
+    ld_factors_one = k * mode_participation_factors_one
+    eta_x = calc_single_ion_ld_factors(wx, k, mass_yb_amu)
+    eta_y = calc_single_ion_ld_factors(wy, k, mass_yb_amu)
+    eta_z = calc_single_ion_ld_factors(wz, k, mass_yb_amu)
+    ld_factors_one_analytical = np.zeros((3, 1, 3), dtype = np.complex128)
+    ld_factors_one_analytical[0, 0, 2] = eta_x
+    ld_factors_one_analytical[1, 0, 1] = eta_y
+    ld_factors_one_analytical[2, 0, 0] = eta_z
+    print(f"\nLD factors: {ld_factors_one_analytical}")
+    print("\nRatio of single ion LD factors from mode participation calculation to analytical calculation: \n", ld_factors_one / ld_factors_one_analytical)
+ 
+def check_two_ion_case():
+    wz = 2 * np.pi * .5e6  # axial trap frequency
+    wy = 2 * np.pi * 1.5e6  # radial trap frequency
+    wx = 2 * np.pi * 2e6  # radial trap frequency
+    wy_tilt = np.sqrt(wy**2 - wz**2)
+    wx_tilt = np.sqrt(wx**2 - wz**2)
+    mass_yb_amu = 170.936
+    k = 2 * np.pi / 355e-9
+    mode_analysis_two= mode_analyzer(N =2, wz = wz, wy = wy, wx = wx, ionmass_amu= mass_yb_amu)
+    mode_analysis_two.run()
+    mode_participation_factors_two= calculate_mode_participation_factors(mode_analysis_two)
+    ld_factors_two = k * mode_participation_factors_two
+    eta_x = calc_single_ion_ld_factors(wx, k, mass_yb_amu)
+    eta_y = calc_single_ion_ld_factors(wy, k, mass_yb_amu)
+    eta_z = calc_single_ion_ld_factors(wz, k, mass_yb_amu)
+    eta_COM_x = eta_x / np.sqrt(2)
+    eta_COM_y = eta_y / np.sqrt(2)
+    eta_COM_z = eta_z / np.sqrt(2)
+    eta_stretch_z = eta_z /np.sqrt(2) /3**(1/4)
+    # I don't know the analytical expressions for the tilt modes off the top of my head... let's assume its the square root of the normalized mode frequency
+    eta_tilt_x = eta_x / np.sqrt(2) / np.sqrt(wx_tilt / wx)
+    eta_tilt_y = eta_y / np.sqrt(2) / np.sqrt(wy_tilt / wy)
+    ld_factors_two_analytical = np.zeros((3, 2, 6), dtype = np.complex128)
+    ld_factors_two_analytical[0, 0, 5] = eta_COM_x
+    ld_factors_two_analytical[0, 1, 5] = eta_COM_x
+    ld_factors_two_analytical[0, 0, 4] = eta_tilt_x
+    ld_factors_two_analytical[0, 1, 4] = -eta_tilt_x
+    ld_factors_two_analytical[1, 0, 3] = eta_COM_y
+    ld_factors_two_analytical[1, 1, 3] = eta_COM_y
+    ld_factors_two_analytical[1, 0, 2] = eta_tilt_y
+    ld_factors_two_analytical[1, 1, 2] = -eta_tilt_y
+    ld_factors_two_analytical[2, 0, 0] = eta_COM_z
+    ld_factors_two_analytical[2, 1, 0] = eta_COM_z
+    ld_factors_two_analytical[2, 0, 1] = eta_stretch_z
+    ld_factors_two_analytical[2, 1, 1] = -eta_stretch_z
+    # these seem to work out!
+    print(f"\nLD factors: {ld_factors_two_analytical}")
+    print("Ratio of two ion LD factors from mode participation calculation to analytical calculation: \n", ld_factors_two / ld_factors_two_analytical)
+ 
+ 
+
+ 
+### Example usage  ###  
+if __name__ == '__main__':
+    # Test analysis: 
+    #run()
+    #check_two_ion_case() 
+    check_single_ion_case()
+

From 241605a7ee2993227403fead47cc7459a2e0c1a2 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 13:26:59 -0600
Subject: [PATCH 02/31]   add args to dimensionless parameter function in mode
 analysis

---
 src/ionsim/trapped_ion_mode_analysis.py | 229 +++++++++++++-----------
 1 file changed, 129 insertions(+), 100 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 8a42d45..6bf73a0 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -1,5 +1,5 @@
 import numpy as np
-from generalized_mode_analysis import GeneralizedModeAnalysis as mode_analyzer
+#from generalized_mode_analysis import GeneralizedModeAnalysis as mode_analyzer
 #plt.style.use('ionsim')
 from time import time as timer
 import scipy.constants as const
@@ -8,12 +8,16 @@
 
 
 
-def characteristic_length(q, m, wz):
-    k_e = 1 / (4 * np.pi * const.epsilon_0)  # Coulomb constant
-    l0 = ((k_e * q ** 2) / (.5 * m * wz ** 2)) ** (1 / 3)
-    return l0
+
+########## List of substantive changes by Ethan: 
+## 1. Made dimensionless variable function take in inputs so the user can choose the charge, mass, and trap freq scales to use.   
 
 
+def characteristic_length(q: float, mass: float, omega: float):
+    """ Computes characteristic length in trapped ion system """
+    k_e = 1 / (4 * np.pi * const.epsilon_0)  # Coulomb constant
+    l0 = ((k_e * q ** 2) / (.5 * mass * omega ** 2)) ** (1 / 3)
+    return l0
 
 def get_norm(en,H):
     """
@@ -23,7 +27,6 @@ def get_norm(en,H):
     return norm
 
 
-
 def normalize_eigen_vectors(ens,H,evs=None): 
     """
     Rescale the eigen vectors ens w.r.t. the Hamiltonian H.
@@ -74,69 +77,88 @@ def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float
             - omega_x: harmonic trap frequency in the x direction in units of rad/s. 
             - omega_y: harmonic trap frequency in the y direction in units of rad/s. 
             - omega_z: harmonic trap frequency in the z direction. This is the "axial" direction used for 1D chains.  
-            - atomic masses: array of atomic masses in amu or a single number => same mass for each ion. 
+            - atomic masses: array of atomic masses or a single number => same mass for each ion. Units are amu (atomic mass units) 
             - atomic numbers: array of (Z) atomic numbers (# of protons of an element) or a single number => all ions are the same.  
         """ 
         # TODO: Decide how to handle units and how much of pint we should use.  
         self.num_ions = num_ions
 
-        self.atomic_masses = convert_num_to_array(atomic_masses)
+        self.atomic_masses = convert_num_to_array(atomic_masses) * const.u # kg from amu 
         self.atomic_numbers = convert_num_to_array(atomic_numbers)
 
         # Safety checks: 
         assert len(self.atomic_masses) == num_ions
         assert len(self.atomic_numbers) == num_ions
 
-        # Proton charge q:
-        #self.q = z * const.e # const.e ==> elementary charge in Coulombs 
-        self.proton_charges = self.atomic_numbers * const.e # Z * e, const.e ==> elementary charge in Coulombs 
+        # Charge of all the protons -> total nuclear charge:
+        self.nuclear_charges = self.atomic_numbers * const.e # Z * e, const.e ==> elementary charge in Coulombs 
 
         # Trapping frequencies for each ion 
-        self.wx_E = self.calculate_species_trap_frequencies(self.m_E, q_E, omega_x)   
-        self.wy_E = self.calculate_species_trap_frequencies(self.m_E, self.q_E, omega_y)
-        self.wz_E = self.calculate_species_trap_frequencies(self.m_E, self.q_E, omega_z)    
-
-        self.z = ensure_numpy_array(Z,N)
-        self.ionmass_amu = ensure_numpy_array(ionmass_amu,N)
-        self.q_E = self.Z * const.e
-        self.atomic_mass = self.ionmass_amu * const.atomic_mass # amu -> kilograms  
+        # TODO: loop and vectorize trap frequencies?  
+        self.omega_x = self.calculate_species_trap_frequencies(omega_x)   
+        self.omega_y = self.calculate_species_trap_frequencies(omega_y)
+        self.omega_z = self.calculate_species_trap_frequencies(omega_z)    
+
+        # Check that the trap is stable: 
+        if not self.trap_is_stable():
+            raise IonSimError(f"Error: Trap is unstable from negative trap frequencies.")
+
         self.initial_equilibrium_guess = None
         self.hasrun = False 
     
+
     def calculate_species_trap_frequencies(self, omega: float) -> Vector[float]: 
+        """ Computes relative trap frequencies for each species:
+
+            w_i = sqrt(q_{i} m_0 / m_{i} q_0) omega 
+
+        """ 
+        # TODO: Does Wes have a ref. for this? 
         # assume that the trapping frequency given corresponds to the first ion species 
-        u_r_sqr = np.sqrt(self.atomic_masses / self.proton_charges[0]) * omega
-        omega_E = np.sqrt(self.proton_charges / self.atomic_masses) * u_r_sqr    
-        return omega_E 
+        q0 = self.nuclear_charges[0] # charge of first ion 
+        m0 = self.atomic_masses[0] # charge of first ion 
+        species_trap_frequencies = np.sqrt(self.nuclear_charges * m0 /(q0 * self.atomic_masses)) * omega 
+        return species_trap_frequencies 
    
-    def dimensionless_parameters(self):
-        # normalize to the first ion species    
-        q0 = self.q_E[0]
-        m0 = self.m_E[0]    
-        w0 = self.wz_E[0]   
-        self.q0 = q0
-        self.m0 = m0
-        self.w0 = w0
-        # ion properties    
-        self.m = self.m_E / m0
-        self.q = self.q_E / q0  
-        # trap frequencies
-        self.wz = self.wz_E / w0 
-        self.wy = self.wy_E / w0 
-        self.wx = self.wx_E / w0 
-        # system parameters
-        self.l0 = characteristic_length(q0, m0, w0)  # characteristic length 
-        self.t0 = 1 / w0  # characteristic time
-        self.v0 = self.l0 / self.t0  # characteristic velocity
-        self.E0 = 0.5 * m0 * self.v0 ** 2  # characteristic energy  
 
-    
-    
     def trap_is_stable(self):
         # check that all trap frequencies are positive  
-        return np.all(self.wz_E > 0) and np.all(self.wy_E > 0) and np.all(self.wx_E > 0)
+        return np.all(self.omega_z > 0) and np.all(self.omega_y > 0) and np.all(self.omega_x > 0)
+
+
+    #def dimensionless_parameters(self):
+    #def nondimensionalize_parameters(self):
+    def convert_parameters_to_dimensionless(self, charge_scale: float=1., mass_scale: float=1., trap_freq_scale: float=1.):
+        """ Compute dimensionless parameters, using first ion's properties and axial (z) trapping frequency. 
 
+            - characteristic length
+            - characteristic time 
+            - characteristic velocity 
+            - characteristic energy 
 
+        """
+        # Store the scales so they are retrievable by the user 
+        self.charge_scale = charge_scale
+        self.mass_scale = mass_scale
+        self.trap_freq_scale = trap_freq_scale
+
+        # Normalize to the first ion species and axial (z) trap frequency 
+        # ion properties    
+        self.m = self.atomic_masses / self.mass_scale 
+        self.q = self.nuclear_charges / self.charge_scale 
+
+        # trap frequencies
+        self.wz = self.omega_z / self.trap_freq_scale 
+        self.wy = self.omega_y / self.trap_freq_scale 
+        self.wx = self.omega_x / self.trap_freq_scale  
+
+        # TODO: Is hbar set to 1? Reconcile hbar somewhere  
+        # Compute characteristic length, time, velocity, and energy scales and store in a dictionary  
+        self.characteristic_parameters = {}
+        self.characteristic_parameters['length'] = characteristic_length(self.charge_scale, self.mass_scale, self.trap_freq_scale)  
+        self.characteristic_parameters['time'] = 1 / trap_freq_scale # characteristic time
+        self.characteristic_parameters['velocity'] = self.characteristic_parameters['length'] * trap_freq_scale  # characteristic velocity
+        self.characteristic_parameters['energy'] = 0.5 * mass_scale * self.characteristic_parameters['velocity'] ** 2  # characteristic energy  
 
     def check_for_zero_modes(self):
         assert np.all(self.evals > 0), "All eigenvalues must be positive"   
@@ -147,13 +169,13 @@ def check_outer_relation(self):
         Eval, Evec = np.linalg.eig(D)
         _, en = self.sort_modes(Eval,Evec)
         en = self.normalize_eigen_vectors(en,H)
-        Outers = np.zeros((6*self.N,6*self.N),dtype=complex)
-        for i in range(6*self.N):
+        Outers = np.zeros((6*self.num_ions,6*self.num_ions),dtype=complex)
+        for i in range(6*self.num_ions):
             norm = get_norm(en[:,i],H)
             Outers = Outers + np.outer(en[:,i],en[:,i].conj())/norm 
         I_right = H @ Outers
         I_left  = Outers @ H
-        eye = np.eye(6*self.N,dtype=complex)
+        eye = np.eye(6*self.num_ions,dtype=complex)
         np.set_printoptions(precision=2, suppress=True) 
         try:
             assert np.allclose(I_left,eye) 
@@ -187,8 +209,10 @@ def checks(self):
         self.check_diagnolization()
 
     def run(self):
-        self.dimensionless_parameters()
-        assert self.trap_is_stable()    
+        #self.dimensionless_parameters()
+        # Convert to dimensionless units using axial trap frequency and first ion's mass and charge  
+        self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
+        #assert self.trap_is_stable()    
         
         self.u = self.calculate_equilibrium_positions()
         #self.reindex_ions(self.u)
@@ -247,17 +271,17 @@ def calculate_normal_modes(self, H_matrix):
 
 
     def get_initial_equilibrium_guess(self):
-        self.u0 = np.zeros(3*self.N)
-        self.u0[:] = (np.random.rand(3*self.N) * 2 - 1) * self.N 
+        self.u0 = np.zeros(3*self.num_ions)
+        self.u0[:] = (np.random.rand(3*self.num_ions) * 2 - 1) * self.num_ions 
         return self.u0
 
 
 
     def reindex_ions(self, u):  
         # based on the distance from the center of the trap, reindex the ions, smallest index is closest to the center
-        x = u[0:self.N]
-        y = u[self.N:2*self.N]
-        z = u[2*self.N:]
+        x = u[0:self.num_ions]
+        y = u[self.num_ions:2*self.num_ions]
+        z = u[2*self.num_ions:]
         r = np.sqrt(x**2 + y**2 + z**2)
         idx = np.argsort(r)
         u = np.hstack((x[idx], y[idx], z[idx]))
@@ -277,9 +301,9 @@ def find_equilibrium_positions(self, u0):
         return out.x
 
     def potential_trap(self, pos_array):
-        x = pos_array[0:self.N]
-        y = pos_array[self.N:2*self.N]
-        z = pos_array[2*self.N:]
+        x = pos_array[0:self.num_ions]
+        y = pos_array[self.num_ions:2*self.num_ions]
+        z = pos_array[2*self.num_ions:]
         V_trap = 0.5 * np.sum((self.m * self.wx ** 2) * x ** 2) + \
             0.5 * np.sum((self.m * self.wy ** 2) * y ** 2) + \
                 0.5 * np.sum((self.m * self.wz ** 2) * z ** 2)
@@ -288,9 +312,9 @@ def potential_trap(self, pos_array):
 
 
     def potential_coulomb(self, pos_array):
-        x = pos_array[0:self.N]
-        y = pos_array[self.N:2*self.N]
-        z = pos_array[2*self.N:]
+        x = pos_array[0:self.num_ions]
+        y = pos_array[self.num_ions:2*self.num_ions]
+        z = pos_array[2*self.num_ions:]
 
         dx = x[:, np.newaxis] - x
         dy = y[:, np.newaxis] - y
@@ -309,9 +333,9 @@ def potential(self, pos_array):
 
 
     def force_trap(self, pos_array):
-        x = pos_array[0:self.N]
-        y = pos_array[self.N:2*self.N]
-        z = pos_array[2*self.N:]
+        x = pos_array[0:self.num_ions]
+        y = pos_array[self.num_ions:2*self.num_ions]
+        z = pos_array[2*self.num_ions:]
 
         Ftrapx = self.m * self.wx**2 * x
         Ftrapy = self.m * self.wy**2 * y
@@ -323,9 +347,9 @@ def force_trap(self, pos_array):
 
 
     def force_coulomb(self, pos_array): 
-        x = pos_array[0:self.N]
-        y = pos_array[self.N:2*self.N]
-        z = pos_array[2*self.N:]
+        x = pos_array[0:self.num_ions]
+        y = pos_array[self.num_ions:2*self.num_ions]
+        z = pos_array[2*self.num_ions:]
 
         dx = x[:, np.newaxis] - x
         dy = y[:, np.newaxis] - y
@@ -361,19 +385,19 @@ def force(self, pos_array):
 
 
     def hessian_trap(self, pos_array):
-        Hxx = np.diag(self.m * (self.wx**2) * np.ones(self.N))
-        Hyy = np.diag(self.m * (self.wy**2) * np.ones(self.N))  
-        Hzz = np.diag(self.m * (self.wz**2) * np.ones(self.N))  
-        zeros = np.zeros((self.N, self.N))  
+        Hxx = np.diag(self.m * (self.wx**2) * np.ones(self.num_ions))
+        Hyy = np.diag(self.m * (self.wy**2) * np.ones(self.num_ions))  
+        Hzz = np.diag(self.m * (self.wz**2) * np.ones(self.num_ions))  
+        zeros = np.zeros((self.num_ions, self.num_ions))  
         H = np.block([[Hxx, zeros, zeros], [zeros, Hyy, zeros], [zeros, zeros, Hzz]])
         return H
 
 
 
     def hessian_coulomb(self, pos_array):
-        x = pos_array[0:self.N]
-        y = pos_array[self.N:2*self.N]
-        z = pos_array[2*self.N:]
+        x = pos_array[0:self.num_ions]
+        y = pos_array[self.num_ions:2*self.num_ions]
+        z = pos_array[2*self.num_ions:]
         
         dx = x[:, np.newaxis] - x
         dy = y[:, np.newaxis] - y
@@ -395,16 +419,16 @@ def hessian_coulomb(self, pos_array):
         Hzz = (rsep2 - 3 * dzsq) * rsep5
 
         # Above, for alpha == beta
-        Hxx[np.diag_indices(self.N)] = -np.sum(Hxx, axis=0)
-        Hyy[np.diag_indices(self.N)] = -np.sum(Hyy, axis=0)
-        Hzz[np.diag_indices(self.N)] = -np.sum(Hzz, axis=0)
+        Hxx[np.diag_indices(self.num_ions)] = -np.sum(Hxx, axis=0)
+        Hyy[np.diag_indices(self.num_ions)] = -np.sum(Hyy, axis=0)
+        Hzz[np.diag_indices(self.num_ions)] = -np.sum(Hzz, axis=0)
 
         Hxy = -3 * dx * dy * rsep5
-        Hxy[np.diag_indices(self.N)] = 3 * np.sum(dx * dy * rsep5, axis=0)
+        Hxy[np.diag_indices(self.num_ions)] = 3 * np.sum(dx * dy * rsep5, axis=0)
         Hxz = -3 * dx * dz * rsep5
-        Hxz[np.diag_indices(self.N)] = 3 * np.sum(dx * dz * rsep5, axis=0)
+        Hxz[np.diag_indices(self.num_ions)] = 3 * np.sum(dx * dz * rsep5, axis=0)
         Hyz = -3 * dy * dz * rsep5
-        Hyz[np.diag_indices(self.N)] = 3 * np.sum(dy * dz * rsep5, axis=0)
+        Hyz[np.diag_indices(self.num_ions)] = 3 * np.sum(dy * dz * rsep5, axis=0)
         
         Hxx *= qq
         Hyy *= qq
@@ -426,13 +450,13 @@ def get_mass_matrix(self,m):
         return np.diag(np.tile(m, 3)) 
 
     def get_E_matrix(self,u): 
-        PE_matrix = np.zeros((3*self.N, 3*self.N), dtype=np.complex128)
-        KE_matrix = np.zeros((3*self.N, 3*self.N), dtype=np.complex128)
-        E_matrix = np.zeros((6*self.N, 6*self.N), dtype=np.complex128)
+        PE_matrix = np.zeros((3*self.num_ions, 3*self.num_ions), dtype=np.complex128)
+        KE_matrix = np.zeros((3*self.num_ions, 3*self.num_ions), dtype=np.complex128)
+        E_matrix = np.zeros((6*self.num_ions, 6*self.num_ions), dtype=np.complex128)
 
         PE_matrix = self.hessian(u)
         KE_matrix = self.get_mass_matrix(self.m) 
-        zeros = np.zeros((3*self.N, 3*self.N)) 
+        zeros = np.zeros((3*self.num_ions, 3*self.num_ions)) 
         E_matrix = np.block([[PE_matrix, zeros], [zeros, KE_matrix]])
         return E_matrix
 
@@ -448,16 +472,16 @@ def get_H_matrix(self, T_matrix, E_matrix):
     def get_momentum_transform(self):
         # assuming no magnetic field
         mass_matrix = self.get_mass_matrix(self.m)  
-        eye = np.eye(3*self.N)  
-        zeros = np.zeros((3*self.N, 3*self.N))
+        eye = np.eye(3*self.num_ions)  
+        zeros = np.zeros((3*self.num_ions, 3*self.num_ions))
         T = np.block([[eye, zeros], [zeros, mass_matrix]])  
         return T    
 
 
 
     def get_symplectic_matrix(self):
-        zeros = np.zeros((3*self.N, 3*self.N), dtype=np.complex128)
-        I = np.eye(3*self.N, dtype=np.complex128)
+        zeros = np.zeros((3*self.num_ions, 3*self.num_ions), dtype=np.complex128)
+        I = np.eye(3*self.num_ions, dtype=np.complex128)
         J = np.block([[zeros, I], [-I, zeros]])
         return J    
 
@@ -487,7 +511,7 @@ def organize_modes(self,evals, evecs):
  
  
 #classes
-class GeneralizedModeAnalysisWithBranchSortedModes(mode_analyzer):
+class GeneralizedModeAnalysisWithBranchSortedModes(TrappedIonModeAnalysis):
  
     def _xyz_classify_modes(self, evecs):
         N_coords, N_modes = np.shape(evecs)
@@ -527,9 +551,9 @@ def organize_modes(self, evals, evecs):
  
     def reindex_ions_by_z(self, u):  
         # based on the position along z, lowest i, is 0, up to N - 1
-        x = u[0:self.N]
-        y = u[self.N:2*self.N]
-        z = u[2*self.N:]
+        x = u[0:self.num_ions]
+        y = u[self.num_ions:2*self.num_ions]
+        z = u[2*self.num_ions:]
  
         idx = np.argsort(z)
         self.u = np.hstack((x[idx], y[idx], z[idx]))
@@ -542,15 +566,17 @@ def reindex_ions_by_z(self, u):
         self.Z = self.Z[idx]
         self.ionmass_amu = self.ionmass_amu[idx]
         # trapping frequencies
-        self.wz_E = self.wz_E[idx]
-        self.wy_E = self.wy_E[idx]
-        self.wx_E = self.wx_E[idx]
+        self.omega_z = self.omega_z[idx]
+        self.omega_y = self.omega_y[idx]
+        self.omega_x = self.omega_x[idx]
         # all ions are the same so this is safe
-        self.dimensionless_parameters()
+        #self.dimensionless_parameters()
+        self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
  
  
     def run(self):
-        self.dimensionless_parameters()
+        #self.dimensionless_parameters()
+        self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
         assert self.trap_is_stable()    
         
         self.u = self.calculate_equilibrium_positions()
@@ -577,7 +603,7 @@ def two_central_ion_separation(wz_Hz, l_2_cent):
     mass_yb_amu = 170.936 # TODO: hardcoded for now
     four_ion_analysis = GeneralizedModeAnalysisWithBranchSortedModes(N=4, wz=2*np.pi*wz_Hz, wy=2*np.pi*10e6, wx=2*np.pi*11e6, ionmass_amu=mass_yb_amu)
     four_ion_analysis.run()
-    positions_z = four_ion_analysis.u[2*4:] * four_ion_analysis.l0
+    positions_z = four_ion_analysis.u[2*4:] * four_ion_analysis.characteristic_parameters['length']
     central_ions = np.argsort(np.abs(positions_z))[:2]
     dl = np.abs(positions_z[central_ions[0]] - positions_z[central_ions[1]])
     return dl - l_2_cent
@@ -612,7 +638,7 @@ def calculate_mode_participation_factors(mode_analysis):
             direction_index = pos_coord // num_ions
             ion_index = pos_coord % num_ions
             factor = np.sqrt(2* mode_analysis.m[ion_index] * mode_analysis.evals[mode_index] )
-            zpm_dimensionful = np.sqrt(const.hbar / (2* mode_analysis.m0 * mode_analysis.w0))
+            zpm_dimensionful = np.sqrt(const.hbar / (2* mode_analysis.mass_scale * mode_analysis.trap_freq_scale))
             mode_participation_factors[direction_index, ion_index, mode_index] = zpm_dimensionful * factor * evecs[pos_coord, mode_index]
     return mode_participation_factors
  
@@ -630,7 +656,10 @@ def check_single_ion_case():
     wx = 2 * np.pi * 10e6  # radial trap frequency, something high to avoid any issues with mode ordering
     mass_yb_amu = 170.936
     k = 2 * np.pi / 355e-9
-    mode_analysis_one = mode_analyzer(N =1, wz = wz, wy = wy, wx = wx, ionmass_amu= mass_yb_amu)
+    atomic_nums = np.array([70]) 
+    num_ions = 1
+
+    mode_analysis_one = TrappedIonModeAnalysis(num_ions, wx, wy, wz, np.ones(num_ions)*mass_yb_amu, atomic_nums) 
     mode_analysis_one.run()
     mode_participation_factors_one = calculate_mode_participation_factors(mode_analysis_one)
     ld_factors_one = k * mode_participation_factors_one

From f233d4308baf5787b14c4a8b1b839ad1ee7d1a11 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 13:44:23 -0600
Subject: [PATCH 03/31]  + helper function to unpack flat-1D-arr of ion
 coordinates

---
 src/ionsim/trapped_ion_mode_analysis.py | 68 +++++++++++++++++--------
 1 file changed, 47 insertions(+), 21 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 6bf73a0..8dff9aa 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -11,6 +11,12 @@
 
 ########## List of substantive changes by Ethan: 
 ## 1. Made dimensionless variable function take in inputs so the user can choose the charge, mass, and trap freq scales to use.   
+## 2. Added a helper function to pack/unpack the ion coordinates to reduce code. 
+
+########## List of non-substantive changes by Ethan: 
+## 1. Variable changes for readable (e.g. omega_x instead of wx) 
+## 2. Type-hinting in functions  
+
 
 
 def characteristic_length(q: float, mass: float, omega: float):
@@ -270,7 +276,21 @@ def calculate_normal_modes(self, H_matrix):
 
 
 
+    def ion_coordinates_from_flattened(self, flattened_coordinate_vector: Vector) -> tuple(Vector, Vector, Vector):
+        """ Returns arrays corresponding to an ion's x, y, and z coordinates, respectively, from a flattened vector u:
+                x = u[0:N], y = u[N:2N], z[2N:]
+
+            -e.g. x[1], y[1], z[1] is the x, y, z coordinate values for ion 2. 
+        """
+        x = flattened_coordinate_vector[0:self.num_ions]
+        y = flattened_coordinate_vector[self.num_ions:2*self.num_ions]
+        z = flattened_coordinate_vector[2*self.num_ions:]
+        return x, y, z 
+        
+
     def get_initial_equilibrium_guess(self):
+        """ Initialize the set of ion coordinates {u} as a flattened 1D array. """
+        # TODO: option for random vs. equally spaced along one axis? 
         self.u0 = np.zeros(3*self.num_ions)
         self.u0[:] = (np.random.rand(3*self.num_ions) * 2 - 1) * self.num_ions 
         return self.u0
@@ -279,9 +299,7 @@ def get_initial_equilibrium_guess(self):
 
     def reindex_ions(self, u):  
         # based on the distance from the center of the trap, reindex the ions, smallest index is closest to the center
-        x = u[0:self.num_ions]
-        y = u[self.num_ions:2*self.num_ions]
-        z = u[2*self.num_ions:]
+        x,y,z = self.ion_coordinates_from_flattened(u)
         r = np.sqrt(x**2 + y**2 + z**2)
         idx = np.argsort(r)
         u = np.hstack((x[idx], y[idx], z[idx]))
@@ -294,16 +312,19 @@ def reindex_ions(self, u):
 
 
 
-    def find_equilibrium_positions(self, u0):
+    def find_equilibrium_positions(self, u0: Vector):
+        """ Solves for equilibrium position vector: u0, which is a flattened spatial grid. """
         bfgs_tolerance = 1e-34
         out = opt.minimize(self.potential, u0, method='BFGS', jac=self.force,
                                     options={'gtol': bfgs_tolerance, 'disp': False})
         return out.x
 
-    def potential_trap(self, pos_array):
-        x = pos_array[0:self.num_ions]
-        y = pos_array[self.num_ions:2*self.num_ions]
-        z = pos_array[2*self.num_ions:]
+    def potential_trap(self, pos_array: Vector):
+        """ Computes trapping potential, takes in a flattened (1D) array of all position coordinates """
+ #        x = pos_array[0:self.num_ions]
+ #        y = pos_array[self.num_ions:2*self.num_ions]
+ #        z = pos_array[2*self.num_ions:]
+        x,y,z = self.ion_coordinates_from_flattened(pos_array)
         V_trap = 0.5 * np.sum((self.m * self.wx ** 2) * x ** 2) + \
             0.5 * np.sum((self.m * self.wy ** 2) * y ** 2) + \
                 0.5 * np.sum((self.m * self.wz ** 2) * z ** 2)
@@ -311,10 +332,12 @@ def potential_trap(self, pos_array):
     
 
 
-    def potential_coulomb(self, pos_array):
-        x = pos_array[0:self.num_ions]
-        y = pos_array[self.num_ions:2*self.num_ions]
-        z = pos_array[2*self.num_ions:]
+    def potential_coulomb(self, pos_array: Vector):
+        """ Computes Coulomb potential, takes in a flattened (1D) array of all position coordinates """
+ #        x = pos_array[0:self.num_ions]
+ #        y = pos_array[self.num_ions:2*self.num_ions]
+ #        z = pos_array[2*self.num_ions:]
+        x,y,z = self.ion_coordinates_from_flattened(pos_array)
 
         dx = x[:, np.newaxis] - x
         dy = y[:, np.newaxis] - y
@@ -333,9 +356,10 @@ def potential(self, pos_array):
 
 
     def force_trap(self, pos_array):
-        x = pos_array[0:self.num_ions]
-        y = pos_array[self.num_ions:2*self.num_ions]
-        z = pos_array[2*self.num_ions:]
+#         x = pos_array[0:self.num_ions]
+#         y = pos_array[self.num_ions:2*self.num_ions]
+#         z = pos_array[2*self.num_ions:]
+        x,y,z = self.ion_coordinates_from_flattened(pos_array)
 
         Ftrapx = self.m * self.wx**2 * x
         Ftrapy = self.m * self.wy**2 * y
@@ -347,9 +371,10 @@ def force_trap(self, pos_array):
 
 
     def force_coulomb(self, pos_array): 
-        x = pos_array[0:self.num_ions]
-        y = pos_array[self.num_ions:2*self.num_ions]
-        z = pos_array[2*self.num_ions:]
+        #x = pos_array[0:self.num_ions]
+        #y = pos_array[self.num_ions:2*self.num_ions]
+        #z = pos_array[2*self.num_ions:]
+        x,y,z = self.ion_coordinates_from_flattened(pos_array)
 
         dx = x[:, np.newaxis] - x
         dy = y[:, np.newaxis] - y
@@ -395,9 +420,10 @@ def hessian_trap(self, pos_array):
 
 
     def hessian_coulomb(self, pos_array):
-        x = pos_array[0:self.num_ions]
-        y = pos_array[self.num_ions:2*self.num_ions]
-        z = pos_array[2*self.num_ions:]
+ #        x = pos_array[0:self.num_ions]
+ #        y = pos_array[self.num_ions:2*self.num_ions]
+ #        z = pos_array[2*self.num_ions:]
+        x,y,z = self.ion_coordinates_from_flattened(pos_array)
         
         dx = x[:, np.newaxis] - x
         dy = y[:, np.newaxis] - y

From f8b6153f35ec4ec55b8b578b6f8c736d4787a4e5 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 14:50:07 -0600
Subject: [PATCH 04/31]   rename eval, evec -> eigvals, eigvecs

---
 src/ionsim/trapped_ion_mode_analysis.py | 247 ++++++++++++------------
 1 file changed, 123 insertions(+), 124 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 8dff9aa..59853e6 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -2,6 +2,7 @@
 #from generalized_mode_analysis import GeneralizedModeAnalysis as mode_analyzer
 #plt.style.use('ionsim')
 from time import time as timer
+from numpy.typing import NDArray
 import scipy.constants as const
 import warnings
 import scipy.optimize as opt    
@@ -67,7 +68,7 @@ def get_canonical_transformation(H,ens,evs=None):
     return T
 
 
-def convert_num_to_array(x):
+def convert_to_array(x: float | int | list | NDArray):
     """ Converts a scalar to a vector or returns the vector """ 
     if not hasattr(x, "__len__"):
         x = np.ones(N) * x 
@@ -89,8 +90,8 @@ def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float
         # TODO: Decide how to handle units and how much of pint we should use.  
         self.num_ions = num_ions
 
-        self.atomic_masses = convert_num_to_array(atomic_masses) * const.u # kg from amu 
-        self.atomic_numbers = convert_num_to_array(atomic_numbers)
+        self.atomic_masses = convert_to_array(atomic_masses) * const.u # kg from amu 
+        self.atomic_numbers = convert_to_array(atomic_numbers)
 
         # Safety checks: 
         assert len(self.atomic_masses) == num_ions
@@ -143,6 +144,14 @@ def convert_parameters_to_dimensionless(self, charge_scale: float=1., mass_scale
             - characteristic energy 
 
         """
+        # Check positivity:
+        if trap_freq_scale <= 0:
+            raise IonSimError(f"Trap frequency scale should be positive. Received {trap_freq_scale}")
+        if mass_scale <= 0:
+            raise IonSimError(f"Mass scale should be positive. Received {mass_scale}")
+ #        if charge_scale <= 0:
+ #            raise IonSimError(f"Charge scale should be positive. Received {charge_scale}")
+
         # Store the scales so they are retrievable by the user 
         self.charge_scale = charge_scale
         self.mass_scale = mass_scale
@@ -167,7 +176,7 @@ def convert_parameters_to_dimensionless(self, charge_scale: float=1., mass_scale
         self.characteristic_parameters['energy'] = 0.5 * mass_scale * self.characteristic_parameters['velocity'] ** 2  # characteristic energy  
 
     def check_for_zero_modes(self):
-        assert np.all(self.evals > 0), "All eigenvalues must be positive"   
+        assert np.all(self.eigvals > 0), "All eigenvalues must be positive"   
 
     def check_outer_relation(self): 
         H = self.H_matrix.copy()       
@@ -191,8 +200,8 @@ def check_outer_relation(self):
 
 
 
-    def has_duplicate_evals(self,evals):
-        evs = evals.copy()    
+    def has_duplicate_eigvals(eigvals):
+        evs = eigvals.copy()    
         return np.any(np.triu(np.isclose(evs[:, None], evs[None, :], atol=1e-6), k=1))  
     
 
@@ -200,13 +209,13 @@ def has_duplicate_evals(self,evals):
     def check_diagnolization(self):
         M = np.linalg.inv(self.T_matrix) @ self.S_matrix    
         H_diag = M.T @ self.E_matrix @ M    
-        H_diag_check = np.diag(np.tile(self.evals,2)) 
+        H_diag_check = np.diag(np.tile(self.eigvals,2)) 
         np.set_printoptions(precision=2, suppress=True) 
         try:
             assert np.allclose(H_diag,H_diag_check)
         except AssertionError:
             warnings.warn("Diagnolization check failed")    
-            print("has duplicate evals: ", self.has_duplicate_evals(self.evals))
+            print("has duplicate eigenvalues: ", self.has_duplicate_eigvals(self.eigvals))
 
 
 
@@ -218,15 +227,16 @@ def run(self):
         #self.dimensionless_parameters()
         # Convert to dimensionless units using axial trap frequency and first ion's mass and charge  
         self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
-        #assert self.trap_is_stable()    
+        #assert self.trap_is_stable()   # checked in the constructor  
         
         self.u = self.calculate_equilibrium_positions()
-        #self.reindex_ions(self.u)
+        self.reindex_ions()
         self.E_matrix = self.get_E_matrix(self.u)  
         self.T_matrix = self.get_momentum_transform() 
         self.H_matrix = self.get_H_matrix(self.T_matrix, self.E_matrix)   
-        self.evals, self.evecs = self.calculate_normal_modes(self.H_matrix)
-        self.evecs_vel = self.get_eigen_vectors_xv_coords(self.T_matrix,self.evecs)    
+
+        self.eigvals, self.eigvecs = self.calculate_normal_modes(self.H_matrix)
+        self.eigvecs_vel = self.get_eigen_vectors_xv_coords(self.T_matrix,self.eigvecs)    
         self.check_for_zero_modes() 
         self.S_matrix = self.get_canonical_transformation() 
         self.checks() 
@@ -248,12 +258,12 @@ def calculate_equilibrium_positions(self):
 
 
     def get_canonical_transformation(self):
-        return get_canonical_transformation(self.H_matrix,self.evecs,evs=self.evals)    
+        return get_canonical_transformation(self.H_matrix,self.eigvecs,evs=self.eigvals)    
 
 
 
-    def normalize_eigen_vectors(self, evecs, H_matrix,evs=None):
-        return normalize_eigen_vectors(evecs,H_matrix,evs=evs) 
+    def normalize_eigen_vectors(self, eigvecs, H_matrix,evs=None):
+        return normalize_eigen_vectors(eigvecs,H_matrix,evs=evs) 
 
 
 
@@ -263,16 +273,22 @@ def get_eigen_vectors_xv_coords(self,T,ens):
         return ens_vel 
 
 
-    def calculate_normal_modes(self, H_matrix):
+    def calculate_normal_modes(self, H_matrix: Matrix) -> (Vector, list[Vector]):
+        """ Compute normal mode frequencies from eigenvalues of the matrix D. Returns eigenvalues and eigenvectors.
+
+            Solves det( D + i*omega*Identity) == 0 for omega.
+            
+            The matrix "D" is defind as D = JH, where J = (0 , I ; I, 0) (eq. 1.67 of thesis)
 
+        """
         J = self.get_symplectic_matrix()
         D_matrix = J @ H_matrix  
         # u_n(t) = exp(-i w_n t) u_n(0), minus by convention
-        evals, evecs = np.linalg.eig(-D_matrix)   
-        evals, evecs = self.organize_modes(evals, evecs)
-        evecs = self.normalize_eigen_vectors(evecs, H_matrix) 
+        eigvals, eigvecs = np.linalg.eig(-D_matrix)   
+        eigvals, eigvecs = self.organize_modes(eigvals, eigvecs)
+        eigvecs = self.normalize_eigen_vectors(eigvecs, H_matrix) 
 
-        return evals, evecs 
+        return eigvals, eigvecs 
 
 
 
@@ -280,7 +296,7 @@ def ion_coordinates_from_flattened(self, flattened_coordinate_vector: Vector) ->
         """ Returns arrays corresponding to an ion's x, y, and z coordinates, respectively, from a flattened vector u:
                 x = u[0:N], y = u[N:2N], z[2N:]
 
-            -e.g. x[1], y[1], z[1] is the x, y, z coordinate values for ion 2. 
+            e.g. x[1], y[1], z[1] is the x, y, z coordinate values for ion 2. 
         """
         x = flattened_coordinate_vector[0:self.num_ions]
         y = flattened_coordinate_vector[self.num_ions:2*self.num_ions]
@@ -296,20 +312,20 @@ def get_initial_equilibrium_guess(self):
         return self.u0
 
 
-
-    def reindex_ions(self, u):  
-        # based on the distance from the center of the trap, reindex the ions, smallest index is closest to the center
-        x,y,z = self.ion_coordinates_from_flattened(u)
+    def reindex_ions(self):  
+        """ Re-indexes the ions to order based on the distance from the center of the trap, where smallest index is closest to the center """
+        # TODO: How are even/odd cases is handled? 
+        x,y,z = self.ion_coordinates_from_flattened(self.u)
         r = np.sqrt(x**2 + y**2 + z**2)
         idx = np.argsort(r)
-        u = np.hstack((x[idx], y[idx], z[idx]))
-        self.u = u 
+
+        reindexed_positions = np.hstack((x[idx], y[idx], z[idx]))
+        self.u = reindexed_positions 
         self.m = self.m[idx]
-        self.m_E = self.m_E[idx]
         self.q = self.q[idx]
-        self.q_E = self.q_E[idx]
-        self.Z = self.Z[idx]
-
+        self.atomic_masses = self.atomic_masses[idx]
+        self.nuclear_charges = self.nuclear_charges[idx]
+        self.atomic_numbers = self.atomic_numbers[idx]
 
 
     def find_equilibrium_positions(self, u0: Vector):
@@ -319,12 +335,12 @@ def find_equilibrium_positions(self, u0: Vector):
                                     options={'gtol': bfgs_tolerance, 'disp': False})
         return out.x
 
-    def potential_trap(self, pos_array: Vector):
+    def potential_trap(self, positions: Vector):
         """ Computes trapping potential, takes in a flattened (1D) array of all position coordinates """
- #        x = pos_array[0:self.num_ions]
- #        y = pos_array[self.num_ions:2*self.num_ions]
- #        z = pos_array[2*self.num_ions:]
-        x,y,z = self.ion_coordinates_from_flattened(pos_array)
+ #        x = positions[0:self.num_ions]
+ #        y = positions[self.num_ions:2*self.num_ions]
+ #        z = positions[2*self.num_ions:]
+        x,y,z = self.ion_coordinates_from_flattened(positions)
         V_trap = 0.5 * np.sum((self.m * self.wx ** 2) * x ** 2) + \
             0.5 * np.sum((self.m * self.wy ** 2) * y ** 2) + \
                 0.5 * np.sum((self.m * self.wz ** 2) * z ** 2)
@@ -332,12 +348,12 @@ def potential_trap(self, pos_array: Vector):
     
 
 
-    def potential_coulomb(self, pos_array: Vector):
+    def potential_coulomb(self, positions: Vector):
         """ Computes Coulomb potential, takes in a flattened (1D) array of all position coordinates """
- #        x = pos_array[0:self.num_ions]
- #        y = pos_array[self.num_ions:2*self.num_ions]
- #        z = pos_array[2*self.num_ions:]
-        x,y,z = self.ion_coordinates_from_flattened(pos_array)
+ #        x = positions[0:self.num_ions]
+ #        y = positions[self.num_ions:2*self.num_ions]
+ #        z = positions[2*self.num_ions:]
+        x,y,z = self.ion_coordinates_from_flattened(positions)
 
         dx = x[:, np.newaxis] - x
         dy = y[:, np.newaxis] - y
@@ -350,16 +366,16 @@ def potential_coulomb(self, pos_array: Vector):
         V_Coulomb *= .5 
         return V_Coulomb
 
-    def potential(self, pos_array):
-        return self.potential_trap(pos_array) + self.potential_coulomb(pos_array)   
+    def potential(self, positions):
+        return self.potential_trap(positions) + self.potential_coulomb(positions)   
 
 
 
-    def force_trap(self, pos_array):
-#         x = pos_array[0:self.num_ions]
-#         y = pos_array[self.num_ions:2*self.num_ions]
-#         z = pos_array[2*self.num_ions:]
-        x,y,z = self.ion_coordinates_from_flattened(pos_array)
+    def force_trap(self, positions):
+#         x = positions[0:self.num_ions]
+#         y = positions[self.num_ions:2*self.num_ions]
+#         z = positions[2*self.num_ions:]
+        x,y,z = self.ion_coordinates_from_flattened(positions)
 
         Ftrapx = self.m * self.wx**2 * x
         Ftrapy = self.m * self.wy**2 * y
@@ -370,11 +386,11 @@ def force_trap(self, pos_array):
 
 
 
-    def force_coulomb(self, pos_array): 
-        #x = pos_array[0:self.num_ions]
-        #y = pos_array[self.num_ions:2*self.num_ions]
-        #z = pos_array[2*self.num_ions:]
-        x,y,z = self.ion_coordinates_from_flattened(pos_array)
+    def force_coulomb(self, positions): 
+        #x = positions[0:self.num_ions]
+        #y = positions[self.num_ions:2*self.num_ions]
+        #z = positions[2*self.num_ions:]
+        x,y,z = self.ion_coordinates_from_flattened(positions)
 
         dx = x[:, np.newaxis] - x
         dy = y[:, np.newaxis] - y
@@ -399,17 +415,17 @@ def force_coulomb(self, pos_array):
 
 
 
-    def force(self, pos_array):
-        Force = self.force_coulomb(pos_array) + self.force_trap(pos_array)  
+    def force(self, positions):
+        Force = self.force_coulomb(positions) + self.force_trap(positions)  
         return Force
 
-    def force(self, pos_array):
-        Force = self.force_coulomb(pos_array) + self.force_trap(pos_array)  
+    def force(self, positions):
+        Force = self.force_coulomb(positions) + self.force_trap(positions)  
         return Force
 
 
-
-    def hessian_trap(self, pos_array):
+    def hessian_trap(self, positions: Vector):
+        """ Computes the Hessian of the trap  """ 
         Hxx = np.diag(self.m * (self.wx**2) * np.ones(self.num_ions))
         Hyy = np.diag(self.m * (self.wy**2) * np.ones(self.num_ions))  
         Hzz = np.diag(self.m * (self.wz**2) * np.ones(self.num_ions))  
@@ -417,13 +433,9 @@ def hessian_trap(self, pos_array):
         H = np.block([[Hxx, zeros, zeros], [zeros, Hyy, zeros], [zeros, zeros, Hzz]])
         return H
 
-
-
-    def hessian_coulomb(self, pos_array):
- #        x = pos_array[0:self.num_ions]
- #        y = pos_array[self.num_ions:2*self.num_ions]
- #        z = pos_array[2*self.num_ions:]
-        x,y,z = self.ion_coordinates_from_flattened(pos_array)
+    def hessian_coulomb(self, positions: Vector):
+        """ Computes the Hessian of the Coulomb interaction """ 
+        x,y,z = self.ion_coordinates_from_flattened(positions)
         
         dx = x[:, np.newaxis] - x
         dy = y[:, np.newaxis] - y
@@ -468,33 +480,31 @@ def hessian_coulomb(self, pos_array):
         return H_coulomb
 
 
-    def hessian(self, pos_array):
-        H = self.hessian_coulomb(pos_array) + self.hessian_trap(pos_array)  
+    def hessian(self, positions: Vector):
+        """ Computes total Hessian (trap + Coulomb interaction) """
+        H = self.hessian_coulomb(positions) + self.hessian_trap(positions)  
         return H
 
-    def get_mass_matrix(self,m):    
+    def get_mass_matrix(self, m): 
         return np.diag(np.tile(m, 3)) 
 
-    def get_E_matrix(self,u): 
+    def get_E_matrix(self, positions: Vector): 
+        """ 6N x 6N energy matrix for N ions """ 
         PE_matrix = np.zeros((3*self.num_ions, 3*self.num_ions), dtype=np.complex128)
         KE_matrix = np.zeros((3*self.num_ions, 3*self.num_ions), dtype=np.complex128)
         E_matrix = np.zeros((6*self.num_ions, 6*self.num_ions), dtype=np.complex128)
 
-        PE_matrix = self.hessian(u)
+        PE_matrix = self.hessian(positions)
         KE_matrix = self.get_mass_matrix(self.m) 
         zeros = np.zeros((3*self.num_ions, 3*self.num_ions)) 
         E_matrix = np.block([[PE_matrix, zeros], [zeros, KE_matrix]])
         return E_matrix
 
-
-
     def get_H_matrix(self, T_matrix, E_matrix):  
         T_matrix_inv = np.linalg.inv(T_matrix)  
         H_matrix = T_matrix_inv.T @ E_matrix @ T_matrix_inv
         return H_matrix 
 
-
-
     def get_momentum_transform(self):
         # assuming no magnetic field
         mass_matrix = self.get_mass_matrix(self.m)  
@@ -503,35 +513,30 @@ def get_momentum_transform(self):
         T = np.block([[eye, zeros], [zeros, mass_matrix]])  
         return T    
 
-
-
     def get_symplectic_matrix(self):
         zeros = np.zeros((3*self.num_ions, 3*self.num_ions), dtype=np.complex128)
         I = np.eye(3*self.num_ions, dtype=np.complex128)
         J = np.block([[zeros, I], [-I, zeros]])
         return J    
 
-    def sort_modes(self,evals, evecs):
-       evals = np.imag(evals)
-       sort_dex = np.argsort(evals)
-       evals = evals[sort_dex]
-       evecs = evecs[:,sort_dex]
-       return evals, evecs  
+    def sort_modes(self, eigvals, eigvecs):
+       eigvals = np.imag(eigvals)
+       sort_dex = np.argsort(eigvals)
+       eigvals = eigvals[sort_dex]
+       eigvecs = eigvecs[:,sort_dex]
+       return eigvals, eigvecs  
     
-
-
-    def split_modes(self,evals, evecs): 
-       half = len(evals) // 2
-       evals = evals[half:]
-       evecs = evecs[:,half:]
-       return evals, evecs
+    def split_modes(self, eigvals, eigvecs): 
+       half = len(eigvals) // 2
+       eigvals = eigvals[half:]
+       eigvecs = eigvecs[:,half:]
+       return eigvals, eigvecs
     
 
-
-    def organize_modes(self,evals, evecs):  
-        evals, evecs = self.sort_modes(evals, evecs)    
-        evals, evecs = self.split_modes(evals, evecs)
-        return evals, evecs 
+    def organize_modes(self, eigvals, eigvecs):  
+        eigvals, eigvecs = self.sort_modes(eigvals, eigvecs)    
+        eigvals, eigvecs = self.split_modes(eigvals, eigvecs)
+        return eigvals, eigvecs 
 
 
  
@@ -568,34 +573,32 @@ def sort_by_branch(self, evals, evecs):
             sorted_by_branch_evecs[:, direction*N_ions:(direction+1)*N_ions] = evecs[:, direction_indices]
         return sorted_by_branch_evals, sorted_by_branch_evecs
  
-    def organize_modes(self, evals, evecs):        
-        evals, evecs = self.sort_modes(evals, evecs)
-        evals, evecs = self.split_modes(evals, evecs)
-        evals, evecs = self.sort_by_branch(evals, evecs)
-        return evals, evecs
+    def organize_modes(self, eigvals, eigvecs):        
+        eigvals, eigvecs = self.sort_modes(eigvals, eigvecs)
+        eigvals, eigvecs = self.split_modes(eigvals, eigvecs)
+        eigvals, eigvecs = self.sort_by_branch(eigvals, eigvecs)
+        return eigvals, eigvecs
  
  
-    def reindex_ions_by_z(self, u):  
+    def reindex_ions_by_z(self): 
         # based on the position along z, lowest i, is 0, up to N - 1
-        x = u[0:self.num_ions]
-        y = u[self.num_ions:2*self.num_ions]
-        z = u[2*self.num_ions:]
+        x,y,z = self.ion_coordinates_from_flattened(self.u)
  
         idx = np.argsort(z)
         self.u = np.hstack((x[idx], y[idx], z[idx]))
  
-        # reindex the other arrays accordingly
+        # Reindex all other arrays accordingly
         self.m = self.m[idx]
-        self.m_E = self.m_E[idx]
         self.q = self.q[idx]
-        self.q_E = self.q_E[idx]
-        self.Z = self.Z[idx]
-        self.ionmass_amu = self.ionmass_amu[idx]
+        self.atomic_masses = self.atomic_masses[idx]
+        self.nuclear_charges = self.nuclear_charges[idx]
+        self.atomic_numbers = self.atomic_numbers[idx]
+
         # trapping frequencies
-        self.omega_z = self.omega_z[idx]
-        self.omega_y = self.omega_y[idx]
         self.omega_x = self.omega_x[idx]
-        # all ions are the same so this is safe
+        self.omega_y = self.omega_y[idx]
+        self.omega_z = self.omega_z[idx]
+        # all ions are the same so this is safe. TODO: When does this change? 
         #self.dimensionless_parameters()
         self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
  
@@ -603,15 +606,15 @@ def reindex_ions_by_z(self, u):
     def run(self):
         #self.dimensionless_parameters()
         self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
-        assert self.trap_is_stable()    
+        #assert self.trap_is_stable()    
         
         self.u = self.calculate_equilibrium_positions()
-        self.reindex_ions_by_z(self.u)
+        self.reindex_ions_by_z()
         self.E_matrix = self.get_E_matrix(self.u)  
         self.T_matrix = self.get_momentum_transform()
         self.H_matrix = self.get_H_matrix(self.T_matrix, self.E_matrix)   
-        self.evals, self.evecs = self.calculate_normal_modes(self.H_matrix)
-        self.evecs_vel = self.get_eigen_vectors_xv_coords(self.T_matrix,self.evecs)    
+        self.eigvals, self.eigvecs = self.calculate_normal_modes(self.H_matrix)
+        self.eigvecs_vel = self.get_eigen_vectors_xv_coords(self.T_matrix,self.eigvecs)    
         self.check_for_zero_modes()
         self.S_matrix = self.get_canonical_transformation()
         self.checks()
@@ -621,10 +624,6 @@ def run(self):
 #You could redefine the equilibrium finding function to assert that the equilibrium is linear. For example, make a wrapper inside the function for the potential, Jacobian, and Hessian that forces x_i and y_i = 0. 
  
 #This is the code for the Lamb-Dicke parameters: (not that overall phases don't matter here, but the relative phase does. We could pin down the phase with some convention like, "each mode's first non-negative LD value is defined positive." I also include some extra helper functions. 
- 
-
-## helper ? 
-
 def two_central_ion_separation(wz_Hz, l_2_cent):
     mass_yb_amu = 170.936 # TODO: hardcoded for now
     four_ion_analysis = GeneralizedModeAnalysisWithBranchSortedModes(N=4, wz=2*np.pi*wz_Hz, wy=2*np.pi*10e6, wx=2*np.pi*11e6, ionmass_amu=mass_yb_amu)
@@ -655,17 +654,17 @@ def calc_mode_energies(res, Fock_cutoffs):
     return energies_m1, energies_m2
  
 def calculate_mode_participation_factors(mode_analysis):
-    evecs = mode_analysis.evecs
-    num_coords, num_modes = np.shape(evecs)
+    eigvecs = mode_analysis.eigvecs
+    num_coords, num_modes = np.shape(eigvecs)
     num_ions = num_modes // 3
     mode_participation_factors = np.zeros((3, num_ions, num_modes), dtype = np.complex128)
     for mode_index in range(num_modes):
         for pos_coord in range(num_coords//2):
             direction_index = pos_coord // num_ions
             ion_index = pos_coord % num_ions
-            factor = np.sqrt(2* mode_analysis.m[ion_index] * mode_analysis.evals[mode_index] )
+            factor = np.sqrt(2* mode_analysis.m[ion_index] * mode_analysis.eigvals[mode_index] )
             zpm_dimensionful = np.sqrt(const.hbar / (2* mode_analysis.mass_scale * mode_analysis.trap_freq_scale))
-            mode_participation_factors[direction_index, ion_index, mode_index] = zpm_dimensionful * factor * evecs[pos_coord, mode_index]
+            mode_participation_factors[direction_index, ion_index, mode_index] = zpm_dimensionful * factor * eigvecs[pos_coord, mode_index]
     return mode_participation_factors
  
  

From b259fc4035534b42a719dc3917e1a457a9a89408 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 17:36:46 -0600
Subject: [PATCH 05/31]  consolidate ion equilibrium position solving functions

---
 src/ionsim/trapped_ion_mode_analysis.py | 80 ++++++++++---------------
 1 file changed, 32 insertions(+), 48 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 59853e6..f0c680b 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -17,6 +17,7 @@
 ########## List of non-substantive changes by Ethan: 
 ## 1. Variable changes for readable (e.g. omega_x instead of wx) 
 ## 2. Type-hinting in functions  
+## 3. Consolidated the functions for solving for ion positions equilibria 
 
 
 
@@ -110,7 +111,7 @@ def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float
         if not self.trap_is_stable():
             raise IonSimError(f"Error: Trap is unstable from negative trap frequencies.")
 
-        self.initial_equilibrium_guess = None
+        #self.initial_equilibrium_guess = None
         self.hasrun = False 
     
 
@@ -129,7 +130,7 @@ def calculate_species_trap_frequencies(self, omega: float) -> Vector[float]:
    
 
     def trap_is_stable(self):
-        # check that all trap frequencies are positive  
+        # Checks that all trap frequencies are positive  
         return np.all(self.omega_z > 0) and np.all(self.omega_y > 0) and np.all(self.omega_x > 0)
 
 
@@ -229,7 +230,8 @@ def run(self):
         self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
         #assert self.trap_is_stable()   # checked in the constructor  
         
-        self.u = self.calculate_equilibrium_positions()
+        #self.u = self.calculate_equilibrium_positions()
+        self.u = self.solve_for_equilibrium_positions()
         self.reindex_ions()
         self.E_matrix = self.get_E_matrix(self.u)  
         self.T_matrix = self.get_momentum_transform() 
@@ -243,20 +245,6 @@ def run(self):
         self.hasrun = True  
     
 
-
-    def calculate_equilibrium_positions(self):
-
-        if self.initial_equilibrium_guess is None:
-            self.initial_equilibrium_guess = self.get_initial_equilibrium_guess()
-        else:
-            self.u0 = self.initial_equilibrium_guess
-
-        u = self.find_equilibrium_positions(self.u0)
-        self.p0 = self.potential(u)
-        return u   
-
-
-
     def get_canonical_transformation(self):
         return get_canonical_transformation(self.H_matrix,self.eigvecs,evs=self.eigvals)    
 
@@ -291,7 +279,6 @@ def calculate_normal_modes(self, H_matrix: Matrix) -> (Vector, list[Vector]):
         return eigvals, eigvecs 
 
 
-
     def ion_coordinates_from_flattened(self, flattened_coordinate_vector: Vector) -> tuple(Vector, Vector, Vector):
         """ Returns arrays corresponding to an ion's x, y, and z coordinates, respectively, from a flattened vector u:
                 x = u[0:N], y = u[N:2N], z[2N:]
@@ -304,14 +291,6 @@ def ion_coordinates_from_flattened(self, flattened_coordinate_vector: Vector) ->
         return x, y, z 
         
 
-    def get_initial_equilibrium_guess(self):
-        """ Initialize the set of ion coordinates {u} as a flattened 1D array. """
-        # TODO: option for random vs. equally spaced along one axis? 
-        self.u0 = np.zeros(3*self.num_ions)
-        self.u0[:] = (np.random.rand(3*self.num_ions) * 2 - 1) * self.num_ions 
-        return self.u0
-
-
     def reindex_ions(self):  
         """ Re-indexes the ions to order based on the distance from the center of the trap, where smallest index is closest to the center """
         # TODO: How are even/odd cases is handled? 
@@ -328,31 +307,39 @@ def reindex_ions(self):
         self.atomic_numbers = self.atomic_numbers[idx]
 
 
-    def find_equilibrium_positions(self, u0: Vector):
-        """ Solves for equilibrium position vector: u0, which is a flattened spatial grid. """
+    def solve_for_equilibrium_positions(self, positions_guess: Vector | None=None):
+        """ Solves for equilibrium position vector: u, which represents a flattened spatial grid. """
+
+        # TODO: option for an initial guess choice? 
+        # Set the initial guess for the solver  
+        if positions_guess == None:
+            # Initialize a random guess  
+            u0 = np.zeros(3*self.num_ions)
+            u0[:] = (np.random.rand(3*self.num_ions) * 2 - 1) * self.num_ions 
+        else:
+            u0 = positions_guess
+
+        # Solve for the equilibrium positions by minimizing the potential energy (trap + Coulomb) 
         bfgs_tolerance = 1e-34
-        out = opt.minimize(self.potential, u0, method='BFGS', jac=self.force,
+        solver_output = opt.minimize(self.potential_energy, u0, method='BFGS', jac=self.force,
                                     options={'gtol': bfgs_tolerance, 'disp': False})
-        return out.x
+        equilibrium_positions = solver_output.x 
+
+        # Get potential energy at equilibrium positions  
+        self.p0 = self.potential_energy(equilibrium_positions)
+        return equilibrium_positions 
+
 
     def potential_trap(self, positions: Vector):
         """ Computes trapping potential, takes in a flattened (1D) array of all position coordinates """
- #        x = positions[0:self.num_ions]
- #        y = positions[self.num_ions:2*self.num_ions]
- #        z = positions[2*self.num_ions:]
         x,y,z = self.ion_coordinates_from_flattened(positions)
         V_trap = 0.5 * np.sum((self.m * self.wx ** 2) * x ** 2) + \
             0.5 * np.sum((self.m * self.wy ** 2) * y ** 2) + \
                 0.5 * np.sum((self.m * self.wz ** 2) * z ** 2)
         return V_trap
     
-
-
     def potential_coulomb(self, positions: Vector):
         """ Computes Coulomb potential, takes in a flattened (1D) array of all position coordinates """
- #        x = positions[0:self.num_ions]
- #        y = positions[self.num_ions:2*self.num_ions]
- #        z = positions[2*self.num_ions:]
         x,y,z = self.ion_coordinates_from_flattened(positions)
 
         dx = x[:, np.newaxis] - x
@@ -366,15 +353,11 @@ def potential_coulomb(self, positions: Vector):
         V_Coulomb *= .5 
         return V_Coulomb
 
-    def potential(self, positions):
+    def potential_energy(self, positions):
         return self.potential_trap(positions) + self.potential_coulomb(positions)   
 
 
-
     def force_trap(self, positions):
-#         x = positions[0:self.num_ions]
-#         y = positions[self.num_ions:2*self.num_ions]
-#         z = positions[2*self.num_ions:]
         x,y,z = self.ion_coordinates_from_flattened(positions)
 
         Ftrapx = self.m * self.wx**2 * x
@@ -387,15 +370,12 @@ def force_trap(self, positions):
 
 
     def force_coulomb(self, positions): 
-        #x = positions[0:self.num_ions]
-        #y = positions[self.num_ions:2*self.num_ions]
-        #z = positions[2*self.num_ions:]
         x,y,z = self.ion_coordinates_from_flattened(positions)
 
         dx = x[:, np.newaxis] - x
         dy = y[:, np.newaxis] - y
         dz = z[:, np.newaxis] - z
-        rsep = np.sqrt(dx ** 2 + dy ** 2 + dz ** 2).astype(np.float64)
+        rsep = np.sqrt(dx**2 + dy**2 + dz**2).astype(np.float64)
         qq = (self.q * self.q[:, np.newaxis]).astype(np.float64)    
 
         with np.errstate(divide='ignore', invalid='ignore'):
@@ -424,6 +404,7 @@ def force(self, positions):
         return Force
 
 
+    #### Matrix methods 
     def hessian_trap(self, positions: Vector):
         """ Computes the Hessian of the trap  """ 
         Hxx = np.diag(self.m * (self.wx**2) * np.ones(self.num_ions))
@@ -501,6 +482,7 @@ def get_E_matrix(self, positions: Vector):
         return E_matrix
 
     def get_H_matrix(self, T_matrix, E_matrix):  
+        """ Computes 6N x 6N Hermitian Hamiltonian matrix in (q, p) canonical coordinates """ 
         T_matrix_inv = np.linalg.inv(T_matrix)  
         H_matrix = T_matrix_inv.T @ E_matrix @ T_matrix_inv
         return H_matrix 
@@ -514,11 +496,13 @@ def get_momentum_transform(self):
         return T    
 
     def get_symplectic_matrix(self):
+        """ Computes J a 6N x 6N symplectic matrix defined as J = (0 , I ; I, 0) (eq. 1.67 of thesis) """
         zeros = np.zeros((3*self.num_ions, 3*self.num_ions), dtype=np.complex128)
         I = np.eye(3*self.num_ions, dtype=np.complex128)
         J = np.block([[zeros, I], [-I, zeros]])
         return J    
 
+    ### Mode organizing helper methods  
     def sort_modes(self, eigvals, eigvecs):
        eigvals = np.imag(eigvals)
        sort_dex = np.argsort(eigvals)
@@ -608,7 +592,7 @@ def run(self):
         self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
         #assert self.trap_is_stable()    
         
-        self.u = self.calculate_equilibrium_positions()
+        self.u = self.solve_for_equilibrium_positions()
         self.reindex_ions_by_z()
         self.E_matrix = self.get_E_matrix(self.u)  
         self.T_matrix = self.get_momentum_transform()

From 1e8d4387ab44dd75d424c9a174ee56e1bc5d1608 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 18:16:10 -0600
Subject: [PATCH 06/31]  mv eigvec helper fxns into class, expose matrix
 function args

---
 src/ionsim/trapped_ion_mode_analysis.py | 187 +++++++++++-------------
 1 file changed, 88 insertions(+), 99 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index f0c680b..670a467 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -18,6 +18,8 @@
 ## 1. Variable changes for readable (e.g. omega_x instead of wx) 
 ## 2. Type-hinting in functions  
 ## 3. Consolidated the functions for solving for ion positions equilibria 
+## 4. Moved eigenvector helper functions into the class. 
+## 5. Reduced number of matrices that we store as class attributes: Need to check whether we want to be storing those. 
 
 
 
@@ -27,48 +29,6 @@ def characteristic_length(q: float, mass: float, omega: float):
     l0 = ((k_e * q ** 2) / (.5 * mass * omega ** 2)) ** (1 / 3)
     return l0
 
-def get_norm(en,H):
-    """
-    Get the norm of the eigen vector en w.r.t. the Hamiltonian H.
-    """
-    norm = np.sqrt(en.T.conj() @ H @ en)
-    return norm
-
-
-def normalize_eigen_vectors(ens,H,evs=None): 
-    """
-    Rescale the eigen vectors ens w.r.t. the Hamiltonian H.
-    """
-    ens_rescaled = np.zeros_like(ens,dtype=complex)
-    num_coords,num_evs = np.shape(ens)
-    if evs is None:
-        evs = np.ones(num_evs)
-    for i in range(num_evs):
-        en = ens[:,i].reshape(num_coords,1)
-        norm = get_norm(en,H)
-        ens_rescaled[:,i] = en[:,0]/norm
-        ens_rescaled[:,i] *= np.sqrt(evs[i])
-    return ens_rescaled 
-
-
-
-def get_canonical_transformation(H,ens,evs=None):
-    """
-    Given the eigen-solve of the dynamical matrix, get the transform matrix 
-    to the canonical coordinates.
-    X = S X', where X' = (Q,P)^T and X = (q,p)^T
-    """
-    ## TODO: understand the sign in the transformation matrix
-    sign = -1
-    num_coords, num_evs = np.shape(ens)
-    assert num_coords //2 == num_evs    
-    T = np.zeros((num_coords,num_coords),dtype=complex) 
-    ens = normalize_eigen_vectors(ens,H,evs=evs)
-    T = np.sqrt(2)*np.concatenate((np.real(ens), sign*np.imag(ens)), axis=1)
-    # the sign is likely due to the convention of writing the time evolution as exp(-iwt)
-    return T
-
-
 def convert_to_array(x: float | int | list | NDArray):
     """ Converts a scalar to a vector or returns the vector """ 
     if not hasattr(x, "__len__"):
@@ -87,6 +47,7 @@ def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float
             - omega_z: harmonic trap frequency in the z direction. This is the "axial" direction used for 1D chains.  
             - atomic masses: array of atomic masses or a single number => same mass for each ion. Units are amu (atomic mass units) 
             - atomic numbers: array of (Z) atomic numbers (# of protons of an element) or a single number => all ions are the same.  
+
         """ 
         # TODO: Decide how to handle units and how much of pint we should use.  
         self.num_ions = num_ions
@@ -176,18 +137,41 @@ def convert_parameters_to_dimensionless(self, charge_scale: float=1., mass_scale
         self.characteristic_parameters['velocity'] = self.characteristic_parameters['length'] * trap_freq_scale  # characteristic velocity
         self.characteristic_parameters['energy'] = 0.5 * mass_scale * self.characteristic_parameters['velocity'] ** 2  # characteristic energy  
 
+
+    def get_norm(self, eigenvector: Vector, H: Matrix) -> float:
+        """ Computesthe norm of the eigenvector en w.r.t. the Hamiltonian H. """
+        norm = np.sqrt(eigenvector.T.conj() @ H @ eigenvector)
+        return norm
+    
+    
+    def get_canonical_transformation(self, H, eigenvectors, eigenvalues: Vector | None=None):
+        """ Given the eigensolve of the dynamical matrix, get the transform matrix to the canonical coordinates.
+            X = S X', where X' = (Q,P)^T and X = (q,p)^T
+        """
+        ## TODO: understand the sign in the transformation matrix
+        sign = -1
+        num_coords, num_eigenvalues = np.shape(eigenvectors)
+        assert num_coords //2 == num_eigenvalues 
+        T = np.zeros((num_coords,num_coords),dtype=complex) 
+        eigenvectors = self.normalize_eigenvectors(eigenvectors, H, eigenvalues)
+        T = np.sqrt(2)*np.concatenate((np.real(eigenvectors), sign*np.imag(eigenvectors)), axis=1)
+        # the sign is likely due to the convention of writing the time evolution as exp(-iwt)
+        return T
+
+
     def check_for_zero_modes(self):
         assert np.all(self.eigvals > 0), "All eigenvalues must be positive"   
 
-    def check_outer_relation(self): 
-        H = self.H_matrix.copy()       
-        D = self.get_symplectic_matrix() @ H
+    def check_outer_relation(self, H: Matrix): 
+        # TODO: Should this return something (True/False)? 
+        #H = self.H_matrix.copy()       
+        D = self.build_symplectic_matrix() @ H
         Eval, Evec = np.linalg.eig(D)
         _, en = self.sort_modes(Eval,Evec)
-        en = self.normalize_eigen_vectors(en,H)
+        en = self.normalize_eigenvectors(en,H)
         Outers = np.zeros((6*self.num_ions,6*self.num_ions),dtype=complex)
         for i in range(6*self.num_ions):
-            norm = get_norm(en[:,i],H)
+            norm = self.get_norm(en[:,i],H)
             Outers = Outers + np.outer(en[:,i],en[:,i].conj())/norm 
         I_right = H @ Outers
         I_left  = Outers @ H
@@ -199,17 +183,13 @@ def check_outer_relation(self):
         except AssertionError:
             warnings.warn("Outer relation check failed")
 
-
-
     def has_duplicate_eigvals(eigvals):
         evs = eigvals.copy()    
         return np.any(np.triu(np.isclose(evs[:, None], evs[None, :], atol=1e-6), k=1))  
     
-
-
-    def check_diagnolization(self):
-        M = np.linalg.inv(self.T_matrix) @ self.S_matrix    
-        H_diag = M.T @ self.E_matrix @ M    
+    def check_diagonalization(self, T: Matrix, S: Matrix, E: Matrix) -> bool:
+        M = np.linalg.inv(T) @ S
+        H_diag = M.T @ E @ M    
         H_diag_check = np.diag(np.tile(self.eigvals,2)) 
         np.set_printoptions(precision=2, suppress=True) 
         try:
@@ -218,41 +198,45 @@ def check_diagnolization(self):
             warnings.warn("Diagnolization check failed")    
             print("has duplicate eigenvalues: ", self.has_duplicate_eigvals(self.eigvals))
 
-
-
-    def checks(self):   
-        self.check_outer_relation()
-        self.check_diagnolization()
-
-    def run(self):
-        #self.dimensionless_parameters()
+    def solve_ion_trap_equilibria(self): 
+    #def run(self):
         # Convert to dimensionless units using axial trap frequency and first ion's mass and charge  
         self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
         #assert self.trap_is_stable()   # checked in the constructor  
         
-        #self.u = self.calculate_equilibrium_positions()
         self.u = self.solve_for_equilibrium_positions()
         self.reindex_ions()
-        self.E_matrix = self.get_E_matrix(self.u)  
-        self.T_matrix = self.get_momentum_transform() 
-        self.H_matrix = self.get_H_matrix(self.T_matrix, self.E_matrix)   
-
-        self.eigvals, self.eigvecs = self.calculate_normal_modes(self.H_matrix)
-        self.eigvecs_vel = self.get_eigen_vectors_xv_coords(self.T_matrix,self.eigvecs)    
+        # Compute matrices 
+        mass_matrix = self.build_mass_matrix(self.m)  
+        E_matrix = self.build_E_matrix(self.u, mass_matrix)  
+        T_matrix = self.build_momentum_transform_matrix(mass_matrix)
+        H_matrix = self.compute_H_matrix(T_matrix, E_matrix)   
+
+        self.eigvals, self.eigvecs = self.calculate_normal_modes(H_matrix)
+        self.eigvecs_vel = self.get_eigen_vectors_xv_coords(T_matrix, self.eigvecs)    
         self.check_for_zero_modes() 
-        self.S_matrix = self.get_canonical_transformation() 
-        self.checks() 
+        S_matrix = self.get_canonical_transformation(H_matrix, self.eigvecs, self.eigvals) 
+
+        # Perform checks 
+        self.check_outer_relation(H_matrix)
+        self.check_diagonalization(T_matrix, S_matrix, E_matrix)
         self.hasrun = True  
     
 
-    def get_canonical_transformation(self):
-        return get_canonical_transformation(self.H_matrix,self.eigvecs,evs=self.eigvals)    
-
-
-
-    def normalize_eigen_vectors(self, eigvecs, H_matrix,evs=None):
-        return normalize_eigen_vectors(eigvecs,H_matrix,evs=evs) 
-
+    def normalize_eigenvectors(self, eigvecs, H: Matrix, eigenvalues: Vector | None=None): 
+        """
+        Rescale the eigenvectors ens w.r.t. the Hamiltonian H.
+        """
+        eigvecs_rescaled = np.zeros_like(eigvecs,dtype=complex)
+        num_coords, num_eigenvalues = np.shape(eigvecs)
+        if eigenvalues is None:
+            eigenvalues = np.ones(num_eigenvalues)
+        for i in range(num_eigenvalues):
+            en = eigvecs[:,i].reshape(num_coords,1)
+            norm = self.get_norm(en,H)
+            eigvecs_rescaled[:,i] = en[:,0]/norm
+            eigvecs_rescaled[:,i] *= np.sqrt(eigenvalues[i])
+        return eigvecs_rescaled 
 
 
     def get_eigen_vectors_xv_coords(self,T,ens): 
@@ -269,12 +253,12 @@ def calculate_normal_modes(self, H_matrix: Matrix) -> (Vector, list[Vector]):
             The matrix "D" is defind as D = JH, where J = (0 , I ; I, 0) (eq. 1.67 of thesis)
 
         """
-        J = self.get_symplectic_matrix()
+        J = self.build_symplectic_matrix()
         D_matrix = J @ H_matrix  
         # u_n(t) = exp(-i w_n t) u_n(0), minus by convention
         eigvals, eigvecs = np.linalg.eig(-D_matrix)   
         eigvals, eigvecs = self.organize_modes(eigvals, eigvecs)
-        eigvecs = self.normalize_eigen_vectors(eigvecs, H_matrix) 
+        eigvecs = self.normalize_eigenvectors(eigvecs, H_matrix) 
 
         return eigvals, eigvecs 
 
@@ -461,41 +445,40 @@ def hessian_coulomb(self, positions: Vector):
         return H_coulomb
 
 
-    def hessian(self, positions: Vector):
+    def hessian(self, positions: Vector) -> Matrix:
         """ Computes total Hessian (trap + Coulomb interaction) """
         H = self.hessian_coulomb(positions) + self.hessian_trap(positions)  
         return H
 
-    def get_mass_matrix(self, m): 
-        return np.diag(np.tile(m, 3)) 
+    def build_mass_matrix(self, masses: Vector) -> Matrix: 
+        return np.diag(np.tile(masses, 3)) 
 
-    def get_E_matrix(self, positions: Vector): 
+    def build_E_matrix(self, positions: Vector, mass_matrix: Matrix) -> Matrix: 
         """ 6N x 6N energy matrix for N ions """ 
         PE_matrix = np.zeros((3*self.num_ions, 3*self.num_ions), dtype=np.complex128)
         KE_matrix = np.zeros((3*self.num_ions, 3*self.num_ions), dtype=np.complex128)
         E_matrix = np.zeros((6*self.num_ions, 6*self.num_ions), dtype=np.complex128)
 
         PE_matrix = self.hessian(positions)
-        KE_matrix = self.get_mass_matrix(self.m) 
+        KE_matrix = mass_matrix 
         zeros = np.zeros((3*self.num_ions, 3*self.num_ions)) 
         E_matrix = np.block([[PE_matrix, zeros], [zeros, KE_matrix]])
         return E_matrix
 
-    def get_H_matrix(self, T_matrix, E_matrix):  
+    def compute_H_matrix(self, T_matrix, E_matrix):  
         """ Computes 6N x 6N Hermitian Hamiltonian matrix in (q, p) canonical coordinates """ 
         T_matrix_inv = np.linalg.inv(T_matrix)  
         H_matrix = T_matrix_inv.T @ E_matrix @ T_matrix_inv
         return H_matrix 
 
-    def get_momentum_transform(self):
+    def build_momentum_transform_matrix(self, mass_matrix: Matrix) -> Matrix:
         # assuming no magnetic field
-        mass_matrix = self.get_mass_matrix(self.m)  
         eye = np.eye(3*self.num_ions)  
         zeros = np.zeros((3*self.num_ions, 3*self.num_ions))
         T = np.block([[eye, zeros], [zeros, mass_matrix]])  
         return T    
 
-    def get_symplectic_matrix(self):
+    def build_symplectic_matrix(self) -> Matrix:
         """ Computes J a 6N x 6N symplectic matrix defined as J = (0 , I ; I, 0) (eq. 1.67 of thesis) """
         zeros = np.zeros((3*self.num_ions, 3*self.num_ions), dtype=np.complex128)
         I = np.eye(3*self.num_ions, dtype=np.complex128)
@@ -587,21 +570,27 @@ def reindex_ions_by_z(self):
         self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
  
  
-    def run(self):
+    #def run(self):
+    def solve_ion_trap_equilibria(self): 
         #self.dimensionless_parameters()
         self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
         #assert self.trap_is_stable()    
         
         self.u = self.solve_for_equilibrium_positions()
         self.reindex_ions_by_z()
-        self.E_matrix = self.get_E_matrix(self.u)  
-        self.T_matrix = self.get_momentum_transform()
-        self.H_matrix = self.get_H_matrix(self.T_matrix, self.E_matrix)   
-        self.eigvals, self.eigvecs = self.calculate_normal_modes(self.H_matrix)
-        self.eigvecs_vel = self.get_eigen_vectors_xv_coords(self.T_matrix,self.eigvecs)    
+        #self.E_matrix = self.build_E_matrix(self.u)  
+        mass_matrix = self.build_mass_matrix(self.m)  
+        E_matrix = self.build_E_matrix(self.u, mass_matrix)  
+        T_matrix = self.build_momentum_transform_matrix(mass_matrix)
+        H_matrix = self.compute_H_matrix(T_matrix, E_matrix)   
+
+        self.eigvals, self.eigvecs = self.calculate_normal_modes(H_matrix)
+        self.eigvecs_vel = self.get_eigen_vectors_xv_coords(T_matrix, self.eigvecs)    
         self.check_for_zero_modes()
-        self.S_matrix = self.get_canonical_transformation()
-        self.checks()
+        #S_matrix = self.get_canonical_transformation(H_matrix) 
+        S_matrix = self.get_canonical_transformation(H_matrix, self.eigvecs, self.eigvals) 
+        self.check_outer_relation(H_matrix)
+        self.check_diagonalization(T_matrix, S_matrix, E_matrix)
         self.hasrun = True  
     
  
@@ -669,7 +658,7 @@ def check_single_ion_case():
     num_ions = 1
 
     mode_analysis_one = TrappedIonModeAnalysis(num_ions, wx, wy, wz, np.ones(num_ions)*mass_yb_amu, atomic_nums) 
-    mode_analysis_one.run()
+    mode_analysis_one.solve_ion_trap_equilibria()
     mode_participation_factors_one = calculate_mode_participation_factors(mode_analysis_one)
     ld_factors_one = k * mode_participation_factors_one
     eta_x = calc_single_ion_ld_factors(wx, k, mass_yb_amu)

From c7305eeb8e158d9c30fc02ef43b3ae63fdc036f3 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 18:36:19 -0600
Subject: [PATCH 07/31]  mv num_to_arr conversion into class, consolidate 2nd
 test

---
 src/ionsim/trapped_ion_mode_analysis.py | 81 ++++++++++---------------
 1 file changed, 33 insertions(+), 48 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 670a467..7c32d88 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -1,15 +1,9 @@
 import numpy as np
-#from generalized_mode_analysis import GeneralizedModeAnalysis as mode_analyzer
-#plt.style.use('ionsim')
-from time import time as timer
 from numpy.typing import NDArray
 import scipy.constants as const
 import warnings
 import scipy.optimize as opt    
 
-
-
-
 ########## List of substantive changes by Ethan: 
 ## 1. Made dimensionless variable function take in inputs so the user can choose the charge, mass, and trap freq scales to use.   
 ## 2. Added a helper function to pack/unpack the ion coordinates to reduce code. 
@@ -21,6 +15,10 @@
 ## 4. Moved eigenvector helper functions into the class. 
 ## 5. Reduced number of matrices that we store as class attributes: Need to check whether we want to be storing those. 
 
+########## Questions: 
+# 1. Is there a need for the "has run" boolean? 
+# 2. Should the dimensionless parameters have a naming convention, e.g. omega_x --> omega_x_ND 
+# 3. See TODO's 
 
 
 def characteristic_length(q: float, mass: float, omega: float):
@@ -29,12 +27,6 @@ def characteristic_length(q: float, mass: float, omega: float):
     l0 = ((k_e * q ** 2) / (.5 * mass * omega ** 2)) ** (1 / 3)
     return l0
 
-def convert_to_array(x: float | int | list | NDArray):
-    """ Converts a scalar to a vector or returns the vector """ 
-    if not hasattr(x, "__len__"):
-        x = np.ones(N) * x 
-        return x
-    return np.array(x) 
 
 
 class TrappedIonModeAnalysis:
@@ -52,8 +44,8 @@ def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float
         # TODO: Decide how to handle units and how much of pint we should use.  
         self.num_ions = num_ions
 
-        self.atomic_masses = convert_to_array(atomic_masses) * const.u # kg from amu 
-        self.atomic_numbers = convert_to_array(atomic_numbers)
+        self.atomic_masses = self.convert_to_array(atomic_masses) * const.u # kg from amu 
+        self.atomic_numbers = self.convert_to_array(atomic_numbers)
 
         # Safety checks: 
         assert len(self.atomic_masses) == num_ions
@@ -70,10 +62,10 @@ def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float
 
         # Check that the trap is stable: 
         if not self.trap_is_stable():
-            raise IonSimError(f"Error: Trap is unstable from negative trap frequencies.")
+            raise IonSimError(f"Trap is unstable from negative trap frequencies.")
 
         #self.initial_equilibrium_guess = None
-        self.hasrun = False 
+        #self.hasrun = False 
     
 
     def calculate_species_trap_frequencies(self, omega: float) -> Vector[float]: 
@@ -94,6 +86,13 @@ def trap_is_stable(self):
         # Checks that all trap frequencies are positive  
         return np.all(self.omega_z > 0) and np.all(self.omega_y > 0) and np.all(self.omega_x > 0)
 
+    def convert_to_array(self, x: float | int | list | NDArray) -> NDArray:
+        """ Converts a scalar to a vector or returns the vector """ 
+        if not hasattr(x, "__len__"):
+            x = np.ones(self.num_ions) * x 
+            return x
+        return np.array(x) 
+
 
     #def dimensionless_parameters(self):
     #def nondimensionalize_parameters(self):
@@ -139,7 +138,7 @@ def convert_parameters_to_dimensionless(self, charge_scale: float=1., mass_scale
 
 
     def get_norm(self, eigenvector: Vector, H: Matrix) -> float:
-        """ Computesthe norm of the eigenvector en w.r.t. the Hamiltonian H. """
+        """ Computes the norm of the eigenvector en w.r.t. the Hamiltonian H. """
         norm = np.sqrt(eigenvector.T.conj() @ H @ eigenvector)
         return norm
     
@@ -198,14 +197,14 @@ def check_diagonalization(self, T: Matrix, S: Matrix, E: Matrix) -> bool:
             warnings.warn("Diagnolization check failed")    
             print("has duplicate eigenvalues: ", self.has_duplicate_eigvals(self.eigvals))
 
-    def solve_ion_trap_equilibria(self): 
     #def run(self):
+    def solve_ion_trap_equilibrium(self): 
         # Convert to dimensionless units using axial trap frequency and first ion's mass and charge  
         self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
-        #assert self.trap_is_stable()   # checked in the constructor  
         
         self.u = self.solve_for_equilibrium_positions()
         self.reindex_ions()
+
         # Compute matrices 
         mass_matrix = self.build_mass_matrix(self.m)  
         E_matrix = self.build_E_matrix(self.u, mass_matrix)  
@@ -220,7 +219,7 @@ def solve_ion_trap_equilibria(self):
         # Perform checks 
         self.check_outer_relation(H_matrix)
         self.check_diagonalization(T_matrix, S_matrix, E_matrix)
-        self.hasrun = True  
+        #self.hasrun = True  
     
 
     def normalize_eigenvectors(self, eigvecs, H: Matrix, eigenvalues: Vector | None=None): 
@@ -340,7 +339,6 @@ def potential_coulomb(self, positions: Vector):
     def potential_energy(self, positions):
         return self.potential_trap(positions) + self.potential_coulomb(positions)   
 
-
     def force_trap(self, positions):
         x,y,z = self.ion_coordinates_from_flattened(positions)
 
@@ -351,8 +349,6 @@ def force_trap(self, positions):
         force_trap = np.hstack((Ftrapx, Ftrapy, Ftrapz))
         return force_trap
 
-
-
     def force_coulomb(self, positions): 
         x,y,z = self.ion_coordinates_from_flattened(positions)
 
@@ -377,12 +373,6 @@ def force_coulomb(self, positions):
         force_coulomb *= 0.5    
         return force_coulomb
 
-
-
-    def force(self, positions):
-        Force = self.force_coulomb(positions) + self.force_trap(positions)  
-        return Force
-
     def force(self, positions):
         Force = self.force_coulomb(positions) + self.force_trap(positions)  
         return Force
@@ -506,8 +496,6 @@ def organize_modes(self, eigvals, eigvecs):
         return eigvals, eigvecs 
 
 
- 
- 
 #classes
 class GeneralizedModeAnalysisWithBranchSortedModes(TrappedIonModeAnalysis):
  
@@ -566,19 +554,14 @@ def reindex_ions_by_z(self):
         self.omega_y = self.omega_y[idx]
         self.omega_z = self.omega_z[idx]
         # all ions are the same so this is safe. TODO: When does this change? 
-        #self.dimensionless_parameters()
         self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
  
- 
     #def run(self):
-    def solve_ion_trap_equilibria(self): 
-        #self.dimensionless_parameters()
+    def solve_ion_trap_equilibrium(self): 
         self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
-        #assert self.trap_is_stable()    
         
         self.u = self.solve_for_equilibrium_positions()
         self.reindex_ions_by_z()
-        #self.E_matrix = self.build_E_matrix(self.u)  
         mass_matrix = self.build_mass_matrix(self.m)  
         E_matrix = self.build_E_matrix(self.u, mass_matrix)  
         T_matrix = self.build_momentum_transform_matrix(mass_matrix)
@@ -587,11 +570,10 @@ def solve_ion_trap_equilibria(self):
         self.eigvals, self.eigvecs = self.calculate_normal_modes(H_matrix)
         self.eigvecs_vel = self.get_eigen_vectors_xv_coords(T_matrix, self.eigvecs)    
         self.check_for_zero_modes()
-        #S_matrix = self.get_canonical_transformation(H_matrix) 
         S_matrix = self.get_canonical_transformation(H_matrix, self.eigvecs, self.eigvals) 
         self.check_outer_relation(H_matrix)
         self.check_diagonalization(T_matrix, S_matrix, E_matrix)
-        self.hasrun = True  
+        #self.hasrun = True  
     
  
 #You could redefine the equilibrium finding function to assert that the equilibrium is linear. For example, make a wrapper inside the function for the potential, Jacobian, and Hessian that forces x_i and y_i = 0. 
@@ -654,11 +636,12 @@ def check_single_ion_case():
     wx = 2 * np.pi * 10e6  # radial trap frequency, something high to avoid any issues with mode ordering
     mass_yb_amu = 170.936
     k = 2 * np.pi / 355e-9
-    atomic_nums = np.array([70]) 
+    #atomic_nums = np.array([70]) 
+    Z = 70 
     num_ions = 1
 
-    mode_analysis_one = TrappedIonModeAnalysis(num_ions, wx, wy, wz, np.ones(num_ions)*mass_yb_amu, atomic_nums) 
-    mode_analysis_one.solve_ion_trap_equilibria()
+    mode_analysis_one = TrappedIonModeAnalysis(num_ions, wx, wy, wz, np.ones(num_ions)*mass_yb_amu, Z) 
+    mode_analysis_one.solve_ion_trap_equilibrium()
     mode_participation_factors_one = calculate_mode_participation_factors(mode_analysis_one)
     ld_factors_one = k * mode_participation_factors_one
     eta_x = calc_single_ion_ld_factors(wx, k, mass_yb_amu)
@@ -679,8 +662,11 @@ def check_two_ion_case():
     wx_tilt = np.sqrt(wx**2 - wz**2)
     mass_yb_amu = 170.936
     k = 2 * np.pi / 355e-9
-    mode_analysis_two= mode_analyzer(N =2, wz = wz, wy = wy, wx = wx, ionmass_amu= mass_yb_amu)
-    mode_analysis_two.run()
+    num_ions = 2
+    Z = 70
+    #mode_analysis_two= mode_analyzer(N =2, wz = wz, wy = wy, wx = wx, ionmass_amu= mass_yb_amu)
+    mode_analysis_two = TrappedIonModeAnalysis(num_ions, wx, wy, wz, np.ones(num_ions)*mass_yb_amu, Z) 
+    mode_analysis_two.solve_ion_trap_equilibrium()
     mode_participation_factors_two= calculate_mode_participation_factors(mode_analysis_two)
     ld_factors_two = k * mode_participation_factors_two
     eta_x = calc_single_ion_ld_factors(wx, k, mass_yb_amu)
@@ -716,7 +702,6 @@ def check_two_ion_case():
 ### Example usage  ###  
 if __name__ == '__main__':
     # Test analysis: 
-    #run()
-    #check_two_ion_case() 
-    check_single_ion_case()
-
+    # TODO: convert to tests 
+    check_two_ion_case() 
+    #check_single_ion_case()

From da1b89695fdf88345437bdcafc0637bd518aa325 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 18:44:31 -0600
Subject: [PATCH 08/31]  mv Lamb-Dicke parameter calcs into class

---
 src/ionsim/trapped_ion_mode_analysis.py | 42 +++++++++++++------------
 1 file changed, 22 insertions(+), 20 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 7c32d88..bfb978c 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -496,6 +496,22 @@ def organize_modes(self, eigvals, eigvecs):
         return eigvals, eigvecs 
 
 
+    def calculate_mode_participation_factors(self):
+        """ Computes ion-mode participation factors, related to the Lamb-Dicke parameters """ 
+        eigvecs = self.eigvecs
+        num_coords, num_modes = np.shape(eigvecs)
+        num_ions = num_modes // 3
+        mode_participation_factors = np.zeros((3, num_ions, num_modes), dtype = np.complex128)
+        for mode_index in range(num_modes):
+            for pos_coord in range(num_coords//2):
+                direction_index = pos_coord // num_ions
+                ion_index = pos_coord % num_ions
+                prefactor = np.sqrt(2 * self.m[ion_index] * self.eigvals[mode_index] )
+                zpm_dimensionful = np.sqrt(const.hbar / (2 * self.mass_scale * self.trap_freq_scale))
+                mode_participation_factors[direction_index, ion_index, mode_index] = zpm_dimensionful * prefactor * eigvecs[pos_coord, mode_index]
+        return mode_participation_factors
+ 
+
 #classes
 class GeneralizedModeAnalysisWithBranchSortedModes(TrappedIonModeAnalysis):
  
@@ -608,22 +624,6 @@ def calc_mode_energies(res, Fock_cutoffs):
         energies_m2[k] = qt.expect(E2_op, state)
     return energies_m1, energies_m2
  
-def calculate_mode_participation_factors(mode_analysis):
-    eigvecs = mode_analysis.eigvecs
-    num_coords, num_modes = np.shape(eigvecs)
-    num_ions = num_modes // 3
-    mode_participation_factors = np.zeros((3, num_ions, num_modes), dtype = np.complex128)
-    for mode_index in range(num_modes):
-        for pos_coord in range(num_coords//2):
-            direction_index = pos_coord // num_ions
-            ion_index = pos_coord % num_ions
-            factor = np.sqrt(2* mode_analysis.m[ion_index] * mode_analysis.eigvals[mode_index] )
-            zpm_dimensionful = np.sqrt(const.hbar / (2* mode_analysis.mass_scale * mode_analysis.trap_freq_scale))
-            mode_participation_factors[direction_index, ion_index, mode_index] = zpm_dimensionful * factor * eigvecs[pos_coord, mode_index]
-    return mode_participation_factors
- 
- 
- 
 def calc_single_ion_ld_factors(omega, k, mass_amu):
     z0 = np.sqrt(const.hbar / (2 * mass_amu * const.u * omega))
     return k * z0 # ignore the phase
@@ -642,7 +642,8 @@ def check_single_ion_case():
 
     mode_analysis_one = TrappedIonModeAnalysis(num_ions, wx, wy, wz, np.ones(num_ions)*mass_yb_amu, Z) 
     mode_analysis_one.solve_ion_trap_equilibrium()
-    mode_participation_factors_one = calculate_mode_participation_factors(mode_analysis_one)
+    #mode_participation_factors_one = calculate_mode_participation_factors(mode_analysis_one)
+    mode_participation_factors_one = mode_analysis_one.calculate_mode_participation_factors()
     ld_factors_one = k * mode_participation_factors_one
     eta_x = calc_single_ion_ld_factors(wx, k, mass_yb_amu)
     eta_y = calc_single_ion_ld_factors(wy, k, mass_yb_amu)
@@ -667,7 +668,8 @@ def check_two_ion_case():
     #mode_analysis_two= mode_analyzer(N =2, wz = wz, wy = wy, wx = wx, ionmass_amu= mass_yb_amu)
     mode_analysis_two = TrappedIonModeAnalysis(num_ions, wx, wy, wz, np.ones(num_ions)*mass_yb_amu, Z) 
     mode_analysis_two.solve_ion_trap_equilibrium()
-    mode_participation_factors_two= calculate_mode_participation_factors(mode_analysis_two)
+    #mode_participation_factors_two= calculate_mode_participation_factors(mode_analysis_two)
+    mode_participation_factors_two = mode_analysis_two.calculate_mode_participation_factors()
     ld_factors_two = k * mode_participation_factors_two
     eta_x = calc_single_ion_ld_factors(wx, k, mass_yb_amu)
     eta_y = calc_single_ion_ld_factors(wy, k, mass_yb_amu)
@@ -703,5 +705,5 @@ def check_two_ion_case():
 if __name__ == '__main__':
     # Test analysis: 
     # TODO: convert to tests 
-    check_two_ion_case() 
-    #check_single_ion_case()
+    #check_two_ion_case() 
+    check_single_ion_case()

From b414cbdb720c0f409b5179ecc26bda17ba202e39 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 19:15:42 -0600
Subject: [PATCH 09/31]  consolidate check for dupe eigvals, consolidate LD ref

---
 src/ionsim/trapped_ion_mode_analysis.py | 95 ++++++++++++++-----------
 1 file changed, 52 insertions(+), 43 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index bfb978c..780a91b 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -18,17 +18,17 @@
 ########## Questions: 
 # 1. Is there a need for the "has run" boolean? 
 # 2. Should the dimensionless parameters have a naming convention, e.g. omega_x --> omega_x_ND 
-# 3. See TODO's 
+# 3. What sets the phases of the Lamb-Dicke parameters?  
+# 4. See TODO's 
 
 
-def characteristic_length(q: float, mass: float, omega: float):
+def characteristic_length(q: float, mass: float, omega: float) -> float:
     """ Computes characteristic length in trapped ion system """
     k_e = 1 / (4 * np.pi * const.epsilon_0)  # Coulomb constant
     l0 = ((k_e * q ** 2) / (.5 * mass * omega ** 2)) ** (1 / 3)
     return l0
 
 
-
 class TrappedIonModeAnalysis:
     def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float, atomic_masses: Vector[float] | float, atomic_numbers: Vector[int] | int): 
         """ Class for determining properties of trapped-ion phonon modes, requiring the following system parameters:
@@ -74,14 +74,13 @@ def calculate_species_trap_frequencies(self, omega: float) -> Vector[float]:
             w_i = sqrt(q_{i} m_0 / m_{i} q_0) omega 
 
         """ 
-        # TODO: Does Wes have a ref. for this? 
+        # TODO: Does Wes have a ref. for this? Why is omega species-dependent?  
         # assume that the trapping frequency given corresponds to the first ion species 
         q0 = self.nuclear_charges[0] # charge of first ion 
         m0 = self.atomic_masses[0] # charge of first ion 
         species_trap_frequencies = np.sqrt(self.nuclear_charges * m0 /(q0 * self.atomic_masses)) * omega 
         return species_trap_frequencies 
    
-
     def trap_is_stable(self):
         # Checks that all trap frequencies are positive  
         return np.all(self.omega_z > 0) and np.all(self.omega_y > 0) and np.all(self.omega_x > 0)
@@ -93,7 +92,6 @@ def convert_to_array(self, x: float | int | list | NDArray) -> NDArray:
             return x
         return np.array(x) 
 
-
     #def dimensionless_parameters(self):
     #def nondimensionalize_parameters(self):
     def convert_parameters_to_dimensionless(self, charge_scale: float=1., mass_scale: float=1., trap_freq_scale: float=1.):
@@ -162,8 +160,6 @@ def check_for_zero_modes(self):
         assert np.all(self.eigvals > 0), "All eigenvalues must be positive"   
 
     def check_outer_relation(self, H: Matrix): 
-        # TODO: Should this return something (True/False)? 
-        #H = self.H_matrix.copy()       
         D = self.build_symplectic_matrix() @ H
         Eval, Evec = np.linalg.eig(D)
         _, en = self.sort_modes(Eval,Evec)
@@ -182,10 +178,6 @@ def check_outer_relation(self, H: Matrix):
         except AssertionError:
             warnings.warn("Outer relation check failed")
 
-    def has_duplicate_eigvals(eigvals):
-        evs = eigvals.copy()    
-        return np.any(np.triu(np.isclose(evs[:, None], evs[None, :], atol=1e-6), k=1))  
-    
     def check_diagonalization(self, T: Matrix, S: Matrix, E: Matrix) -> bool:
         M = np.linalg.inv(T) @ S
         H_diag = M.T @ E @ M    
@@ -194,8 +186,10 @@ def check_diagonalization(self, T: Matrix, S: Matrix, E: Matrix) -> bool:
         try:
             assert np.allclose(H_diag,H_diag_check)
         except AssertionError:
+            has_duplicate_eigenvalues = np.any(np.triu(np.isclose(eigenvalues[:, None], eigenvalues[None, :], atol = 1E-6), k=1))
             warnings.warn("Diagnolization check failed")    
-            print("has duplicate eigenvalues: ", self.has_duplicate_eigvals(self.eigvals))
+            print("has duplicate eigenvalues: ", has_duplicate_eigvals)
+
 
     #def run(self):
     def solve_ion_trap_equilibrium(self): 
@@ -212,7 +206,7 @@ def solve_ion_trap_equilibrium(self):
         H_matrix = self.compute_H_matrix(T_matrix, E_matrix)   
 
         self.eigvals, self.eigvecs = self.calculate_normal_modes(H_matrix)
-        self.eigvecs_vel = self.get_eigen_vectors_xv_coords(T_matrix, self.eigvecs)    
+        self.eigvecs_vel = self.get_eigenvectors_xv_coords(T_matrix, self.eigvecs)    
         self.check_for_zero_modes() 
         S_matrix = self.get_canonical_transformation(H_matrix, self.eigvecs, self.eigvals) 
 
@@ -238,7 +232,7 @@ def normalize_eigenvectors(self, eigvecs, H: Matrix, eigenvalues: Vector | None=
         return eigvecs_rescaled 
 
 
-    def get_eigen_vectors_xv_coords(self,T,ens): 
+    def get_eigenvectors_xv_coords(self,T,ens): 
         ens_vel = np.zeros_like(ens,dtype=complex)
         ens_vel[:,:] = np.linalg.inv(T) @ ens
         return ens_vel 
@@ -313,7 +307,7 @@ def solve_for_equilibrium_positions(self, positions_guess: Vector | None=None):
         return equilibrium_positions 
 
 
-    def potential_trap(self, positions: Vector):
+    def potential_trap(self, positions: Vector) -> Vector:
         """ Computes trapping potential, takes in a flattened (1D) array of all position coordinates """
         x,y,z = self.ion_coordinates_from_flattened(positions)
         V_trap = 0.5 * np.sum((self.m * self.wx ** 2) * x ** 2) + \
@@ -321,7 +315,7 @@ def potential_trap(self, positions: Vector):
                 0.5 * np.sum((self.m * self.wz ** 2) * z ** 2)
         return V_trap
     
-    def potential_coulomb(self, positions: Vector):
+    def potential_coulomb(self, positions: Vector) -> Vector:
         """ Computes Coulomb potential, takes in a flattened (1D) array of all position coordinates """
         x,y,z = self.ion_coordinates_from_flattened(positions)
 
@@ -336,10 +330,11 @@ def potential_coulomb(self, positions: Vector):
         V_Coulomb *= .5 
         return V_Coulomb
 
-    def potential_energy(self, positions):
+    def potential_energy(self, positions: Vector) -> Vector:
+        """ Computes the potential energy of each ion in its coordinate basis """ 
         return self.potential_trap(positions) + self.potential_coulomb(positions)   
 
-    def force_trap(self, positions):
+    def force_trap(self, positions: Vector):
         x,y,z = self.ion_coordinates_from_flattened(positions)
 
         Ftrapx = self.m * self.wx**2 * x
@@ -379,7 +374,7 @@ def force(self, positions):
 
 
     #### Matrix methods 
-    def hessian_trap(self, positions: Vector):
+    def hessian_trap(self, positions: Vector) -> Matrix:
         """ Computes the Hessian of the trap  """ 
         Hxx = np.diag(self.m * (self.wx**2) * np.ones(self.num_ions))
         Hyy = np.diag(self.m * (self.wy**2) * np.ones(self.num_ions))  
@@ -388,7 +383,7 @@ def hessian_trap(self, positions: Vector):
         H = np.block([[Hxx, zeros, zeros], [zeros, Hyy, zeros], [zeros, zeros, Hzz]])
         return H
 
-    def hessian_coulomb(self, positions: Vector):
+    def hessian_coulomb(self, positions: Vector) -> Matrix:
         """ Computes the Hessian of the Coulomb interaction """ 
         x,y,z = self.ion_coordinates_from_flattened(positions)
         
@@ -441,6 +436,7 @@ def hessian(self, positions: Vector) -> Matrix:
         return H
 
     def build_mass_matrix(self, masses: Vector) -> Matrix: 
+        """ Builds a diagonal mass matrix representing the ions """ 
         return np.diag(np.tile(masses, 3)) 
 
     def build_E_matrix(self, positions: Vector, mass_matrix: Matrix) -> Matrix: 
@@ -510,8 +506,22 @@ def calculate_mode_participation_factors(self):
                 zpm_dimensionful = np.sqrt(const.hbar / (2 * self.mass_scale * self.trap_freq_scale))
                 mode_participation_factors[direction_index, ion_index, mode_index] = zpm_dimensionful * prefactor * eigvecs[pos_coord, mode_index]
         return mode_participation_factors
- 
 
+    # Derived properties 
+    def compute_reference_single_ion_lamb_dicke_factors(self, wavenumber: float) -> (float, float, float):
+        """ Computes analytical Lamb-Dicke parameter by eta = k * sqrt(hbar / m omega) for an ion in a light-field with wavevector |k| 
+
+            - wavenumber: wavevector magnitude |k| in units of 1 / m  
+
+        """
+        # TODO: better way to handle units? 
+        # Convention to use first value of trap frequency arrays (representing one of the ions)  
+        trap_frequencies = np.array([self.omega_x[0], self.omega_y[0], self.omega_z[0]])        
+        mass = self.atomic_masses[0] 
+        eta_x, eta_y, eta_z = wavenumber * np.sqrt(const.hbar / (2 * mass * trap_frequencies))
+        return eta_x, eta_y, eta_z    # ignore the phase
+ 
+    
 #classes
 class GeneralizedModeAnalysisWithBranchSortedModes(TrappedIonModeAnalysis):
  
@@ -532,6 +542,7 @@ def _xyz_classify_modes(self, evecs):
         return classifier
  
     def sort_by_branch(self, evals, evecs):
+        """ Sorts the eigenvalues, eigenvectors by mode branch (radial x, radial y, axial (z)) """
         classifier = self._xyz_classify_modes(evecs)
         # within each branch, sort by frequency
         N_ions = len(evals) // 3
@@ -543,7 +554,8 @@ def sort_by_branch(self, evals, evecs):
             sorted_by_branch_evals[direction*N_ions:(direction+1)*N_ions] = evals[direction_indices]
             sorted_by_branch_evecs[:, direction*N_ions:(direction+1)*N_ions] = evecs[:, direction_indices]
         return sorted_by_branch_evals, sorted_by_branch_evecs
- 
+
+    # Override from parent  
     def organize_modes(self, eigvals, eigvecs):        
         eigvals, eigvecs = self.sort_modes(eigvals, eigvecs)
         eigvals, eigvecs = self.split_modes(eigvals, eigvecs)
@@ -584,13 +596,15 @@ def solve_ion_trap_equilibrium(self):
         H_matrix = self.compute_H_matrix(T_matrix, E_matrix)   
 
         self.eigvals, self.eigvecs = self.calculate_normal_modes(H_matrix)
-        self.eigvecs_vel = self.get_eigen_vectors_xv_coords(T_matrix, self.eigvecs)    
+        self.eigvecs_vel = self.get_eigenvectors_xv_coords(T_matrix, self.eigvecs)    
         self.check_for_zero_modes()
-        S_matrix = self.get_canonical_transformation(H_matrix, self.eigvecs, self.eigvals) 
         self.check_outer_relation(H_matrix)
+
+        S_matrix = self.get_canonical_transformation(H_matrix, self.eigvecs, self.eigvals) 
         self.check_diagonalization(T_matrix, S_matrix, E_matrix)
         #self.hasrun = True  
-    
+
+
  
 #You could redefine the equilibrium finding function to assert that the equilibrium is linear. For example, make a wrapper inside the function for the potential, Jacobian, and Hessian that forces x_i and y_i = 0. 
  
@@ -624,10 +638,6 @@ def calc_mode_energies(res, Fock_cutoffs):
         energies_m2[k] = qt.expect(E2_op, state)
     return energies_m1, energies_m2
  
-def calc_single_ion_ld_factors(omega, k, mass_amu):
-    z0 = np.sqrt(const.hbar / (2 * mass_amu * const.u * omega))
-    return k * z0 # ignore the phase
- 
 def check_single_ion_case():
     # for a single ion, the mode participation factors should just be the LD factors for each mode and direction.
     # this is a good sanity check to make sure the mode participation factor calculation is correct.
@@ -642,18 +652,18 @@ def check_single_ion_case():
 
     mode_analysis_one = TrappedIonModeAnalysis(num_ions, wx, wy, wz, np.ones(num_ions)*mass_yb_amu, Z) 
     mode_analysis_one.solve_ion_trap_equilibrium()
-    #mode_participation_factors_one = calculate_mode_participation_factors(mode_analysis_one)
     mode_participation_factors_one = mode_analysis_one.calculate_mode_participation_factors()
-    ld_factors_one = k * mode_participation_factors_one
-    eta_x = calc_single_ion_ld_factors(wx, k, mass_yb_amu)
-    eta_y = calc_single_ion_ld_factors(wy, k, mass_yb_amu)
-    eta_z = calc_single_ion_ld_factors(wz, k, mass_yb_amu)
+    simulated_LD_parameters = k * mode_participation_factors_one
+
+    # Compute analytical reference 
+    eta_x, eta_y, eta_z = mode_analysis_one.compute_reference_single_ion_lamb_dicke_factors(k) 
     ld_factors_one_analytical = np.zeros((3, 1, 3), dtype = np.complex128)
     ld_factors_one_analytical[0, 0, 2] = eta_x
     ld_factors_one_analytical[1, 0, 1] = eta_y
     ld_factors_one_analytical[2, 0, 0] = eta_z
     print(f"\nLD factors: {ld_factors_one_analytical}")
-    print("\nRatio of single ion LD factors from mode participation calculation to analytical calculation: \n", ld_factors_one / ld_factors_one_analytical)
+    #print("\nRatio of single ion LD factors from mode participation calculation to analytical calculation: \n", simulated_LD_parameters / ld_factors_one_analytical)
+    print(f"\nLD factors (via diagonalization):\n {simulated_LD_parameters}")
  
 def check_two_ion_case():
     wz = 2 * np.pi * .5e6  # axial trap frequency
@@ -671,9 +681,7 @@ def check_two_ion_case():
     #mode_participation_factors_two= calculate_mode_participation_factors(mode_analysis_two)
     mode_participation_factors_two = mode_analysis_two.calculate_mode_participation_factors()
     ld_factors_two = k * mode_participation_factors_two
-    eta_x = calc_single_ion_ld_factors(wx, k, mass_yb_amu)
-    eta_y = calc_single_ion_ld_factors(wy, k, mass_yb_amu)
-    eta_z = calc_single_ion_ld_factors(wz, k, mass_yb_amu)
+    eta_x, eta_y, eta_z = mode_analysis_two.compute_reference_single_ion_lamb_dicke_factors(k) 
     eta_COM_x = eta_x / np.sqrt(2)
     eta_COM_y = eta_y / np.sqrt(2)
     eta_COM_z = eta_z / np.sqrt(2)
@@ -696,8 +704,9 @@ def check_two_ion_case():
     ld_factors_two_analytical[2, 1, 1] = -eta_stretch_z
     # these seem to work out!
     print(f"\nLD factors: {ld_factors_two_analytical}")
-    print("Ratio of two ion LD factors from mode participation calculation to analytical calculation: \n", ld_factors_two / ld_factors_two_analytical)
- 
+    #print("Ratio of two ion LD factors from mode participation calculation to analytical calculation: \n", ld_factors_two / ld_factors_two_analytical)
+    print(f"\nLD factors (via diagonalization):\n {ld_factors_two}")
+
  
 
  
@@ -705,5 +714,5 @@ def check_two_ion_case():
 if __name__ == '__main__':
     # Test analysis: 
     # TODO: convert to tests 
-    #check_two_ion_case() 
-    check_single_ion_case()
+    check_two_ion_case() 
+    #check_single_ion_case()

From c24465d2ce7dec424dd40746788089a5bdc2c046 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 19:38:27 -0600
Subject: [PATCH 10/31]   add draft of test from mode analysis, add mode
 analysis to init

---
 src/ionsim/__init__.py                  |  1 +
 src/ionsim/trapped_ion_mode_analysis.py |  3 +-
 tests/test_mode_analysis.py             | 97 +++++++++++++++++++++++++
 3 files changed, 100 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_mode_analysis.py

diff --git a/src/ionsim/__init__.py b/src/ionsim/__init__.py
index 3adc264..e83f074 100644
--- a/src/ionsim/__init__.py
+++ b/src/ionsim/__init__.py
@@ -30,3 +30,4 @@ def patched_kron(a, b, *args, **kwargs):
 from .zeeman_solver import ZeemanHyperfineSolver
 from .composite_operator import CompositeOperator
 from .dissipator import Dissipator, DissipatorSpontaneousEmission, Lindbladian 
+from .trapped_ion_mode_analysis import TrappedIonModeAnalysis 
diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 780a91b..08c2b73 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -676,11 +676,12 @@ def check_two_ion_case():
     num_ions = 2
     Z = 70
     #mode_analysis_two= mode_analyzer(N =2, wz = wz, wy = wy, wx = wx, ionmass_amu= mass_yb_amu)
-    mode_analysis_two = TrappedIonModeAnalysis(num_ions, wx, wy, wz, np.ones(num_ions)*mass_yb_amu, Z) 
+    mode_analysis_two = TrappedIonModeAnalysis(num_ions, wx, wy, wz, mass_yb_amu, Z) 
     mode_analysis_two.solve_ion_trap_equilibrium()
     #mode_participation_factors_two= calculate_mode_participation_factors(mode_analysis_two)
     mode_participation_factors_two = mode_analysis_two.calculate_mode_participation_factors()
     ld_factors_two = k * mode_participation_factors_two
+
     eta_x, eta_y, eta_z = mode_analysis_two.compute_reference_single_ion_lamb_dicke_factors(k) 
     eta_COM_x = eta_x / np.sqrt(2)
     eta_COM_y = eta_y / np.sqrt(2)
diff --git a/tests/test_mode_analysis.py b/tests/test_mode_analysis.py
new file mode 100644
index 0000000..1ade8c2
--- /dev/null
+++ b/tests/test_mode_analysis.py
@@ -0,0 +1,97 @@
+import unittest
+import numpy as np
+from scipy.linalg import expm
+#from ionsim.zeeman_solver import ZeemanHyperfineSolver 
+from ionsim.trapped_ion_mode_analysis import TrappedIonModeAnalysis 
+from ionsim.testing import assert_array_close
+from ionsim.degree_of_freedom import AtomicSpin
+
+class TestModeAnalysis(unittest.TestCase):
+
+    def setUp(self):
+        TPI = 2*np.pi
+        """Set up the necessary objects for testing."""
+        # Set up tests with a list of dictionaries: 
+        self.test_cases = [
+            {
+                'test name' : 'single ion 171Yb',
+                'mass' : 170.936, # amu  
+                'omega x' : 10. * TPI * 1E6,  # rad/s 
+                'omega y' : 1.9 * TPI * 1E6,
+                'omega z' : 0.5 * TPI * 1E6,
+                'N' : 1, # number of ions 
+                'wavevector' : TPI / (355.0*1E-9), # 1/m
+                'Z' : 70 
+            },
+            {
+                'test name' : 'two ion 171Yb',
+                'mass' : 170.936, # amu  
+                'omega x' : 2. * TPI * 1E6,  # rad/s 
+                'omega y' : 1.5 * TPI * 1E6,
+                'omega z' : 0.5 * TPI * 1E6,
+                'N' : 2, # number of ions 
+                'wavevector' : TPI / (355.0*1E-9), # 1/m
+                'Z' : 70 
+            }
+        ]
+        self.mode_analyzers = {
+            case['test name'] : TrappedIonModeAnalysis(case['N'], case['omega x'], case['omega y'], case['omega z'], case['mass'], case['Z']) 
+                for case in self.test_cases
+            }
+
+        # Test 3: Use AtomicSpin from_species()  
+        self.spin = AtomicSpin.from_species(species='171Yb', term_symbols=['S0'], level_names=['S0,1/2,1/2', 'S0,1/2,-1/2'], magnetic_field = 400.) 
+
+    def test_mode_analysis_solvers(self):
+        # Test functionality of mode analysis for computing Lamb-Dicke parameters as compared to a reference  
+        # Create tests for each solver. Each atom case has a list of tests:  
+        #tests = {case['test name'] : [] for case in self.test_cases}
+
+
+        # Compute Lamb-Dicke parameters for both test cases 
+        for case, mode_analyzer in zip(self.test_cases, self.mode_analyzers):
+            # Solve the ion-trap equilibrium problem
+            mode_analyzer.solve_ion_trap_equilibrium()
+            # Compute and store Lamb-Dicke parameters: 
+            case['Lamb Dicke parameters'] = case['wavevector'] * mode_analyzer.calculate_mode_participation_factors()
+            print(f"Lamb Dicke parameters: {case['Lamb Dicke parameters']}")
+
+
+        # For magnetic field strength, verify energy shift for a state: 
+ #        tests['171Yb'].append({'Magnetic field' : 400, # Gauss
+ #                'state type' : 'hyperfine',
+ #                'F,mF' : (0.5,-0.5),
+ #                'Zeeman shift' : 149.98214416459987 # kHz
+ #            })    
+ #
+ #        tests['171Yb'].append({'Magnetic field' : 400, # Gauss
+ #                'state type' : 'hyperfine',
+ #                'F,mF' : (0.5,0.5),
+ #                'Zeeman shift' : -149.98214416459987 # kHz
+ #            })        
+ #
+ #        # Loop through each solver and loop through each tests
+ #        for name, solver in (self.solvers).items():
+ #            test_list = tests[name] 
+ #            for test in test_list:
+ #                shifts, eigenvecs = solver.solve_at_field(test['Magnetic field'])
+ #                if test['state type'] == 'hyperfine':
+ #                    f, mF = test['F,mF']
+ #                    calculated_value = solver.get_state_energy(shifts, eigenvecs, f = f, mf = mF)
+ #                else:
+ #                    mj, mi = test['mJ,mI']
+ #                    calculated_value = solver.get_state_energy_from_mjmi_pair(shifts, eigenvecs, mj = mj, mi = mi)
+ #
+ #                self.assertAlmostEqual(calculated_value, test['Zeeman shift'], places=6) 
+
+
+ #    def test_AtomicSpin_ZeemanShift(self): 
+ #        """Test the Zeemaen shift functionality within AtomicSpin class."""
+ #        expected_frequency = 149.98214416459987*2. # kHz
+ #        qubit_frequency = self.spin.energy_levels[1].energy - self.spin.energy_levels[0].energy
+ #        qubit_frequency /= (2.* np.pi) # convert from rad/s to Hz 
+ #        self.assertAlmostEqual(expected_frequency, np.abs(qubit_frequency)*1E-3, places=6)
+
+
+if __name__ == '__main__':
+    unittest.main()

From 0a96a39d7eded84d4085bcfe2a58ef1e04836be6 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 21:48:40 -0600
Subject: [PATCH 11/31]   add tests for Lamb-Dicke parameters via mode analysis

---
 tests/test_dissipators.py   |  1 -
 tests/test_mode_analysis.py | 94 ++++++++++++++++++-------------------
 tests/test_zeeman_solver.py |  1 -
 3 files changed, 46 insertions(+), 50 deletions(-)

diff --git a/tests/test_dissipators.py b/tests/test_dissipators.py
index 24592f1..7494bd1 100644
--- a/tests/test_dissipators.py
+++ b/tests/test_dissipators.py
@@ -5,7 +5,6 @@
 from scipy.sparse import kron as skron
 
 from ionsim.zeeman_solver import ZeemanHyperfineSolver 
-from ionsim.testing import assert_array_close
 from ionsim.basis import StandardBasis
 from ionsim.operator import CouplingOperator, EnergyShiftOperator
 from ionsim.named_operators import Pauli, Fock
diff --git a/tests/test_mode_analysis.py b/tests/test_mode_analysis.py
index 1ade8c2..aa87b52 100644
--- a/tests/test_mode_analysis.py
+++ b/tests/test_mode_analysis.py
@@ -1,10 +1,8 @@
 import unittest
 import numpy as np
 from scipy.linalg import expm
-#from ionsim.zeeman_solver import ZeemanHyperfineSolver 
 from ionsim.trapped_ion_mode_analysis import TrappedIonModeAnalysis 
 from ionsim.testing import assert_array_close
-from ionsim.degree_of_freedom import AtomicSpin
 
 class TestModeAnalysis(unittest.TestCase):
 
@@ -39,59 +37,59 @@ def setUp(self):
                 for case in self.test_cases
             }
 
-        # Test 3: Use AtomicSpin from_species()  
-        self.spin = AtomicSpin.from_species(species='171Yb', term_symbols=['S0'], level_names=['S0,1/2,1/2', 'S0,1/2,-1/2'], magnetic_field = 400.) 
+        # Compute Lamb-Dicke parameter analytical reference for each case  
+        self.references = {}
 
-    def test_mode_analysis_solvers(self):
-        # Test functionality of mode analysis for computing Lamb-Dicke parameters as compared to a reference  
-        # Create tests for each solver. Each atom case has a list of tests:  
-        #tests = {case['test name'] : [] for case in self.test_cases}
+        # 1. Single ion case: 
+        k = self.test_cases[0]['wavevector']
+        eta_x, eta_y, eta_z = self.mode_analyzers['single ion 171Yb'].compute_reference_single_ion_lamb_dicke_factors(k)
+        etas_analytical = np.zeros((3, 1, 3), dtype = np.complex128)
+        etas_analytical[0, 0, 2] = eta_x
+        etas_analytical[1, 0, 1] = eta_y
+        etas_analytical[2, 0, 0] = eta_z
+        self.references['single ion 171Yb'] = etas_analytical
+
+        # 2. Two-ion case: 
+        k = self.test_cases[1]['wavevector']
+        eta_x, eta_y, eta_z = self.mode_analyzers['two ion 171Yb'].compute_reference_single_ion_lamb_dicke_factors(k)
+        wx = self.test_cases[1]['omega x']
+        wy = self.test_cases[1]['omega y']
+        wz = self.test_cases[1]['omega z']
 
+        wy_tilt = np.sqrt(wy**2 - wz**2)
+        wx_tilt = np.sqrt(wx**2 - wz**2)
+        eta_COM_x = eta_x / np.sqrt(2)
+        eta_COM_y = eta_y / np.sqrt(2)
+        eta_COM_z = eta_z / np.sqrt(2)
+        eta_stretch_z = eta_z /np.sqrt(2) /3**(1/4)
+        # I don't know the analytical expressions for the tilt modes off the top of my head... let's assume its the square root of the normalized mode frequency
+        eta_tilt_x = eta_x / np.sqrt(2) / np.sqrt(wx_tilt / wx) 
+        eta_tilt_y = eta_y / np.sqrt(2) / np.sqrt(wy_tilt / wy) 
+        etas_analytical = np.zeros((3, 2, 6), dtype = np.complex128)
+        etas_analytical[0, 0, 5] = eta_COM_x
+        etas_analytical[0, 1, 5] = eta_COM_x
+        etas_analytical[0, 0, 4] = eta_tilt_x
+        etas_analytical[0, 1, 4] = -eta_tilt_x
+        etas_analytical[1, 0, 3] = eta_COM_y
+        etas_analytical[1, 1, 3] = eta_COM_y
+        etas_analytical[1, 0, 2] = eta_tilt_y
+        etas_analytical[1, 1, 2] = -eta_tilt_y
+        etas_analytical[2, 0, 0] = eta_COM_z
+        etas_analytical[2, 1, 0] = eta_COM_z
+        etas_analytical[2, 0, 1] = eta_stretch_z
+        etas_analytical[2, 1, 1] = -eta_stretch_z
+        self.references['two ion 171Yb'] = etas_analytical
 
-        # Compute Lamb-Dicke parameters for both test cases 
-        for case, mode_analyzer in zip(self.test_cases, self.mode_analyzers):
+
+    def test_mode_analysis_solvers(self):
+        # Test functionality of mode analysis for computing Lamb-Dicke parameters as compared to a reference  
+        # Compute Lamb-Dicke parameters for all test cases 
+        for case, mode_analyzer, reference_values in zip(self.test_cases, self.mode_analyzers.values(), self.references.values()):
             # Solve the ion-trap equilibrium problem
             mode_analyzer.solve_ion_trap_equilibrium()
             # Compute and store Lamb-Dicke parameters: 
             case['Lamb Dicke parameters'] = case['wavevector'] * mode_analyzer.calculate_mode_participation_factors()
-            print(f"Lamb Dicke parameters: {case['Lamb Dicke parameters']}")
-
-
-        # For magnetic field strength, verify energy shift for a state: 
- #        tests['171Yb'].append({'Magnetic field' : 400, # Gauss
- #                'state type' : 'hyperfine',
- #                'F,mF' : (0.5,-0.5),
- #                'Zeeman shift' : 149.98214416459987 # kHz
- #            })    
- #
- #        tests['171Yb'].append({'Magnetic field' : 400, # Gauss
- #                'state type' : 'hyperfine',
- #                'F,mF' : (0.5,0.5),
- #                'Zeeman shift' : -149.98214416459987 # kHz
- #            })        
- #
- #        # Loop through each solver and loop through each tests
- #        for name, solver in (self.solvers).items():
- #            test_list = tests[name] 
- #            for test in test_list:
- #                shifts, eigenvecs = solver.solve_at_field(test['Magnetic field'])
- #                if test['state type'] == 'hyperfine':
- #                    f, mF = test['F,mF']
- #                    calculated_value = solver.get_state_energy(shifts, eigenvecs, f = f, mf = mF)
- #                else:
- #                    mj, mi = test['mJ,mI']
- #                    calculated_value = solver.get_state_energy_from_mjmi_pair(shifts, eigenvecs, mj = mj, mi = mi)
- #
- #                self.assertAlmostEqual(calculated_value, test['Zeeman shift'], places=6) 
-
-
- #    def test_AtomicSpin_ZeemanShift(self): 
- #        """Test the Zeemaen shift functionality within AtomicSpin class."""
- #        expected_frequency = 149.98214416459987*2. # kHz
- #        qubit_frequency = self.spin.energy_levels[1].energy - self.spin.energy_levels[0].energy
- #        qubit_frequency /= (2.* np.pi) # convert from rad/s to Hz 
- #        self.assertAlmostEqual(expected_frequency, np.abs(qubit_frequency)*1E-3, places=6)
-
+            assert_array_close(np.abs(case['Lamb Dicke parameters']), np.abs(reference_values), atol = 1E-10, rtol=None)
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/test_zeeman_solver.py b/tests/test_zeeman_solver.py
index 68d3d6b..ca06fa6 100644
--- a/tests/test_zeeman_solver.py
+++ b/tests/test_zeeman_solver.py
@@ -2,7 +2,6 @@
 import numpy as np
 from scipy.linalg import expm
 from ionsim.zeeman_solver import ZeemanHyperfineSolver 
-from ionsim.testing import assert_array_close
 from ionsim.degree_of_freedom import AtomicSpin
 
 class TestZeemanSolver(unittest.TestCase):

From b26c2fe696b138d46743299cbd87c59e504a7e6f Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 21:54:13 -0600
Subject: [PATCH 12/31]   fix type hinting bugs with custom Vector object

---
 src/ionsim/trapped_ion_mode_analysis.py | 162 ++++++++++++------------
 1 file changed, 81 insertions(+), 81 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 08c2b73..976067c 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -30,7 +30,7 @@ def characteristic_length(q: float, mass: float, omega: float) -> float:
 
 
 class TrappedIonModeAnalysis:
-    def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float, atomic_masses: Vector[float] | float, atomic_numbers: Vector[int] | int): 
+    def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float, atomic_masses: Vector | float, atomic_numbers: Vector | int): 
         """ Class for determining properties of trapped-ion phonon modes, requiring the following system parameters:
 
             - num_ions: the number of ions
@@ -68,7 +68,7 @@ def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float
         #self.hasrun = False 
     
 
-    def calculate_species_trap_frequencies(self, omega: float) -> Vector[float]: 
+    def calculate_species_trap_frequencies(self, omega: float) -> Vector: 
         """ Computes relative trap frequencies for each species:
 
             w_i = sqrt(q_{i} m_0 / m_{i} q_0) omega 
@@ -638,82 +638,82 @@ def calc_mode_energies(res, Fock_cutoffs):
         energies_m2[k] = qt.expect(E2_op, state)
     return energies_m1, energies_m2
  
-def check_single_ion_case():
-    # for a single ion, the mode participation factors should just be the LD factors for each mode and direction.
-    # this is a good sanity check to make sure the mode participation factor calculation is correct.
-    wz = 2 * np.pi * .5e6  # axial trap frequency
-    wy = 2 * np.pi * 1.9e6  # radial trap frequency
-    wx = 2 * np.pi * 10e6  # radial trap frequency, something high to avoid any issues with mode ordering
-    mass_yb_amu = 170.936
-    k = 2 * np.pi / 355e-9
-    #atomic_nums = np.array([70]) 
-    Z = 70 
-    num_ions = 1
-
-    mode_analysis_one = TrappedIonModeAnalysis(num_ions, wx, wy, wz, np.ones(num_ions)*mass_yb_amu, Z) 
-    mode_analysis_one.solve_ion_trap_equilibrium()
-    mode_participation_factors_one = mode_analysis_one.calculate_mode_participation_factors()
-    simulated_LD_parameters = k * mode_participation_factors_one
-
-    # Compute analytical reference 
-    eta_x, eta_y, eta_z = mode_analysis_one.compute_reference_single_ion_lamb_dicke_factors(k) 
-    ld_factors_one_analytical = np.zeros((3, 1, 3), dtype = np.complex128)
-    ld_factors_one_analytical[0, 0, 2] = eta_x
-    ld_factors_one_analytical[1, 0, 1] = eta_y
-    ld_factors_one_analytical[2, 0, 0] = eta_z
-    print(f"\nLD factors: {ld_factors_one_analytical}")
-    #print("\nRatio of single ion LD factors from mode participation calculation to analytical calculation: \n", simulated_LD_parameters / ld_factors_one_analytical)
-    print(f"\nLD factors (via diagonalization):\n {simulated_LD_parameters}")
- 
-def check_two_ion_case():
-    wz = 2 * np.pi * .5e6  # axial trap frequency
-    wy = 2 * np.pi * 1.5e6  # radial trap frequency
-    wx = 2 * np.pi * 2e6  # radial trap frequency
-    wy_tilt = np.sqrt(wy**2 - wz**2)
-    wx_tilt = np.sqrt(wx**2 - wz**2)
-    mass_yb_amu = 170.936
-    k = 2 * np.pi / 355e-9
-    num_ions = 2
-    Z = 70
-    #mode_analysis_two= mode_analyzer(N =2, wz = wz, wy = wy, wx = wx, ionmass_amu= mass_yb_amu)
-    mode_analysis_two = TrappedIonModeAnalysis(num_ions, wx, wy, wz, mass_yb_amu, Z) 
-    mode_analysis_two.solve_ion_trap_equilibrium()
-    #mode_participation_factors_two= calculate_mode_participation_factors(mode_analysis_two)
-    mode_participation_factors_two = mode_analysis_two.calculate_mode_participation_factors()
-    ld_factors_two = k * mode_participation_factors_two
-
-    eta_x, eta_y, eta_z = mode_analysis_two.compute_reference_single_ion_lamb_dicke_factors(k) 
-    eta_COM_x = eta_x / np.sqrt(2)
-    eta_COM_y = eta_y / np.sqrt(2)
-    eta_COM_z = eta_z / np.sqrt(2)
-    eta_stretch_z = eta_z /np.sqrt(2) /3**(1/4)
-    # I don't know the analytical expressions for the tilt modes off the top of my head... let's assume its the square root of the normalized mode frequency
-    eta_tilt_x = eta_x / np.sqrt(2) / np.sqrt(wx_tilt / wx)
-    eta_tilt_y = eta_y / np.sqrt(2) / np.sqrt(wy_tilt / wy)
-    ld_factors_two_analytical = np.zeros((3, 2, 6), dtype = np.complex128)
-    ld_factors_two_analytical[0, 0, 5] = eta_COM_x
-    ld_factors_two_analytical[0, 1, 5] = eta_COM_x
-    ld_factors_two_analytical[0, 0, 4] = eta_tilt_x
-    ld_factors_two_analytical[0, 1, 4] = -eta_tilt_x
-    ld_factors_two_analytical[1, 0, 3] = eta_COM_y
-    ld_factors_two_analytical[1, 1, 3] = eta_COM_y
-    ld_factors_two_analytical[1, 0, 2] = eta_tilt_y
-    ld_factors_two_analytical[1, 1, 2] = -eta_tilt_y
-    ld_factors_two_analytical[2, 0, 0] = eta_COM_z
-    ld_factors_two_analytical[2, 1, 0] = eta_COM_z
-    ld_factors_two_analytical[2, 0, 1] = eta_stretch_z
-    ld_factors_two_analytical[2, 1, 1] = -eta_stretch_z
-    # these seem to work out!
-    print(f"\nLD factors: {ld_factors_two_analytical}")
-    #print("Ratio of two ion LD factors from mode participation calculation to analytical calculation: \n", ld_factors_two / ld_factors_two_analytical)
-    print(f"\nLD factors (via diagonalization):\n {ld_factors_two}")
-
- 
-
- 
-### Example usage  ###  
-if __name__ == '__main__':
-    # Test analysis: 
-    # TODO: convert to tests 
-    check_two_ion_case() 
-    #check_single_ion_case()
+ #def check_single_ion_case():
+ #    # for a single ion, the mode participation factors should just be the LD factors for each mode and direction.
+ #    # this is a good sanity check to make sure the mode participation factor calculation is correct.
+ #    wz = 2 * np.pi * .5e6  # axial trap frequency
+ #    wy = 2 * np.pi * 1.9e6  # radial trap frequency
+ #    wx = 2 * np.pi * 10e6  # radial trap frequency, something high to avoid any issues with mode ordering
+ #    mass_yb_amu = 170.936
+ #    k = 2 * np.pi / 355e-9
+ #    #atomic_nums = np.array([70]) 
+ #    Z = 70 
+ #    num_ions = 1
+ #
+ #    mode_analysis_one = TrappedIonModeAnalysis(num_ions, wx, wy, wz, np.ones(num_ions)*mass_yb_amu, Z) 
+ #    mode_analysis_one.solve_ion_trap_equilibrium()
+ #    mode_participation_factors_one = mode_analysis_one.calculate_mode_participation_factors()
+ #    simulated_LD_parameters = k * mode_participation_factors_one
+ #
+ #    # Compute analytical reference 
+ #    eta_x, eta_y, eta_z = mode_analysis_one.compute_reference_single_ion_lamb_dicke_factors(k) 
+ #    ld_factors_one_analytical = np.zeros((3, 1, 3), dtype = np.complex128)
+ #    ld_factors_one_analytical[0, 0, 2] = eta_x
+ #    ld_factors_one_analytical[1, 0, 1] = eta_y
+ #    ld_factors_one_analytical[2, 0, 0] = eta_z
+ #    print(f"\nLD factors: {ld_factors_one_analytical}")
+ #    #print("\nRatio of single ion LD factors from mode participation calculation to analytical calculation: \n", simulated_LD_parameters / ld_factors_one_analytical)
+ #    print(f"\nLD factors (via diagonalization):\n {simulated_LD_parameters}")
+ # 
+ #def check_two_ion_case():
+ #    wz = 2 * np.pi * .5e6  # axial trap frequency
+ #    wy = 2 * np.pi * 1.5e6  # radial trap frequency
+ #    wx = 2 * np.pi * 2e6  # radial trap frequency
+ #    wy_tilt = np.sqrt(wy**2 - wz**2)
+ #    wx_tilt = np.sqrt(wx**2 - wz**2)
+ #    mass_yb_amu = 170.936
+ #    k = 2 * np.pi / 355e-9
+ #    num_ions = 2
+ #    Z = 70
+ #    #mode_analysis_two= mode_analyzer(N =2, wz = wz, wy = wy, wx = wx, ionmass_amu= mass_yb_amu)
+ #    mode_analysis_two = TrappedIonModeAnalysis(num_ions, wx, wy, wz, mass_yb_amu, Z) 
+ #    mode_analysis_two.solve_ion_trap_equilibrium()
+ #    #mode_participation_factors_two= calculate_mode_participation_factors(mode_analysis_two)
+ #    mode_participation_factors_two = mode_analysis_two.calculate_mode_participation_factors()
+ #    ld_factors_two = k * mode_participation_factors_two
+ #
+ #    eta_x, eta_y, eta_z = mode_analysis_two.compute_reference_single_ion_lamb_dicke_factors(k) 
+ #    eta_COM_x = eta_x / np.sqrt(2)
+ #    eta_COM_y = eta_y / np.sqrt(2)
+ #    eta_COM_z = eta_z / np.sqrt(2)
+ #    eta_stretch_z = eta_z /np.sqrt(2) /3**(1/4)
+ #    # I don't know the analytical expressions for the tilt modes off the top of my head... let's assume its the square root of the normalized mode frequency
+ #    eta_tilt_x = eta_x / np.sqrt(2) / np.sqrt(wx_tilt / wx)
+ #    eta_tilt_y = eta_y / np.sqrt(2) / np.sqrt(wy_tilt / wy)
+ #    ld_factors_two_analytical = np.zeros((3, 2, 6), dtype = np.complex128)
+ #    ld_factors_two_analytical[0, 0, 5] = eta_COM_x
+ #    ld_factors_two_analytical[0, 1, 5] = eta_COM_x
+ #    ld_factors_two_analytical[0, 0, 4] = eta_tilt_x
+ #    ld_factors_two_analytical[0, 1, 4] = -eta_tilt_x
+ #    ld_factors_two_analytical[1, 0, 3] = eta_COM_y
+ #    ld_factors_two_analytical[1, 1, 3] = eta_COM_y
+ #    ld_factors_two_analytical[1, 0, 2] = eta_tilt_y
+ #    ld_factors_two_analytical[1, 1, 2] = -eta_tilt_y
+ #    ld_factors_two_analytical[2, 0, 0] = eta_COM_z
+ #    ld_factors_two_analytical[2, 1, 0] = eta_COM_z
+ #    ld_factors_two_analytical[2, 0, 1] = eta_stretch_z
+ #    ld_factors_two_analytical[2, 1, 1] = -eta_stretch_z
+ #    # these seem to work out!
+ #    print(f"\nLD factors: {ld_factors_two_analytical}")
+ #    #print("Ratio of two ion LD factors from mode participation calculation to analytical calculation: \n", ld_factors_two / ld_factors_two_analytical)
+ #    print(f"\nLD factors (via diagonalization):\n {ld_factors_two}")
+ #
+ # 
+ #
+ # 
+ #### Example usage  ###  
+ #if __name__ == '__main__':
+ #    # Test analysis: 
+ #    # TODO: convert to tests 
+ #    check_two_ion_case() 
+ #    #check_single_ion_case()

From 93abb734ba6371ddaeb8a2294e523e485c81c529 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 21:57:18 -0600
Subject: [PATCH 13/31]  add custom type Vector,Matrix import to class

---
 src/ionsim/trapped_ion_mode_analysis.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 976067c..1af1948 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -3,6 +3,7 @@
 import scipy.constants as const
 import warnings
 import scipy.optimize as opt    
+from ionsim.custom_types import Matrix, Vector
 
 ########## List of substantive changes by Ethan: 
 ## 1. Made dimensionless variable function take in inputs so the user can choose the charge, mass, and trap freq scales to use.   

From bd829462b4263259faf0234ba2f6ce5c80dfb013 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 22:00:54 -0600
Subject: [PATCH 14/31]   fix tuple type-hinting

---
 src/ionsim/trapped_ion_mode_analysis.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 1af1948..6bf536c 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -257,7 +257,7 @@ def calculate_normal_modes(self, H_matrix: Matrix) -> (Vector, list[Vector]):
         return eigvals, eigvecs 
 
 
-    def ion_coordinates_from_flattened(self, flattened_coordinate_vector: Vector) -> tuple(Vector, Vector, Vector):
+    def ion_coordinates_from_flattened(self, flattened_coordinate_vector: Vector) -> tuple[Vector, Vector, Vector]:
         """ Returns arrays corresponding to an ion's x, y, and z coordinates, respectively, from a flattened vector u:
                 x = u[0:N], y = u[N:2N], z[2N:]
 
@@ -287,7 +287,6 @@ def reindex_ions(self):
 
     def solve_for_equilibrium_positions(self, positions_guess: Vector | None=None):
         """ Solves for equilibrium position vector: u, which represents a flattened spatial grid. """
-
         # TODO: option for an initial guess choice? 
         # Set the initial guess for the solver  
         if positions_guess == None:
@@ -345,7 +344,7 @@ def force_trap(self, positions: Vector):
         force_trap = np.hstack((Ftrapx, Ftrapy, Ftrapz))
         return force_trap
 
-    def force_coulomb(self, positions): 
+    def force_coulomb(self, positions: Vector): 
         x,y,z = self.ion_coordinates_from_flattened(positions)
 
         dx = x[:, np.newaxis] - x
@@ -369,7 +368,7 @@ def force_coulomb(self, positions):
         force_coulomb *= 0.5    
         return force_coulomb
 
-    def force(self, positions):
+    def force(self, positions: Vector):
         Force = self.force_coulomb(positions) + self.force_trap(positions)  
         return Force
 

From 287a7ac6873fef02d0f0713409a3f3499636e5e9 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 22:06:06 -0600
Subject: [PATCH 15/31]   rm code -> moved to tests

---
 src/ionsim/trapped_ion_mode_analysis.py | 79 -------------------------
 1 file changed, 79 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 6bf536c..07438eb 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -638,82 +638,3 @@ def calc_mode_energies(res, Fock_cutoffs):
         energies_m2[k] = qt.expect(E2_op, state)
     return energies_m1, energies_m2
  
- #def check_single_ion_case():
- #    # for a single ion, the mode participation factors should just be the LD factors for each mode and direction.
- #    # this is a good sanity check to make sure the mode participation factor calculation is correct.
- #    wz = 2 * np.pi * .5e6  # axial trap frequency
- #    wy = 2 * np.pi * 1.9e6  # radial trap frequency
- #    wx = 2 * np.pi * 10e6  # radial trap frequency, something high to avoid any issues with mode ordering
- #    mass_yb_amu = 170.936
- #    k = 2 * np.pi / 355e-9
- #    #atomic_nums = np.array([70]) 
- #    Z = 70 
- #    num_ions = 1
- #
- #    mode_analysis_one = TrappedIonModeAnalysis(num_ions, wx, wy, wz, np.ones(num_ions)*mass_yb_amu, Z) 
- #    mode_analysis_one.solve_ion_trap_equilibrium()
- #    mode_participation_factors_one = mode_analysis_one.calculate_mode_participation_factors()
- #    simulated_LD_parameters = k * mode_participation_factors_one
- #
- #    # Compute analytical reference 
- #    eta_x, eta_y, eta_z = mode_analysis_one.compute_reference_single_ion_lamb_dicke_factors(k) 
- #    ld_factors_one_analytical = np.zeros((3, 1, 3), dtype = np.complex128)
- #    ld_factors_one_analytical[0, 0, 2] = eta_x
- #    ld_factors_one_analytical[1, 0, 1] = eta_y
- #    ld_factors_one_analytical[2, 0, 0] = eta_z
- #    print(f"\nLD factors: {ld_factors_one_analytical}")
- #    #print("\nRatio of single ion LD factors from mode participation calculation to analytical calculation: \n", simulated_LD_parameters / ld_factors_one_analytical)
- #    print(f"\nLD factors (via diagonalization):\n {simulated_LD_parameters}")
- # 
- #def check_two_ion_case():
- #    wz = 2 * np.pi * .5e6  # axial trap frequency
- #    wy = 2 * np.pi * 1.5e6  # radial trap frequency
- #    wx = 2 * np.pi * 2e6  # radial trap frequency
- #    wy_tilt = np.sqrt(wy**2 - wz**2)
- #    wx_tilt = np.sqrt(wx**2 - wz**2)
- #    mass_yb_amu = 170.936
- #    k = 2 * np.pi / 355e-9
- #    num_ions = 2
- #    Z = 70
- #    #mode_analysis_two= mode_analyzer(N =2, wz = wz, wy = wy, wx = wx, ionmass_amu= mass_yb_amu)
- #    mode_analysis_two = TrappedIonModeAnalysis(num_ions, wx, wy, wz, mass_yb_amu, Z) 
- #    mode_analysis_two.solve_ion_trap_equilibrium()
- #    #mode_participation_factors_two= calculate_mode_participation_factors(mode_analysis_two)
- #    mode_participation_factors_two = mode_analysis_two.calculate_mode_participation_factors()
- #    ld_factors_two = k * mode_participation_factors_two
- #
- #    eta_x, eta_y, eta_z = mode_analysis_two.compute_reference_single_ion_lamb_dicke_factors(k) 
- #    eta_COM_x = eta_x / np.sqrt(2)
- #    eta_COM_y = eta_y / np.sqrt(2)
- #    eta_COM_z = eta_z / np.sqrt(2)
- #    eta_stretch_z = eta_z /np.sqrt(2) /3**(1/4)
- #    # I don't know the analytical expressions for the tilt modes off the top of my head... let's assume its the square root of the normalized mode frequency
- #    eta_tilt_x = eta_x / np.sqrt(2) / np.sqrt(wx_tilt / wx)
- #    eta_tilt_y = eta_y / np.sqrt(2) / np.sqrt(wy_tilt / wy)
- #    ld_factors_two_analytical = np.zeros((3, 2, 6), dtype = np.complex128)
- #    ld_factors_two_analytical[0, 0, 5] = eta_COM_x
- #    ld_factors_two_analytical[0, 1, 5] = eta_COM_x
- #    ld_factors_two_analytical[0, 0, 4] = eta_tilt_x
- #    ld_factors_two_analytical[0, 1, 4] = -eta_tilt_x
- #    ld_factors_two_analytical[1, 0, 3] = eta_COM_y
- #    ld_factors_two_analytical[1, 1, 3] = eta_COM_y
- #    ld_factors_two_analytical[1, 0, 2] = eta_tilt_y
- #    ld_factors_two_analytical[1, 1, 2] = -eta_tilt_y
- #    ld_factors_two_analytical[2, 0, 0] = eta_COM_z
- #    ld_factors_two_analytical[2, 1, 0] = eta_COM_z
- #    ld_factors_two_analytical[2, 0, 1] = eta_stretch_z
- #    ld_factors_two_analytical[2, 1, 1] = -eta_stretch_z
- #    # these seem to work out!
- #    print(f"\nLD factors: {ld_factors_two_analytical}")
- #    #print("Ratio of two ion LD factors from mode participation calculation to analytical calculation: \n", ld_factors_two / ld_factors_two_analytical)
- #    print(f"\nLD factors (via diagonalization):\n {ld_factors_two}")
- #
- # 
- #
- # 
- #### Example usage  ###  
- #if __name__ == '__main__':
- #    # Test analysis: 
- #    # TODO: convert to tests 
- #    check_two_ion_case() 
- #    #check_single_ion_case()

From 87314346f5e646e2e516bfaed5ca3eb38af3700c Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 22:22:40 -0600
Subject: [PATCH 16/31]   vectorize Coulomb force

---
 src/ionsim/trapped_ion_mode_analysis.py | 39 ++++++++++++++-----------
 1 file changed, 22 insertions(+), 17 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 07438eb..02affd9 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -5,19 +5,19 @@
 import scipy.optimize as opt    
 from ionsim.custom_types import Matrix, Vector
 
-########## List of substantive changes by Ethan: 
+########## List of substantive changes: 
 ## 1. Made dimensionless variable function take in inputs so the user can choose the charge, mass, and trap freq scales to use.   
 ## 2. Added a helper function to pack/unpack the ion coordinates to reduce code. 
 
-########## List of non-substantive changes by Ethan: 
-## 1. Variable changes for readable (e.g. omega_x instead of wx) 
+########## List of non-substantive changes: 
+## 1. Variable changes for readability (e.g. omega_x instead of wx) 
 ## 2. Type-hinting in functions  
 ## 3. Consolidated the functions for solving for ion positions equilibria 
 ## 4. Moved eigenvector helper functions into the class. 
 ## 5. Reduced number of matrices that we store as class attributes: Need to check whether we want to be storing those. 
 
 ########## Questions: 
-# 1. Is there a need for the "has run" boolean? 
+# 1. Is there a need for the "has run" boolean? e.g. is there a time where it stays False?  
 # 2. Should the dimensionless parameters have a naming convention, e.g. omega_x --> omega_x_ND 
 # 3. What sets the phases of the Lamb-Dicke parameters?  
 # 4. See TODO's 
@@ -345,30 +345,35 @@ def force_trap(self, positions: Vector):
         return force_trap
 
     def force_coulomb(self, positions: Vector): 
+        """ Computes the Coulomb force by derivative w.r.t. ion coordinates """
         x,y,z = self.ion_coordinates_from_flattened(positions)
+        r = [x,y,z]
+        dr = []
 
-        dx = x[:, np.newaxis] - x
-        dy = y[:, np.newaxis] - y
-        dz = z[:, np.newaxis] - z
-        rsep = np.sqrt(dx**2 + dy**2 + dz**2).astype(np.float64)
+        for coord in r:
+            dr.append(coord[:, np.newaxis] - coord)
+        #rsep = np.sqrt(np.sum( [dx**2 for dx in dr])).astype(float)
+        rsep = np.sqrt(dr[0]**2 + dr[1]**2 + dr[2]**2).astype(np.float64)
         qq = (self.q * self.q[:, np.newaxis]).astype(np.float64)    
 
         with np.errstate(divide='ignore', invalid='ignore'):
             rsep3 = np.where(rsep != 0., rsep ** (-3), 0)
 
-        fx = dx * rsep3 * qq
-        fy = dy * rsep3 * qq
-        fz = dz * rsep3 * qq    
-
-        Fx = -np.sum(fx, axis=1)
-        Fy = -np.sum(fy, axis=1)
-        Fz = -np.sum(fz, axis=1)
+        F_r = [] # Force in each direction 
+        for coord in dr:
+            f_r = coord * rsep3 * qq
+            F_r.append(-np.sum(np.array(f_r), axis=1))
 
-        force_coulomb = np.hstack((Fx, Fy, Fz))
+        force_coulomb = np.hstack(tuple(F_r))
         force_coulomb *= 0.5    
         return force_coulomb
 
-    def force(self, positions: Vector):
+    def force(self, positions: Vector) -> Vector:
+        """ Computes the force vector, e.g. derivative of the potential energy w.r.t. ion position: dU/dx, dU/dy, dU/dz. 
+
+            - returns a one-dimensional length-3N array characterizing forces w.r.t. ion coordinates for N ions. 
+        
+        """
         Force = self.force_coulomb(positions) + self.force_trap(positions)  
         return Force
 

From 55a2a7464ddd2db7bca29bb207c7497f517d2f6c Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 22:30:48 -0600
Subject: [PATCH 17/31]  add loops for force & potential calcs

---
 src/ionsim/trapped_ion_mode_analysis.py | 37 ++++++++++++-------------
 1 file changed, 18 insertions(+), 19 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 02affd9..e425d84 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -126,6 +126,7 @@ def convert_parameters_to_dimensionless(self, charge_scale: float=1., mass_scale
         self.wz = self.omega_z / self.trap_freq_scale 
         self.wy = self.omega_y / self.trap_freq_scale 
         self.wx = self.omega_x / self.trap_freq_scale  
+        self.trap_frequencies_ND = [self.wx, self.wy, self.wz]
 
         # TODO: Is hbar set to 1? Reconcile hbar somewhere  
         # Compute characteristic length, time, velocity, and energy scales and store in a dictionary  
@@ -309,22 +310,21 @@ def solve_for_equilibrium_positions(self, positions_guess: Vector | None=None):
 
     def potential_trap(self, positions: Vector) -> Vector:
         """ Computes trapping potential, takes in a flattened (1D) array of all position coordinates """
-        x,y,z = self.ion_coordinates_from_flattened(positions)
-        V_trap = 0.5 * np.sum((self.m * self.wx ** 2) * x ** 2) + \
-            0.5 * np.sum((self.m * self.wy ** 2) * y ** 2) + \
-                0.5 * np.sum((self.m * self.wz ** 2) * z ** 2)
+        r = self.ion_coordinates_from_flattened(positions)
+        V_trap = 0.
+        for coord, omega in zip(r, self.trap_frequencies_ND): 
+            V_trap += 0.5 * np.sum((self.m * omega**2) * coord**2)
         return V_trap
     
     def potential_coulomb(self, positions: Vector) -> Vector:
         """ Computes Coulomb potential, takes in a flattened (1D) array of all position coordinates """
-        x,y,z = self.ion_coordinates_from_flattened(positions)
+        r = self.ion_coordinates_from_flattened(positions)
+        dr = []
+        for coord in r:
+            dr.append(coord[:, np.newaxis] - coord)
 
-        dx = x[:, np.newaxis] - x
-        dy = y[:, np.newaxis] - y
-        dz = z[:, np.newaxis] - z
-        rsep = np.sqrt(dx ** 2 + dy ** 2 + dz ** 2).astype(np.float64)
+        rsep = np.sqrt(dr[0]**2 + dr[1]**2 + dr[2]**2).astype(np.float64)
         qq = (self.q * self.q[:, np.newaxis]).astype(np.float64)    
-
         with np.errstate(divide='ignore'):
             V_Coulomb = np.sum( np.where(rsep != 0., qq / rsep, 0) ) / 2 # divide by 2 to avoid double counting
         V_Coulomb *= .5 
@@ -335,13 +335,13 @@ def potential_energy(self, positions: Vector) -> Vector:
         return self.potential_trap(positions) + self.potential_coulomb(positions)   
 
     def force_trap(self, positions: Vector):
-        x,y,z = self.ion_coordinates_from_flattened(positions)
+        r = self.ion_coordinates_from_flattened(positions)
 
-        Ftrapx = self.m * self.wx**2 * x
-        Ftrapy = self.m * self.wy**2 * y
-        Ftrapz = self.m * self.wz**2 * z
+        F_r = []
+        for coord, omega in zip(r, self.trap_frequencies_ND): 
+            F_r.append(self.m * omega**2 * coord)
 
-        force_trap = np.hstack((Ftrapx, Ftrapy, Ftrapz))
+        force_trap = np.hstack(F_r)
         return force_trap
 
     def force_coulomb(self, positions: Vector): 
@@ -352,7 +352,6 @@ def force_coulomb(self, positions: Vector):
 
         for coord in r:
             dr.append(coord[:, np.newaxis] - coord)
-        #rsep = np.sqrt(np.sum( [dx**2 for dx in dr])).astype(float)
         rsep = np.sqrt(dr[0]**2 + dr[1]**2 + dr[2]**2).astype(np.float64)
         qq = (self.q * self.q[:, np.newaxis]).astype(np.float64)    
 
@@ -360,11 +359,11 @@ def force_coulomb(self, positions: Vector):
             rsep3 = np.where(rsep != 0., rsep ** (-3), 0)
 
         F_r = [] # Force in each direction 
-        for coord in dr:
-            f_r = coord * rsep3 * qq
+        for dx in dr:
+            f_r = dx * rsep3 * qq
             F_r.append(-np.sum(np.array(f_r), axis=1))
 
-        force_coulomb = np.hstack(tuple(F_r))
+        force_coulomb = np.hstack(F_r)
         force_coulomb *= 0.5    
         return force_coulomb
 

From dbe60cd6728c9d1c9e2f6f964bb67cc10377e619 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 21 Apr 2026 22:44:16 -0600
Subject: [PATCH 18/31]  convert many fxn's to loops over dimensions

---
 src/ionsim/trapped_ion_mode_analysis.py | 52 +++++++++++--------------
 1 file changed, 23 insertions(+), 29 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index e425d84..8ea8805 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -157,7 +157,6 @@ def get_canonical_transformation(self, H, eigenvectors, eigenvalues: Vector | No
         # the sign is likely due to the convention of writing the time evolution as exp(-iwt)
         return T
 
-
     def check_for_zero_modes(self):
         assert np.all(self.eigvals > 0), "All eigenvalues must be positive"   
 
@@ -346,8 +345,7 @@ def force_trap(self, positions: Vector):
 
     def force_coulomb(self, positions: Vector): 
         """ Computes the Coulomb force by derivative w.r.t. ion coordinates """
-        x,y,z = self.ion_coordinates_from_flattened(positions)
-        r = [x,y,z]
+        r = self.ion_coordinates_from_flattened(positions)
         dr = []
 
         for coord in r:
@@ -376,45 +374,43 @@ def force(self, positions: Vector) -> Vector:
         Force = self.force_coulomb(positions) + self.force_trap(positions)  
         return Force
 
-
-    #### Matrix methods 
     def hessian_trap(self, positions: Vector) -> Matrix:
         """ Computes the Hessian of the trap  """ 
-        Hxx = np.diag(self.m * (self.wx**2) * np.ones(self.num_ions))
-        Hyy = np.diag(self.m * (self.wy**2) * np.ones(self.num_ions))  
-        Hzz = np.diag(self.m * (self.wz**2) * np.ones(self.num_ions))  
+        H_rr = []
+        for i, omega in enumerate(self.trap_frequencies_ND): 
+            H_rr.append(np.diag(self.m * omega**2 * np.ones(self.num_ions)))
         zeros = np.zeros((self.num_ions, self.num_ions))  
-        H = np.block([[Hxx, zeros, zeros], [zeros, Hyy, zeros], [zeros, zeros, Hzz]])
+        H = np.block([[H_rr[0], zeros, zeros], [zeros, H_rr[1], zeros], [zeros, zeros, H_rr[2]]])
         return H
 
     def hessian_coulomb(self, positions: Vector) -> Matrix:
         """ Computes the Hessian of the Coulomb interaction """ 
-        x,y,z = self.ion_coordinates_from_flattened(positions)
-        
-        dx = x[:, np.newaxis] - x
-        dy = y[:, np.newaxis] - y
-        dz = z[:, np.newaxis] - z
-        rsep = np.sqrt(dx ** 2 + dy ** 2 + dz ** 2).astype(np.float64)  
+        r = self.ion_coordinates_from_flattened(positions)
+        dr = []
+
+        for coord in r:
+            dr.append(coord[:, np.newaxis] - coord)
+        rsep = np.sqrt(dr[0]**2 + dr[1]**2 + dr[2]**2).astype(np.float64)
         qq = (self.q * self.q[:, np.newaxis]).astype(np.float64)    
 
         with np.errstate(divide='ignore'):
             rsep5 = np.where(rsep != 0., rsep ** (-5), 0)
 
-        dxsq = dx ** 2
-        dysq = dy ** 2
-        dzsq = dz ** 2
+        dr_sq = np.array(dr)**2
         rsep2 = rsep ** 2
 
         # X derivatives, Y derivatives for alpha != beta
-        Hxx = (rsep2 - 3 * dxsq) * rsep5
-        Hyy = (rsep2 - 3 * dysq) * rsep5
-        Hzz = (rsep2 - 3 * dzsq) * rsep5
+        H_rr = (rsep2 - 3 * dr_sq) * rsep5 # form: [Hxx, Hyy, Hzz] 
 
         # Above, for alpha == beta
-        Hxx[np.diag_indices(self.num_ions)] = -np.sum(Hxx, axis=0)
-        Hyy[np.diag_indices(self.num_ions)] = -np.sum(Hyy, axis=0)
-        Hzz[np.diag_indices(self.num_ions)] = -np.sum(Hzz, axis=0)
-
+        # Compute diagonals: 
+        for i in range(3):
+            H_rr[i][np.diag_indices(self.num_ions)] = -np.sum(H_rr[i], axis=0)
+        dx = dr[0]
+        dy = dr[1]
+        dz = dr[2]
+
+        # Off-diagonal elements:
         Hxy = -3 * dx * dy * rsep5
         Hxy[np.diag_indices(self.num_ions)] = 3 * np.sum(dx * dy * rsep5, axis=0)
         Hxz = -3 * dx * dz * rsep5
@@ -422,14 +418,12 @@ def hessian_coulomb(self, positions: Vector) -> Matrix:
         Hyz = -3 * dy * dz * rsep5
         Hyz[np.diag_indices(self.num_ions)] = 3 * np.sum(dy * dz * rsep5, axis=0)
         
-        Hxx *= qq
-        Hyy *= qq
-        Hzz *= qq
+        H_rr *= qq
         Hxy *= qq
         Hxz *= qq
         Hyz *= qq
 
-        H_coulomb = np.block([[Hxx, Hxy, Hxz], [Hxy, Hyy, Hyz], [Hxz, Hyz, Hzz]])
+        H_coulomb = np.block([[H_rr[0], Hxy, Hxz], [Hxy, H_rr[1], Hyz], [Hxz, Hyz, H_rr[2]]])
         H_coulomb /= 2
         return H_coulomb
 

From 0cdfb57df592c8386fae3b8f28b20f6001521b81 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Wed, 22 Apr 2026 08:37:45 -0600
Subject: [PATCH 19/31]  .

---
 src/ionsim/trapped_ion_mode_analysis.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 8ea8805..a212c25 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -137,7 +137,7 @@ def convert_parameters_to_dimensionless(self, charge_scale: float=1., mass_scale
         self.characteristic_parameters['energy'] = 0.5 * mass_scale * self.characteristic_parameters['velocity'] ** 2  # characteristic energy  
 
 
-    def get_norm(self, eigenvector: Vector, H: Matrix) -> float:
+    def get_eigenvector_norm(self, eigenvector: Vector, H: Matrix) -> float:
         """ Computes the norm of the eigenvector en w.r.t. the Hamiltonian H. """
         norm = np.sqrt(eigenvector.T.conj() @ H @ eigenvector)
         return norm
@@ -167,7 +167,7 @@ def check_outer_relation(self, H: Matrix):
         en = self.normalize_eigenvectors(en,H)
         Outers = np.zeros((6*self.num_ions,6*self.num_ions),dtype=complex)
         for i in range(6*self.num_ions):
-            norm = self.get_norm(en[:,i],H)
+            norm = self.get_eigenvector_norm(en[:,i],H)
             Outers = Outers + np.outer(en[:,i],en[:,i].conj())/norm 
         I_right = H @ Outers
         I_left  = Outers @ H
@@ -194,7 +194,9 @@ def check_diagonalization(self, T: Matrix, S: Matrix, E: Matrix) -> bool:
 
     #def run(self):
     def solve_ion_trap_equilibrium(self): 
+        """ Diagonalizes the Coulomb + harmonic trap Hamiltonian for a system of ions. """
         # Convert to dimensionless units using axial trap frequency and first ion's mass and charge  
+        # TODO: take in input here or in class constructor for mass, charge, trap scales. 
         self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
         
         self.u = self.solve_for_equilibrium_positions()
@@ -227,7 +229,7 @@ def normalize_eigenvectors(self, eigvecs, H: Matrix, eigenvalues: Vector | None=
             eigenvalues = np.ones(num_eigenvalues)
         for i in range(num_eigenvalues):
             en = eigvecs[:,i].reshape(num_coords,1)
-            norm = self.get_norm(en,H)
+            norm = self.get_eigenvector_norm(en,H)
             eigvecs_rescaled[:,i] = en[:,0]/norm
             eigvecs_rescaled[:,i] *= np.sqrt(eigenvalues[i])
         return eigvecs_rescaled 
@@ -584,6 +586,7 @@ def reindex_ions_by_z(self):
  
     #def run(self):
     def solve_ion_trap_equilibrium(self): 
+        """ Diagonalizes the Coulomb + harmonic trap Hamiltonian for a system of ions. """
         self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
         
         self.u = self.solve_for_equilibrium_positions()

From 52fa9645af142ba5d99eb105d3662caf71711dd1 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Wed, 22 Apr 2026 10:30:13 -0600
Subject: [PATCH 20/31]   add documentation for Lamb-Dicke parameter matrix
 form

---
 src/ionsim/trapped_ion_mode_analysis.py | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index a212c25..96c14f5 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -20,7 +20,9 @@
 # 1. Is there a need for the "has run" boolean? e.g. is there a time where it stays False?  
 # 2. Should the dimensionless parameters have a naming convention, e.g. omega_x --> omega_x_ND 
 # 3. What sets the phases of the Lamb-Dicke parameters?  
-# 4. See TODO's 
+# 4. Does the LD calculation properly do k dot r?
+# 5. Can we vectorize this to convert loops into vector operations?  
+# 6. See TODO's 
 
 
 def characteristic_length(q: float, mass: float, omega: float) -> float:
@@ -492,8 +494,18 @@ def organize_modes(self, eigvals, eigvecs):
         return eigvals, eigvecs 
 
 
-    def calculate_mode_participation_factors(self):
-        """ Computes ion-mode participation factors, related to the Lamb-Dicke parameters """ 
+    def calculate_mode_participation_factors(self) -> Matrix:
+        """ Computes ion-mode participation factors, related to the Lamb-Dicke parameters.
+
+            Mode participation factors (eta) take the following form:
+             
+            shape of eta matrix: (dimension, ion, mode) 
+            e.g. eta[1, 2, 3] is eta in the "y" direction, ion 1, and the 2nd mode. 
+
+            For N ions, this form organizes the 3N modes into N modes per direction "d", where d = x, y, z. 
+
+        """ 
+
         eigvecs = self.eigvecs
         num_coords, num_modes = np.shape(eigvecs)
         num_ions = num_modes // 3

From 3b710dfebc5f7cb401d32c9072f19f94c584e6f9 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Wed, 22 Apr 2026 12:22:48 -0600
Subject: [PATCH 21/31]  + draft class methods for from_species,
 from_spin_basis construction

---
 src/ionsim/degree_of_freedom.py         |  9 ++-
 src/ionsim/trapped_ion_mode_analysis.py | 96 +++++++++++++++++--------
 tests/test_mode_analysis.py             |  4 +-
 3 files changed, 77 insertions(+), 32 deletions(-)

diff --git a/src/ionsim/degree_of_freedom.py b/src/ionsim/degree_of_freedom.py
index f8633ee..7fad0a9 100644
--- a/src/ionsim/degree_of_freedom.py
+++ b/src/ionsim/degree_of_freedom.py
@@ -19,12 +19,16 @@
 class DegreeOfFreedom(ABC):
     """A degree of freedom in a basis of states."""
     energy_levels: Sequence[EnergyLevel]
-    name: str | None = None # TODO: will we use these names?
+    #name: str | None = None # TODO: will we use these names?
 
 @dataclass(frozen=True, eq=False)
 class AtomicSpin(DegreeOfFreedom):
     """An atomic spin degree of freedom, i.e. its coupled spin and orbital angular momentum."""
     energy_levels: list[AtomicInternalEnergyLevel]
+    # Store mass and atomic number 
+    atomic_mass: float 
+    atomic_number: int # should we allowe for effective Z's which are non-integer?) 
+    name: str | None = None # TODO: will we use these names?
 
     @classmethod
     def from_species(cls, species: str, term_symbols: list[str] | None = None, level_names: list[str] | None = None,
@@ -102,7 +106,7 @@ def from_species(cls, species: str, term_symbols: list[str] | None = None, level
                         level = HyperfineLevel(**fine_data, i=nuclear_spin, f=f, mf=mf, external_energy_shift = zeeman_shift_energy)
                         if level_names is None or level.name in level_names:
                             levels.append(level)
-        return cls(levels, name)
+        return cls(levels, mass, z, name)
 
     @classmethod
     def get_level_factory(cls, coupling_scheme: str):
@@ -185,6 +189,7 @@ def check_uniqueness_of_term_symbol(term_symbol: str, levels_data: list[dict]):
 class MotionalMode(DegreeOfFreedom):
     """An normal mode of motion for a linear chain of ions."""
     energy_levels: list[CollectiveMotionalEnergyLevel]
+    name: str | None = None # TODO: will we use these names?
 
     @classmethod
     def from_frequency(cls, frequency: float, fock_dimension: int, name: str | None = None):
diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 96c14f5..fd80c76 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -4,6 +4,8 @@
 import warnings
 import scipy.optimize as opt    
 from ionsim.custom_types import Matrix, Vector
+from ionsim.degree_of_freedom import AtomicSpin
+from ionsim.basis import Basis 
 
 ########## List of substantive changes: 
 ## 1. Made dimensionless variable function take in inputs so the user can choose the charge, mass, and trap freq scales to use.   
@@ -17,7 +19,9 @@
 ## 5. Reduced number of matrices that we store as class attributes: Need to check whether we want to be storing those. 
 
 ########## Questions: 
-# 1. Is there a need for the "has run" boolean? e.g. is there a time where it stays False?  
+# -1. Should the branch sorting be the default option? This seems like what we would want to do for our problems  
+# 0. LD parameter matrix shape? 
+# 1. Is there a need for the "has run" boolean? e.g. is there a time where it stays False? Ah I think it refers to whether the eqb solver has been successfully executed. 
 # 2. Should the dimensionless parameters have a naming convention, e.g. omega_x --> omega_x_ND 
 # 3. What sets the phases of the Lamb-Dicke parameters?  
 # 4. Does the LD calculation properly do k dot r?
@@ -67,17 +71,14 @@ def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float
         if not self.trap_is_stable():
             raise IonSimError(f"Trap is unstable from negative trap frequencies.")
 
-        #self.initial_equilibrium_guess = None
-        #self.hasrun = False 
-    
-
     def calculate_species_trap_frequencies(self, omega: float) -> Vector: 
         """ Computes relative trap frequencies for each species:
 
-            w_i = sqrt(q_{i} m_0 / m_{i} q_0) omega 
+            w_i = sqrt(q_{i} m_0 / m_{i} q_0) omega
+
+            from omega_(secular, i) = sqrt(q_i V_0 / (2 m_i d^2) ) for the i'th ion. 
 
         """ 
-        # TODO: Does Wes have a ref. for this? Why is omega species-dependent?  
         # assume that the trapping frequency given corresponds to the first ion species 
         q0 = self.nuclear_charges[0] # charge of first ion 
         m0 = self.atomic_masses[0] # charge of first ion 
@@ -193,20 +194,20 @@ def check_diagonalization(self, T: Matrix, S: Matrix, E: Matrix) -> bool:
             warnings.warn("Diagnolization check failed")    
             print("has duplicate eigenvalues: ", has_duplicate_eigvals)
 
-
     #def run(self):
+
     def solve_ion_trap_equilibrium(self): 
         """ Diagonalizes the Coulomb + harmonic trap Hamiltonian for a system of ions. """
         # Convert to dimensionless units using axial trap frequency and first ion's mass and charge  
         # TODO: take in input here or in class constructor for mass, charge, trap scales. 
         self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
         
-        self.u = self.solve_for_equilibrium_positions()
+        self.equilibrium_positions = self.solve_for_equilibrium_positions()
         self.reindex_ions()
 
-        # Compute matrices 
+        # Compute helper matrices for computing normal mode properties 
         mass_matrix = self.build_mass_matrix(self.m)  
-        E_matrix = self.build_E_matrix(self.u, mass_matrix)  
+        E_matrix = self.build_E_matrix(self.equilibrium_positions, mass_matrix)  
         T_matrix = self.build_momentum_transform_matrix(mass_matrix)
         H_matrix = self.compute_H_matrix(T_matrix, E_matrix)   
 
@@ -219,12 +220,10 @@ def solve_ion_trap_equilibrium(self):
         self.check_outer_relation(H_matrix)
         self.check_diagonalization(T_matrix, S_matrix, E_matrix)
         #self.hasrun = True  
-    
+
 
     def normalize_eigenvectors(self, eigvecs, H: Matrix, eigenvalues: Vector | None=None): 
-        """
-        Rescale the eigenvectors ens w.r.t. the Hamiltonian H.
-        """
+        """ Rescale the eigenvectors ens w.r.t. the Hamiltonian H. """
         eigvecs_rescaled = np.zeros_like(eigvecs,dtype=complex)
         num_coords, num_eigenvalues = np.shape(eigvecs)
         if eigenvalues is None:
@@ -276,12 +275,12 @@ def ion_coordinates_from_flattened(self, flattened_coordinate_vector: Vector) ->
     def reindex_ions(self):  
         """ Re-indexes the ions to order based on the distance from the center of the trap, where smallest index is closest to the center """
         # TODO: How are even/odd cases is handled? 
-        x,y,z = self.ion_coordinates_from_flattened(self.u)
+        x,y,z = self.ion_coordinates_from_flattened(self.equilibrium_positions)
         r = np.sqrt(x**2 + y**2 + z**2)
         idx = np.argsort(r)
 
         reindexed_positions = np.hstack((x[idx], y[idx], z[idx]))
-        self.u = reindexed_positions 
+        self.equilibrium_positions = reindexed_positions 
         self.m = self.m[idx]
         self.q = self.q[idx]
         self.atomic_masses = self.atomic_masses[idx]
@@ -497,17 +496,18 @@ def organize_modes(self, eigvals, eigvecs):
     def calculate_mode_participation_factors(self) -> Matrix:
         """ Computes ion-mode participation factors, related to the Lamb-Dicke parameters.
 
-            Mode participation factors (eta) take the following form:
+            Mode participation factors (eta) take the following matrix form:
              
-            shape of eta matrix: (dimension, ion, mode) 
+            shape: (dimension, ion, mode), for N ions this is (3, N, N) 
             e.g. eta[1, 2, 3] is eta in the "y" direction, ion 1, and the 2nd mode. 
 
             For N ions, this form organizes the 3N modes into N modes per direction "d", where d = x, y, z. 
 
         """ 
 
+        # TODO: Question: shouldn't the shape be (d, N, N) for d dimensions (3), N ions, and N modes per direction. 
         eigvecs = self.eigvecs
-        num_coords, num_modes = np.shape(eigvecs)
+        num_coords, num_modes = np.shape(eigvecs) 
         num_ions = num_modes // 3
         mode_participation_factors = np.zeros((3, num_ions, num_modes), dtype = np.complex128)
         for mode_index in range(num_modes):
@@ -577,10 +577,10 @@ def organize_modes(self, eigvals, eigvecs):
  
     def reindex_ions_by_z(self): 
         # based on the position along z, lowest i, is 0, up to N - 1
-        x,y,z = self.ion_coordinates_from_flattened(self.u)
+        x,y,z = self.ion_coordinates_from_flattened(self.equilibrium_positions)
  
         idx = np.argsort(z)
-        self.u = np.hstack((x[idx], y[idx], z[idx]))
+        self.equilibrium_positions = np.hstack((x[idx], y[idx], z[idx]))
  
         # Reindex all other arrays accordingly
         self.m = self.m[idx]
@@ -601,10 +601,10 @@ def solve_ion_trap_equilibrium(self):
         """ Diagonalizes the Coulomb + harmonic trap Hamiltonian for a system of ions. """
         self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
         
-        self.u = self.solve_for_equilibrium_positions()
+        self.equilibrium_positions = self.solve_for_equilibrium_positions()
         self.reindex_ions_by_z()
         mass_matrix = self.build_mass_matrix(self.m)  
-        E_matrix = self.build_E_matrix(self.u, mass_matrix)  
+        E_matrix = self.build_E_matrix(self.equilibrium_positions, mass_matrix)  
         T_matrix = self.build_momentum_transform_matrix(mass_matrix)
         H_matrix = self.compute_H_matrix(T_matrix, E_matrix)   
 
@@ -618,15 +618,52 @@ def solve_ion_trap_equilibrium(self):
         #self.hasrun = True  
 
 
+# Extra methods?
+    def return_equilibrium_positions(self, dimensionless: bool) -> Vector:
+        """ Return equilibrium positions in either dimensionless or SI units (inverse meters) """
+        if dimensionless:
+            return self.equilibrium_positions
+        else:
+            return self.equilibrium_positions * self.characteristic_parameters['length']
+
+
+    @classmethod
+    def from_species(cls, species_name: str, num_ions: int, omega_x: float, omega_y: float, omega_z: float): 
+        """ Build the mode analysis class from a species name, corresponding to a system of N ions under harmonic trapping. """ 
+        # Import necessary data for the species 
+        species_data = AtomicSpin.get_config_data(species_name)
+        atomic_mass = species_data['mass']
+        atomic_number = species_data['Z']
+        # Construct the class 
+        return cls(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_number)
+
+    @classmethod
+    def from_atomic_spin_basis(cls, atomic_structure_basis: StandardBasis, omega_x: float, omega_y: float, omega_z: float): 
+        """ Build the mode analysis class from a basis of AtomicSpin degrees of freedom under harmonic trapping. """ 
+        # Extract number of ions and the mass and atomic number from the DOF in the basis 
+        DOFs = atomic_structure_basis.degrees_of_freedom
+        num_ions = len(DOFs)
+        atomic_masses = []
+        atomic_numbers = []
+
+        for DOF in DOFs:
+            if not isinstance(DOF, AtomicSpin):
+                raise IonSimError("Atomic structure basis should only contain AtomicSpin or AtomicStructure objects. No motional modes should be included.")
+                atomic_masses.append(DOF.atomic_mass) 
+                atomic_numbers.append(DOF.atomic_number) 
+        # Construct the class 
+        return cls(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_number)
+
+
  
 #You could redefine the equilibrium finding function to assert that the equilibrium is linear. For example, make a wrapper inside the function for the potential, Jacobian, and Hessian that forces x_i and y_i = 0. 
  
 #This is the code for the Lamb-Dicke parameters: (not that overall phases don't matter here, but the relative phase does. We could pin down the phase with some convention like, "each mode's first non-negative LD value is defined positive." I also include some extra helper functions. 
 def two_central_ion_separation(wz_Hz, l_2_cent):
     mass_yb_amu = 170.936 # TODO: hardcoded for now
-    four_ion_analysis = GeneralizedModeAnalysisWithBranchSortedModes(N=4, wz=2*np.pi*wz_Hz, wy=2*np.pi*10e6, wx=2*np.pi*11e6, ionmass_amu=mass_yb_amu)
-    four_ion_analysis.run()
-    positions_z = four_ion_analysis.u[2*4:] * four_ion_analysis.characteristic_parameters['length']
+    four_ion_analysis = GeneralizedModeAnalysisWithBranchSortedModes(num_ions=4, omega_x=2*np.pi*11e6, omega_y=2*np.pi*10e6, omega_z=2*np.pi*wz_Hz, atomic_masses=mass_yb_amu, atomic_numbers = 70)
+    four_ion_analysis.solve_ion_trap_equilibrium()
+    positions_z = four_ion_analysis.equilibrium_positions[2*4:] * four_ion_analysis.characteristic_parameters['length']
     central_ions = np.argsort(np.abs(positions_z))[:2]
     dl = np.abs(positions_z[central_ions[0]] - positions_z[central_ions[1]])
     return dl - l_2_cent
@@ -650,4 +687,7 @@ def calc_mode_energies(res, Fock_cutoffs):
         energies_m1[k] = qt.expect(E1_op, state)
         energies_m2[k] = qt.expect(E2_op, state)
     return energies_m1, energies_m2
- 
+
+
+if __name__ == '__main__':
+    print(two_central_ion_separation(10. * 1E6 * 2*np.pi, 0.05)) 
diff --git a/tests/test_mode_analysis.py b/tests/test_mode_analysis.py
index aa87b52..eb29f79 100644
--- a/tests/test_mode_analysis.py
+++ b/tests/test_mode_analysis.py
@@ -88,8 +88,8 @@ def test_mode_analysis_solvers(self):
             # Solve the ion-trap equilibrium problem
             mode_analyzer.solve_ion_trap_equilibrium()
             # Compute and store Lamb-Dicke parameters: 
-            case['Lamb Dicke parameters'] = case['wavevector'] * mode_analyzer.calculate_mode_participation_factors()
-            assert_array_close(np.abs(case['Lamb Dicke parameters']), np.abs(reference_values), atol = 1E-10, rtol=None)
+            computed_Lamb_Dicke_parameters = case['wavevector'] * mode_analyzer.calculate_mode_participation_factors()
+            assert_array_close(np.abs(computed_Lamb_Dicke_parameters), np.abs(reference_values), atol = 1E-10, rtol=None)
 
 if __name__ == '__main__':
     unittest.main()

From c4016939c1850c1042a06849cffd47b17b7ca927 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Wed, 22 Apr 2026 12:37:01 -0600
Subject: [PATCH 22/31]  fix basis import, add draft method for creating
 motional modes

---
 src/ionsim/trapped_ion_mode_analysis.py | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index fd80c76..0c10ad0 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -4,8 +4,8 @@
 import warnings
 import scipy.optimize as opt    
 from ionsim.custom_types import Matrix, Vector
-from ionsim.degree_of_freedom import AtomicSpin
-from ionsim.basis import Basis 
+from ionsim.degree_of_freedom import AtomicSpin, MotionalMode
+from ionsim.basis import StandardBasis 
 
 ########## List of substantive changes: 
 ## 1. Made dimensionless variable function take in inputs so the user can choose the charge, mass, and trap freq scales to use.   
@@ -655,6 +655,21 @@ def from_atomic_spin_basis(cls, atomic_structure_basis: StandardBasis, omega_x:
         return cls(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_number)
 
 
+
+    def build_mode_DOFs(self, mode_indices: list[int], fock_dimensions: Vector | int) -> MotionalMode:
+        """ Builds and returns an IonSim Motional Degree of Freedom.
+
+            - Applies each fock dimension to each mode, or applies the same fock dimension to all the modes  
+        """
+        modes = []
+        # Convert list of Fock dimensions to an array if it's not already an array 
+        fock_dimensions = self.convert_to_array(fock_dimensions).astype(int)
+        for idx, fock_dim in zip(mode_indices, fock_dimensions):
+            mode_index = idx # or some function of this index 
+            modes.append(MotionalMode.from_frequency(self.eigvals[mode_index], fock_dim))
+
+        return modes 
+
  
 #You could redefine the equilibrium finding function to assert that the equilibrium is linear. For example, make a wrapper inside the function for the potential, Jacobian, and Hessian that forces x_i and y_i = 0. 
  

From 370cdc5751d72315dc9422681ed5aa25dfd13e74 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Wed, 22 Apr 2026 19:10:08 -0600
Subject: [PATCH 23/31]   save draft w notes from group

---
 src/ionsim/trapped_ion_mode_analysis.py | 75 ++++++++++++++++---------
 1 file changed, 47 insertions(+), 28 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 0c10ad0..e71862d 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -7,28 +7,23 @@
 from ionsim.degree_of_freedom import AtomicSpin, MotionalMode
 from ionsim.basis import StandardBasis 
 
-########## List of substantive changes: 
-## 1. Made dimensionless variable function take in inputs so the user can choose the charge, mass, and trap freq scales to use.   
-## 2. Added a helper function to pack/unpack the ion coordinates to reduce code. 
-
-########## List of non-substantive changes: 
-## 1. Variable changes for readability (e.g. omega_x instead of wx) 
-## 2. Type-hinting in functions  
-## 3. Consolidated the functions for solving for ion positions equilibria 
-## 4. Moved eigenvector helper functions into the class. 
-## 5. Reduced number of matrices that we store as class attributes: Need to check whether we want to be storing those. 
-
 ########## Questions: 
 # -1. Should the branch sorting be the default option? This seems like what we would want to do for our problems  
 # 0. LD parameter matrix shape? 
-# 1. Is there a need for the "has run" boolean? e.g. is there a time where it stays False? Ah I think it refers to whether the eqb solver has been successfully executed. 
+# 1.  "has run" boolean? e.g. is there a time where it stays False? Ah I think it refers to whether the eqb solver has been successfully executed. 
 # 2. Should the dimensionless parameters have a naming convention, e.g. omega_x --> omega_x_ND 
-# 3. What sets the phases of the Lamb-Dicke parameters?  
-# 4. Does the LD calculation properly do k dot r?
-# 5. Can we vectorize this to convert loops into vector operations?  
-# 6. See TODO's 
+
+### 3. What sets the phases of the Lamb-Dicke parameters?  
+#       - need to define a convention and stick with it 
+
+### Notes:
+# - won't get orthonormal eigenvectors if there's degeneracies; the orthonormality is w.r.t the H matrix 
 
 
+## References: 
+# https://arxiv.org/abs/2007.12725
+# https://arxiv.org/abs/quant-ph/9702053 
+
 def characteristic_length(q: float, mass: float, omega: float) -> float:
     """ Computes characteristic length in trapped ion system """
     k_e = 1 / (4 * np.pi * const.epsilon_0)  # Coulomb constant
@@ -78,6 +73,8 @@ def calculate_species_trap_frequencies(self, omega: float) -> Vector:
 
             from omega_(secular, i) = sqrt(q_i V_0 / (2 m_i d^2) ) for the i'th ion. 
 
+            - An atom will experience a potentially different pseudopotential based on its mass and charge.
+            - This is a valid approach when micromotion is small compared to secular motion. 
         """ 
         # assume that the trapping frequency given corresponds to the first ion species 
         q0 = self.nuclear_charges[0] # charge of first ion 
@@ -98,7 +95,8 @@ def convert_to_array(self, x: float | int | list | NDArray) -> NDArray:
 
     #def dimensionless_parameters(self):
     #def nondimensionalize_parameters(self):
-    def convert_parameters_to_dimensionless(self, charge_scale: float=1., mass_scale: float=1., trap_freq_scale: float=1.):
+    #def convert_parameters_to_dimensionless(self, charge_scale: float=1., mass_scale: float=1., trap_freq_scale: float=1.):
+    def set_up_dimensionless_parameeters(self, charge_scale: float=1., mass_scale: float=1., trap_freq_scale: float=1.):
         """ Compute dimensionless parameters, using first ion's properties and axial (z) trapping frequency. 
 
             - characteristic length
@@ -131,7 +129,6 @@ def convert_parameters_to_dimensionless(self, charge_scale: float=1., mass_scale
         self.wx = self.omega_x / self.trap_freq_scale  
         self.trap_frequencies_ND = [self.wx, self.wy, self.wz]
 
-        # TODO: Is hbar set to 1? Reconcile hbar somewhere  
         # Compute characteristic length, time, velocity, and energy scales and store in a dictionary  
         self.characteristic_parameters = {}
         self.characteristic_parameters['length'] = characteristic_length(self.charge_scale, self.mass_scale, self.trap_freq_scale)  
@@ -200,7 +197,8 @@ def solve_ion_trap_equilibrium(self):
         """ Diagonalizes the Coulomb + harmonic trap Hamiltonian for a system of ions. """
         # Convert to dimensionless units using axial trap frequency and first ion's mass and charge  
         # TODO: take in input here or in class constructor for mass, charge, trap scales. 
-        self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
+        #self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
+        self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
         
         self.equilibrium_positions = self.solve_for_equilibrium_positions()
         self.reindex_ions()
@@ -221,6 +219,10 @@ def solve_ion_trap_equilibrium(self):
         self.check_diagonalization(T_matrix, S_matrix, E_matrix)
         #self.hasrun = True  
 
+    @property
+    def normal_mode_frequencies(self):
+        return self.eigvals * self.trap_freq_scale 
+
 
     def normalize_eigenvectors(self, eigvecs, H: Matrix, eigenvalues: Vector | None=None): 
         """ Rescale the eigenvectors ens w.r.t. the Hamiltonian H. """
@@ -290,7 +292,6 @@ def reindex_ions(self):
 
     def solve_for_equilibrium_positions(self, positions_guess: Vector | None=None):
         """ Solves for equilibrium position vector: u, which represents a flattened spatial grid. """
-        # TODO: option for an initial guess choice? 
         # Set the initial guess for the solver  
         if positions_guess == None:
             # Initialize a random guess  
@@ -498,14 +499,14 @@ def calculate_mode_participation_factors(self) -> Matrix:
 
             Mode participation factors (eta) take the following matrix form:
              
-            shape: (dimension, ion, mode), for N ions this is (3, N, N) 
+            shape: (dimension, ion, mode), for N ions this is (3, N, 3N). This is the most general case.  
             e.g. eta[1, 2, 3] is eta in the "y" direction, ion 1, and the 2nd mode. 
 
             For N ions, this form organizes the 3N modes into N modes per direction "d", where d = x, y, z. 
 
         """ 
 
-        # TODO: Question: shouldn't the shape be (d, N, N) for d dimensions (3), N ions, and N modes per direction. 
+        # TODO: For the linear case: the shape should be (d, N, N) for d dimensions (3), N ions, and N modes per direction. 
         eigvecs = self.eigvecs
         num_coords, num_modes = np.shape(eigvecs) 
         num_ions = num_modes // 3
@@ -519,6 +520,13 @@ def calculate_mode_participation_factors(self) -> Matrix:
                 mode_participation_factors[direction_index, ion_index, mode_index] = zpm_dimensionful * prefactor * eigvecs[pos_coord, mode_index]
         return mode_participation_factors
 
+
+    # Can get lamb-dicke parameters from wavevctor \dot mode_participation_factors[:, i, m]
+    ### -- Coordinate systems must be the same / correspond. 
+
+
+
+
     # Derived properties 
     def compute_reference_single_ion_lamb_dicke_factors(self, wavenumber: float) -> (float, float, float):
         """ Computes analytical Lamb-Dicke parameter by eta = k * sqrt(hbar / m omega) for an ion in a light-field with wavevector |k| 
@@ -555,6 +563,8 @@ def _xyz_classify_modes(self, evecs):
  
     def sort_by_branch(self, evals, evecs):
         """ Sorts the eigenvalues, eigenvectors by mode branch (radial x, radial y, axial (z)) """
+        # Hessian block diagonal H_xx, H_yy, H_zz non-zero for linear chain 
+        ## only makes sense for linear chain 
         classifier = self._xyz_classify_modes(evecs)
         # within each branch, sort by frequency
         N_ions = len(evals) // 3
@@ -594,12 +604,14 @@ def reindex_ions_by_z(self):
         self.omega_y = self.omega_y[idx]
         self.omega_z = self.omega_z[idx]
         # all ions are the same so this is safe. TODO: When does this change? 
-        self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
+        self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
  
     #def run(self):
+    #def run_normal_mode_analysis()
     def solve_ion_trap_equilibrium(self): 
         """ Diagonalizes the Coulomb + harmonic trap Hamiltonian for a system of ions. """
-        self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
+        #self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
+        self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
         
         self.equilibrium_positions = self.solve_for_equilibrium_positions()
         self.reindex_ions_by_z()
@@ -613,6 +625,7 @@ def solve_ion_trap_equilibrium(self):
         self.check_for_zero_modes()
         self.check_outer_relation(H_matrix)
 
+        # S matrix may be important; it's how you convert from x, p coordinates to normal mode coordinates 
         S_matrix = self.get_canonical_transformation(H_matrix, self.eigvecs, self.eigvals) 
         self.check_diagonalization(T_matrix, S_matrix, E_matrix)
         #self.hasrun = True  
@@ -637,11 +650,16 @@ def from_species(cls, species_name: str, num_ions: int, omega_x: float, omega_y:
         # Construct the class 
         return cls(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_number)
 
+
+    basis = StandardBasis([*spins])
+
     @classmethod
-    def from_atomic_spin_basis(cls, atomic_structure_basis: StandardBasis, omega_x: float, omega_y: float, omega_z: float): 
+    #def from_atomic_spin_basis(cls, atomic_structure_basis: StandardBasis, omega_x: float, omega_y: float, omega_z: float): 
+    def from_atomic_spin_basis(cls, spins: list[degree_of_freedom], omega_x: float, omega_y: float, omega_z: float): 
         """ Build the mode analysis class from a basis of AtomicSpin degrees of freedom under harmonic trapping. """ 
         # Extract number of ions and the mass and atomic number from the DOF in the basis 
-        DOFs = atomic_structure_basis.degrees_of_freedom
+        #DOFs = atomic_structure_basis.degrees_of_freedom
+        DOFs = spins 
         num_ions = len(DOFs)
         atomic_masses = []
         atomic_numbers = []
@@ -656,7 +674,7 @@ def from_atomic_spin_basis(cls, atomic_structure_basis: StandardBasis, omega_x:
 
 
 
-    def build_mode_DOFs(self, mode_indices: list[int], fock_dimensions: Vector | int) -> MotionalMode:
+    def build_mode_DOFs(self, mode_indices: list[int], fock_dimensions: Vector | int) -> list[MotionalMode]:
         """ Builds and returns an IonSim Motional Degree of Freedom.
 
             - Applies each fock dimension to each mode, or applies the same fock dimension to all the modes  
@@ -670,7 +688,8 @@ def build_mode_DOFs(self, mode_indices: list[int], fock_dimensions: Vector | int
 
         return modes 
 
- 
+
+
 #You could redefine the equilibrium finding function to assert that the equilibrium is linear. For example, make a wrapper inside the function for the potential, Jacobian, and Hessian that forces x_i and y_i = 0. 
  
 #This is the code for the Lamb-Dicke parameters: (not that overall phases don't matter here, but the relative phase does. We could pin down the phase with some convention like, "each mode's first non-negative LD value is defined positive." I also include some extra helper functions. 

From 12d5631aed2dc114edfa0d134631a153447504c6 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 16 Jun 2026 17:56:07 -0600
Subject: [PATCH 24/31]   draft mode analysis with linear chain class + tests

---
 .gitignore                                 |   4 +
 src/ionsim/trapped_ion_mode_analysis.py    | 483 +++++++++++++++++++--
 tests/__pycache__/__init__.cpython-312.pyc | Bin 153 -> 124 bytes
 tests/test_mode_analysis.py                | 132 +++++-
 4 files changed, 582 insertions(+), 37 deletions(-)

diff --git a/.gitignore b/.gitignore
index 93b1f4e..955cfce 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,6 +4,10 @@
 # MacOS
 .DS_Store
 
+# AI Agent Scripts
+sandia-claudecode
+
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index e71862d..a8ec419 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -2,7 +2,7 @@
 from numpy.typing import NDArray
 import scipy.constants as const
 import warnings
-import scipy.optimize as opt    
+import scipy.optimize as opt
 from ionsim.custom_types import Matrix, Vector
 from ionsim.degree_of_freedom import AtomicSpin, MotionalMode
 from ionsim.basis import StandardBasis 
@@ -692,36 +692,453 @@ def build_mode_DOFs(self, mode_indices: list[int], fock_dimensions: Vector | int
 
 #You could redefine the equilibrium finding function to assert that the equilibrium is linear. For example, make a wrapper inside the function for the potential, Jacobian, and Hessian that forces x_i and y_i = 0. 
  
-#This is the code for the Lamb-Dicke parameters: (not that overall phases don't matter here, but the relative phase does. We could pin down the phase with some convention like, "each mode's first non-negative LD value is defined positive." I also include some extra helper functions. 
-def two_central_ion_separation(wz_Hz, l_2_cent):
-    mass_yb_amu = 170.936 # TODO: hardcoded for now
-    four_ion_analysis = GeneralizedModeAnalysisWithBranchSortedModes(num_ions=4, omega_x=2*np.pi*11e6, omega_y=2*np.pi*10e6, omega_z=2*np.pi*wz_Hz, atomic_masses=mass_yb_amu, atomic_numbers = 70)
-    four_ion_analysis.solve_ion_trap_equilibrium()
-    positions_z = four_ion_analysis.equilibrium_positions[2*4:] * four_ion_analysis.characteristic_parameters['length']
-    central_ions = np.argsort(np.abs(positions_z))[:2]
-    dl = np.abs(positions_z[central_ions[0]] - positions_z[central_ions[1]])
-    return dl - l_2_cent
-    
-def find_wz_for_desired_central_ion_separation(l_2_cent, bounds=(0.1e6, 0.5e6)):
-    result = root_scalar(two_central_ion_separation, args=(l_2_cent,), bracket=bounds, method='bisect')
-    if not result.converged:
-        raise ValueError("Root finding did not converge. Try adjusting the bounds or check the function behavior.")
-    return result.root
-    
+#This is the code for the Lamb-Dicke parameters: (not that overall phases don't matter here, but the relative phase does. We could pin down the phase with some convention like, "each mode's first non-negative LD value is defined positive.")
+        
+ 
+
+
+class LinearIonChainAnalysis(GeneralizedModeAnalysisWithBranchSortedModes):
+    """
+    Subclass for setting up and analyzing linear ion chains.
+
+    In a linear ion chain, ions are typically aligned along the z-axis (axial direction),
+    with minimal displacement in the x and y directions (radial directions).
+
+    This class extends GeneralizedModeAnalysisWithBranchSortedModes to:
+    1. Enforce linear chain configuration
+    2. Provide specialized methods for linear chain analysis
+    3. Offer convenience methods for common linear chain scenarios
+    """
+
+    def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float,
+                 atomic_masses: np.ndarray | float, atomic_numbers: np.ndarray | int):
+        """
+        Initialize a linear ion chain analysis.
+
+        Parameters:
+        -----------
+        num_ions : int
+            Number of ions in the chain
+        omega_x, omega_y, omega_z : float
+            Trap frequencies in x, y, and z directions (rad/s)
+        atomic_masses : array or float
+            Atomic masses in amu (atomic mass units)
+        atomic_numbers : array or int
+            Atomic numbers (number of protons)
+        """
+        super().__init__(num_ions, omega_x, omega_y, omega_z, atomic_masses, atomic_numbers)
+
+    def solve_ion_trap_equilibrium(self):
+        """
+        Solve for equilibrium positions and analyze normal modes for a linear chain.
+
+        This method:
+        1. Sets up dimensionless parameters
+        2. Solves for equilibrium positions
+        3. Reindexes ions by their z-position (closest to center first)
+        4. Computes normal modes with branch sorting
+        5. Performs validation checks
+        """
+        # Set up dimensionless parameters using first ion's properties and axial trap frequency
+        self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
+
+        # Solve for equilibrium positions
+        self.equilibrium_positions = self.solve_for_equilibrium_positions()
+
+        # Reindex ions by their z-position (important for linear chains)
+        self.reindex_ions_by_z()
+
+        # Build necessary matrices
+        mass_matrix = self.build_mass_matrix(self.m)
+        E_matrix = self.build_E_matrix(self.equilibrium_positions, mass_matrix)
+        T_matrix = self.build_momentum_transform_matrix(mass_matrix)
+        H_matrix = self.compute_H_matrix(T_matrix, E_matrix)
+
+        # Calculate normal modes with branch sorting (x, y, z branches)
+        self.eigvals, self.eigvecs = self.calculate_normal_modes(H_matrix)
+        self.eigvecs_vel = self.get_eigenvectors_xv_coords(T_matrix, self.eigvecs)
+
+        # Validate results
+        self.check_for_zero_modes()
+        self.check_outer_relation(H_matrix)
+
+        # Compute canonical transformation matrix
+        S_matrix = self.get_canonical_transformation(H_matrix, self.eigvecs, self.eigvals)
+        self.check_diagonalization(T_matrix, S_matrix, E_matrix)
+
+    def get_axial_modes(self):
+        """
+        Get the axial (z-direction) normal modes.
+
+        Returns:
+        --------
+        tuple
+            (axial_eigenvalues, axial_eigenvectors) for the axial modes
+        """
+        # Axial modes are the last third of the modes (after branch sorting)
+        n_modes_per_branch = self.num_ions
+        axial_start = 2 * n_modes_per_branch
+        axial_end = 3 * n_modes_per_branch
+
+        axial_eigvals = self.eigvals[axial_start:axial_end]
+        axial_eigvecs = self.eigvecs[:, axial_start:axial_end]
+
+        return axial_eigvals, axial_eigvecs
+
+    def get_radial_modes(self, direction='x'):
+        """
+        Get the radial normal modes in the specified direction.
+
+        Parameters:
+        -----------
+        direction : str
+            'x' or 'y' for radial direction
+
+        Returns:
+        --------
+        tuple
+            (radial_eigenvalues, radial_eigenvectors) for the specified radial modes
+        """
+        if direction not in ['x', 'y']:
+            raise ValueError("Direction must be 'x' or 'y'")
+
+        n_modes_per_branch = self.num_ions
+        if direction == 'x':
+            start_idx = 0
+        else:  # direction == 'y'
+            start_idx = n_modes_per_branch
+
+        end_idx = start_idx + n_modes_per_branch
+
+        radial_eigvals = self.eigvals[start_idx:end_idx]
+        radial_eigvecs = self.eigvecs[:, start_idx:end_idx]
+
+        return radial_eigvals, radial_eigvecs
+
+    def get_ion_spacing(self):
+        """
+        Get the spacing between consecutive ions along the z-axis.
+
+        Returns:
+        --------
+        np.ndarray
+            Array of distances between consecutive ions (in meters)
+        """
+        # Get z positions (dimensionful)
+        x, y, z = self.ion_coordinates_from_flattened(self.equilibrium_positions)
+        z_dimensionful = z * self.characteristic_parameters['length']
+
+        # Sort positions and calculate spacing
+        sorted_z = np.sort(z_dimensionful)
+        ion_spacing = np.diff(sorted_z)
+
+        return ion_spacing
+
+    def get_center_of_mass_position(self):
+        """
+        Get the center of mass position of the ion chain.
+
+        Returns:
+        --------
+        tuple
+            (com_x, com_y, com_z) center of mass coordinates in meters
+        """
+        # Get positions (dimensionful)
+        x, y, z = self.ion_coordinates_from_flattened(self.equilibrium_positions)
+        x_dim = x * self.characteristic_parameters['length']
+        y_dim = y * self.characteristic_parameters['length']
+        z_dim = z * self.characteristic_parameters['length']
+
+        # Calculate center of mass
+        total_mass = np.sum(self.atomic_masses)
+        com_x = np.sum(x_dim * self.atomic_masses) / total_mass
+        com_y = np.sum(y_dim * self.atomic_masses) / total_mass
+        com_z = np.sum(z_dim * self.atomic_masses) / total_mass
+
+        return com_x, com_y, com_z
+
+    def get_axial_mode_frequencies(self):
+        """
+        Get the frequencies of the axial modes.
+
+        Returns:
+        --------
+        np.ndarray
+            Array of axial mode frequencies in rad/s
+        """
+        axial_eigvals, _ = self.get_axial_modes()
+        return axial_eigvals * self.trap_freq_scale
+
+    def get_radial_mode_frequencies(self, direction='x'):
+        """
+        Get the frequencies of the radial modes in the specified direction.
+
+        Parameters:
+        -----------
+        direction : str
+            'x' or 'y' for radial direction
+
+        Returns:
+        --------
+        np.ndarray
+            Array of radial mode frequencies in rad/s
+        """
+        radial_eigvals, _ = self.get_radial_modes(direction)
+        return radial_eigvals * self.trap_freq_scale
+
+    def get_mode_participation_factors_by_branch(self):
+        """
+        Get mode participation factors organized by branch (x, y, z).
+
+        Returns:
+        --------
+        dict
+            Dictionary with keys 'x', 'y', 'z' containing mode participation factors
+            for each branch. Each value is a 2D array of shape (num_ions, num_modes_per_branch)
+        """
+        # Get full mode participation factors
+        mode_pf = self.calculate_mode_participation_factors()
+
+        # Organize by branch
+        n_modes_per_branch = self.num_ions
+
+        result = {
+            'x': mode_pf[0, :, :n_modes_per_branch],
+            'y': mode_pf[1, :, n_modes_per_branch:2*n_modes_per_branch],
+            'z': mode_pf[2, :, 2*n_modes_per_branch:3*n_modes_per_branch]
+        }
+
+        return result
+
+    def calculate_lamb_dicke_parameters(self, wavevector: np.ndarray | float) -> dict:
+        """
+        Calculate Lamb-Dicke parameters for all ions and modes, organized by branch.
+
+        The Lamb-Dicke parameter η = k · Δr_0 represents the ratio of the spatial
+        extent of the ion's zero-point motion to the wavelength of the laser.
+
+        Parameters:
+        -----------
+        wavevector : np.ndarray or float
+            - If array: Full wavevector (kx, ky, kz) as a 3-element array in units of 1/m
+            - If float: Magnitude of the laser wavevector |k| in units of 1/m (for backward compatibility)
+
+        Returns:
+        --------
+        dict
+            Dictionary with keys 'x', 'y', 'z' containing Lamb-Dicke parameters
+            for each branch. Each value is a 2D array of shape (num_ions, num_modes_per_branch)
+            representing η_{direction}[ion_index, mode_index]
+        """
+        # Get mode participation factors (which are proportional to zero-point motion)
+        mode_pf_by_branch = self.get_mode_participation_factors_by_branch()
+
+        # Handle both full wavevector (preferred) and scalar wavenumber (backward compatibility)
+        if isinstance(wavevector, (float, int)):
+            # Scalar case: multiply each component by the same wavenumber
+            # This assumes the laser is equally coupled to all directions
+            lamb_dicke_parameters = {}
+            for direction in ['x', 'y', 'z']:
+                lamb_dicke_parameters[direction] = wavevector * mode_pf_by_branch[direction]
+        else:
+            # Vector case: compute dot product k · Δr_0 for proper directionality
+            wavevector = np.asarray(wavevector)
+            if wavevector.shape != (3,):
+                raise ValueError("Wavevector must be a 3-element array (kx, ky, kz)")
+
+            # Get full mode participation factors
+            full_mode_pf = self.calculate_mode_participation_factors()
+
+            # Compute Lamb-Dicke parameters as dot product: η = k · Δr_0
+            # This gives us shape (num_ions, num_modes) for the total LD parameter
+            num_ions = self.num_ions
+            num_modes = 3 * num_ions
+
+            # Reshape for dot product: (3, num_ions, num_modes) · (3,) -> (num_ions, num_modes)
+            total_ld_params = np.zeros((num_ions, num_modes), dtype=complex)
+            for i in range(3):
+                total_ld_params += wavevector[i] * full_mode_pf[i, :, :]
+
+            # Organize by branch for consistency with scalar case
+            n_modes_per_branch = num_ions
+            lamb_dicke_parameters = {
+                'x': total_ld_params[:, :n_modes_per_branch],
+                'y': total_ld_params[:, n_modes_per_branch:2*n_modes_per_branch],
+                'z': total_ld_params[:, 2*n_modes_per_branch:3*n_modes_per_branch]
+            }
+
+        return lamb_dicke_parameters
+
+    def calculate_lamb_dicke_parameters_full(self, wavevector: np.ndarray | float) -> np.ndarray:
+        """
+        Calculate full Lamb-Dicke parameter matrix.
+
+        Parameters:
+        -----------
+        wavevector : np.ndarray or float
+            - If array: Full wavevector (kx, ky, kz) as a 3-element array in units of 1/m
+            - If float: Magnitude of the laser wavevector |k| in units of 1/m (for backward compatibility)
+
+        Returns:
+        --------
+        np.ndarray
+            - If scalar wavevector: Full matrix of shape (3, num_ions, num_modes)
+              where eta[direction, ion, mode] gives the component-wise LD parameter
+            - If vector wavevector: Total LD parameter matrix of shape (num_ions, num_modes)
+              where eta[ion, mode] gives the total LD parameter k · Δr_0
+        """
+        # Get full mode participation factors
+        mode_pf = self.calculate_mode_participation_factors()
+
+        if isinstance(wavevector, (float, int)):
+            # Scalar case: component-wise multiplication (backward compatibility)
+            return wavevector * mode_pf
+        else:
+            # Vector case: compute dot product k · Δr_0 (physically accurate)
+            wavevector = np.asarray(wavevector)
+            if wavevector.shape != (3,):
+                raise ValueError("Wavevector must be a 3-element array (kx, ky, kz)")
+
+            # Compute total Lamb-Dicke parameter as dot product
+            # Shape: (3, num_ions, num_modes) · (3,) -> (num_ions, num_modes)
+            num_ions = self.num_ions
+            num_modes = 3 * num_ions
+            total_ld_params = np.zeros((num_ions, num_modes), dtype=complex)
+
+            for i in range(3):
+                total_ld_params += wavevector[i] * mode_pf[i, :, :]
+
+            return total_ld_params
+
+    def get_axial_lamb_dicke_parameters(self, wavevector: np.ndarray | float) -> np.ndarray:
+        """
+        Get Lamb-Dicke parameters for axial (z) modes only.
+
+        Parameters:
+        -----------
+        wavevector : np.ndarray or float
+            - If array: Full wavevector (kx, ky, kz) as a 3-element array in units of 1/m
+            - If float: Magnitude of the laser wavevector |k| in units of 1/m (for backward compatibility)
+
+        Returns:
+        --------
+        np.ndarray
+            Lamb-Dicke parameters for axial modes, shape (num_ions, num_axial_modes)
+        """
+        lamb_dicke_by_branch = self.calculate_lamb_dicke_parameters(wavevector)
+        return lamb_dicke_by_branch['z']
+
+    def get_radial_lamb_dicke_parameters(self, wavevector: np.ndarray | float, direction: str = 'x') -> np.ndarray:
+        """
+        Get Lamb-Dicke parameters for radial modes in specified direction.
+
+        Parameters:
+        -----------
+        wavevector : np.ndarray or float
+            - If array: Full wavevector (kx, ky, kz) as a 3-element array in units of 1/m
+            - If float: Magnitude of the laser wavevector |k| in units of 1/m (for backward compatibility)
+        direction : str
+            'x' or 'y' for radial direction
+
+        Returns:
+        --------
+        np.ndarray
+            Lamb-Dicke parameters for radial modes, shape (num_ions, num_radial_modes)
+        """
+        if direction not in ['x', 'y']:
+            raise ValueError("Direction must be 'x' or 'y'")
+
+        lamb_dicke_by_branch = self.calculate_lamb_dicke_parameters(wavevector)
+        return lamb_dicke_by_branch[direction]
+
+    def get_two_central_ion_separation(self) -> float:
+        """
+        Calculate the separation between the two central ions in the chain.
+
+        Returns:
+        --------
+        float
+            Distance between the two central ions in meters
+        """
+        # Get z positions (dimensionful)
+        x, y, z = self.ion_coordinates_from_flattened(self.equilibrium_positions)
+        z_dimensionful = z * self.characteristic_parameters['length']
+
+        # Find the two central ions (closest to the center)
+        central_ions = np.argsort(np.abs(z_dimensionful))[:2]
+
+        # Calculate separation
+        dl = np.abs(z_dimensionful[central_ions[0]] - z_dimensionful[central_ions[1]])
+        return dl
+
+    def find_axial_frequency_for_desired_central_ion_separation(self, target_separation: float,
+                                                               bounds: tuple = (0.1e6, 0.5e6)) -> float:
+        """
+        Find the axial frequency (wz) required to achieve a desired separation between central ions.
+
+        Parameters:
+        -----------
+        target_separation : float
+            Desired separation between central ions in meters
+        bounds : tuple
+            Tuple of (min_freq, max_freq) in Hz for the root finding algorithm
+
+        Returns:
+        --------
+        float
+            Axial frequency in rad/s that achieves the target separation
+        """
+        from scipy.optimize import root_scalar
+
+        def separation_error(wz_Hz):
+            # Create a temporary analysis object with the current wz
+            temp_analysis = GeneralizedModeAnalysisWithBranchSortedModes(
+                num_ions=self.num_ions,
+                omega_x=self.omega_x[0],
+                omega_y=self.omega_y[0],
+                omega_z=2*np.pi*wz_Hz,
+                atomic_masses=self.atomic_masses[0],
+                atomic_numbers=self.atomic_numbers[0]
+            )
+            temp_analysis.solve_ion_trap_equilibrium()
+            current_separation = temp_analysis.get_two_central_ion_separation()
+            return current_separation - target_separation
+
+        # Convert bounds from Hz to rad/s for the root finding
+        bounds_rad_s = (2*np.pi*bounds[0], 2*np.pi*bounds[1])
+
+        result = root_scalar(separation_error, bracket=bounds_rad_s, method='bisect')
+        if not result.converged:
+            raise ValueError("Root finding did not converge. Try adjusting the bounds or check the function behavior.")
+
+        return result.root
+
+    def print_chain_summary(self):
+        """
+        Print a summary of the linear ion chain configuration.
+        """
+        print(f"Linear Ion Chain Analysis Summary")
+        print(f"Number of ions: {self.num_ions}")
+        print(f"Trap frequencies: x={self.omega_x[0]/(2*np.pi)/1e6:.2f} MHz, "
+              f"y={self.omega_y[0]/(2*np.pi)/1e6:.2f} MHz, "
+              f"z={self.omega_z[0]/(2*np.pi)/1e6:.2f} MHz")
+
+        # Get ion spacing
+        spacing = self.get_ion_spacing()
+        print(f"Ion spacing (z-axis): {spacing*1e6:.2f} µm")
+
+        # Get COM position
+        com_x, com_y, com_z = self.get_center_of_mass_position()
+        print(f"Center of mass position: ({com_x*1e6:.2f} µm, {com_y*1e6:.2f} µm, {com_z*1e6:.2f} µm)")
+
+        # Get mode frequencies
+        axial_freqs = self.get_axial_mode_frequencies()
+        radial_x_freqs = self.get_radial_mode_frequencies('x')
+        radial_y_freqs = self.get_radial_mode_frequencies('y')
+
+        print(f"\nAxial mode frequencies: {axial_freqs/(2*np.pi)/1e6:.2f} MHz")
+        print(f"Radial X mode frequencies: {radial_x_freqs/(2*np.pi)/1e6:.2f} MHz")
+        print(f"Radial Y mode frequencies: {radial_y_freqs/(2*np.pi)/1e6:.2f} MHz")
+
+
  
-def calc_mode_energies(res, Fock_cutoffs):
-    energies_m1 = np.empty(len(res.states))
-    energies_m2 = np.empty(len(res.states))
-    N_op = qt.num(Fock_cutoffs[0])
-    eye = qt.qeye(Fock_cutoffs[0])
-    spin_eye = qt.qeye(2)
-    E1_op = qt.tensor([spin_eye, N_op, eye])
-    E2_op = qt.tensor([spin_eye, eye, N_op])
-    for k, state in enumerate(res.states):
-        energies_m1[k] = qt.expect(E1_op, state)
-        energies_m2[k] = qt.expect(E2_op, state)
-    return energies_m1, energies_m2
-
-
-if __name__ == '__main__':
-    print(two_central_ion_separation(10. * 1E6 * 2*np.pi, 0.05)) 
diff --git a/tests/__pycache__/__init__.cpython-312.pyc b/tests/__pycache__/__init__.cpython-312.pyc
index fe3ce6e0a002e2c8c389d9e7d12ed9ce7f41ebb0..265d7e66b451e6f62f04190a67eeb4a85c870647 100644
GIT binary patch
delta 33
ncmbQqSi^IgmzRqH2>ACIW=-TV;*!=cD9X=DO)e>(m}&w5h#d(r

delta 62
zcmb=K$;fk>mzRqH2s*Uxr%mKBQn%0#Elw>e)=w%bt;$T+4=BpdN=+^)*7wZM3(m}q
QFG@{IOfJbUn&@Ey097g!?*IS*

diff --git a/tests/test_mode_analysis.py b/tests/test_mode_analysis.py
index eb29f79..07c5ad2 100644
--- a/tests/test_mode_analysis.py
+++ b/tests/test_mode_analysis.py
@@ -1,7 +1,7 @@
 import unittest
 import numpy as np
 from scipy.linalg import expm
-from ionsim.trapped_ion_mode_analysis import TrappedIonModeAnalysis 
+from ionsim.trapped_ion_mode_analysis import TrappedIonModeAnalysis, LinearIonChainAnalysis, GeneralizedModeAnalysisWithBranchSortedModes 
 from ionsim.testing import assert_array_close
 
 class TestModeAnalysis(unittest.TestCase):
@@ -82,14 +82,138 @@ def setUp(self):
 
 
     def test_mode_analysis_solvers(self):
-        # Test functionality of mode analysis for computing Lamb-Dicke parameters as compared to a reference  
-        # Compute Lamb-Dicke parameters for all test cases 
+        # Test functionality of mode analysis for computing Lamb-Dicke parameters as compared to a reference
+        # Compute Lamb-Dicke parameters for all test cases
         for case, mode_analyzer, reference_values in zip(self.test_cases, self.mode_analyzers.values(), self.references.values()):
             # Solve the ion-trap equilibrium problem
             mode_analyzer.solve_ion_trap_equilibrium()
-            # Compute and store Lamb-Dicke parameters: 
+            # Compute and store Lamb-Dicke parameters:
             computed_Lamb_Dicke_parameters = case['wavevector'] * mode_analyzer.calculate_mode_participation_factors()
             assert_array_close(np.abs(computed_Lamb_Dicke_parameters), np.abs(reference_values), atol = 1E-10, rtol=None)
 
+    def test_linear_ion_chain_analysis(self):
+        """Test the LinearIonChainAnalysis class with a 5-ion Yb+ chain."""
+        # Create a 5-ion linear chain of Yb+ ions
+        num_ions = 5
+        omega_x = 2 * np.pi * 10e6  # 10 MHz
+        omega_y = 2 * np.pi * 10e6  # 10 MHz
+        omega_z = 2 * np.pi * 1e6   # 1 MHz (axial frequency typically lower)
+        atomic_mass = 170.936  # Yb+ mass in amu
+        atomic_number = 70      # Yb atomic number
+
+        # Create and analyze the linear chain
+        chain = LinearIonChainAnalysis(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_number)
+        chain.solve_ion_trap_equilibrium()
+
+        # Test that we can get axial and radial modes
+        axial_eigvals, axial_eigvecs = chain.get_axial_modes()
+        radial_x_eigvals, radial_x_eigvecs = chain.get_radial_modes('x')
+        radial_y_eigvals, radial_y_eigvecs = chain.get_radial_modes('y')
+
+        # Verify shapes
+        self.assertEqual(axial_eigvals.shape, (num_ions,))
+        self.assertEqual(axial_eigvecs.shape, (6*num_ions, num_ions))
+        self.assertEqual(radial_x_eigvals.shape, (num_ions,))
+        self.assertEqual(radial_x_eigvecs.shape, (6*num_ions, num_ions))
+
+        # Test that we can get ion spacing
+        spacing = chain.get_ion_spacing()
+        self.assertEqual(spacing.shape, (num_ions-1,))
+        self.assertTrue(np.all(spacing > 0))  # All spacings should be positive
+
+        # Test that we can get center of mass
+        com_x, com_y, com_z = chain.get_center_of_mass_position()
+        self.assertIsInstance(com_x, (float, np.floating))
+        self.assertIsInstance(com_y, (float, np.floating))
+        self.assertIsInstance(com_z, (float, np.floating))
+
+        # Test that we can get mode frequencies
+        axial_freqs = chain.get_axial_mode_frequencies()
+        radial_x_freqs = chain.get_radial_mode_frequencies('x')
+        radial_y_freqs = chain.get_radial_mode_frequencies('y')
+
+        self.assertEqual(axial_freqs.shape, (num_ions,))
+        self.assertEqual(radial_x_freqs.shape, (num_ions,))
+        self.assertEqual(radial_y_freqs.shape, (num_ions,))
+
+        # Test mode participation factors by branch
+        mode_pf_by_branch = chain.get_mode_participation_factors_by_branch()
+        self.assertIn('x', mode_pf_by_branch)
+        self.assertIn('y', mode_pf_by_branch)
+        self.assertIn('z', mode_pf_by_branch)
+
+        # Test that frequencies are positive
+        self.assertTrue(np.all(axial_freqs > 0))
+        self.assertTrue(np.all(radial_x_freqs > 0))
+        self.assertTrue(np.all(radial_y_freqs > 0))
+
+        # Test Lamb-Dicke parameter calculations
+        wavenumber = 2 * np.pi / (355.0 * 1e-9)  # Typical laser wavelength for Yb+
+
+        # Test full Lamb-Dicke parameter matrix
+        full_ld_params = chain.calculate_lamb_dicke_parameters_full(wavenumber)
+        self.assertEqual(full_ld_params.shape, (3, num_ions, 3*num_ions))
+
+        # Test branch-organized Lamb-Dicke parameters
+        ld_by_branch = chain.calculate_lamb_dicke_parameters(wavenumber)
+        self.assertIn('x', ld_by_branch)
+        self.assertIn('y', ld_by_branch)
+        self.assertIn('z', ld_by_branch)
+        self.assertEqual(ld_by_branch['x'].shape, (num_ions, num_ions))
+        self.assertEqual(ld_by_branch['y'].shape, (num_ions, num_ions))
+        self.assertEqual(ld_by_branch['z'].shape, (num_ions, num_ions))
+
+        # Test axial Lamb-Dicke parameters
+        axial_ld = chain.get_axial_lamb_dicke_parameters(wavenumber)
+        self.assertEqual(axial_ld.shape, (num_ions, num_ions))
+
+        # Test radial Lamb-Dicke parameters
+        radial_x_ld = chain.get_radial_lamb_dicke_parameters(wavenumber, 'x')
+        radial_y_ld = chain.get_radial_lamb_dicke_parameters(wavenumber, 'y')
+        self.assertEqual(radial_x_ld.shape, (num_ions, num_ions))
+        self.assertEqual(radial_y_ld.shape, (num_ions, num_ions))
+
+        # Verify that branch-organized LD params match the corresponding parts of full matrix
+        n_modes_per_branch = num_ions
+        np.testing.assert_array_almost_equal(ld_by_branch['x'], full_ld_params[0, :, :n_modes_per_branch])
+        np.testing.assert_array_almost_equal(ld_by_branch['y'], full_ld_params[1, :, n_modes_per_branch:2*n_modes_per_branch])
+        np.testing.assert_array_almost_equal(ld_by_branch['z'], full_ld_params[2, :, 2*n_modes_per_branch:3*n_modes_per_branch])
+
+        # Test that Lamb-Dicke parameters are related to mode participation factors
+        mode_pf = chain.calculate_mode_participation_factors()
+        expected_full_ld = wavenumber * mode_pf
+        np.testing.assert_array_almost_equal(full_ld_params, expected_full_ld)
+
 if __name__ == '__main__':
+    # Run the unit tests
     unittest.main()
+
+    # Example usage of the new methods (from the original main section)
+    print("\n" + "="*60)
+    print("Example usage of LinearIonChainAnalysis methods")
+    print("="*60)
+
+    # Create a 4-ion analysis
+    four_ion_analysis = GeneralizedModeAnalysisWithBranchSortedModes(
+        num_ions=4,
+        omega_x=2*np.pi*11e6,
+        omega_y=2*np.pi*10e6,
+        omega_z=2*np.pi*10e6,  # Start with 10 MHz axial frequency
+        atomic_masses=170.936,
+        atomic_numbers=70
+    )
+    four_ion_analysis.solve_ion_trap_equilibrium()
+
+    # Get current central ion separation
+    current_separation = four_ion_analysis.get_two_central_ion_separation()
+    print(f"Current central ion separation: {current_separation*1e6:.3f} µm")
+
+    # Find axial frequency for desired separation (e.g., 50 µm)
+    target_separation = 50e-6  # 50 micrometers
+    wz_for_target = four_ion_analysis.find_axial_frequency_for_desired_central_ion_separation(
+        target_separation,
+        bounds=(0.1e6, 0.5e6)
+    )
+    print(f"Required axial frequency for {target_separation*1e6:.1f} µm separation: {wz_for_target/(2*np.pi)/1e6:.3f} MHz")
+
+    print("\nExample completed successfully!")

From 6e78fa6befdcbb3b65bbc5cf93a1db7a7d27521c Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 16 Jun 2026 19:08:06 -0600
Subject: [PATCH 25/31]   refactor of mode analysis + linear chain class

---
 src/ionsim/trapped_ion_mode_analysis.py | 688 ++++++++++++------------
 1 file changed, 335 insertions(+), 353 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index a8ec419..dd14e13 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -32,20 +32,33 @@ def characteristic_length(q: float, mass: float, omega: float) -> float:
 
 
 class TrappedIonModeAnalysis:
-    def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float, atomic_masses: Vector | float, atomic_numbers: Vector | int): 
+    def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float, atomic_masses: Vector | float, atomic_numbers: Vector | int,
+                 mode_organization: str = 'frequency_only', reindexing_strategy: str = 'distance'):
         """ Class for determining properties of trapped-ion phonon modes, requiring the following system parameters:
 
             - num_ions: the number of ions
-            - omega_x: harmonic trap frequency in the x direction in units of rad/s. 
-            - omega_y: harmonic trap frequency in the y direction in units of rad/s. 
-            - omega_z: harmonic trap frequency in the z direction. This is the "axial" direction used for 1D chains.  
-            - atomic masses: array of atomic masses or a single number => same mass for each ion. Units are amu (atomic mass units) 
-            - atomic numbers: array of (Z) atomic numbers (# of protons of an element) or a single number => all ions are the same.  
+            - omega_x: harmonic trap frequency in the x direction in units of rad/s.
+            - omega_y: harmonic trap frequency in the y direction in units of rad/s.
+            - omega_z: harmonic trap frequency in the z direction. This is the "axial" direction used for 1D chains.
+            - atomic masses: array of atomic masses or a single number => same mass for each ion. Units are amu (atomic mass units)
+            - atomic numbers: array of (Z) atomic numbers (# of protons of an element) or a single number => all ions are the same.
+            - mode_organization: strategy for organizing modes ('frequency_only' or 'branch_sorted')
+            - reindexing_strategy: strategy for reindexing ions ('distance' or 'z_axis')
 
-        """ 
-        # TODO: Decide how to handle units and how much of pint we should use.  
+        """
+        # TODO: Decide how to handle units and how much of pint we should use.
         self.num_ions = num_ions
 
+        # Store configuration strategies
+        self.mode_organization = mode_organization
+        self.reindexing_strategy = reindexing_strategy
+
+        # Validate configuration
+        if mode_organization not in ['frequency_only', 'branch_sorted']:
+            raise ValueError(f"Invalid mode_organization: {mode_organization}. Must be 'frequency_only' or 'branch_sorted'")
+        if reindexing_strategy not in ['distance', 'z_axis']:
+            raise ValueError(f"Invalid reindexing_strategy: {reindexing_strategy}. Must be 'distance' or 'z_axis'")
+
         self.atomic_masses = self.convert_to_array(atomic_masses) * const.u # kg from amu 
         self.atomic_numbers = self.convert_to_array(atomic_numbers)
 
@@ -191,30 +204,33 @@ def check_diagonalization(self, T: Matrix, S: Matrix, E: Matrix) -> bool:
             warnings.warn("Diagnolization check failed")    
             print("has duplicate eigenvalues: ", has_duplicate_eigvals)
 
-    #def run(self):
-
-    def solve_ion_trap_equilibrium(self): 
+    def solve_ion_trap_equilibrium(self):
         """ Diagonalizes the Coulomb + harmonic trap Hamiltonian for a system of ions. """
-        # Convert to dimensionless units using axial trap frequency and first ion's mass and charge  
-        # TODO: take in input here or in class constructor for mass, charge, trap scales. 
+        # Convert to dimensionless units using axial trap frequency and first ion's mass and charge
+        # TODO: take in input here or in class constructor for mass, charge, trap scales.
         #self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
         self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
-        
+
         self.equilibrium_positions = self.solve_for_equilibrium_positions()
-        self.reindex_ions()
 
-        # Compute helper matrices for computing normal mode properties 
-        mass_matrix = self.build_mass_matrix(self.m)  
-        E_matrix = self.build_E_matrix(self.equilibrium_positions, mass_matrix)  
+        # Use configurable reindexing strategy
+        if self.reindexing_strategy == 'z_axis':
+            self.reindex_ions_by_z()
+        else:  # 'distance'
+            self.reindex_ions()
+
+        # Compute helper matrices for computing normal mode properties
+        mass_matrix = self.build_mass_matrix(self.m)
+        E_matrix = self.build_E_matrix(self.equilibrium_positions, mass_matrix)
         T_matrix = self.build_momentum_transform_matrix(mass_matrix)
-        H_matrix = self.compute_H_matrix(T_matrix, E_matrix)   
+        H_matrix = self.compute_H_matrix(T_matrix, E_matrix)
 
         self.eigvals, self.eigvecs = self.calculate_normal_modes(H_matrix)
-        self.eigvecs_vel = self.get_eigenvectors_xv_coords(T_matrix, self.eigvecs)    
-        self.check_for_zero_modes() 
-        S_matrix = self.get_canonical_transformation(H_matrix, self.eigvecs, self.eigvals) 
+        self.eigvecs_vel = self.get_eigenvectors_xv_coords(T_matrix, self.eigvecs)
+        self.check_for_zero_modes()
+        S_matrix = self.get_canonical_transformation(H_matrix, self.eigvecs, self.eigvals)
 
-        # Perform checks 
+        # Perform checks
         self.check_outer_relation(H_matrix)
         self.check_diagonalization(T_matrix, S_matrix, E_matrix)
         #self.hasrun = True  
@@ -274,21 +290,42 @@ def ion_coordinates_from_flattened(self, flattened_coordinate_vector: Vector) ->
         return x, y, z 
         
 
-    def reindex_ions(self):  
-        """ Re-indexes the ions to order based on the distance from the center of the trap, where smallest index is closest to the center """
-        # TODO: How are even/odd cases is handled? 
+    def reindex_ions(self):
+        """Re-indexes the ions to order based on the distance from the center of the trap, where smallest index is closest to the center."""
+        # TODO: How are even/odd cases is handled?
         x,y,z = self.ion_coordinates_from_flattened(self.equilibrium_positions)
         r = np.sqrt(x**2 + y**2 + z**2)
         idx = np.argsort(r)
 
         reindexed_positions = np.hstack((x[idx], y[idx], z[idx]))
-        self.equilibrium_positions = reindexed_positions 
+        self.equilibrium_positions = reindexed_positions
+        self.m = self.m[idx]
+        self.q = self.q[idx]
+        self.atomic_masses = self.atomic_masses[idx]
+        self.nuclear_charges = self.nuclear_charges[idx]
+        self.atomic_numbers = self.atomic_numbers[idx]
+
+    def reindex_ions_by_z(self):
+        """Re-indexes ions based on their position along the z-axis, with the ion closest to the center having index 0."""
+        x,y,z = self.ion_coordinates_from_flattened(self.equilibrium_positions)
+        idx = np.argsort(z)
+
+        self.equilibrium_positions = np.hstack((x[idx], y[idx], z[idx]))
+
+        # Reindex all other arrays accordingly
         self.m = self.m[idx]
         self.q = self.q[idx]
         self.atomic_masses = self.atomic_masses[idx]
         self.nuclear_charges = self.nuclear_charges[idx]
         self.atomic_numbers = self.atomic_numbers[idx]
 
+        # trapping frequencies
+        self.omega_x = self.omega_x[idx]
+        self.omega_y = self.omega_y[idx]
+        self.omega_z = self.omega_z[idx]
+        # all ions are the same so this is safe. TODO: When does this change?
+        self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
+
 
     def solve_for_equilibrium_positions(self, positions_guess: Vector | None=None):
         """ Solves for equilibrium position vector: u, which represents a flattened spatial grid. """
@@ -488,10 +525,49 @@ def split_modes(self, eigvals, eigvecs):
        return eigvals, eigvecs
     
 
-    def organize_modes(self, eigvals, eigvecs):  
-        eigvals, eigvecs = self.sort_modes(eigvals, eigvecs)    
+    def organize_modes(self, eigvals, eigvecs):
+        eigvals, eigvecs = self.sort_modes(eigvals, eigvecs)
         eigvals, eigvecs = self.split_modes(eigvals, eigvecs)
-        return eigvals, eigvecs 
+
+        # Apply branch sorting if configured
+        if self.mode_organization == 'branch_sorted':
+            eigvals, eigvecs = self.sort_by_branch(eigvals, eigvecs)
+
+        return eigvals, eigvecs
+
+    def _xyz_classify_modes(self, evecs):
+        """Classify modes by their dominant spatial direction (x, y, z)."""
+        N_coords, N_modes = np.shape(evecs)
+        N_ions = N_coords // 6
+        classifier = np.zeros(N_modes, dtype=int)
+        for mode_index in range(N_modes):
+            x_score = np.sum(np.abs(evecs[0:N_ions, mode_index])) + np.sum(np.abs(evecs[3*N_ions:4*N_ions, mode_index]))
+            y_score = np.sum(np.abs(evecs[N_ions:2*N_ions, mode_index])) + np.sum(np.abs(evecs[4*N_ions:5*N_ions, mode_index]))
+            z_score = np.sum(np.abs(evecs[2*N_ions:3*N_ions, mode_index])) + np.sum(np.abs(evecs[5*N_ions:6*N_ions, mode_index]))
+            scores = [x_score, y_score, z_score]
+            max_index = np.argmax(scores)
+            classifier[mode_index] = max_index
+        for direction in range(3):
+            count = np.sum(classifier == direction)
+            assert count == N_modes // 3, f"Classification error: direction {direction} has {count} modes, expected {N_modes // 3}"
+        return classifier
+
+    def sort_by_branch(self, evals, evecs):
+        """Sorts the eigenvalues, eigenvectors by mode branch (radial x, radial y, axial (z)).
+
+        Only makes sense for linear chain configurations where Hessian is block diagonal.
+        """
+        classifier = self._xyz_classify_modes(evecs)
+        # within each branch, sort by frequency
+        N_ions = len(evals) // 3
+        sorted_by_branch_evals = np.zeros_like(evals)
+        sorted_by_branch_evecs = np.zeros_like(evecs)
+        for direction in range(3):
+            direction_indices = np.where(classifier == direction)[0]
+            # they are already sorted by frequency from the original mode analysis code, so we can just take them in order
+            sorted_by_branch_evals[direction*N_ions:(direction+1)*N_ions] = evals[direction_indices]
+            sorted_by_branch_evecs[:, direction*N_ions:(direction+1)*N_ions] = evecs[:, direction_indices]
+        return sorted_by_branch_evals, sorted_by_branch_evecs 
 
 
     def calculate_mode_participation_factors(self) -> Matrix:
@@ -524,9 +600,6 @@ def calculate_mode_participation_factors(self) -> Matrix:
     # Can get lamb-dicke parameters from wavevctor \dot mode_participation_factors[:, i, m]
     ### -- Coordinate systems must be the same / correspond. 
 
-
-
-
     # Derived properties 
     def compute_reference_single_ion_lamb_dicke_factors(self, wavenumber: float) -> (float, float, float):
         """ Computes analytical Lamb-Dicke parameter by eta = k * sqrt(hbar / m omega) for an ion in a light-field with wavevector |k| 
@@ -540,121 +613,25 @@ def compute_reference_single_ion_lamb_dicke_factors(self, wavenumber: float) ->
         mass = self.atomic_masses[0] 
         eta_x, eta_y, eta_z = wavenumber * np.sqrt(const.hbar / (2 * mass * trap_frequencies))
         return eta_x, eta_y, eta_z    # ignore the phase
- 
-    
-#classes
-class GeneralizedModeAnalysisWithBranchSortedModes(TrappedIonModeAnalysis):
- 
-    def _xyz_classify_modes(self, evecs):
-        N_coords, N_modes = np.shape(evecs)
-        N_ions = N_coords // 6
-        classifier = np.zeros(N_modes, dtype=int)
-        for mode_index in range(N_modes):
-            x_score = np.sum(np.abs(evecs[0:N_ions, mode_index])) + np.sum(np.abs(evecs[3*N_ions:4*N_ions, mode_index]))
-            y_score = np.sum(np.abs(evecs[N_ions:2*N_ions, mode_index])) + np.sum(np.abs(evecs[4*N_ions:5*N_ions, mode_index]))
-            z_score = np.sum(np.abs(evecs[2*N_ions:3*N_ions, mode_index])) + np.sum(np.abs(evecs[5*N_ions:6*N_ions, mode_index]))
-            scores = [x_score, y_score, z_score]
-            max_index = np.argmax(scores)
-            classifier[mode_index] = max_index
-        for direction in range(3):
-            count = np.sum(classifier == direction)
-            assert count == N_modes // 3, f"Classification error: direction {direction} has {count} modes, expected {N_modes // 3}"
-        return classifier
- 
-    def sort_by_branch(self, evals, evecs):
-        """ Sorts the eigenvalues, eigenvectors by mode branch (radial x, radial y, axial (z)) """
-        # Hessian block diagonal H_xx, H_yy, H_zz non-zero for linear chain 
-        ## only makes sense for linear chain 
-        classifier = self._xyz_classify_modes(evecs)
-        # within each branch, sort by frequency
-        N_ions = len(evals) // 3
-        sorted_by_branch_evals = np.zeros_like(evals)
-        sorted_by_branch_evecs = np.zeros_like(evecs)
-        for direction in range(3):
-            direction_indices = np.where(classifier == direction)[0]
-            # they are already sorted by frequency from the original mode analysis code, so we can just take them in order
-            sorted_by_branch_evals[direction*N_ions:(direction+1)*N_ions] = evals[direction_indices]
-            sorted_by_branch_evecs[:, direction*N_ions:(direction+1)*N_ions] = evecs[:, direction_indices]
-        return sorted_by_branch_evals, sorted_by_branch_evecs
 
-    # Override from parent  
-    def organize_modes(self, eigvals, eigvecs):        
-        eigvals, eigvecs = self.sort_modes(eigvals, eigvecs)
-        eigvals, eigvecs = self.split_modes(eigvals, eigvecs)
-        eigvals, eigvecs = self.sort_by_branch(eigvals, eigvecs)
-        return eigvals, eigvecs
- 
- 
-    def reindex_ions_by_z(self): 
-        # based on the position along z, lowest i, is 0, up to N - 1
-        x,y,z = self.ion_coordinates_from_flattened(self.equilibrium_positions)
- 
-        idx = np.argsort(z)
-        self.equilibrium_positions = np.hstack((x[idx], y[idx], z[idx]))
- 
-        # Reindex all other arrays accordingly
-        self.m = self.m[idx]
-        self.q = self.q[idx]
-        self.atomic_masses = self.atomic_masses[idx]
-        self.nuclear_charges = self.nuclear_charges[idx]
-        self.atomic_numbers = self.atomic_numbers[idx]
-
-        # trapping frequencies
-        self.omega_x = self.omega_x[idx]
-        self.omega_y = self.omega_y[idx]
-        self.omega_z = self.omega_z[idx]
-        # all ions are the same so this is safe. TODO: When does this change? 
-        self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
- 
-    #def run(self):
-    #def run_normal_mode_analysis()
-    def solve_ion_trap_equilibrium(self): 
-        """ Diagonalizes the Coulomb + harmonic trap Hamiltonian for a system of ions. """
-        #self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
-        self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
-        
-        self.equilibrium_positions = self.solve_for_equilibrium_positions()
-        self.reindex_ions_by_z()
-        mass_matrix = self.build_mass_matrix(self.m)  
-        E_matrix = self.build_E_matrix(self.equilibrium_positions, mass_matrix)  
-        T_matrix = self.build_momentum_transform_matrix(mass_matrix)
-        H_matrix = self.compute_H_matrix(T_matrix, E_matrix)   
-
-        self.eigvals, self.eigvecs = self.calculate_normal_modes(H_matrix)
-        self.eigvecs_vel = self.get_eigenvectors_xv_coords(T_matrix, self.eigvecs)    
-        self.check_for_zero_modes()
-        self.check_outer_relation(H_matrix)
-
-        # S matrix may be important; it's how you convert from x, p coordinates to normal mode coordinates 
-        S_matrix = self.get_canonical_transformation(H_matrix, self.eigvecs, self.eigvals) 
-        self.check_diagonalization(T_matrix, S_matrix, E_matrix)
-        #self.hasrun = True  
-
-
-# Extra methods?
     def return_equilibrium_positions(self, dimensionless: bool) -> Vector:
-        """ Return equilibrium positions in either dimensionless or SI units (inverse meters) """
+        """Return equilibrium positions in either dimensionless or SI units (inverse meters)."""
         if dimensionless:
             return self.equilibrium_positions
         else:
             return self.equilibrium_positions * self.characteristic_parameters['length']
 
-
     @classmethod
-    def from_species(cls, species_name: str, num_ions: int, omega_x: float, omega_y: float, omega_z: float): 
-        """ Build the mode analysis class from a species name, corresponding to a system of N ions under harmonic trapping. """ 
-        # Import necessary data for the species 
+    def from_species(cls, species_name: str, num_ions: int, omega_x: float, omega_y: float, omega_z: float):
+        """Build the mode analysis class from a species name, corresponding to a system of N ions under harmonic trapping."""
+        # Import necessary data for the species
         species_data = AtomicSpin.get_config_data(species_name)
         atomic_mass = species_data['mass']
         atomic_number = species_data['Z']
-        # Construct the class 
+        # Construct the class
         return cls(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_number)
 
-
-    basis = StandardBasis([*spins])
-
     @classmethod
-    #def from_atomic_spin_basis(cls, atomic_structure_basis: StandardBasis, omega_x: float, omega_y: float, omega_z: float): 
     def from_atomic_spin_basis(cls, spins: list[degree_of_freedom], omega_x: float, omega_y: float, omega_z: float): 
         """ Build the mode analysis class from a basis of AtomicSpin degrees of freedom under harmonic trapping. """ 
         # Extract number of ions and the mass and atomic number from the DOF in the basis 
@@ -675,69 +652,188 @@ def from_atomic_spin_basis(cls, spins: list[degree_of_freedom], omega_x: float,
 
 
     def build_mode_DOFs(self, mode_indices: list[int], fock_dimensions: Vector | int) -> list[MotionalMode]:
-        """ Builds and returns an IonSim Motional Degree of Freedom.
+        """Builds and returns an IonSim Motional Degree of Freedom.
 
-            - Applies each fock dimension to each mode, or applies the same fock dimension to all the modes  
+            - Applies each fock dimension to each mode, or applies the same fock dimension to all the modes
         """
         modes = []
-        # Convert list of Fock dimensions to an array if it's not already an array 
+        # Convert list of Fock dimensions to an array if it's not already an array
         fock_dimensions = self.convert_to_array(fock_dimensions).astype(int)
         for idx, fock_dim in zip(mode_indices, fock_dimensions):
-            mode_index = idx # or some function of this index 
+            mode_index = idx # or some function of this index
             modes.append(MotionalMode.from_frequency(self.eigvals[mode_index], fock_dim))
 
-        return modes 
+        return modes
+
+    
+#classes
+# GeneralizedModeAnalysisWithBranchSortedModes has been consolidated into TrappedIonModeAnalysis
+# with configurable mode_organization and reindexing_strategy parameters.
+
+ #class LinearIonChainAnalysis(TrappedIonModeAnalysis):
+ # 
+ #    def _xyz_classify_modes(self, evecs):
+ #        N_coords, N_modes = np.shape(evecs)
+ #        N_ions = N_coords // 6
+ #        classifier = np.zeros(N_modes, dtype=int)
+ #        for mode_index in range(N_modes):
+ #            x_score = np.sum(np.abs(evecs[0:N_ions, mode_index])) + np.sum(np.abs(evecs[3*N_ions:4*N_ions, mode_index]))
+ #            y_score = np.sum(np.abs(evecs[N_ions:2*N_ions, mode_index])) + np.sum(np.abs(evecs[4*N_ions:5*N_ions, mode_index]))
+ #            z_score = np.sum(np.abs(evecs[2*N_ions:3*N_ions, mode_index])) + np.sum(np.abs(evecs[5*N_ions:6*N_ions, mode_index]))
+ #            scores = [x_score, y_score, z_score]
+ #            max_index = np.argmax(scores)
+ #            classifier[mode_index] = max_index
+ #        for direction in range(3):
+ #            count = np.sum(classifier == direction)
+ #            assert count == N_modes // 3, f"Classification error: direction {direction} has {count} modes, expected {N_modes // 3}"
+ #        return classifier
+ #
+ #    def sort_by_branch(self, evals, evecs):
+ #        """ Sorts the eigenvalues, eigenvectors by mode branch (radial x, radial y, axial (z)) """
+ #        # Hessian block diagonal H_xx, H_yy, H_zz non-zero for linear chain 
+ #        ## only makes sense for linear chain 
+ #        classifier = self._xyz_classify_modes(evecs)
+ #        # within each branch, sort by frequency
+ #        N_ions = len(evals) // 3
+ #        sorted_by_branch_evals = np.zeros_like(evals)
+ #        sorted_by_branch_evecs = np.zeros_like(evecs)
+ #        for direction in range(3):
+ #            direction_indices = np.where(classifier == direction)[0]
+ #            # they are already sorted by frequency from the original mode analysis code, so we can just take them in order
+ #            sorted_by_branch_evals[direction*N_ions:(direction+1)*N_ions] = evals[direction_indices]
+ #            sorted_by_branch_evecs[:, direction*N_ions:(direction+1)*N_ions] = evecs[:, direction_indices]
+ #        return sorted_by_branch_evals, sorted_by_branch_evecs
+ #
+ #    # Override from parent  
+ #    def organize_modes(self, eigvals, eigvecs):        
+ #        eigvals, eigvecs = self.sort_modes(eigvals, eigvecs)
+ #        eigvals, eigvecs = self.split_modes(eigvals, eigvecs)
+ #        eigvals, eigvecs = self.sort_by_branch(eigvals, eigvecs)
+ #        return eigvals, eigvecs
+ # 
+ # 
+ #    def reindex_ions_by_z(self): 
+ #        # based on the position along z, lowest i, is 0, up to N - 1
+ #        x,y,z = self.ion_coordinates_from_flattened(self.equilibrium_positions)
+ # 
+ #        idx = np.argsort(z)
+ #        self.equilibrium_positions = np.hstack((x[idx], y[idx], z[idx]))
+ # 
+ #        # Reindex all other arrays accordingly
+ #        self.m = self.m[idx]
+ #        self.q = self.q[idx]
+ #        self.atomic_masses = self.atomic_masses[idx]
+ #        self.nuclear_charges = self.nuclear_charges[idx]
+ #        self.atomic_numbers = self.atomic_numbers[idx]
+ #
+ #        # trapping frequencies
+ #        self.omega_x = self.omega_x[idx]
+ #        self.omega_y = self.omega_y[idx]
+ #        self.omega_z = self.omega_z[idx]
+ #        # all ions are the same so this is safe. TODO: When does this change? 
+ #        self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
+ # 
+ #    def solve_ion_trap_equilibrium(self): 
+ #        """ Diagonalizes the Coulomb + harmonic trap Hamiltonian for a system of ions. """
+ #        self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
+ #        
+ #        self.equilibrium_positions = self.solve_for_equilibrium_positions()
+ #        self.reindex_ions_by_z()
+ #        mass_matrix = self.build_mass_matrix(self.m)  
+ #        E_matrix = self.build_E_matrix(self.equilibrium_positions, mass_matrix)  
+ #        T_matrix = self.build_momentum_transform_matrix(mass_matrix)
+ #        H_matrix = self.compute_H_matrix(T_matrix, E_matrix)   
+ #
+ #        self.eigvals, self.eigvecs = self.calculate_normal_modes(H_matrix)
+ #        self.eigvecs_vel = self.get_eigenvectors_xv_coords(T_matrix, self.eigvecs)    
+ #        self.check_for_zero_modes()
+ #        self.check_outer_relation(H_matrix)
+ #
+ #        # S matrix may be important; it's how you convert from x, p coordinates to normal mode coordinates 
+ #        S_matrix = self.get_canonical_transformation(H_matrix, self.eigvecs, self.eigvals) 
+ #        self.check_diagonalization(T_matrix, S_matrix, E_matrix)
+ #        #self.hasrun = True  
+
 
+# Extra methods?
+ #    def return_equilibrium_positions(self, dimensionless: bool) -> Vector:
+ #        """ Return equilibrium positions in either dimensionless or SI units (inverse meters) """
+ #        if dimensionless:
+ #            return self.equilibrium_positions
+ #        else:
+ #            return self.equilibrium_positions * self.characteristic_parameters['length']
+ #
+
+ #    @classmethod
+ #    def from_species(cls, species_name: str, num_ions: int, omega_x: float, omega_y: float, omega_z: float): 
+ #        """ Build the mode analysis class from a species name, corresponding to a system of N ions under harmonic trapping. """ 
+ #        # Import necessary data for the species 
+ #        species_data = AtomicSpin.get_config_data(species_name)
+ #        atomic_mass = species_data['mass']
+ #        atomic_number = species_data['Z']
+ #        # Construct the class 
+ #        return cls(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_number)
+
+
+ #    def build_mode_DOFs(self, mode_indices: list[int], fock_dimensions: Vector | int) -> list[MotionalMode]:
+ #        """ Builds and returns an IonSim Motional Degree of Freedom.
+ #
+ #            - Applies each fock dimension to each mode, or applies the same fock dimension to all the modes  
+ #        """
+ #        modes = []
+ #        # Convert list of Fock dimensions to an array if it's not already an array 
+ #        fock_dimensions = self.convert_to_array(fock_dimensions).astype(int)
+ #        for idx, fock_dim in zip(mode_indices, fock_dimensions):
+ #            mode_index = idx # or some function of this index 
+ #            modes.append(MotionalMode.from_frequency(self.eigvals[mode_index], fock_dim))
+ #
+ #        return modes 
+ #
 
 
 #You could redefine the equilibrium finding function to assert that the equilibrium is linear. For example, make a wrapper inside the function for the potential, Jacobian, and Hessian that forces x_i and y_i = 0. 
  
 #This is the code for the Lamb-Dicke parameters: (not that overall phases don't matter here, but the relative phase does. We could pin down the phase with some convention like, "each mode's first non-negative LD value is defined positive.")
-        
- 
-
 
-class LinearIonChainAnalysis(GeneralizedModeAnalysisWithBranchSortedModes):
-    """
-    Subclass for setting up and analyzing linear ion chains.
+class LinearIonChainAnalysis(TrappedIonModeAnalysis):
+    """ Subclass for setting up and analyzing linear ion chains.
 
-    In a linear ion chain, ions are typically aligned along the z-axis (axial direction),
-    with minimal displacement in the x and y directions (radial directions).
+        In a linear ion chain, ions are typically aligned along the z-axis (axial direction),
+        with minimal displacement in the x and y directions (radial directions).
 
-    This class extends GeneralizedModeAnalysisWithBranchSortedModes to:
-    1. Enforce linear chain configuration
-    2. Provide specialized methods for linear chain analysis
-    3. Offer convenience methods for common linear chain scenarios
+        This class accomplishes:     
+            1. Enforce linear chain configuration
+            2. Provide specialized methods for linear chain analysis
+            3. Offer convenience methods for common linear chain scenarios
     """
 
-    def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float,
-                 atomic_masses: np.ndarray | float, atomic_numbers: np.ndarray | int):
-        """
-        Initialize a linear ion chain analysis.
-
-        Parameters:
-        -----------
-        num_ions : int
-            Number of ions in the chain
-        omega_x, omega_y, omega_z : float
-            Trap frequencies in x, y, and z directions (rad/s)
-        atomic_masses : array or float
-            Atomic masses in amu (atomic mass units)
-        atomic_numbers : array or int
-            Atomic numbers (number of protons)
-        """
-        super().__init__(num_ions, omega_x, omega_y, omega_z, atomic_masses, atomic_numbers)
+ #    def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float,
+ #                 atomic_masses: np.ndarray | float, atomic_numbers: np.ndarray | int):
+ #        """
+ #        Initialize a linear ion chain analysis.
+ #
+ #        Parameters:
+ #        -----------
+ #        num_ions : int
+ #            Number of ions in the chain
+ #        omega_x, omega_y, omega_z : float
+ #            Trap frequencies in x, y, and z directions (rad/s)
+ #        atomic_masses : array or float
+ #            Atomic masses in amu (atomic mass units)
+ #        atomic_numbers : array or int
+ #            Atomic numbers (number of protons)
+ #        """
+ #        super().__init__(num_ions, omega_x, omega_y, omega_z, atomic_masses, atomic_numbers)
 
     def solve_ion_trap_equilibrium(self):
-        """
-        Solve for equilibrium positions and analyze normal modes for a linear chain.
-
-        This method:
-        1. Sets up dimensionless parameters
-        2. Solves for equilibrium positions
-        3. Reindexes ions by their z-position (closest to center first)
-        4. Computes normal modes with branch sorting
-        5. Performs validation checks
+        """ Solve for equilibrium positions and analyze normal modes for a linear chain.
+
+            This method:
+            1. Sets up dimensionless parameters
+            2. Solves for equilibrium positions
+            3. Reindexes ions by their z-position (closest to center first)
+            4. Computes normal modes with branch sorting
+            5. Performs validation checks
         """
         # Set up dimensionless parameters using first ion's properties and axial trap frequency
         self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
@@ -766,15 +862,8 @@ def solve_ion_trap_equilibrium(self):
         S_matrix = self.get_canonical_transformation(H_matrix, self.eigvecs, self.eigvals)
         self.check_diagonalization(T_matrix, S_matrix, E_matrix)
 
-    def get_axial_modes(self):
-        """
-        Get the axial (z-direction) normal modes.
-
-        Returns:
-        --------
-        tuple
-            (axial_eigenvalues, axial_eigenvectors) for the axial modes
-        """
+    def get_axial_modes(self) -> tuple:
+        """ Returns eigenvalues, eigenvectors for the axial normal modes """  
         # Axial modes are the last third of the modes (after branch sorting)
         n_modes_per_branch = self.num_ions
         axial_start = 2 * n_modes_per_branch
@@ -785,19 +874,11 @@ def get_axial_modes(self):
 
         return axial_eigvals, axial_eigvecs
 
-    def get_radial_modes(self, direction='x'):
-        """
-        Get the radial normal modes in the specified direction.
-
-        Parameters:
-        -----------
-        direction : str
-            'x' or 'y' for radial direction
+    def get_radial_modes(self, direction='x') -> tuple:
+        """ Returns eigenvalues, eigenvectors for the radial normal modes in the specified direction.
 
-        Returns:
-        --------
-        tuple
-            (radial_eigenvalues, radial_eigenvectors) for the specified radial modes
+            Parameters:
+                direction : str; 'x' or 'y' for radial direction
         """
         if direction not in ['x', 'y']:
             raise ValueError("Direction must be 'x' or 'y'")
@@ -816,14 +897,7 @@ def get_radial_modes(self, direction='x'):
         return radial_eigvals, radial_eigvecs
 
     def get_ion_spacing(self):
-        """
-        Get the spacing between consecutive ions along the z-axis.
-
-        Returns:
-        --------
-        np.ndarray
-            Array of distances between consecutive ions (in meters)
-        """
+        """ Returns an array of ion spacings between consecutive ions along the z-axis. """
         # Get z positions (dimensionful)
         x, y, z = self.ion_coordinates_from_flattened(self.equilibrium_positions)
         z_dimensionful = z * self.characteristic_parameters['length']
@@ -834,15 +908,8 @@ def get_ion_spacing(self):
 
         return ion_spacing
 
-    def get_center_of_mass_position(self):
-        """
-        Get the center of mass position of the ion chain.
-
-        Returns:
-        --------
-        tuple
-            (com_x, com_y, com_z) center of mass coordinates in meters
-        """
+    def get_center_of_mass_position(self) -> tuple:
+        """ Get the center of mass position of the ion chain. """
         # Get positions (dimensionful)
         x, y, z = self.ion_coordinates_from_flattened(self.equilibrium_positions)
         x_dim = x * self.characteristic_parameters['length']
@@ -858,42 +925,21 @@ def get_center_of_mass_position(self):
         return com_x, com_y, com_z
 
     def get_axial_mode_frequencies(self):
-        """
-        Get the frequencies of the axial modes.
-
-        Returns:
-        --------
-        np.ndarray
-            Array of axial mode frequencies in rad/s
-        """
+        """ Returns the frequencies [rad/s] of the axial modes. """
         axial_eigvals, _ = self.get_axial_modes()
         return axial_eigvals * self.trap_freq_scale
 
     def get_radial_mode_frequencies(self, direction='x'):
-        """
-        Get the frequencies of the radial modes in the specified direction.
-
-        Parameters:
-        -----------
-        direction : str
-            'x' or 'y' for radial direction
-
-        Returns:
-        --------
-        np.ndarray
-            Array of radial mode frequencies in rad/s
+        """ Returns the frequencies of the radial modes in rad/s for the specified direction.
+            direction = 'x' or 'y' for radial direction
         """
         radial_eigvals, _ = self.get_radial_modes(direction)
         return radial_eigvals * self.trap_freq_scale
 
     def get_mode_participation_factors_by_branch(self):
-        """
-        Get mode participation factors organized by branch (x, y, z).
+        """ Returns a dictionary of mode participation factors organized by branch (x, y, z).
 
-        Returns:
-        --------
-        dict
-            Dictionary with keys 'x', 'y', 'z' containing mode participation factors
+            Output dictionary with keys 'x', 'y', 'z' containing mode participation factors
             for each branch. Each value is a 2D array of shape (num_ions, num_modes_per_branch)
         """
         # Get full mode participation factors
@@ -910,25 +956,19 @@ def get_mode_participation_factors_by_branch(self):
 
         return result
 
-    def calculate_lamb_dicke_parameters(self, wavevector: np.ndarray | float) -> dict:
-        """
-        Calculate Lamb-Dicke parameters for all ions and modes, organized by branch.
-
-        The Lamb-Dicke parameter η = k · Δr_0 represents the ratio of the spatial
-        extent of the ion's zero-point motion to the wavelength of the laser.
-
-        Parameters:
-        -----------
-        wavevector : np.ndarray or float
-            - If array: Full wavevector (kx, ky, kz) as a 3-element array in units of 1/m
-            - If float: Magnitude of the laser wavevector |k| in units of 1/m (for backward compatibility)
+    def calculate_lamb_dicke_parameters(self, wavevector: Vector) -> dict:
+        """ Calculate Lamb-Dicke parameters for all ions and modes, organized by branch.
 
-        Returns:
-        --------
-        dict
-            Dictionary with keys 'x', 'y', 'z' containing Lamb-Dicke parameters
-            for each branch. Each value is a 2D array of shape (num_ions, num_modes_per_branch)
-            representing η_{direction}[ion_index, mode_index]
+            The Lamb-Dicke parameter η = k · Δr_0 represents the ratio of the spatial
+            extent of the ion's zero-point motion to the wavelength of the laser.
+    
+            Parameters:
+                wavevector : k = (kx, ky, kz) as a 3-element array in units of 1/m
+    
+            Returns:
+                Dictionary with keys 'x', 'y', 'z' containing Lamb-Dicke parameters
+                for each branch. Each value is a 2D array of shape (num_ions, num_modes_per_branch)
+                representing η_{direction}[ion_index, mode_index]
         """
         # Get mode participation factors (which are proportional to zero-point motion)
         mode_pf_by_branch = self.get_mode_participation_factors_by_branch()
@@ -951,16 +991,15 @@ def calculate_lamb_dicke_parameters(self, wavevector: np.ndarray | float) -> dic
 
             # Compute Lamb-Dicke parameters as dot product: η = k · Δr_0
             # This gives us shape (num_ions, num_modes) for the total LD parameter
-            num_ions = self.num_ions
-            num_modes = 3 * num_ions
+            num_modes = 3 * self. num_ions
 
             # Reshape for dot product: (3, num_ions, num_modes) · (3,) -> (num_ions, num_modes)
-            total_ld_params = np.zeros((num_ions, num_modes), dtype=complex)
+            total_ld_params = np.zeros((self.num_ions, num_modes), dtype=complex)
             for i in range(3):
                 total_ld_params += wavevector[i] * full_mode_pf[i, :, :]
 
             # Organize by branch for consistency with scalar case
-            n_modes_per_branch = num_ions
+            n_modes_per_branch = self.num_ions
             lamb_dicke_parameters = {
                 'x': total_ld_params[:, :n_modes_per_branch],
                 'y': total_ld_params[:, n_modes_per_branch:2*n_modes_per_branch],
@@ -969,81 +1008,43 @@ def calculate_lamb_dicke_parameters(self, wavevector: np.ndarray | float) -> dic
 
         return lamb_dicke_parameters
 
-    def calculate_lamb_dicke_parameters_full(self, wavevector: np.ndarray | float) -> np.ndarray:
-        """
-        Calculate full Lamb-Dicke parameter matrix.
-
-        Parameters:
-        -----------
-        wavevector : np.ndarray or float
-            - If array: Full wavevector (kx, ky, kz) as a 3-element array in units of 1/m
-            - If float: Magnitude of the laser wavevector |k| in units of 1/m (for backward compatibility)
-
-        Returns:
-        --------
-        np.ndarray
-            - If scalar wavevector: Full matrix of shape (3, num_ions, num_modes)
-              where eta[direction, ion, mode] gives the component-wise LD parameter
-            - If vector wavevector: Total LD parameter matrix of shape (num_ions, num_modes)
-              where eta[ion, mode] gives the total LD parameter k · Δr_0
+    def calculate_lamb_dicke_parameters_full(self, wavevector: Vector) -> Matrix:
+        """ Calculate full Lamb-Dicke parameter matrix from a laser wavevector k = (kx, ky, kz) 
+
+            Returns: 
+                - Total LD parameter matrix of shape (num_ions, num_modes) where eta[ion, mode] 
+                    gives the total LD parameter k · Δr_0.
         """
         # Get full mode participation factors
         mode_pf = self.calculate_mode_participation_factors()
 
-        if isinstance(wavevector, (float, int)):
-            # Scalar case: component-wise multiplication (backward compatibility)
-            return wavevector * mode_pf
-        else:
-            # Vector case: compute dot product k · Δr_0 (physically accurate)
-            wavevector = np.asarray(wavevector)
-            if wavevector.shape != (3,):
-                raise ValueError("Wavevector must be a 3-element array (kx, ky, kz)")
+        # Vector case: compute dot product k · Δr_0 (physically accurate)
+        wavevector = np.asarray(wavevector)
+        if wavevector.shape != (3,):
+            raise ValueError(f"Wavevector must be a 3-element array (kx, ky, kz). Received {wavevector}")
 
-            # Compute total Lamb-Dicke parameter as dot product
-            # Shape: (3, num_ions, num_modes) · (3,) -> (num_ions, num_modes)
-            num_ions = self.num_ions
-            num_modes = 3 * num_ions
-            total_ld_params = np.zeros((num_ions, num_modes), dtype=complex)
+        # Compute total Lamb-Dicke parameter as dot product
+        # Shape: (3, num_ions, num_modes) · (3,) -> (num_ions, num_modes)
+        num_modes = 3 * self.num_ions
+        total_ld_params = np.zeros((self.num_ions, num_modes), dtype=complex)
 
-            for i in range(3):
-                total_ld_params += wavevector[i] * mode_pf[i, :, :]
+        for i in range(3):
+            total_ld_params += wavevector[i] * mode_pf[i, :, :]
 
-            return total_ld_params
+        return total_ld_params
 
-    def get_axial_lamb_dicke_parameters(self, wavevector: np.ndarray | float) -> np.ndarray:
-        """
-        Get Lamb-Dicke parameters for axial (z) modes only.
-
-        Parameters:
-        -----------
-        wavevector : np.ndarray or float
-            - If array: Full wavevector (kx, ky, kz) as a 3-element array in units of 1/m
-            - If float: Magnitude of the laser wavevector |k| in units of 1/m (for backward compatibility)
-
-        Returns:
-        --------
-        np.ndarray
-            Lamb-Dicke parameters for axial modes, shape (num_ions, num_axial_modes)
+    def get_axial_lamb_dicke_parameters(self, wavevector: Vector) -> Vector:
+        """ Get Lamb-Dicke parameters for axial (z) modes only.
+
+            Returns: Lamb-Dicke parameters for axial modes, NDArray of shape (num_ions, num_axial_modes)
         """
         lamb_dicke_by_branch = self.calculate_lamb_dicke_parameters(wavevector)
         return lamb_dicke_by_branch['z']
 
-    def get_radial_lamb_dicke_parameters(self, wavevector: np.ndarray | float, direction: str = 'x') -> np.ndarray:
-        """
-        Get Lamb-Dicke parameters for radial modes in specified direction.
-
-        Parameters:
-        -----------
-        wavevector : np.ndarray or float
-            - If array: Full wavevector (kx, ky, kz) as a 3-element array in units of 1/m
-            - If float: Magnitude of the laser wavevector |k| in units of 1/m (for backward compatibility)
-        direction : str
-            'x' or 'y' for radial direction
-
-        Returns:
-        --------
-        np.ndarray
-            Lamb-Dicke parameters for radial modes, shape (num_ions, num_radial_modes)
+    def get_radial_lamb_dicke_parameters(self, wavevector: Vector | float, direction: str = 'x') -> Vector:
+        """ Get Lamb-Dicke parameters for radial modes modes in a specified direction.
+
+            Returns: Lamb-Dicke parameters for radial modes, array of shape (num_ions, num_radial_modes)
         """
         if direction not in ['x', 'y']:
             raise ValueError("Direction must be 'x' or 'y'")
@@ -1052,14 +1053,7 @@ def get_radial_lamb_dicke_parameters(self, wavevector: np.ndarray | float, direc
         return lamb_dicke_by_branch[direction]
 
     def get_two_central_ion_separation(self) -> float:
-        """
-        Calculate the separation between the two central ions in the chain.
-
-        Returns:
-        --------
-        float
-            Distance between the two central ions in meters
-        """
+        """ Calculate the separation between the two central ions in the chain. Returns: Distance between the two central ions in meters """
         # Get z positions (dimensionful)
         x, y, z = self.ion_coordinates_from_flattened(self.equilibrium_positions)
         z_dimensionful = z * self.characteristic_parameters['length']
@@ -1073,20 +1067,12 @@ def get_two_central_ion_separation(self) -> float:
 
     def find_axial_frequency_for_desired_central_ion_separation(self, target_separation: float,
                                                                bounds: tuple = (0.1e6, 0.5e6)) -> float:
-        """
-        Find the axial frequency (wz) required to achieve a desired separation between central ions.
-
-        Parameters:
-        -----------
-        target_separation : float
-            Desired separation between central ions in meters
-        bounds : tuple
-            Tuple of (min_freq, max_freq) in Hz for the root finding algorithm
-
-        Returns:
-        --------
-        float
-            Axial frequency in rad/s that achieves the target separation
+        """ Computes the axial frequency (wz) [rad/s] required to achieve a desired separation between central ions.
+
+            Parameters:
+                - target_separation : float; Desired separation between central ions in meters
+                - bounds : tuple; (min_freq, max_freq) in Hz for the root finding algorithm
+
         """
         from scipy.optimize import root_scalar
 
@@ -1114,9 +1100,7 @@ def separation_error(wz_Hz):
         return result.root
 
     def print_chain_summary(self):
-        """
-        Print a summary of the linear ion chain configuration.
-        """
+        """ Print a summary of the linear ion chain configuration. """
         print(f"Linear Ion Chain Analysis Summary")
         print(f"Number of ions: {self.num_ions}")
         print(f"Trap frequencies: x={self.omega_x[0]/(2*np.pi)/1e6:.2f} MHz, "
@@ -1140,5 +1124,3 @@ def print_chain_summary(self):
         print(f"Radial X mode frequencies: {radial_x_freqs/(2*np.pi)/1e6:.2f} MHz")
         print(f"Radial Y mode frequencies: {radial_y_freqs/(2*np.pi)/1e6:.2f} MHz")
 
-
- 

From 89512c6cd61b01fbbb2a169d18b01b4b96324f26 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Tue, 16 Jun 2026 19:13:41 -0600
Subject: [PATCH 26/31]   add default values for linear chain class constructor

---
 src/ionsim/trapped_ion_mode_analysis.py | 23 ++++++-----------------
 1 file changed, 6 insertions(+), 17 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index dd14e13..7bf574e 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -807,23 +807,12 @@ class LinearIonChainAnalysis(TrappedIonModeAnalysis):
             3. Offer convenience methods for common linear chain scenarios
     """
 
- #    def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float,
- #                 atomic_masses: np.ndarray | float, atomic_numbers: np.ndarray | int):
- #        """
- #        Initialize a linear ion chain analysis.
- #
- #        Parameters:
- #        -----------
- #        num_ions : int
- #            Number of ions in the chain
- #        omega_x, omega_y, omega_z : float
- #            Trap frequencies in x, y, and z directions (rad/s)
- #        atomic_masses : array or float
- #            Atomic masses in amu (atomic mass units)
- #        atomic_numbers : array or int
- #            Atomic numbers (number of protons)
- #        """
- #        super().__init__(num_ions, omega_x, omega_y, omega_z, atomic_masses, atomic_numbers)
+    def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float,
+                 atomic_masses: np.ndarray | float, atomic_numbers: np.ndarray | int):
+        """ Initialize a linear ion chain analysis, defaulting to branch sorted modes. """
+        mode_organization = 'branched'
+        reindexing_strategy = 'z_axis'
+        super().__init__(num_ions, omega_x, omega_y, omega_z, atomic_masses, atomic_numbers, mode_organization)
 
     def solve_ion_trap_equilibrium(self):
         """ Solve for equilibrium positions and analyze normal modes for a linear chain.

From 17a0106321f22f92f15923245aa0bb254e68fcbe Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Wed, 17 Jun 2026 08:44:55 -0600
Subject: [PATCH 27/31]  update normal mode analysis with refactor and tests

---
 src/ionsim/__init__.py                  |  2 +-
 src/ionsim/trapped_ion_mode_analysis.py | 38 +++++++++----------------
 tests/test_mode_analysis.py             |  4 +--
 3 files changed, 17 insertions(+), 27 deletions(-)

diff --git a/src/ionsim/__init__.py b/src/ionsim/__init__.py
index e83f074..14a9a3c 100644
--- a/src/ionsim/__init__.py
+++ b/src/ionsim/__init__.py
@@ -30,4 +30,4 @@ def patched_kron(a, b, *args, **kwargs):
 from .zeeman_solver import ZeemanHyperfineSolver
 from .composite_operator import CompositeOperator
 from .dissipator import Dissipator, DissipatorSpontaneousEmission, Lindbladian 
-from .trapped_ion_mode_analysis import TrappedIonModeAnalysis 
+from .trapped_ion_mode_analysis import TrappedIonModeAnalysis, LinearIonChainAnalysis 
diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 7bf574e..08712d7 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -106,9 +106,6 @@ def convert_to_array(self, x: float | int | list | NDArray) -> NDArray:
             return x
         return np.array(x) 
 
-    #def dimensionless_parameters(self):
-    #def nondimensionalize_parameters(self):
-    #def convert_parameters_to_dimensionless(self, charge_scale: float=1., mass_scale: float=1., trap_freq_scale: float=1.):
     def set_up_dimensionless_parameeters(self, charge_scale: float=1., mass_scale: float=1., trap_freq_scale: float=1.):
         """ Compute dimensionless parameters, using first ion's properties and axial (z) trapping frequency. 
 
@@ -205,7 +202,17 @@ def check_diagonalization(self, T: Matrix, S: Matrix, E: Matrix) -> bool:
             print("has duplicate eigenvalues: ", has_duplicate_eigvals)
 
     def solve_ion_trap_equilibrium(self):
-        """ Diagonalizes the Coulomb + harmonic trap Hamiltonian for a system of ions. """
+        """ Solve for equilibrium positions and analyze normal modes for a linear chain by minizing the Coulomb + 
+                harmonic trap energy functional for a system of ions
+
+            This method:
+            1. Sets up dimensionless parameters
+            2. Solves for equilibrium positions
+            3. Reindexes ions by their z-position (closest to center first)
+            4. Computes normal modes with branch sorting
+            5. Performs validation checks
+
+        """
         # Convert to dimensionless units using axial trap frequency and first ion's mass and charge
         # TODO: take in input here or in class constructor for mass, charge, trap scales.
         #self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
@@ -603,9 +610,7 @@ def calculate_mode_participation_factors(self) -> Matrix:
     # Derived properties 
     def compute_reference_single_ion_lamb_dicke_factors(self, wavenumber: float) -> (float, float, float):
         """ Computes analytical Lamb-Dicke parameter by eta = k * sqrt(hbar / m omega) for an ion in a light-field with wavevector |k| 
-
             - wavenumber: wavevector magnitude |k| in units of 1 / m  
-
         """
         # TODO: better way to handle units? 
         # Convention to use first value of trap frequency arrays (representing one of the ions)  
@@ -649,8 +654,6 @@ def from_atomic_spin_basis(cls, spins: list[degree_of_freedom], omega_x: float,
         # Construct the class 
         return cls(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_number)
 
-
-
     def build_mode_DOFs(self, mode_indices: list[int], fock_dimensions: Vector | int) -> list[MotionalMode]:
         """Builds and returns an IonSim Motional Degree of Freedom.
 
@@ -815,15 +818,7 @@ def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float
         super().__init__(num_ions, omega_x, omega_y, omega_z, atomic_masses, atomic_numbers, mode_organization)
 
     def solve_ion_trap_equilibrium(self):
-        """ Solve for equilibrium positions and analyze normal modes for a linear chain.
-
-            This method:
-            1. Sets up dimensionless parameters
-            2. Solves for equilibrium positions
-            3. Reindexes ions by their z-position (closest to center first)
-            4. Computes normal modes with branch sorting
-            5. Performs validation checks
-        """
+        """ Solve for equilibrium positions and analyze normal modes for a linear chain. It reindexes ions by their axial position. """
         # Set up dimensionless parameters using first ion's properties and axial trap frequency
         self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
 
@@ -865,10 +860,7 @@ def get_axial_modes(self) -> tuple:
 
     def get_radial_modes(self, direction='x') -> tuple:
         """ Returns eigenvalues, eigenvectors for the radial normal modes in the specified direction.
-
-            Parameters:
-                direction : str; 'x' or 'y' for radial direction
-        """
+                - direction : str; 'x' or 'y' for radial direction """
         if direction not in ['x', 'y']:
             raise ValueError("Direction must be 'x' or 'y'")
 
@@ -919,9 +911,7 @@ def get_axial_mode_frequencies(self):
         return axial_eigvals * self.trap_freq_scale
 
     def get_radial_mode_frequencies(self, direction='x'):
-        """ Returns the frequencies of the radial modes in rad/s for the specified direction.
-            direction = 'x' or 'y' for radial direction
-        """
+        """ Returns the frequencies of the radial modes in rad/s for the specified direction ('x' or 'y'). """
         radial_eigvals, _ = self.get_radial_modes(direction)
         return radial_eigvals * self.trap_freq_scale
 
diff --git a/tests/test_mode_analysis.py b/tests/test_mode_analysis.py
index 07c5ad2..60c9c7b 100644
--- a/tests/test_mode_analysis.py
+++ b/tests/test_mode_analysis.py
@@ -1,7 +1,7 @@
 import unittest
 import numpy as np
 from scipy.linalg import expm
-from ionsim.trapped_ion_mode_analysis import TrappedIonModeAnalysis, LinearIonChainAnalysis, GeneralizedModeAnalysisWithBranchSortedModes 
+from ionsim.trapped_ion_mode_analysis import TrappedIonModeAnalysis, LinearIonChainAnalysis
 from ionsim.testing import assert_array_close
 
 class TestModeAnalysis(unittest.TestCase):
@@ -194,7 +194,7 @@ def test_linear_ion_chain_analysis(self):
     print("="*60)
 
     # Create a 4-ion analysis
-    four_ion_analysis = GeneralizedModeAnalysisWithBranchSortedModes(
+    four_ion_analysis = TrappedIonModeAnalysis(
         num_ions=4,
         omega_x=2*np.pi*11e6,
         omega_y=2*np.pi*10e6,

From 3ee867e56f4ffb2faad25c991560d829907500a3 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Wed, 17 Jun 2026 11:43:52 -0600
Subject: [PATCH 28/31]   rm old code that's been incorporated into parent
 class

---
 src/ionsim/trapped_ion_mode_analysis.py | 328 ++++++++++++------------
 1 file changed, 159 insertions(+), 169 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 08712d7..97af312 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -215,9 +215,7 @@ def solve_ion_trap_equilibrium(self):
         """
         # Convert to dimensionless units using axial trap frequency and first ion's mass and charge
         # TODO: take in input here or in class constructor for mass, charge, trap scales.
-        #self.convert_parameters_to_dimensionless(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
         self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
-
         self.equilibrium_positions = self.solve_for_equilibrium_positions()
 
         # Use configurable reindexing strategy
@@ -560,10 +558,8 @@ def _xyz_classify_modes(self, evecs):
         return classifier
 
     def sort_by_branch(self, evals, evecs):
-        """Sorts the eigenvalues, eigenvectors by mode branch (radial x, radial y, axial (z)).
-
-        Only makes sense for linear chain configurations where Hessian is block diagonal.
-        """
+        """ Sorts the eigenvalues, eigenvectors by mode branch (radial x, radial y, axial (z)).
+            Only makes sense for linear chain configurations where Hessian is block diagonal. """
         classifier = self._xyz_classify_modes(evecs)
         # within each branch, sort by frequency
         N_ions = len(evals) // 3
@@ -640,7 +636,6 @@ def from_species(cls, species_name: str, num_ions: int, omega_x: float, omega_y:
     def from_atomic_spin_basis(cls, spins: list[degree_of_freedom], omega_x: float, omega_y: float, omega_z: float): 
         """ Build the mode analysis class from a basis of AtomicSpin degrees of freedom under harmonic trapping. """ 
         # Extract number of ions and the mass and atomic number from the DOF in the basis 
-        #DOFs = atomic_structure_basis.degrees_of_freedom
         DOFs = spins 
         num_ions = len(DOFs)
         atomic_masses = []
@@ -668,134 +663,80 @@ def build_mode_DOFs(self, mode_indices: list[int], fock_dimensions: Vector | int
 
         return modes
 
-    
-#classes
-# GeneralizedModeAnalysisWithBranchSortedModes has been consolidated into TrappedIonModeAnalysis
-# with configurable mode_organization and reindexing_strategy parameters.
-
- #class LinearIonChainAnalysis(TrappedIonModeAnalysis):
- # 
- #    def _xyz_classify_modes(self, evecs):
- #        N_coords, N_modes = np.shape(evecs)
- #        N_ions = N_coords // 6
- #        classifier = np.zeros(N_modes, dtype=int)
- #        for mode_index in range(N_modes):
- #            x_score = np.sum(np.abs(evecs[0:N_ions, mode_index])) + np.sum(np.abs(evecs[3*N_ions:4*N_ions, mode_index]))
- #            y_score = np.sum(np.abs(evecs[N_ions:2*N_ions, mode_index])) + np.sum(np.abs(evecs[4*N_ions:5*N_ions, mode_index]))
- #            z_score = np.sum(np.abs(evecs[2*N_ions:3*N_ions, mode_index])) + np.sum(np.abs(evecs[5*N_ions:6*N_ions, mode_index]))
- #            scores = [x_score, y_score, z_score]
- #            max_index = np.argmax(scores)
- #            classifier[mode_index] = max_index
- #        for direction in range(3):
- #            count = np.sum(classifier == direction)
- #            assert count == N_modes // 3, f"Classification error: direction {direction} has {count} modes, expected {N_modes // 3}"
- #        return classifier
- #
- #    def sort_by_branch(self, evals, evecs):
- #        """ Sorts the eigenvalues, eigenvectors by mode branch (radial x, radial y, axial (z)) """
- #        # Hessian block diagonal H_xx, H_yy, H_zz non-zero for linear chain 
- #        ## only makes sense for linear chain 
- #        classifier = self._xyz_classify_modes(evecs)
- #        # within each branch, sort by frequency
- #        N_ions = len(evals) // 3
- #        sorted_by_branch_evals = np.zeros_like(evals)
- #        sorted_by_branch_evecs = np.zeros_like(evecs)
- #        for direction in range(3):
- #            direction_indices = np.where(classifier == direction)[0]
- #            # they are already sorted by frequency from the original mode analysis code, so we can just take them in order
- #            sorted_by_branch_evals[direction*N_ions:(direction+1)*N_ions] = evals[direction_indices]
- #            sorted_by_branch_evecs[:, direction*N_ions:(direction+1)*N_ions] = evecs[:, direction_indices]
- #        return sorted_by_branch_evals, sorted_by_branch_evecs
- #
- #    # Override from parent  
- #    def organize_modes(self, eigvals, eigvecs):        
- #        eigvals, eigvecs = self.sort_modes(eigvals, eigvecs)
- #        eigvals, eigvecs = self.split_modes(eigvals, eigvecs)
- #        eigvals, eigvecs = self.sort_by_branch(eigvals, eigvecs)
- #        return eigvals, eigvecs
- # 
- # 
- #    def reindex_ions_by_z(self): 
- #        # based on the position along z, lowest i, is 0, up to N - 1
- #        x,y,z = self.ion_coordinates_from_flattened(self.equilibrium_positions)
- # 
- #        idx = np.argsort(z)
- #        self.equilibrium_positions = np.hstack((x[idx], y[idx], z[idx]))
- # 
- #        # Reindex all other arrays accordingly
- #        self.m = self.m[idx]
- #        self.q = self.q[idx]
- #        self.atomic_masses = self.atomic_masses[idx]
- #        self.nuclear_charges = self.nuclear_charges[idx]
- #        self.atomic_numbers = self.atomic_numbers[idx]
- #
- #        # trapping frequencies
- #        self.omega_x = self.omega_x[idx]
- #        self.omega_y = self.omega_y[idx]
- #        self.omega_z = self.omega_z[idx]
- #        # all ions are the same so this is safe. TODO: When does this change? 
- #        self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
- # 
- #    def solve_ion_trap_equilibrium(self): 
- #        """ Diagonalizes the Coulomb + harmonic trap Hamiltonian for a system of ions. """
- #        self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
- #        
- #        self.equilibrium_positions = self.solve_for_equilibrium_positions()
- #        self.reindex_ions_by_z()
- #        mass_matrix = self.build_mass_matrix(self.m)  
- #        E_matrix = self.build_E_matrix(self.equilibrium_positions, mass_matrix)  
- #        T_matrix = self.build_momentum_transform_matrix(mass_matrix)
- #        H_matrix = self.compute_H_matrix(T_matrix, E_matrix)   
- #
- #        self.eigvals, self.eigvecs = self.calculate_normal_modes(H_matrix)
- #        self.eigvecs_vel = self.get_eigenvectors_xv_coords(T_matrix, self.eigvecs)    
- #        self.check_for_zero_modes()
- #        self.check_outer_relation(H_matrix)
- #
- #        # S matrix may be important; it's how you convert from x, p coordinates to normal mode coordinates 
- #        S_matrix = self.get_canonical_transformation(H_matrix, self.eigvecs, self.eigvals) 
- #        self.check_diagonalization(T_matrix, S_matrix, E_matrix)
- #        #self.hasrun = True  
-
-
-# Extra methods?
- #    def return_equilibrium_positions(self, dimensionless: bool) -> Vector:
- #        """ Return equilibrium positions in either dimensionless or SI units (inverse meters) """
- #        if dimensionless:
- #            return self.equilibrium_positions
- #        else:
- #            return self.equilibrium_positions * self.characteristic_parameters['length']
- #
-
- #    @classmethod
- #    def from_species(cls, species_name: str, num_ions: int, omega_x: float, omega_y: float, omega_z: float): 
- #        """ Build the mode analysis class from a species name, corresponding to a system of N ions under harmonic trapping. """ 
- #        # Import necessary data for the species 
- #        species_data = AtomicSpin.get_config_data(species_name)
- #        atomic_mass = species_data['mass']
- #        atomic_number = species_data['Z']
- #        # Construct the class 
- #        return cls(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_number)
-
-
- #    def build_mode_DOFs(self, mode_indices: list[int], fock_dimensions: Vector | int) -> list[MotionalMode]:
- #        """ Builds and returns an IonSim Motional Degree of Freedom.
- #
- #            - Applies each fock dimension to each mode, or applies the same fock dimension to all the modes  
- #        """
- #        modes = []
- #        # Convert list of Fock dimensions to an array if it's not already an array 
- #        fock_dimensions = self.convert_to_array(fock_dimensions).astype(int)
- #        for idx, fock_dim in zip(mode_indices, fock_dimensions):
- #            mode_index = idx # or some function of this index 
- #            modes.append(MotionalMode.from_frequency(self.eigvals[mode_index], fock_dim))
- #
- #        return modes 
- #
+    def _ensure_mode_analysis_run(self):
+        """Ensure that mode analysis has ran. If not, run it automatically."""
+        if not hasattr(self, 'eigvals') or not hasattr(self, 'eigvecs'):
+            self.solve_ion_trap_equilibrium()
+
+    def get_mode_participation_factors_for_mode(self, mode_index: int) -> Matrix:
+        """ Returns mode participation factors for a specific mode by its index.
+
+            Outpout: Array of shape (3, num_ions) containing participation factors
+                where result[dimension, ion] gives the participation factor
+                for that dimension and ion in the specified mode.
+        """
+        self._ensure_mode_analysis_run()
+
+        # Get full participation factor matrix
+        full_participation = self.calculate_mode_participation_factors()
+
+        # Extract participation factors for the specific mode
+        mode_participation = full_participation[:, :, mode_index]
+
+        return mode_participation
+
+    def get_lamb_dicke_parameters_for_mode(self, mode_index: int, wavevector: Vector) -> Matrix:
+        """ Return Lamb-Dicke parameters for a specific mode by its index.
+
+                - mode_index: Index of the mode of interest
+                - wavevector: Wavevector k as a 3-element array (kx, ky, kz) or scalar wavenumber
+
+                Returns an array of shape (num_ions,) containing Lamb-Dicke parameters
+                where result[ion] gives the LD parameter for that ion in the specified mode.
+        """
+        self._ensure_mode_analysis_run()
+
+        # Get participation factors for this mode
+        mode_participation = self.get_mode_participation_factors_for_mode(mode_index)
+
+        # Calculate Lamb-Dicke parameters
+        if wavevector.shape != (3,):
+            raise ValueError("Wavevector must be a 3-element array (kx, ky, kz)")
+
+        # Compute dot product: (3, num_ions) · (3,) -> (num_ions,)
+        lamb_dicke_params = np.zeros(self.num_ions, dtype=complex)
+        for i in range(3):
+            lamb_dicke_params += wavevector[i] * mode_participation[i, :]
+
+        return lamb_dicke_params
+
+    def get_mode_properties(self, mode_index: int) -> dict:
+        """Get comprehensive properties for a specific mode. Returns a dictionary containing frequency,
+                participation factors, and other properties for the specified mode.
+
+                - mode_index: Index of the mode of interest
+
+            Returns:
+                Dictionary containing:
+                - 'frequency': Mode frequency in rad/s
+                - 'participation_factors': Array of shape (3, num_ions)
+                - 'frequency_MHz': Mode frequency in MHz (convenience)
+        """
+        self._ensure_mode_analysis_run()
+
+        # Get basic mode properties
+        frequency = self.eigvals[mode_index] * self.trap_freq_scale
+        participation_factors = self.get_mode_participation_factors_for_mode(mode_index)
+
+        return {
+            'frequency': frequency,  # rad/s
+            'frequency_MHz': frequency / (2 * np.pi * 1e6),  # MHz
+            'participation_factors': participation_factors,
+            'mode_index': mode_index
+        }
 
 
 #You could redefine the equilibrium finding function to assert that the equilibrium is linear. For example, make a wrapper inside the function for the potential, Jacobian, and Hessian that forces x_i and y_i = 0. 
- 
 #This is the code for the Lamb-Dicke parameters: (not that overall phases don't matter here, but the relative phase does. We could pin down the phase with some convention like, "each mode's first non-negative LD value is defined positive.")
 
 class LinearIonChainAnalysis(TrappedIonModeAnalysis):
@@ -813,9 +754,9 @@ class LinearIonChainAnalysis(TrappedIonModeAnalysis):
     def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float,
                  atomic_masses: np.ndarray | float, atomic_numbers: np.ndarray | int):
         """ Initialize a linear ion chain analysis, defaulting to branch sorted modes. """
-        mode_organization = 'branched'
+        mode_organization = 'branch_sorted'
         reindexing_strategy = 'z_axis'
-        super().__init__(num_ions, omega_x, omega_y, omega_z, atomic_masses, atomic_numbers, mode_organization)
+        super().__init__(num_ions, omega_x, omega_y, omega_z, atomic_masses, atomic_numbers, mode_organization, reindexing_strategy)
 
     def solve_ion_trap_equilibrium(self):
         """ Solve for equilibrium positions and analyze normal modes for a linear chain. It reindexes ions by their axial position. """
@@ -935,6 +876,67 @@ def get_mode_participation_factors_by_branch(self):
 
         return result
 
+    def get_mode_properties_by_branch(self, branch: str, mode_index_in_branch: int) -> dict:
+        """ Returns a dictionary containing properties (freq, participations) for a specific mode within a branch (x, y, or z).
+
+            For linear chains with branch-sorted modes, this provides an intuitive way
+            to access mode properties by branch and mode number within that branch.
+
+                - branch: 'x', 'y', or 'z' for the mode branch
+                - mode_index_in_branch: Index of the mode within the specified branch (0 to num_ions-1)
+        """
+        self._ensure_mode_analysis_run()
+
+        if branch not in ['x', 'y', 'z']:
+            raise ValueError("Branch must be 'x', 'y', or 'z'")
+
+        if mode_index_in_branch < 0 or mode_index_in_branch >= self.num_ions:
+            raise ValueError(f"Mode index in branch must be between 0 and {self.num_ions-1}")
+
+        # Calculate the global mode index
+        n_modes_per_branch = self.num_ions
+        if branch == 'x':
+            global_mode_index = mode_index_in_branch
+        elif branch == 'y':
+            global_mode_index = n_modes_per_branch + mode_index_in_branch
+        else:  # branch == 'z'
+            global_mode_index = 2 * n_modes_per_branch + mode_index_in_branch
+
+        # Get properties using the parent class method
+        properties = super().get_mode_properties(global_mode_index)
+
+        # Add branch-specific information
+        properties['branch'] = branch
+        properties['mode_index_in_branch'] = mode_index_in_branch
+
+        return properties
+
+    def get_lamb_dicke_parameters_by_branch(self, branch: str, mode_index_in_branch: int,
+                                          wavevector: Vector | float) -> Vector:
+        """Get Lamb-Dicke parameters for a specific mode within a branch.
+
+            Parameters:
+                branch: 'x', 'y', or 'z' for the mode branch
+                mode_index_in_branch: Index of the mode within the specified branch
+                wavevector: Wavevector k as a 3-element array or scalar wavenumber
+
+            Returns:
+                Array of Lamb-Dicke parameters for each ion in the specified mode
+        """
+        self._ensure_mode_analysis_run()
+
+        # Get the global mode index
+        n_modes_per_branch = self.num_ions
+        if branch == 'x':
+            global_mode_index = mode_index_in_branch
+        elif branch == 'y':
+            global_mode_index = n_modes_per_branch + mode_index_in_branch
+        else:  # branch == 'z'
+            global_mode_index = 2 * n_modes_per_branch + mode_index_in_branch
+
+        # Use parent class method to get LD parameters
+        return super().get_lamb_dicke_parameters_for_mode(global_mode_index, wavevector)
+
     def calculate_lamb_dicke_parameters(self, wavevector: Vector) -> dict:
         """ Calculate Lamb-Dicke parameters for all ions and modes, organized by branch.
 
@@ -952,38 +954,28 @@ def calculate_lamb_dicke_parameters(self, wavevector: Vector) -> dict:
         # Get mode participation factors (which are proportional to zero-point motion)
         mode_pf_by_branch = self.get_mode_participation_factors_by_branch()
 
-        # Handle both full wavevector (preferred) and scalar wavenumber (backward compatibility)
-        if isinstance(wavevector, (float, int)):
-            # Scalar case: multiply each component by the same wavenumber
-            # This assumes the laser is equally coupled to all directions
-            lamb_dicke_parameters = {}
-            for direction in ['x', 'y', 'z']:
-                lamb_dicke_parameters[direction] = wavevector * mode_pf_by_branch[direction]
-        else:
-            # Vector case: compute dot product k · Δr_0 for proper directionality
-            wavevector = np.asarray(wavevector)
-            if wavevector.shape != (3,):
-                raise ValueError("Wavevector must be a 3-element array (kx, ky, kz)")
-
-            # Get full mode participation factors
-            full_mode_pf = self.calculate_mode_participation_factors()
-
-            # Compute Lamb-Dicke parameters as dot product: η = k · Δr_0
-            # This gives us shape (num_ions, num_modes) for the total LD parameter
-            num_modes = 3 * self. num_ions
-
-            # Reshape for dot product: (3, num_ions, num_modes) · (3,) -> (num_ions, num_modes)
-            total_ld_params = np.zeros((self.num_ions, num_modes), dtype=complex)
-            for i in range(3):
-                total_ld_params += wavevector[i] * full_mode_pf[i, :, :]
-
-            # Organize by branch for consistency with scalar case
-            n_modes_per_branch = self.num_ions
-            lamb_dicke_parameters = {
-                'x': total_ld_params[:, :n_modes_per_branch],
-                'y': total_ld_params[:, n_modes_per_branch:2*n_modes_per_branch],
-                'z': total_ld_params[:, 2*n_modes_per_branch:3*n_modes_per_branch]
-            }
+        if wavevector.shape != (3,):
+            raise ValueError("Wavevector must be a 3-element array (kx, ky, kz)")
+
+        # Get full mode participation factors
+        full_mode_pf = self.calculate_mode_participation_factors()
+
+        # Compute Lamb-Dicke parameters as dot product: η = k · Δr_0
+        # This gives us shape (num_ions, num_modes) for the total LD parameter
+        num_modes = 3 * self. num_ions
+
+        # Reshape for dot product: (3, num_ions, num_modes) · (3,) -> (num_ions, num_modes)
+        total_ld_params = np.zeros((self.num_ions, num_modes), dtype=complex)
+        for i in range(3):
+            total_ld_params += wavevector[i] * full_mode_pf[i, :, :]
+
+        # Organize by branch for consistency with scalar case
+        n_modes_per_branch = self.num_ions
+        lamb_dicke_parameters = {
+            'x': total_ld_params[:, :n_modes_per_branch],
+            'y': total_ld_params[:, n_modes_per_branch:2*n_modes_per_branch],
+            'z': total_ld_params[:, 2*n_modes_per_branch:3*n_modes_per_branch]
+        }
 
         return lamb_dicke_parameters
 
@@ -997,8 +989,6 @@ def calculate_lamb_dicke_parameters_full(self, wavevector: Vector) -> Matrix:
         # Get full mode participation factors
         mode_pf = self.calculate_mode_participation_factors()
 
-        # Vector case: compute dot product k · Δr_0 (physically accurate)
-        wavevector = np.asarray(wavevector)
         if wavevector.shape != (3,):
             raise ValueError(f"Wavevector must be a 3-element array (kx, ky, kz). Received {wavevector}")
 
@@ -1057,7 +1047,7 @@ def find_axial_frequency_for_desired_central_ion_separation(self, target_separat
 
         def separation_error(wz_Hz):
             # Create a temporary analysis object with the current wz
-            temp_analysis = GeneralizedModeAnalysisWithBranchSortedModes(
+            temp_analysis = LinearIonChainAnalysis(
                 num_ions=self.num_ions,
                 omega_x=self.omega_x[0],
                 omega_y=self.omega_y[0],

From 428f7978175648a7e862ac4b405147f9d3f21b6b Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Wed, 1 Jul 2026 09:48:51 -0600
Subject: [PATCH 29/31]   update branch with gitlab version, class finalization
 + example

---
 examples/example_spin_dependent_squeezing.py | 725 +++++++++++++++++++
 examples/style/plot_style_data.txt           |  37 +
 src/ionsim/degree_of_freedom.py              |  13 +-
 src/ionsim/trapped_ion_mode_analysis.py      | 144 ++--
 tests/test_mode_analysis.py                  |  90 +--
 5 files changed, 881 insertions(+), 128 deletions(-)
 create mode 100644 examples/example_spin_dependent_squeezing.py
 create mode 100644 examples/style/plot_style_data.txt

diff --git a/examples/example_spin_dependent_squeezing.py b/examples/example_spin_dependent_squeezing.py
new file mode 100644
index 0000000..8290b25
--- /dev/null
+++ b/examples/example_spin_dependent_squeezing.py
@@ -0,0 +1,725 @@
+import ionsim as ism
+
+import numpy as np
+from scipy.sparse import kron as skron
+import sys
+import matplotlib
+import matplotlib.pyplot as plt
+import platform
+from scipy.special import factorial 
+matplotlib.rcParams['text.usetex']=True 
+style_path_data = 'style/plot_style_data.txt'
+
+TPI = 2*np.pi
+
+def four_tone_hamiltonian(basis, spin_basis, modes, spectator_modes, etas, rabi_rate_x, rabi_rate_y, omega_x_red, omega_x_blue, omega_y_red, omega_y_blue, targetIon: AtomicSpin, 
+                        index_mode_i:int, index_mode_j:int, phi:float, sparse=False, mod=None): 
+    ''' Applies 4-tone MS-like gate to a single ion "k" to produce a spin-dependent motional operation. 
+        Consists of 2 sidebands (red and blue) on mode i and another 2 sidebands (red and blue) on mode j:  
+
+        H(t) = hbar Omega_x sigma_x a_i exp(-i Delta t) + h.c. + hbar Omega_y sigma_y a_j exp(-i n Delta t + phi)  + h.c. 
+
+        - currently sets the first MS phase to zero, accepts input for second MS phase as phi 
+    ''' 
+    # Hamiltonian: H = -eta * Omega/2 sum_{k} [-i sigma_+,k exp(iphi) + i sigma_-,k exp(-iphi)] (a e^i delta t + a† e^{-i delta t} )
+    # Hamiltonian: H = -eta * Omega J_y (a e^i delta t + a† e^{-i delta t} ) # for phi = 0; Jy = 1/2 (sigma_y,1 + sigma_y,2)
+    motional_basis = ism.StandardBasis([*modes])
+
+    DOFs = basis.degrees_of_freedom
+
+    target_ion_index = DOFs.index(targetIon)
+
+    operators = []
+
+    # build matrices for fundamental raising/lowering spin/Fock operators: 
+    # Raising operator for target ion "k" 
+    raise_spin_k = spin_basis.enlarge_matrix(ism.Pauli.plus, [targetIon])
+    
+    fock_dimension = len(modes[index_mode_i].energy_levels)
+    assert fock_dimension == len(modes[index_mode_j].energy_levels)
+
+    # First, include modes that are supposed to be included. (non-spectators)
+    raise_motion = motional_basis.enlarge_matrix(ism.Fock.raising(fock_dimension), [modes[index_mode_i]])
+    lower_motion = motional_basis.enlarge_matrix(ism.Fock.lowering(fock_dimension), [modes[index_mode_j]])
+
+    # Extract Lamb-Dicke parameter for mode i and j, ion k
+    eta_ki = etas[target_ion_index, index_mode_i]
+    eta_kj = etas[target_ion_index, index_mode_j]
+
+    # spin phase = ( phi_R + phi_B )/ 2 
+    #   -------------------
+    #   x-spin MS-like gate: 
+    # For mode i : 
+    # Can achieve sigma_x by setting spin phase of tones 1 and 2 to zero or each as 2pi  
+    phase_x = np.pi/2. 
+
+    x_prefactor = 1j*rabi_rate_x * 0.5 * eta_ki * np.exp(1j*phase_x)
+
+    operator_red_mode_i = x_prefactor * skron(raise_spin_k, lower_motion) 
+    operator_blue_mode_i = x_prefactor * skron(raise_spin_k, raise_motion)
+    # h.c. is already included by ionsim upon providing these coupling operators to the hamiltonian 
+
+    #   -------------------
+    #   y-spin MS-like gate: 
+    # For mode j : 
+    phase_y = 0. 
+    y_prefactor = 1j*rabi_rate_y * 0.5 * eta_kj * np.exp(1j*(phase_y + phi))
+
+    operator_red_mode_j = y_prefactor * skron(raise_spin_k, lower_motion) 
+    operator_blue_mode_j = y_prefactor * skron(raise_spin_k, raise_motion)
+
+    operators.extend([
+        ism.CouplingOperator.from_matrix(basis, operator_red_mode_i, omega_x_red, modulation_function=mod), 
+        ism.CouplingOperator.from_matrix(basis, operator_blue_mode_i, omega_x_blue, modulation_function=mod), 
+        ism.CouplingOperator.from_matrix(basis, operator_red_mode_j, omega_y_red, modulation_function=mod), 
+        ism.CouplingOperator.from_matrix(basis, operator_blue_mode_j, omega_y_blue, modulation_function=mod), 
+    ])
+
+    for mode in spectator_modes:
+        mode_index = modes.index(mode)
+        spectator_raise_motion = motional_basis.enlarge_matrix(ism.Fock.raising(fock_dimension), [mode])
+        spectator_lower_motion = motional_basis.enlarge_matrix(ism.Fock.lowering(fock_dimension), [mode])
+
+        # Extract Lamb-Dicke parameter 
+        eta_ki = etas[target_ion_index, mode_index]
+        eta_kj = etas[target_ion_index, mode_index]
+
+        x_prefactor = 1j*rabi_rate_x * 0.5 * eta_ki * np.exp(1j*phase_x)
+    
+        operator_red_x_spectator = x_prefactor * skron(raise_spin_k, spectator_lower_motion) 
+        operator_blue_x_spectator = x_prefactor * skron(raise_spin_k, spectator_raise_motion)
+
+        y_prefactor = 1j*rabi_rate_y * 0.5 * eta_kj * np.exp(1j*(phase_y + phi))
+    
+        operator_red_y_spectator = y_prefactor * skron(raise_spin_k, spectator_lower_motion) 
+        operator_blue_y_spectator = y_prefactor * skron(raise_spin_k, spectator_raise_motion)
+
+        operators.extend([
+            ism.CouplingOperator.from_matrix(basis, operator_red_x_spectator, omega_x_red, modulation_function=mod), 
+            ism.CouplingOperator.from_matrix(basis, operator_blue_x_spectator, omega_x_blue, modulation_function=mod), 
+            ism.CouplingOperator.from_matrix(basis, operator_red_y_spectator, omega_y_red, modulation_function=mod), 
+            ism.CouplingOperator.from_matrix(basis, operator_blue_y_spectator, omega_y_blue, modulation_function=mod), 
+        ])
+
+    interaction_frame_energies = [-state.energy for state in basis.states]
+    return ism.Hamiltonian(basis, operators, interaction_frame_energies, sparse=sparse)
+
+
+def thermal_state_populations(nbar: float, fock_dimension): 
+    fock_populations = np.zeros(fock_dimension)
+
+    normalization = 1./(1. + nbar) 
+    for n in range(fock_dimension):
+        fock_populations[n] = (nbar/(1. + nbar))**(n) 
+
+    return fock_populations*normalization 
+
+
+def dephasing_dissipator(basis, spin_basis, modes, dephasing_rates: list[float]):
+    ''' Dissipator for ion motional mode dephasing ''' 
+
+    motional_basis = ism.StandardBasis([*modes])
+    spins = spin_basis.degrees_of_freedom
+    
+    # Build spin basis identity 
+    spin_identities = []
+    for spin in spins:
+        # Identity matrix for mode m in Fock space, enlarged to fit dimensionality of M modes 
+        spin_identities.append(spin_basis.enlarge_matrix(ism.Pauli.I, [spin]))
+
+    spin_identity = np.sum(spin_identities, axis=0)
+
+    # Motional mode dephasing is proportional to the phonon number operators 
+    # loop over modes:
+    lindblad_operators = []
+    for mode, dephasing_rate in zip(modes, dephasing_rates):
+        fock_dim = len(mode.energy_levels)
+        mode_number_matrix = motional_basis.enlarge_matrix(ism.Fock.number(fock_dim), [mode]) 
+
+        # One lindblad operator per mode: L = sqrt(Gamma) * (a†_m a_m)
+        dephasing_operator = np.sqrt(dephasing_rate) * skron(spin_identity, mode_number_matrix) 
+
+        # Lindblad operator is diagonal in the Fock occupancy number basis; use EnergyShiftoperator or GeneralOperator 
+        lindblad_operators.append(ism.EnergyShiftOperator.from_matrix(basis, dephasing_operator))
+
+    frame_energies = [-state.energy for state in basis.states] 
+    return ism.Dissipator(basis, lindblad_operators, frame_energies)
+
+
+
+def heating_dissipator(basis, modes, heating_rates: list[float]): 
+    ''' Dissipator for ion mode heating ''' 
+
+    motional_basis = ism.StandardBasis([*modes])
+    
+    # Build spin basis identity 
+    spin_identities = []
+    for spin in spins:
+        # Identity matrix for mode m in Fock space, enlarged to fit dimensionality of M modes 
+        spin_identities.append(spin_basis.enlarge_matrix(ism.Pauli.I, [spin]))
+
+    spin_identity = np.sum(spin_identities, axis=0)
+
+    # Phonon raising and lowering operators 
+    # loop over modes:
+    lindblad_operators = []
+    for mode, heating_rate in zip(modes, heating_rates):
+        fock_dim = len(mode.energy_levels)
+        mode_lowering_matrix = motional_basis.enlarge_matrix(ism.Fock.lowering(fock_dim), [mode]) 
+        mode_raising_matrix = motional_basis.enlarge_matrix(ism.Fock.raising(fock_dim), [mode]) 
+
+        # Two lindblad operators per mode: L1 = sqrt(Gamma) * a_m , L2 = sqrt(Gamma) * a†_m 
+        lowering_operator = np.sqrt(heating_rate) * skron(spin_identity, mode_lowering_matrix) 
+        raising_operator = np.sqrt(heating_rate) * skron(spin_identity, mode_raising_matrix) 
+        #lowering_operator = np.sqrt(decay_rate*(N_thermal+1)) * skron(spin_identity, mode_lowering_matrix) 
+        #raising_operator = np.sqrt(decay_rate*N_thermal) * skron(spin_identity, mode_raising_matrix) 
+
+        lindblad_operators.append(ism.CouplingOperator.from_matrix(basis, lowering_operator, 0.))
+        lindblad_operators.append(ism.CouplingOperator.from_matrix(basis, raising_operator, 0.))
+
+    frame_energies = [-state.energy for state in basis.states] 
+    return ism.Dissipator(basis, lindblad_operators, frame_energies)
+
+
+def carrier_hamiltonian(basis, target_ion, rabi_rate: float, phase: float, omega_carrier: float, sparse: bool=False, mod=None):
+    """ Hamiltonian for carrier spin pulse: (hbar * Omega/2) * sigma_phi e^iphi + h.c. """ 
+    prefactor = rabi_rate * 0.5 * np.exp(-1j*phase)
+    operator_raising = basis.enlarge_matrix(prefactor * ism.Pauli.plus, [target_ion])
+    #raise_spin_k = spin_basis.enlarge_matrix(ism.Pauli.plus, [target_ion])
+
+    operators = [ism.CouplingOperator.from_matrix(basis, operator_raising, omega_carrier, modulation_function=mod)] 
+    interaction_frame_energies = [-state.energy for state in basis.states]
+    return ism.Hamiltonian(basis, operators, interaction_frame_energies, sparse=sparse)
+
+
+def red_sb_pi_pulse_duration(rabi_rate, fock_state, eta, single_ion: bool=True):
+    # Assume either single or 2-ion sysetm
+    if single_ion:
+        t = np.pi / (rabi_rate * np.sqrt(fock_state) * np.abs(eta))
+    else:
+        t = np.pi / (rabi_rate * np.sqrt(fock_state * (fock_state - 1)) * np.abs(eta))
+    return t  
+    
+
+
+def compute_N(rho_motional, motional_basis):
+
+    #for state in motional_basis.states:
+    N_operator = ism.Fock.raising(fock_dimension) @ ism.Fock.lowering(fock_dimension) 
+    return np.trace( N_operator.dot(rho_motional.density_matrix))
+
+def estimate_squeezing_magnitude(motional_rho, mode_basis):
+    ''' Estimate the squeezing parameter from the X and P variances ''' 
+    x, p, x2, p2 = motional_rho.compute_quadratures(True) # includes variance 
+
+    avg_X_tilt = x[0].real
+    avg_X2_tilt = x2[0].real
+    
+    X_variance = avg_X2_tilt - avg_X_tilt**2
+    #print(f"variance of X for the tilt mode: {X_variance}")
+    
+    avg_P_tilt = p[0].real
+    avg_P2_tilt = p2[0].real
+    
+    P_variance = avg_P2_tilt - avg_P_tilt**2
+    #print(f"variance of P for the tilt mode: {P_variance}")
+    if squeezing_phase == 0.: 
+        X_squeezed = False 
+    elif squeezing_phase == np.pi:
+        X_squeezed = True 
+
+    if X_squeezed:
+        r_X = -0.5 * np.log(2. * X_variance)
+        r_P = 0.5 * np.log(2. * P_variance)
+    else:
+        r_X = 0.5 * np.log(2. * X_variance)
+        r_P = -0.5 * np.log(2. * P_variance)
+
+    return r_X, r_P
+
+def squeezed_ground_state_coefficients(r_value, phi, fock_dimension):
+    wave_function = np.zeros(fock_dimension,dtype=complex) 
+    norm = 1./np.sqrt(np.cosh(r_value))
+    for n in range(fock_dimension//2):
+        wave_function[2*n] = np.exp(1j * n * phi) * np.tanh(r_value)**n 
+        #wave_function[2*n] *= norm * np.sqrt(factorial(2 * n))
+        wave_function[2*n] *= norm * np.sqrt(factorial(2 * n))/(2**n) 
+        wave_function[2*n] /= (factorial(n))
+    return wave_function 
+
+
+
+def squeezed_vacuum_populations(fock_dim, r_val):
+    fock_states = np.arange(fock_dim)
+    fock_populations = np.zeros(len(fock_states)) 
+    norm = 1./np.cosh(r_val)
+    for i in range(fock_dim//2):
+        fock_populations[2*i] = norm * (np.tanh(r_val)**(2*i)) * factorial(2 * i)/(2**(2*i) * factorial(i)**2 ) 
+    return fock_populations 
+
+    
+
+def main():
+    import time
+    from matplotlib import pyplot as plt
+
+    print(f"\n --- Simulation of spin-dependent squeezing by 4-tone sideband protocol --- ")
+    
+    include_spectator_mode = False 
+    include_dissipation = False 
+    do_reverse_process = False 
+    thermal_state_IC = False 
+    reverse_phases = False 
+    beam_counter_propagating = True
+
+    num_ions = 2
+
+    # Create 171Yb+ qubits 
+    spins = [
+        ism.AtomicSpin.from_species(species='171Yb+', term_symbols=['S1/2'], level_names=['S1/2,0,0', 'S1/2,1,0'])
+        for _ in range(num_ions)
+    ]
+
+    target_ion = spins[0]    
+    target_ion_index = spins.index(target_ion)
+
+    spin_basis = ism.StandardBasis([*spins])
+
+    # Use mode analysis to get motional mode information 
+    omega_x = 2.1 * TPI * 1E6 # MHz -> rad/s  
+    omega_y = 2.6 * TPI * 1E6 # ""  
+    omega_z = 0.50 * TPI * 1E6 # ""  
+    trap_analysis = ism.LinearIonChainAnalysis.from_atomic_spin_basis(spin_basis, omega_x, omega_y, omega_z) 
+    trap_analysis.solve_ion_trap_equilibrium()
+    print()
+    trap_analysis.print_chain_summary()
+
+    print(trap_analysis.characteristic_parameters)
+
+
+    # Truncate the Hilbert space to a single qubit with the two nearby motional modes 
+    # Choose to do oprerations on the radial x tilt mode; option to include the COM as a spectator mode  
+    #print(trap_analysis.get_mode_properties_by_branch('x', 0))
+    #print(trap_analysis.get_mode_properties_by_branch('x', 1))
+    tilt_mode_index = 0 
+    COM_mode_index = 1 
+
+    fock_dimension = 10
+    branch_dir = 'x'
+    spin_basis = ism.StandardBasis([target_ion]) # Now reset the spin basis to just include the target ion 
+    modes = trap_analysis.build_mode_DOFs_from_branch(branch_dir, [tilt_mode_index, COM_mode_index], fock_dimension)
+    if include_spectator_mode:
+        basis = ism.StandardBasis([spins[target_ion_index], *modes])
+        motional_basis = ism.StandardBasis(modes)
+        spectator_modes = [modes[1]]
+    else:
+        basis = ism.StandardBasis([spins[target_ion_index], modes[tilt_mode_index]])
+        motional_basis = ism.StandardBasis([modes[tilt_mode_index]])
+        modes = [modes[0]]
+        spectator_modes = [] 
+
+
+    target_mode_index = 0 # for tilt mode in our basis 
+
+    # 4-tone protocol targets 2 modes in general with 2 tones each; for squeezing, we target the same mode with all tones 
+    mode_i_index = target_mode_index
+    mode_j_index = target_mode_index
+
+    radial_x_tilt_mode_properties = trap_analysis.get_mode_properties_by_branch('x', tilt_mode_index)
+    mode_i_freq = radial_x_tilt_mode_properties['frequency'] 
+    mode_j_freq = mode_i_freq 
+
+    print(f" - Tilt mode frequency: {mode_i_freq/TPI/1E6} [MHz]\n\n")
+
+    # Laser/driving parameters 
+    print(f" --- Laser parameters --- ")
+    wavelength = 355 # nm 
+    beam_angle = np.pi/4. # relative to axial direction 
+    wavevector = TPI / (355. * 1E-9) * np.array([np.cos(beam_angle), np.sin(beam_angle), 0.]) 
+    if beam_counter_propagating:
+        # Effective wavevector from 
+        wavevector *= 2. # Raman process: \delta k is 2x basic wavevector  
+
+    rabi_rate = 175 * 1E3 * TPI  / np.sqrt(2.) #* np.sqrt(0.5) #/ np.sqrt(4.) # kHz -> Hz  
+    x_rabi_rate = rabi_rate
+    y_rabi_rate = rabi_rate
+    print(f" - Carrier rabi rate: {rabi_rate / TPI / 1E3} [kHz]\n")
+
+    detuning = 40.0 * 1E3 * TPI # 40 kHz -> rad/s 
+    print(f" - Detuning: {detuning/(TPI * 1E3)} [kHz]\n")
+    duration = 60 * 1E-6 
+    print(f" - Gate duration: {duration*1E6:.3f} [µs]\n")
+
+    modulate_amplitude = True 
+
+    print(f" - Estimated number of loops: {duration * detuning / TPI}\n")
+
+    detuning_red = detuning
+    detuning_blue = detuning
+
+    # Blue and red laser frequencies for 1st-sideband tuning:  
+    # For x-spin / operation on mode i, the laser is simply detuned by 1*detuning  
+    offset = 0. * TPI * 1E3 # kHz 
+    omega_qubit = target_ion.energy_levels[1].energy - target_ion.energy_levels[0].energy
+    omega_b_mode_i = (omega_qubit + (mode_i_freq + detuning_blue)) + offset 
+    omega_r_mode_i = (omega_qubit - (mode_i_freq + detuning_red)) - offset 
+    
+    # Second MS pulse 
+    # For y-spin / operation on mode j, it is detuned by n*detuning  
+    n = -1 # for squeezing 
+    omega_b_mode_j = (omega_qubit + (mode_j_freq + n*detuning_blue)) - offset
+    omega_r_mode_j = (omega_qubit - (mode_j_freq + n*detuning_red)) + offset
+
+    ## -- Global parameters for script --- 
+    frequencies = np.array([omega_r_mode_i, omega_r_mode_j, omega_b_mode_i, omega_b_mode_j]) # MHz
+    frequencies_from_carrier = (frequencies - omega_qubit)/TPI/1E6 # MHz
+    print(f"  - laser frequencies (minus carrier): {frequencies_from_carrier}\n\n")
+
+ 
+    def blackman(t: float, magnitude: float = 1.):
+        return magnitude*(0.42 - 0.5*np.cos(TPI * t /duration) + 0.08 * np.cos(2.*TPI*t/duration))
+
+    # Retrieve lamb dicke parameters 
+    lamb_dicke_parameters = trap_analysis.get_radial_lamb_dicke_parameters(wavevector, branch_dir) # matrix of ion, mode 
+    eta = lamb_dicke_parameters[target_ion_index,target_mode_index]
+    print(f" - |eta| for target ion, target mode: {np.abs(eta)}\n")
+    print(f" - Sideband rabi rate (eta*Omega): {np.abs(eta) * rabi_rate / 1E3 / TPI} [kHz]\n")
+    print('\n\n')
+    print(f" - Lamb dicke parameters: \n{lamb_dicke_parameters}")
+    mod_function = None
+    expected_r = duration * ((np.abs(eta) * rabi_rate)**2) / detuning # Square pulse 
+    if modulate_amplitude:
+        print(f" - Using a Blackman pulse shape\n")
+        mod_function = blackman
+        # Area of the blackman pulse shape squared is ~0.3046
+        expected_r = 0.3046 * duration * ((np.abs(eta) * rabi_rate)**2) / detuning
+    else:
+        print(f" - Using a constant (square) pulse shape\n")
+
+    r = expected_r
+    print(f" - Expecting squeezing of r = {expected_r}\n")
+
+    
+    # =============================
+    ## Create Hamiltonians: 
+    squeezing_phase = 0.
+    squeezing_hamiltonian = four_tone_hamiltonian(basis, spin_basis, modes, spectator_modes, lamb_dicke_parameters, x_rabi_rate, y_rabi_rate, omega_r_mode_i, omega_b_mode_i, omega_r_mode_j, 
+                                                    omega_b_mode_j, target_ion, mode_i_index, mode_j_index, squeezing_phase, mod=mod_function)
+
+    reverse_squeezing_hamiltonian = four_tone_hamiltonian(basis, spin_basis, modes, spectator_modes, lamb_dicke_parameters, x_rabi_rate, y_rabi_rate, omega_r_mode_i, omega_b_mode_i, omega_r_mode_j, omega_b_mode_j, 
+                                                        target_ion, mode_i_index, mode_j_index, squeezing_phase + np.pi, mod=mod_function) 
+
+    pi_pulse_hamiltonian = carrier_hamiltonian(basis, target_ion, rabi_rate, 0., np.abs(omega_qubit), False, None)  
+    mode_dephasing_rates = [5.0, 500.] # quanta/second 
+    if include_dissipation : 
+        # Create dissipator and then full Lindbladian: 
+        #ion_heating_dissipator = heating_dissipator(basis, modes, mode_heating_rates) 
+        _dephasing_dissipator = dephasing_dissipator(basis, spin_basis, modes, mode_dephasing_rates) 
+        lindbladian_4tone = ism.Lindbladian(squeezing_hamiltonian, _dephasing_dissipator)
+        #lindbladian_4tone = ism.Lindbladian(squeezing_hamiltonian, ion_heating_dissipator)
+        # Create reverse processes 
+        lindbladian_4tone_reverse = ism.Lindbladian(reverse_squeezing_hamiltonian, _dephasing_dissipator)
+
+    # Initial fock state, all populated in this state:
+    starting_fock_state = 0 
+    excited_qubit1_init_state = False 
+    if excited_qubit1_init_state:
+        qubit1_index = fock_dimension*num_modes 
+    else:
+        qubit1_index = 0
+
+    nbars = [0.1, 0.5]
+
+    IC_fock_pops = [thermal_state_populations(nbar, fock_dimension) for nbar in nbars] # populations  
+    #fock_populations_target_mode_IC = thermal_state(nbar)
+    if thermal_state_IC:
+        print(f"\n--- Starting with thermal ground state ---\n")
+        print(f"Target mode initial populations for nbar = {nbars[0]}, : {IC_fock_pops[0]}")
+        coefs = np.kron(np.array([1., 0.]), IC_fock_pops[0])
+        for m in range(1, len(modes)):
+            print(f"Spectator mode initial populations for nbar = {nbars[m]}, : {IC_fock_pops[m]}")
+            coefs = np.kron(coefs, IC_fock_pops[m])
+        coefs = np.sqrt(coefs)# populations require square root 
+    else:
+        print(f"--- Initializing in the ground state of full Hilbert space ---")
+        IC_fock_pops[0] = np.zeros(fock_dimension) 
+        IC_fock_pops[0][0] = 1. 
+        for m in range(1, len(modes)):
+            IC_fock_pops[m] = np.zeros(fock_dimension) 
+            IC_fock_pops[m][0] = 1.
+        coefs = np.zeros(len(basis.states))
+        starting_fock_state_mode2 = 0
+        coefs[qubit1_index + starting_fock_state + starting_fock_state_mode2] = 1. 
+
+    print_state_names = False 
+    if print_state_names:
+        for indx, state in enumerate(basis.states):
+            print(indx)
+            print(state.name) 
+    init_state = ism.State.from_coefficients(basis, list(coefs)) 
+
+    times = np.linspace(0, duration, 2000) 
+
+    ground_state_level_index = 0 
+    excited_state_level_index = 1 
+    ground_state_level_name = spin_basis.degrees_of_freedom[0].energy_levels[ground_state_level_index].name 
+    excited_state_level_name = spin_basis.degrees_of_freedom[0].energy_levels[excited_state_level_index].name 
+    print("Ground state level name: " + ground_state_level_name)
+    print("Excited state level name: " + excited_state_level_name)
+
+    # Now perform cycling to attempt squeezing operation 
+    print(f"\n\n\n ------- Squeezing protocol ------- ")
+    print(f"Squeezing parameters: r = {r}, phi = {squeezing_phase}")
+
+    # Use rho(t) from previous 
+    # Apply 4 tones with n=-1 detuning relationship to get squeezing  
+    rho_t = init_state 
+
+    ### Evolve Schrodinger or master equation dynamics for spin-dependent squeezing via 4-tone protocol  
+    if include_dissipation:
+        rho_t = rho_t.propagate_using_master_equation(lindbladian_4tone, duration, times)
+    else:
+        rho_t = rho_t.propagate_using_schrodinger_equation(squeezing_hamiltonian, duration, times) # no dissipation  
+    
+    if do_reverse_process: 
+        # Perform pi pulse on qubit:  # Omega * t = pi; t = pi/Omega
+        if reverse_phases: 
+            if include_dissipation:
+                rho_t = rho_t[-1].propagate_using_master_equation(lindbladian_4tone_reverse, duration, times)
+            else:
+                rho_t = rho_t[-1].propagate_using_schrodinger_equation(reverse_squeezing_hamiltonian, duration, times) # no dissipation  
+    
+        else:
+            pi_pulse_duration = np.pi/rabi_rate
+        
+            # Perform squeezing again, which should return you back to the ground state     
+            if include_dissipation:
+                pi_pulse_lindbladian = ism.Lindbladian(hamiltonian = pi_pulse_hamiltonian, dissipator = None)
+                rho_t = rho_t[-1].propagate_using_master_equation(pi_pulse_lindbladian, pi_pulse_duration) 
+                rho_t = rho_t.propagate_using_master_equation(lindbladian_4tone, duration, times)
+            else:
+                rho_t = rho_t[-1].propagate_using_schrodinger_equation(pi_pulse_hamiltonian, pi_pulse_duration) # no dissipation  
+                rho_t = rho_t.propagate_using_schrodinger_equation(squeezing_hamiltonian, duration, times) # no dissipation  
+
+        # Extract final state 
+        final_state = rho_t[-1]
+
+        # Compute state overlap (fidelity of reverse process) 
+        ## TODO: uncomment after merging branch 57 
+        #reversibility_fidelity = final_state.motional_state.compute_state_fidelity(init_state.motional_state.density_matrix)
+        #print(f" - Reversibility fidelity: {reversibility_fidelity}")
+            
+
+    if len(times) != len(rho_t):
+        rho_t = rho_t[:len(times)]
+    print(f"Number of rho(t) timepoints: {len(rho_t)}")
+    print(f"Number of timepoints: {len(times)}")
+    
+    spin_rhos = rho_t.copy()
+    motional_rhos = rho_t.copy()
+    # 1. Trace out motional modes and plot spin-state populations
+    for mode in modes:
+        spin_rhos = [rho.trace_out_degree_of_freedom(mode) for rho in spin_rhos] 
+    
+    new_basis = spin_rhos[0].basis # should be 2-qubit basis 
+    
+    populations = np.array([rho.compute_basis_state_probabilities() for rho in spin_rhos])
+    plt.style.use(style_path_data) 
+    plt.figure(figsize=(6,4))
+    for i,state in enumerate(new_basis.states):
+        plt.plot(times*1E6, populations[:, i], label= state.name)
+    plt.ylabel('Population', fontsize = 16)
+    plt.xlabel('Gate Duration (µs)', fontsize = 20)
+    plt.legend()
+    plt.show()
+
+    # 2. Trace out spin DOFs and plot Fock state populations
+    X_modes = np.zeros((len(times),len(modes)),dtype=complex)
+    P_modes = np.zeros((len(times),len(modes)), dtype=complex)
+    X2_modes = np.zeros((len(times),len(modes)), dtype=complex)
+    P2_modes = np.zeros((len(times),len(modes)), dtype=complex)
+    for spin in [target_ion]:
+        motional_rhos = [rho.trace_out_degree_of_freedom(spin) for rho in motional_rhos]
+        # For each time, compute quadrature expectations w. variances  
+        # TODO: Uncomment after merging Wigner/quadrature helper functions branch 
+ #        for i, t in enumerate(times):
+ #            x, p, x2, p2 = motional_rhos[i].compute_quadratures(True) # includes variance 
+ #            X_modes[i, :] = np.array(x)
+ #            P_modes[i, :] = np.array(p)
+ #            X2_modes[i, :] = np.array(x2)
+ #            P2_modes[i, :] = np.array(p2)
+
+        for mode in spectator_modes: 
+            motional_rhos = [rho.trace_out_degree_of_freedom(mode) for rho in motional_rhos]
+
+    new_basis = motional_rhos[0].basis # should be motional basis  
+
+    fock_populations = np.array([rho.compute_basis_state_probabilities() for rho in motional_rhos])
+ #    print('\nFock populations at end of protocol: \n')
+ #    for state, population in zip(new_basis.states, fock_populations[-1]):
+ #        print(f"\nState: n = {state.name}")
+ #        print(f"Population: {population}")
+
+    fock_states = np.arange(fock_dimension)
+    squeezed_GS_populations = squeezed_vacuum_populations(fock_dimension, r)
+
+    bar_width = 0.35
+    plt.figure(figsize=(6,6))
+    plt.bar(fock_states - bar_width/2, fock_populations[-1], bar_width, color = 'b', label = 'Simulation')
+    if not do_reverse_process: 
+        plt.bar(fock_states + bar_width/2, squeezed_GS_populations, bar_width, color = 'r', label = 'Reference: $r_{target} = ' + str(np.round(r,2)) + '$')
+    else:
+        plt.bar(fock_states + bar_width/2, IC_fock_pops[0], bar_width, color = 'r', label = 'Initial distribution') 
+    plt.title('Target mode Fock distribution', fontsize = 14)
+    plt.ylabel('Population', fontsize = 16)
+    plt.xlabel('Fock state $n$', fontsize = 20)
+    plt.legend()
+    plt.show()
+
+    sys.exit(0)    
+
+
+
+    ############## -------------------------------------------------------------- To include after merging branch 57 --------------------------     ################################# 
+    # Compute wigner distributions: 
+    limit = 4. # phase space limit of x and p 
+    x_grid = np.linspace(-limit, limit, 200)
+    p_grid = np.linspace(-limit, limit, 200)
+
+    last_index = len(times)-1 
+    # Compute Wigner distributions over a few selected indices 
+    sampling_indices = [0, last_index//3, last_index//2, int(2*last_index//3), int(3*last_index//4), int(7*last_index//8), last_index] 
+    #sampling_indices = [0, last_index] 
+
+    sampled_states = [rho_t[i] for i in sampling_indices]
+    Wigner_distributions = [rho.compute_wigner_distribution(x_grid, p_grid) for rho in sampled_states] 
+    #print(len(Wigner_distributions[0]))
+
+    if num_modes == 2:
+        mode_labels = ['Target mode', 'Spectator mode 1']
+    else:
+        mode_labels = ['Target mode']
+
+    for W, t in zip(Wigner_distributions, times[sampling_indices]):
+        for i, mode_label in enumerate(mode_labels):
+            plt.figure(figsize=(6, 5))
+            plt.imshow(W[i], extent=[-limit, limit, -limit, limit], origin='lower', cmap='magma') 
+            if do_reverse_process: 
+                plt.title(f"$W(x,p)$: " + mode_label + f", t = {t*2.*1E6:.2f} [µs]", fontsize = 12)
+            else:
+                plt.title(f"$W(x,p)$, t = {t*1E6:.2f} [µs]", fontsize = 16)
+                #plt.title(f"$W(x,p)$: " + mode_label + f", t = {t*1E6:.2f} [µs]", fontsize = 12)
+            plt.xlabel('Position ($x$)', fontsize = 16)
+            plt.ylabel('Momentum ($p$)', fontsize = 16)
+            cbar = plt.colorbar()
+            cbar.ax.set_ylabel(r'$W(x,p)$', rotation = 0, labelpad = 30, va='center') 
+            if excited_qubit1_init_state: 
+                squeezing_str = 'anti'
+            else:
+                squeezing_str = 'fwd'
+            plt.savefig(squeezing_str + f"_mode_{i}_Wigner_t_{t*1E6:.2f}.pdf", dpi=300)
+            plt.show()
+
+    avg_X_tilt = X_modes[-1,0].real
+    avg_X2_tilt = X2_modes[-1,0].real
+    
+    print(f"<X> for the tilt mode: {avg_X_tilt}")
+    print(f"<X^2> for the tilt mode: {avg_X2_tilt}")
+    
+    X_variance = avg_X2_tilt - avg_X_tilt**2
+    print(f"variance of X for the tilt mode: {X_variance}")
+    
+    avg_P_tilt = P_modes[-1,0].real
+    avg_P2_tilt = P2_modes[-1,0].real
+    
+    print(f"<P> for the tilt mode: {avg_P_tilt}")
+    print(f"<P^2> for the tilt mode: {avg_P2_tilt}")
+    
+    P_variance = avg_P2_tilt - avg_P_tilt**2
+    print(f"variance of P for the tilt mode: {P_variance}")
+
+    print(f"\nVerifying r from variances:")    
+    # If X is squeezed: 
+    if squeezing_phase == 0.:
+        X_squeezed = False 
+    else:
+        X_squeezed = True 
+
+    if X_squeezed:
+        r_X = -0.5 * np.log(2. * X_variance)
+        r_P = 0.5 * np.log(2. * P_variance)
+    else:
+        r_X = 0.5 * np.log(2. * X_variance)
+        r_P = -0.5 * np.log(2. * P_variance)
+
+    print(f"r estimate from X variance: {r_X}") 
+    print(f"r estimate from P variance: {r_P}") 
+
+    target_motional_state = ism.State.from_coefficients(motional_rhos[0].basis, list(squeezed_ground_state_coefficients(r, squeezing_phase, fock_dimension))) 
+    squeezed_state_fidelities = np.array([rho_mot.compute_state_fidelity(target_motional_state.density_matrix) for rho_mot in motional_rhos])
+
+    # Check purity: 
+    print(f"Purity: {np.trace(motional_rhos[-1].density_matrix @ motional_rhos[-1].density_matrix)}")
+    
+    print('\n')
+    print(f'State fidelity with squeezed ground state: {squeezed_state_fidelities[-1]}')
+
+    r_X = np.zeros_like(times)
+    r_P = np.zeros_like(times)
+    n_t = np.zeros_like(times)
+    for i, t in enumerate(times):
+        r_X[i], r_P[i] = estimate_squeezing_magnitude(motional_rhos[i], motional_rhos[i].basis)
+        n_t[i] = compute_N(motional_rhos[i], motional_rhos[i].basis)
+
+    plt.style.use(style_path_data) 
+    plt.figure(figsize=(5,5))
+    plt.plot(times*1E6, r_X, marker = 's', linestyle = 'solid', color = 'k', markersize = 3, linewidth = 1.5, label = r'$r_{X}$')
+    plt.plot(times*1E6, r_P, marker = 'o', linestyle = 'solid', color = 'r', markersize = 3, linewidth = 1.5, label = r'$r_{P}$')
+    plt.axhline(y = r, color = 'k', linestyle='dashed', linewidth = 1.5, label = r'$r_{target} = ' + str(np.round(r, 4)) + '$')
+    plt.xlabel('Time [µs] ', fontsize = 20)
+    plt.ylabel(r'$r$', fontsize = 24, rotation = 0., labelpad=15)
+    plt.legend()
+    plt.show()
+
+    plt.figure(figsize=(5,5))
+    plt.plot(times*1E6, squeezed_state_fidelities, marker = 's', linestyle = 'solid', color = 'k', markersize = 3, linewidth = 1.5, label = r'State fidelity')
+    plt.xlabel('Time [µs] ', fontsize = 20)
+    plt.ylabel(r'$\mathcal{F}$', fontsize = 24, rotation = 0., labelpad=15)
+    plt.legend()
+    plt.show()
+
+    plt.figure(figsize=(5,5))
+    plt.plot(times*1E6, n_t, marker = 's', linestyle = 'solid', color = 'k', markersize = 3, linewidth = 1.5, label = r'Simulation: $N(t)$')
+    plt.axhline(y = np.sinh(r)**2, color = 'k', linestyle='dashed', linewidth = 1.5, label = r'$\langle N \rangle = \sinh(r)^2 $')
+    plt.xlabel('Time [µs] ', fontsize = 20)
+    plt.ylabel(r'$N$', fontsize = 24, rotation = 0., labelpad=15)
+    plt.legend()
+    plt.show()
+
+
+    fock_states = np.arange(fock_dimension)
+    squeezed_GS_populations = squeezed_vacuum_populations(fock_dimension, r)
+
+    bar_width = 0.35
+    plt.figure(figsize=(6,6))
+    plt.bar(fock_states - bar_width/2, fock_populations[-1], bar_width, color = 'b', label = 'Simulation')
+    if not do_reverse_process: 
+        plt.bar(fock_states + bar_width/2, squeezed_GS_populations, bar_width, color = 'r', label = 'Reference: $r_{target} = ' + str(np.round(r,2)) + '$')
+    else:
+        plt.bar(fock_states + bar_width/2, IC_fock_pops[0], bar_width, color = 'r', label = 'Initial distribution') 
+    plt.ylabel('Population', fontsize = 16)
+    plt.xlabel('Fock state $n$', fontsize = 20)
+    plt.legend()
+    plt.show()
+
+    if modulate_amplitude:
+        plt.figure(figsize=(5,5))
+        plt.plot(times*1E6, mod_function(times), marker = 's', linestyle = 'solid', color = 'k', markersize = 2, linewidth = 1.5) 
+        plt.xlabel('Time [µs] ', fontsize = 20)
+        plt.ylabel(r'$\Omega$', fontsize = 24, rotation = 0., labelpad=15)
+        plt.title('Amplitude Modulation')#: $\sigma = ' + str(variance))
+        #plt.legend()
+        plt.show()
+
+if __name__ == '__main__':
+    main()
diff --git a/examples/style/plot_style_data.txt b/examples/style/plot_style_data.txt
new file mode 100644
index 0000000..c4f45fb
--- /dev/null
+++ b/examples/style/plot_style_data.txt
@@ -0,0 +1,37 @@
+xtick.labelsize: 14
+ytick.labelsize: 14
+font.size: 14
+figure.autolayout: True
+figure.figsize: 7.2,4.45
+axes.titlesize : 14
+axes.linewidth : 1.25
+#axes.titleweight : bold
+axes.labelsize : 14
+ytick.right: False
+xtick.top: False
+xtick.minor.visible : False
+ytick.minor.visible : False
+#xtick.major.visible : False
+#ytick.major.visible : False
+ytick.major.right : False
+ytick.minor.right : False
+xtick.major.top : False
+xtick.minor.top : False
+xtick.major.size : 5
+ytick.major.size : 5
+xtick.minor.size : 2.5
+ytick.minor.size : 2.5
+ytick.direction : in
+xtick.direction : in
+ytick.labelright: False
+ytick.labelleft: True
+xtick.labelbottom: True
+xtick.labeltop : False
+lines.linewidth : 1.0
+lines.markersize : 6
+#legend.fontsize: 16
+legend.fontsize: 14
+#legend.loc : upper center
+legend.frameon : False
+mathtext.fontset: stix
+font.family: STIXGeneral
diff --git a/src/ionsim/degree_of_freedom.py b/src/ionsim/degree_of_freedom.py
index 7fad0a9..3718e7c 100644
--- a/src/ionsim/degree_of_freedom.py
+++ b/src/ionsim/degree_of_freedom.py
@@ -27,7 +27,7 @@ class AtomicSpin(DegreeOfFreedom):
     energy_levels: list[AtomicInternalEnergyLevel]
     # Store mass and atomic number 
     atomic_mass: float 
-    atomic_number: int # should we allowe for effective Z's which are non-integer?) 
+    atomic_charge: int 
     name: str | None = None # TODO: will we use these names?
 
     @classmethod
@@ -37,8 +37,9 @@ def from_species(cls, species: str, term_symbols: list[str] | None = None, level
         config_data = cls.get_config_data(species)
         nuclear_spin = config_data['nuclear_spin']
         levels_data = config_data['levels']
-        mass = config_data['mass'] # Daltons
-        z = config_data['Z'] # Atomic number, number of protons 
+        atomic_mass = config_data['mass'] # Daltons
+        atomic_charge = config_data['charge'] # Daltons
+        atomic_number = config_data['Z'] # Atomic number, number of protons 
         magnetic_moment = config_data['magnetic_moment'] # units of \mu_{N}
         structure = 'fine' if nuclear_spin == 0 else 'hyperfine' 
 
@@ -74,11 +75,11 @@ def from_species(cls, species: str, term_symbols: list[str] | None = None, level
                         gj = 2. * (gj1 - 1.) * (k*(k+1) + j1*(j1+1) - l2*(l2 + 1))/((2*j + 1)*(2*k + 1))
                         gj += (3*j*(j+1) - k*(k+1) + s2*(s2+1))/(2.*j*(j+1)) 
                         fine_data['gj'] = gj
-                    Zeeman_solver = ZeemanHyperfineSolver(nuclear_spin, j, None, s2, fine_data['hyperfine_A']*2.*np.pi, mass, magnetic_moment, z, gj = gj)
+                    Zeeman_solver = ZeemanHyperfineSolver(nuclear_spin, j, None, s2, fine_data['hyperfine_A']*2.*np.pi, atomic_mass, magnetic_moment, atomic_number, gj = gj)
                 else:
                     s = fine_data['s']
                     l = fine_data['l']
-                    Zeeman_solver = ZeemanHyperfineSolver(nuclear_spin, j, l, s, fine_data['hyperfine_A']*2.*np.pi, mass, magnetic_moment, z)
+                    Zeeman_solver = ZeemanHyperfineSolver(nuclear_spin, j, l, s, fine_data['hyperfine_A']*2.*np.pi, atomic_mass, magnetic_moment, atomic_number)
                 zeeman_energy_shifts, zeeman_eigenvecs = Zeeman_solver.solve_at_field(magnetic_field)
                 zeeman_energy_shifts *= np.pi*2. # convert to rad/s 
 
@@ -106,7 +107,7 @@ def from_species(cls, species: str, term_symbols: list[str] | None = None, level
                         level = HyperfineLevel(**fine_data, i=nuclear_spin, f=f, mf=mf, external_energy_shift = zeeman_shift_energy)
                         if level_names is None or level.name in level_names:
                             levels.append(level)
-        return cls(levels, mass, z, name)
+        return cls(levels, atomic_mass, atomic_charge, name)
 
     @classmethod
     def get_level_factory(cls, coupling_scheme: str):
diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 97af312..1ec2a07 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -7,19 +7,6 @@
 from ionsim.degree_of_freedom import AtomicSpin, MotionalMode
 from ionsim.basis import StandardBasis 
 
-########## Questions: 
-# -1. Should the branch sorting be the default option? This seems like what we would want to do for our problems  
-# 0. LD parameter matrix shape? 
-# 1.  "has run" boolean? e.g. is there a time where it stays False? Ah I think it refers to whether the eqb solver has been successfully executed. 
-# 2. Should the dimensionless parameters have a naming convention, e.g. omega_x --> omega_x_ND 
-
-### 3. What sets the phases of the Lamb-Dicke parameters?  
-#       - need to define a convention and stick with it 
-
-### Notes:
-# - won't get orthonormal eigenvectors if there's degeneracies; the orthonormality is w.r.t the H matrix 
-
-
 ## References: 
 # https://arxiv.org/abs/2007.12725
 # https://arxiv.org/abs/quant-ph/9702053 
@@ -32,7 +19,7 @@ def characteristic_length(q: float, mass: float, omega: float) -> float:
 
 
 class TrappedIonModeAnalysis:
-    def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float, atomic_masses: Vector | float, atomic_numbers: Vector | int,
+    def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float, atomic_masses: Vector | float, atomic_charges: Vector | int,
                  mode_organization: str = 'frequency_only', reindexing_strategy: str = 'distance'):
         """ Class for determining properties of trapped-ion phonon modes, requiring the following system parameters:
 
@@ -41,12 +28,11 @@ def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float
             - omega_y: harmonic trap frequency in the y direction in units of rad/s.
             - omega_z: harmonic trap frequency in the z direction. This is the "axial" direction used for 1D chains.
             - atomic masses: array of atomic masses or a single number => same mass for each ion. Units are amu (atomic mass units)
-            - atomic numbers: array of (Z) atomic numbers (# of protons of an element) or a single number => all ions are the same.
+            - atomic charges: array of (q) atomic charges (# of protons of an element) or a single number => all ions are the same.
             - mode_organization: strategy for organizing modes ('frequency_only' or 'branch_sorted')
             - reindexing_strategy: strategy for reindexing ions ('distance' or 'z_axis')
 
         """
-        # TODO: Decide how to handle units and how much of pint we should use.
         self.num_ions = num_ions
 
         # Store configuration strategies
@@ -60,14 +46,15 @@ def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float
             raise ValueError(f"Invalid reindexing_strategy: {reindexing_strategy}. Must be 'distance' or 'z_axis'")
 
         self.atomic_masses = self.convert_to_array(atomic_masses) * const.u # kg from amu 
-        self.atomic_numbers = self.convert_to_array(atomic_numbers)
+        self.atomic_charges = self.convert_to_array(atomic_charges)
 
         # Safety checks: 
         assert len(self.atomic_masses) == num_ions
-        assert len(self.atomic_numbers) == num_ions
+        assert len(self.atomic_charges) == num_ions
 
         # Charge of all the protons -> total nuclear charge:
-        self.nuclear_charges = self.atomic_numbers * const.e # Z * e, const.e ==> elementary charge in Coulombs 
+        self.nuclear_charges = np.ones_like(self.atomic_masses) * const.e # Z * e, const.e ==> elementary charge in Coulombs 
+        #self.nuclear_charges = self.atomic_charges * const.e # Z * e, const.e ==> elementary charge in Coulombs 
 
         # Trapping frequencies for each ion 
         # TODO: loop and vectorize trap frequencies?  
@@ -106,7 +93,7 @@ def convert_to_array(self, x: float | int | list | NDArray) -> NDArray:
             return x
         return np.array(x) 
 
-    def set_up_dimensionless_parameeters(self, charge_scale: float=1., mass_scale: float=1., trap_freq_scale: float=1.):
+    def set_up_dimensionless_parameters(self, charge_scale: float=1., mass_scale: float=1., trap_freq_scale: float=1.):
         """ Compute dimensionless parameters, using first ion's properties and axial (z) trapping frequency. 
 
             - characteristic length
@@ -120,8 +107,6 @@ def set_up_dimensionless_parameeters(self, charge_scale: float=1., mass_scale: f
             raise IonSimError(f"Trap frequency scale should be positive. Received {trap_freq_scale}")
         if mass_scale <= 0:
             raise IonSimError(f"Mass scale should be positive. Received {mass_scale}")
- #        if charge_scale <= 0:
- #            raise IonSimError(f"Charge scale should be positive. Received {charge_scale}")
 
         # Store the scales so they are retrievable by the user 
         self.charge_scale = charge_scale
@@ -142,10 +127,9 @@ def set_up_dimensionless_parameeters(self, charge_scale: float=1., mass_scale: f
         # Compute characteristic length, time, velocity, and energy scales and store in a dictionary  
         self.characteristic_parameters = {}
         self.characteristic_parameters['length'] = characteristic_length(self.charge_scale, self.mass_scale, self.trap_freq_scale)  
-        self.characteristic_parameters['time'] = 1 / trap_freq_scale # characteristic time
-        self.characteristic_parameters['velocity'] = self.characteristic_parameters['length'] * trap_freq_scale  # characteristic velocity
-        self.characteristic_parameters['energy'] = 0.5 * mass_scale * self.characteristic_parameters['velocity'] ** 2  # characteristic energy  
-
+        self.characteristic_parameters['time'] = 1 / self.trap_freq_scale # characteristic time
+        self.characteristic_parameters['velocity'] = self.characteristic_parameters['length'] * self.trap_freq_scale  # characteristic velocity
+        self.characteristic_parameters['energy'] = 0.5 * self.mass_scale * self.characteristic_parameters['velocity'] ** 2  # characteristic energy  
 
     def get_eigenvector_norm(self, eigenvector: Vector, H: Matrix) -> float:
         """ Computes the norm of the eigenvector en w.r.t. the Hamiltonian H. """
@@ -182,7 +166,6 @@ def check_outer_relation(self, H: Matrix):
         I_right = H @ Outers
         I_left  = Outers @ H
         eye = np.eye(6*self.num_ions,dtype=complex)
-        np.set_printoptions(precision=2, suppress=True) 
         try:
             assert np.allclose(I_left,eye) 
             assert np.allclose(I_right,eye)
@@ -193,11 +176,10 @@ def check_diagonalization(self, T: Matrix, S: Matrix, E: Matrix) -> bool:
         M = np.linalg.inv(T) @ S
         H_diag = M.T @ E @ M    
         H_diag_check = np.diag(np.tile(self.eigvals,2)) 
-        np.set_printoptions(precision=2, suppress=True) 
         try:
             assert np.allclose(H_diag,H_diag_check)
         except AssertionError:
-            has_duplicate_eigenvalues = np.any(np.triu(np.isclose(eigenvalues[:, None], eigenvalues[None, :], atol = 1E-6), k=1))
+            has_duplicate_eigvals = np.any(np.triu(np.isclose(self.eigvals[:, None], self.eigvals[None, :], atol = 1E-6), k=1))
             warnings.warn("Diagnolization check failed")    
             print("has duplicate eigenvalues: ", has_duplicate_eigvals)
 
@@ -215,7 +197,7 @@ def solve_ion_trap_equilibrium(self):
         """
         # Convert to dimensionless units using axial trap frequency and first ion's mass and charge
         # TODO: take in input here or in class constructor for mass, charge, trap scales.
-        self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
+        self.set_up_dimensionless_parameters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
         self.equilibrium_positions = self.solve_for_equilibrium_positions()
 
         # Use configurable reindexing strategy
@@ -308,7 +290,7 @@ def reindex_ions(self):
         self.q = self.q[idx]
         self.atomic_masses = self.atomic_masses[idx]
         self.nuclear_charges = self.nuclear_charges[idx]
-        self.atomic_numbers = self.atomic_numbers[idx]
+        self.atomic_charges = self.atomic_charges[idx]
 
     def reindex_ions_by_z(self):
         """Re-indexes ions based on their position along the z-axis, with the ion closest to the center having index 0."""
@@ -322,14 +304,14 @@ def reindex_ions_by_z(self):
         self.q = self.q[idx]
         self.atomic_masses = self.atomic_masses[idx]
         self.nuclear_charges = self.nuclear_charges[idx]
-        self.atomic_numbers = self.atomic_numbers[idx]
+        self.atomic_charges = self.atomic_charges[idx]
 
         # trapping frequencies
         self.omega_x = self.omega_x[idx]
         self.omega_y = self.omega_y[idx]
         self.omega_z = self.omega_z[idx]
         # all ions are the same so this is safe. TODO: When does this change?
-        self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
+        self.set_up_dimensionless_parameters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
 
 
     def solve_for_equilibrium_positions(self, positions_guess: Vector | None=None):
@@ -600,15 +582,11 @@ def calculate_mode_participation_factors(self) -> Matrix:
         return mode_participation_factors
 
 
-    # Can get lamb-dicke parameters from wavevctor \dot mode_participation_factors[:, i, m]
-    ### -- Coordinate systems must be the same / correspond. 
-
     # Derived properties 
     def compute_reference_single_ion_lamb_dicke_factors(self, wavenumber: float) -> (float, float, float):
         """ Computes analytical Lamb-Dicke parameter by eta = k * sqrt(hbar / m omega) for an ion in a light-field with wavevector |k| 
             - wavenumber: wavevector magnitude |k| in units of 1 / m  
         """
-        # TODO: better way to handle units? 
         # Convention to use first value of trap frequency arrays (representing one of the ions)  
         trap_frequencies = np.array([self.omega_x[0], self.omega_y[0], self.omega_z[0]])        
         mass = self.atomic_masses[0] 
@@ -628,29 +606,29 @@ def from_species(cls, species_name: str, num_ions: int, omega_x: float, omega_y:
         # Import necessary data for the species
         species_data = AtomicSpin.get_config_data(species_name)
         atomic_mass = species_data['mass']
-        atomic_number = species_data['Z']
+        atomic_charge = species_data['charge']
         # Construct the class
-        return cls(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_number)
+        return cls(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_charge)
 
     @classmethod
-    def from_atomic_spin_basis(cls, spins: list[degree_of_freedom], omega_x: float, omega_y: float, omega_z: float): 
+    def from_atomic_spin_basis(cls, spin_basis: StandardBasis, omega_x: float, omega_y: float, omega_z: float): 
         """ Build the mode analysis class from a basis of AtomicSpin degrees of freedom under harmonic trapping. """ 
         # Extract number of ions and the mass and atomic number from the DOF in the basis 
-        DOFs = spins 
-        num_ions = len(DOFs)
+        spin_DOFs = spin_basis.degrees_of_freedom 
+        num_ions = len(spin_DOFs)
         atomic_masses = []
-        atomic_numbers = []
+        atomic_charges = []
 
-        for DOF in DOFs:
+        for DOF in spin_DOFs:
             if not isinstance(DOF, AtomicSpin):
-                raise IonSimError("Atomic structure basis should only contain AtomicSpin or AtomicStructure objects. No motional modes should be included.")
-                atomic_masses.append(DOF.atomic_mass) 
-                atomic_numbers.append(DOF.atomic_number) 
+                raise IonSimError("Atomic structure basis should only contain AtomicSpin objects. No motional modes should be included.")
+            atomic_masses.append(DOF.atomic_mass) 
+            atomic_charges.append(DOF.atomic_charge) 
         # Construct the class 
-        return cls(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_number)
+        return cls(num_ions, omega_x, omega_y, omega_z, atomic_masses, atomic_charges)
 
     def build_mode_DOFs(self, mode_indices: list[int], fock_dimensions: Vector | int) -> list[MotionalMode]:
-        """Builds and returns an IonSim Motional Degree of Freedom.
+        """ Builds and returns a list of IonSim Motional Degrees of Freedom.
 
             - Applies each fock dimension to each mode, or applies the same fock dimension to all the modes
         """
@@ -693,6 +671,8 @@ def get_lamb_dicke_parameters_for_mode(self, mode_index: int, wavevector: Vector
 
                 Returns an array of shape (num_ions,) containing Lamb-Dicke parameters
                 where result[ion] gives the LD parameter for that ion in the specified mode.
+
+                Note: Lamb Dicke parameters have an overall global phase freedom. 
         """
         self._ensure_mode_analysis_run()
 
@@ -752,16 +732,17 @@ class LinearIonChainAnalysis(TrappedIonModeAnalysis):
     """
 
     def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float,
-                 atomic_masses: np.ndarray | float, atomic_numbers: np.ndarray | int):
+                 atomic_masses: np.ndarray | float, atomic_charges: np.ndarray | int):
         """ Initialize a linear ion chain analysis, defaulting to branch sorted modes. """
         mode_organization = 'branch_sorted'
         reindexing_strategy = 'z_axis'
-        super().__init__(num_ions, omega_x, omega_y, omega_z, atomic_masses, atomic_numbers, mode_organization, reindexing_strategy)
+        super().__init__(num_ions, omega_x, omega_y, omega_z, atomic_masses, atomic_charges, mode_organization, reindexing_strategy)
 
     def solve_ion_trap_equilibrium(self):
         """ Solve for equilibrium positions and analyze normal modes for a linear chain. It reindexes ions by their axial position. """
         # Set up dimensionless parameters using first ion's properties and axial trap frequency
-        self.set_up_dimensionless_parameeters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
+        self.set_up_dimensionless_parameters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
+
 
         # Solve for equilibrium positions
         self.equilibrium_positions = self.solve_for_equilibrium_positions()
@@ -818,15 +799,20 @@ def get_radial_modes(self, direction='x') -> tuple:
 
         return radial_eigvals, radial_eigvecs
 
-    def get_ion_spacing(self):
-        """ Returns an array of ion spacings between consecutive ions along the z-axis. """
+    def get_axial_ion_positions(self):
+        """ Returns an array of ion axial positions along the z-axis. """
         # Get z positions (dimensionful)
         x, y, z = self.ion_coordinates_from_flattened(self.equilibrium_positions)
         z_dimensionful = z * self.characteristic_parameters['length']
-
-        # Sort positions and calculate spacing
         sorted_z = np.sort(z_dimensionful)
-        ion_spacing = np.diff(sorted_z)
+        return sorted_z
+
+
+    def get_ion_spacing(self):
+        """ Returns an array of ion spacings between consecutive ions along the z-axis. """
+        # Get z positions (dimensionful)
+        sorted_z_coords = self.get_axial_ion_positions()
+        ion_spacing = np.diff(sorted_z_coords)
 
         return ion_spacing
 
@@ -937,7 +923,7 @@ def get_lamb_dicke_parameters_by_branch(self, branch: str, mode_index_in_branch:
         # Use parent class method to get LD parameters
         return super().get_lamb_dicke_parameters_for_mode(global_mode_index, wavevector)
 
-    def calculate_lamb_dicke_parameters(self, wavevector: Vector) -> dict:
+    def calculate_lamb_dicke_parameters_by_branch(self, wavevector: Vector) -> dict:
         """ Calculate Lamb-Dicke parameters for all ions and modes, organized by branch.
 
             The Lamb-Dicke parameter η = k · Δr_0 represents the ratio of the spatial
@@ -979,12 +965,32 @@ def calculate_lamb_dicke_parameters(self, wavevector: Vector) -> dict:
 
         return lamb_dicke_parameters
 
+    def build_mode_DOFs_from_branch(self, branch: str, mode_indices: list[int], fock_dimensions: Vector | int) -> list[MotionalMode]:
+        """ Builds and returns a list of IonSim Motional Degrees of Freedom.
+
+            - mode_indices corresponds to a mode's index within a branch (instead of the entire set of 3N modes for N ions)
+            - Applies each fock dimension to each mode, or applies the same fock dimension to all the modes
+            - ex] branch = 'x', mode_indices = [0,1] (tilt and COM), fock_dimensions = 6
+        """
+        modes = []
+
+        # Convert list of Fock dimensions to an array if it's not already an array
+        fock_dimensions = self.convert_to_array(fock_dimensions).astype(int)
+        for mode_idx, fock_dim in zip(mode_indices, fock_dimensions):
+            mode_props = self.get_mode_properties_by_branch(branch, mode_idx)
+            frequency = mode_props['frequency'] # rad/s 
+            modes.append(MotionalMode.from_frequency(frequency, fock_dim))
+
+        return modes
+
     def calculate_lamb_dicke_parameters_full(self, wavevector: Vector) -> Matrix:
         """ Calculate full Lamb-Dicke parameter matrix from a laser wavevector k = (kx, ky, kz) 
 
             Returns: 
                 - Total LD parameter matrix of shape (num_ions, num_modes) where eta[ion, mode] 
                     gives the total LD parameter k · Δr_0.
+
+                Note: Lamb Dicke parameters have an overall global phase freedom. 
         """
         # Get full mode participation factors
         mode_pf = self.calculate_mode_participation_factors()
@@ -997,6 +1003,7 @@ def calculate_lamb_dicke_parameters_full(self, wavevector: Vector) -> Matrix:
         num_modes = 3 * self.num_ions
         total_ld_params = np.zeros((self.num_ions, num_modes), dtype=complex)
 
+        # TODO: Replace with dot product for efficiency 
         for i in range(3):
             total_ld_params += wavevector[i] * mode_pf[i, :, :]
 
@@ -1007,18 +1014,20 @@ def get_axial_lamb_dicke_parameters(self, wavevector: Vector) -> Vector:
 
             Returns: Lamb-Dicke parameters for axial modes, NDArray of shape (num_ions, num_axial_modes)
         """
-        lamb_dicke_by_branch = self.calculate_lamb_dicke_parameters(wavevector)
+        lamb_dicke_by_branch = self.calculate_lamb_dicke_parameters_by_branch(wavevector)
         return lamb_dicke_by_branch['z']
 
     def get_radial_lamb_dicke_parameters(self, wavevector: Vector | float, direction: str = 'x') -> Vector:
         """ Get Lamb-Dicke parameters for radial modes modes in a specified direction.
 
             Returns: Lamb-Dicke parameters for radial modes, array of shape (num_ions, num_radial_modes)
+
+                Note: Lamb Dicke parameters have an overall global phase freedom. 
         """
         if direction not in ['x', 'y']:
             raise ValueError("Direction must be 'x' or 'y'")
 
-        lamb_dicke_by_branch = self.calculate_lamb_dicke_parameters(wavevector)
+        lamb_dicke_by_branch = self.calculate_lamb_dicke_parameters_by_branch(wavevector)
         return lamb_dicke_by_branch[direction]
 
     def get_two_central_ion_separation(self) -> float:
@@ -1053,7 +1062,7 @@ def separation_error(wz_Hz):
                 omega_y=self.omega_y[0],
                 omega_z=2*np.pi*wz_Hz,
                 atomic_masses=self.atomic_masses[0],
-                atomic_numbers=self.atomic_numbers[0]
+                atomic_charges=self.atomic_charges[0]
             )
             temp_analysis.solve_ion_trap_equilibrium()
             current_separation = temp_analysis.get_two_central_ion_separation()
@@ -1070,7 +1079,8 @@ def separation_error(wz_Hz):
 
     def print_chain_summary(self):
         """ Print a summary of the linear ion chain configuration. """
-        print(f"Linear Ion Chain Analysis Summary")
+        print()
+        print(f" ----- Linear Ion Chain Analysis Summary -----")
         print(f"Number of ions: {self.num_ions}")
         print(f"Trap frequencies: x={self.omega_x[0]/(2*np.pi)/1e6:.2f} MHz, "
               f"y={self.omega_y[0]/(2*np.pi)/1e6:.2f} MHz, "
@@ -1078,7 +1088,9 @@ def print_chain_summary(self):
 
         # Get ion spacing
         spacing = self.get_ion_spacing()
-        print(f"Ion spacing (z-axis): {spacing*1e6:.2f} µm")
+        z_coordinates = self.get_axial_ion_positions() 
+        print(f"Ion axial coordinates (z-axis): {z_coordinates*1e6} µm")
+        print(f"Ion spacing (z-axis): {spacing*1e6} µm")
 
         # Get COM position
         com_x, com_y, com_z = self.get_center_of_mass_position()
@@ -1089,7 +1101,7 @@ def print_chain_summary(self):
         radial_x_freqs = self.get_radial_mode_frequencies('x')
         radial_y_freqs = self.get_radial_mode_frequencies('y')
 
-        print(f"\nAxial mode frequencies: {axial_freqs/(2*np.pi)/1e6:.2f} MHz")
-        print(f"Radial X mode frequencies: {radial_x_freqs/(2*np.pi)/1e6:.2f} MHz")
-        print(f"Radial Y mode frequencies: {radial_y_freqs/(2*np.pi)/1e6:.2f} MHz")
+        print(f"\nAxial mode frequencies: {axial_freqs/(2*np.pi)/1e6} MHz")
+        print(f"Radial X mode frequencies: {radial_x_freqs/(2*np.pi)/1e6} MHz")
+        print(f"Radial Y mode frequencies: {radial_y_freqs/(2*np.pi)/1e6} MHz")
 
diff --git a/tests/test_mode_analysis.py b/tests/test_mode_analysis.py
index 60c9c7b..3512436 100644
--- a/tests/test_mode_analysis.py
+++ b/tests/test_mode_analysis.py
@@ -18,8 +18,8 @@ def setUp(self):
                 'omega y' : 1.9 * TPI * 1E6,
                 'omega z' : 0.5 * TPI * 1E6,
                 'N' : 1, # number of ions 
-                'wavevector' : TPI / (355.0*1E-9), # 1/m
-                'Z' : 70 
+                'wavenumber' : TPI / (355.0*1E-9), # 1/m
+                'charge' : 1 
             },
             {
                 'test name' : 'two ion 171Yb',
@@ -28,12 +28,12 @@ def setUp(self):
                 'omega y' : 1.5 * TPI * 1E6,
                 'omega z' : 0.5 * TPI * 1E6,
                 'N' : 2, # number of ions 
-                'wavevector' : TPI / (355.0*1E-9), # 1/m
-                'Z' : 70 
+                'wavenumber' : TPI / (355.0*1E-9), # 1/m
+                'charge' : 1 
             }
         ]
         self.mode_analyzers = {
-            case['test name'] : TrappedIonModeAnalysis(case['N'], case['omega x'], case['omega y'], case['omega z'], case['mass'], case['Z']) 
+            case['test name'] : TrappedIonModeAnalysis(case['N'], case['omega x'], case['omega y'], case['omega z'], case['mass'], case['charge']) 
                 for case in self.test_cases
             }
 
@@ -41,7 +41,7 @@ def setUp(self):
         self.references = {}
 
         # 1. Single ion case: 
-        k = self.test_cases[0]['wavevector']
+        k = self.test_cases[0]['wavenumber']
         eta_x, eta_y, eta_z = self.mode_analyzers['single ion 171Yb'].compute_reference_single_ion_lamb_dicke_factors(k)
         etas_analytical = np.zeros((3, 1, 3), dtype = np.complex128)
         etas_analytical[0, 0, 2] = eta_x
@@ -50,7 +50,7 @@ def setUp(self):
         self.references['single ion 171Yb'] = etas_analytical
 
         # 2. Two-ion case: 
-        k = self.test_cases[1]['wavevector']
+        k = self.test_cases[1]['wavenumber']
         eta_x, eta_y, eta_z = self.mode_analyzers['two ion 171Yb'].compute_reference_single_ion_lamb_dicke_factors(k)
         wx = self.test_cases[1]['omega x']
         wy = self.test_cases[1]['omega y']
@@ -88,22 +88,24 @@ def test_mode_analysis_solvers(self):
             # Solve the ion-trap equilibrium problem
             mode_analyzer.solve_ion_trap_equilibrium()
             # Compute and store Lamb-Dicke parameters:
-            computed_Lamb_Dicke_parameters = case['wavevector'] * mode_analyzer.calculate_mode_participation_factors()
+            computed_Lamb_Dicke_parameters = case['wavenumber'] * mode_analyzer.calculate_mode_participation_factors()
             assert_array_close(np.abs(computed_Lamb_Dicke_parameters), np.abs(reference_values), atol = 1E-10, rtol=None)
 
     def test_linear_ion_chain_analysis(self):
         """Test the LinearIonChainAnalysis class with a 5-ion Yb+ chain."""
         # Create a 5-ion linear chain of Yb+ ions
         num_ions = 5
-        omega_x = 2 * np.pi * 10e6  # 10 MHz
-        omega_y = 2 * np.pi * 10e6  # 10 MHz
-        omega_z = 2 * np.pi * 1e6   # 1 MHz (axial frequency typically lower)
+        TPI = 2*np.pi
+        omega_x = TPI * 9e6  # 10 MHz
+        omega_y = TPI * 10e6  # 10 MHz
+        omega_z = TPI * 1e6   # 1 MHz (axial frequency typically lower)
         atomic_mass = 170.936  # Yb+ mass in amu
-        atomic_number = 70      # Yb atomic number
+        atomic_charge = 1      
 
         # Create and analyze the linear chain
-        chain = LinearIonChainAnalysis(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_number)
+        chain = LinearIonChainAnalysis(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_charge)
         chain.solve_ion_trap_equilibrium()
+        #chain.print_chain_summary()
 
         # Test that we can get axial and radial modes
         axial_eigvals, axial_eigvecs = chain.get_axial_modes()
@@ -148,14 +150,16 @@ def test_linear_ion_chain_analysis(self):
         self.assertTrue(np.all(radial_y_freqs > 0))
 
         # Test Lamb-Dicke parameter calculations
-        wavenumber = 2 * np.pi / (355.0 * 1e-9)  # Typical laser wavelength for Yb+
+        wavenumber = 2 * np.pi / (355.0 * 1e-9)  # example laser wavelength used with Yb+
 
         # Test full Lamb-Dicke parameter matrix
-        full_ld_params = chain.calculate_lamb_dicke_parameters_full(wavenumber)
-        self.assertEqual(full_ld_params.shape, (3, num_ions, 3*num_ions))
+        wavevector = wavenumber * np.array([np.cos(np.pi/4.), np.sin(np.pi/4.), 0.]) # perpendicular to the axial chain 
+        full_ld_params = chain.calculate_lamb_dicke_parameters_full(wavevector)
+        num_modes = 3*num_ions
+        self.assertEqual(full_ld_params.shape, (num_ions, num_modes))
 
         # Test branch-organized Lamb-Dicke parameters
-        ld_by_branch = chain.calculate_lamb_dicke_parameters(wavenumber)
+        ld_by_branch = chain.calculate_lamb_dicke_parameters_by_branch(wavevector)
         self.assertIn('x', ld_by_branch)
         self.assertIn('y', ld_by_branch)
         self.assertIn('z', ld_by_branch)
@@ -164,56 +168,30 @@ def test_linear_ion_chain_analysis(self):
         self.assertEqual(ld_by_branch['z'].shape, (num_ions, num_ions))
 
         # Test axial Lamb-Dicke parameters
-        axial_ld = chain.get_axial_lamb_dicke_parameters(wavenumber)
+        axial_ld = chain.get_axial_lamb_dicke_parameters(wavevector)
         self.assertEqual(axial_ld.shape, (num_ions, num_ions))
 
         # Test radial Lamb-Dicke parameters
-        radial_x_ld = chain.get_radial_lamb_dicke_parameters(wavenumber, 'x')
-        radial_y_ld = chain.get_radial_lamb_dicke_parameters(wavenumber, 'y')
+        radial_x_ld = chain.get_radial_lamb_dicke_parameters(wavevector, 'x')
+        radial_y_ld = chain.get_radial_lamb_dicke_parameters(wavevector, 'y')
         self.assertEqual(radial_x_ld.shape, (num_ions, num_ions))
         self.assertEqual(radial_y_ld.shape, (num_ions, num_ions))
 
         # Verify that branch-organized LD params match the corresponding parts of full matrix
         n_modes_per_branch = num_ions
-        np.testing.assert_array_almost_equal(ld_by_branch['x'], full_ld_params[0, :, :n_modes_per_branch])
-        np.testing.assert_array_almost_equal(ld_by_branch['y'], full_ld_params[1, :, n_modes_per_branch:2*n_modes_per_branch])
-        np.testing.assert_array_almost_equal(ld_by_branch['z'], full_ld_params[2, :, 2*n_modes_per_branch:3*n_modes_per_branch])
+        np.testing.assert_array_almost_equal(ld_by_branch['x'], full_ld_params[:, :n_modes_per_branch])
+        np.testing.assert_array_almost_equal(ld_by_branch['y'], full_ld_params[:, n_modes_per_branch:2*n_modes_per_branch])
+        np.testing.assert_array_almost_equal(ld_by_branch['z'], full_ld_params[:, 2*n_modes_per_branch:3*n_modes_per_branch])
 
         # Test that Lamb-Dicke parameters are related to mode participation factors
-        mode_pf = chain.calculate_mode_participation_factors()
-        expected_full_ld = wavenumber * mode_pf
-        np.testing.assert_array_almost_equal(full_ld_params, expected_full_ld)
+        mode_pf = chain.get_mode_participation_factors_by_branch()
+        #print(mode_pf['x'])
+        #print(ld_by_branch['z'])
+        #mode_pf = chain.get_mode_participation_factors_by_branch()
+        #expected_full_ld = wavenumber * mode_pf
+        #np.testing.assert_array_almost_equal(full_ld_params, expected_full_ld)
+
 
 if __name__ == '__main__':
     # Run the unit tests
     unittest.main()
-
-    # Example usage of the new methods (from the original main section)
-    print("\n" + "="*60)
-    print("Example usage of LinearIonChainAnalysis methods")
-    print("="*60)
-
-    # Create a 4-ion analysis
-    four_ion_analysis = TrappedIonModeAnalysis(
-        num_ions=4,
-        omega_x=2*np.pi*11e6,
-        omega_y=2*np.pi*10e6,
-        omega_z=2*np.pi*10e6,  # Start with 10 MHz axial frequency
-        atomic_masses=170.936,
-        atomic_numbers=70
-    )
-    four_ion_analysis.solve_ion_trap_equilibrium()
-
-    # Get current central ion separation
-    current_separation = four_ion_analysis.get_two_central_ion_separation()
-    print(f"Current central ion separation: {current_separation*1e6:.3f} µm")
-
-    # Find axial frequency for desired separation (e.g., 50 µm)
-    target_separation = 50e-6  # 50 micrometers
-    wz_for_target = four_ion_analysis.find_axial_frequency_for_desired_central_ion_separation(
-        target_separation,
-        bounds=(0.1e6, 0.5e6)
-    )
-    print(f"Required axial frequency for {target_separation*1e6:.1f} µm separation: {wz_for_target/(2*np.pi)/1e6:.3f} MHz")
-
-    print("\nExample completed successfully!")

From e369e210f9b38997f5a8c2d16320ec8e9bb7e1a9 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Thu, 23 Jul 2026 12:31:50 -0600
Subject: [PATCH 30/31]   name change of class for readability, more intuitive
 naming

---
 examples/example_spin_dependent_squeezing.py | 2 +-
 src/ionsim/__init__.py                       | 2 +-
 src/ionsim/trapped_ion_mode_analysis.py      | 9 +++------
 tests/test_mode_analysis.py                  | 8 ++++----
 4 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/examples/example_spin_dependent_squeezing.py b/examples/example_spin_dependent_squeezing.py
index 5b39c95..4c8377b 100644
--- a/examples/example_spin_dependent_squeezing.py
+++ b/examples/example_spin_dependent_squeezing.py
@@ -287,7 +287,7 @@ def main():
     omega_x = 2.1 * TPI * 1E6 # MHz -> rad/s  
     omega_y = 2.6 * TPI * 1E6 # ""  
     omega_z = 0.50 * TPI * 1E6 # ""  
-    trap_analysis = ism.LinearIonChainAnalysis.from_atomic_structure_basis(qubit_basis, omega_x, omega_y, omega_z) 
+    trap_analysis = ism.LinearIonChain.from_atomic_structure_basis(qubit_basis, omega_x, omega_y, omega_z) 
     trap_analysis.solve_ion_trap_equilibrium()
     print()
     trap_analysis.print_chain_summary()
diff --git a/src/ionsim/__init__.py b/src/ionsim/__init__.py
index 9d5b213..b31db74 100644
--- a/src/ionsim/__init__.py
+++ b/src/ionsim/__init__.py
@@ -30,4 +30,4 @@ def patched_kron(a, b, *args, **kwargs):
 from .zeeman_solver import ZeemanHyperfineSolver
 from .composite_operator import CompositeOperator
 from .dissipator import Dissipator, DissipatorSpontaneousEmission, Lindbladian 
-from .trapped_ion_mode_analysis import TrappedIonModeAnalysis, LinearIonChainAnalysis 
+from .trapped_ion_mode_analysis import TrappedIonCrystal, LinearIonChain
diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index a024b40..41ce6ef 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -18,7 +18,7 @@ def characteristic_length(q: float, mass: float, omega: float) -> float:
     return l0
 
 
-class TrappedIonModeAnalysis:
+class TrappedIonCrystal:
     def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float, atomic_masses: Vector | float, atomic_charges: Vector | int,
                  mode_organization: str = 'frequency_only', reindexing_strategy: str = 'distance'):
         """ Class for determining properties of trapped-ion phonon modes, requiring the following system parameters:
@@ -716,10 +716,7 @@ def get_mode_properties(self, mode_index: int) -> dict:
         }
 
 
-#You could redefine the equilibrium finding function to assert that the equilibrium is linear. For example, make a wrapper inside the function for the potential, Jacobian, and Hessian that forces x_i and y_i = 0. 
-#This is the code for the Lamb-Dicke parameters: (not that overall phases don't matter here, but the relative phase does. We could pin down the phase with some convention like, "each mode's first non-negative LD value is defined positive.")
-
-class LinearIonChainAnalysis(TrappedIonModeAnalysis):
+class LinearIonChain(TrappedIonCrystal):
     """ Subclass for setting up and analyzing linear ion chains.
 
         In a linear ion chain, ions are typically aligned along the z-axis (axial direction),
@@ -1056,7 +1053,7 @@ def find_axial_frequency_for_desired_central_ion_separation(self, target_separat
 
         def separation_error(wz_Hz):
             # Create a temporary analysis object with the current wz
-            temp_analysis = LinearIonChainAnalysis(
+            temp_analysis = LinearIonChain(
                 num_ions=self.num_ions,
                 omega_x=self.omega_x[0],
                 omega_y=self.omega_y[0],
diff --git a/tests/test_mode_analysis.py b/tests/test_mode_analysis.py
index 3512436..704a651 100644
--- a/tests/test_mode_analysis.py
+++ b/tests/test_mode_analysis.py
@@ -1,7 +1,7 @@
 import unittest
 import numpy as np
 from scipy.linalg import expm
-from ionsim.trapped_ion_mode_analysis import TrappedIonModeAnalysis, LinearIonChainAnalysis
+from ionsim.trapped_ion_mode_analysis import TrappedIonCrystal, LinearIonChain
 from ionsim.testing import assert_array_close
 
 class TestModeAnalysis(unittest.TestCase):
@@ -33,7 +33,7 @@ def setUp(self):
             }
         ]
         self.mode_analyzers = {
-            case['test name'] : TrappedIonModeAnalysis(case['N'], case['omega x'], case['omega y'], case['omega z'], case['mass'], case['charge']) 
+            case['test name'] : TrappedIonCrystal(case['N'], case['omega x'], case['omega y'], case['omega z'], case['mass'], case['charge']) 
                 for case in self.test_cases
             }
 
@@ -92,7 +92,7 @@ def test_mode_analysis_solvers(self):
             assert_array_close(np.abs(computed_Lamb_Dicke_parameters), np.abs(reference_values), atol = 1E-10, rtol=None)
 
     def test_linear_ion_chain_analysis(self):
-        """Test the LinearIonChainAnalysis class with a 5-ion Yb+ chain."""
+        """Test the LinearIonChain class with a 5-ion Yb+ chain."""
         # Create a 5-ion linear chain of Yb+ ions
         num_ions = 5
         TPI = 2*np.pi
@@ -103,7 +103,7 @@ def test_linear_ion_chain_analysis(self):
         atomic_charge = 1      
 
         # Create and analyze the linear chain
-        chain = LinearIonChainAnalysis(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_charge)
+        chain = LinearIonChain(num_ions, omega_x, omega_y, omega_z, atomic_mass, atomic_charge)
         chain.solve_ion_trap_equilibrium()
         #chain.print_chain_summary()
 

From 12592d275b9a4a16385005e0c54a36fa2a5607a6 Mon Sep 17 00:00:00 2001
From: Ethan McGarrigle <ecmcgar@sandia.gov>
Date: Thu, 23 Jul 2026 13:12:12 -0600
Subject: [PATCH 31/31]   minor clean up of extra comments in mode analysis
 class

---
 src/ionsim/trapped_ion_mode_analysis.py | 23 +++++------------------
 1 file changed, 5 insertions(+), 18 deletions(-)

diff --git a/src/ionsim/trapped_ion_mode_analysis.py b/src/ionsim/trapped_ion_mode_analysis.py
index 41ce6ef..161d6f5 100644
--- a/src/ionsim/trapped_ion_mode_analysis.py
+++ b/src/ionsim/trapped_ion_mode_analysis.py
@@ -57,7 +57,6 @@ def __init__(self, num_ions: int, omega_x: float, omega_y: float, omega_z: float
         #self.nuclear_charges = self.atomic_charges * const.e # Z * e, const.e ==> elementary charge in Coulombs 
 
         # Trapping frequencies for each ion 
-        # TODO: loop and vectorize trap frequencies?  
         self.omega_x = self.calculate_species_trap_frequencies(omega_x)   
         self.omega_y = self.calculate_species_trap_frequencies(omega_y)
         self.omega_z = self.calculate_species_trap_frequencies(omega_z)    
@@ -141,14 +140,13 @@ def get_canonical_transformation(self, H, eigenvectors, eigenvalues: Vector | No
         """ Given the eigensolve of the dynamical matrix, get the transform matrix to the canonical coordinates.
             X = S X', where X' = (Q,P)^T and X = (q,p)^T
         """
-        ## TODO: understand the sign in the transformation matrix
+        # the sign due to the convention of writing the time evolution as exp(-iwt)
         sign = -1
         num_coords, num_eigenvalues = np.shape(eigenvectors)
         assert num_coords //2 == num_eigenvalues 
         T = np.zeros((num_coords,num_coords),dtype=complex) 
         eigenvectors = self.normalize_eigenvectors(eigenvectors, H, eigenvalues)
         T = np.sqrt(2)*np.concatenate((np.real(eigenvectors), sign*np.imag(eigenvectors)), axis=1)
-        # the sign is likely due to the convention of writing the time evolution as exp(-iwt)
         return T
 
     def check_for_zero_modes(self):
@@ -196,14 +194,13 @@ def solve_ion_trap_equilibrium(self):
 
         """
         # Convert to dimensionless units using axial trap frequency and first ion's mass and charge
-        # TODO: take in input here or in class constructor for mass, charge, trap scales.
         self.set_up_dimensionless_parameters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
         self.equilibrium_positions = self.solve_for_equilibrium_positions()
 
         # Use configurable reindexing strategy
         if self.reindexing_strategy == 'z_axis':
             self.reindex_ions_by_z()
-        else:  # 'distance'
+        else:  
             self.reindex_ions()
 
         # Compute helper matrices for computing normal mode properties
@@ -220,7 +217,6 @@ def solve_ion_trap_equilibrium(self):
         # Perform checks
         self.check_outer_relation(H_matrix)
         self.check_diagonalization(T_matrix, S_matrix, E_matrix)
-        #self.hasrun = True  
 
     @property
     def normal_mode_frequencies(self):
@@ -279,7 +275,6 @@ def ion_coordinates_from_flattened(self, flattened_coordinate_vector: Vector) ->
 
     def reindex_ions(self):
         """Re-indexes the ions to order based on the distance from the center of the trap, where smallest index is closest to the center."""
-        # TODO: How are even/odd cases is handled?
         x,y,z = self.ion_coordinates_from_flattened(self.equilibrium_positions)
         r = np.sqrt(x**2 + y**2 + z**2)
         idx = np.argsort(r)
@@ -310,7 +305,7 @@ def reindex_ions_by_z(self):
         self.omega_x = self.omega_x[idx]
         self.omega_y = self.omega_y[idx]
         self.omega_z = self.omega_z[idx]
-        # all ions are the same so this is safe. TODO: When does this change?
+        # uses first ion's properties as the characteristic scalings 
         self.set_up_dimensionless_parameters(self.nuclear_charges[0], self.atomic_masses[0], self.omega_z[0])
 
 
@@ -497,7 +492,6 @@ def build_symplectic_matrix(self) -> Matrix:
         J = np.block([[zeros, I], [-I, zeros]])
         return J    
 
-    ### Mode organizing helper methods  
     def sort_modes(self, eigvals, eigvecs):
        eigvals = np.imag(eigvals)
        sort_dex = np.argsort(eigvals)
@@ -516,7 +510,6 @@ def organize_modes(self, eigvals, eigvecs):
         eigvals, eigvecs = self.sort_modes(eigvals, eigvecs)
         eigvals, eigvecs = self.split_modes(eigvals, eigvecs)
 
-        # Apply branch sorting if configured
         if self.mode_organization == 'branch_sorted':
             eigvals, eigvecs = self.sort_by_branch(eigvals, eigvecs)
 
@@ -549,7 +542,7 @@ def sort_by_branch(self, evals, evecs):
         sorted_by_branch_evecs = np.zeros_like(evecs)
         for direction in range(3):
             direction_indices = np.where(classifier == direction)[0]
-            # they are already sorted by frequency from the original mode analysis code, so we can just take them in order
+            # they are already sorted by frequency, so take them in order  
             sorted_by_branch_evals[direction*N_ions:(direction+1)*N_ions] = evals[direction_indices]
             sorted_by_branch_evecs[:, direction*N_ions:(direction+1)*N_ions] = evecs[:, direction_indices]
         return sorted_by_branch_evals, sorted_by_branch_evecs 
@@ -566,8 +559,6 @@ def calculate_mode_participation_factors(self) -> Matrix:
             For N ions, this form organizes the 3N modes into N modes per direction "d", where d = x, y, z. 
 
         """ 
-
-        # TODO: For the linear case: the shape should be (d, N, N) for d dimensions (3), N ions, and N modes per direction. 
         eigvecs = self.eigvecs
         num_coords, num_modes = np.shape(eigvecs) 
         num_ions = num_modes // 3
@@ -591,7 +582,7 @@ def compute_reference_single_ion_lamb_dicke_factors(self, wavenumber: float) ->
         trap_frequencies = np.array([self.omega_x[0], self.omega_y[0], self.omega_z[0]])        
         mass = self.atomic_masses[0] 
         eta_x, eta_y, eta_z = wavenumber * np.sqrt(const.hbar / (2 * mass * trap_frequencies))
-        return eta_x, eta_y, eta_z    # ignore the phase
+        return eta_x, eta_y, eta_z  
 
     def return_equilibrium_positions(self, dimensionless: bool) -> Vector:
         """Return equilibrium positions in either dimensionless or SI units (inverse meters)."""
@@ -624,7 +615,6 @@ def from_atomic_structure_basis(cls, atomic_structure_basis: StandardBasis, omeg
                 raise IonSimError("Atomic structure basis should only contain AtomicStructure objects. No motional modes should be included.")
             atomic_masses.append(DOF.atomic_mass) 
             atomic_charges.append(DOF.atomic_charge) 
-        # Construct the class 
         return cls(num_ions, omega_x, omega_y, omega_z, atomic_masses, atomic_charges)
 
     def build_mode_DOFs(self, mode_indices: list[int], fock_dimensions: Vector | int) -> list[MotionalMode]:
@@ -633,7 +623,6 @@ def build_mode_DOFs(self, mode_indices: list[int], fock_dimensions: Vector | int
             - Applies each fock dimension to each mode, or applies the same fock dimension to all the modes
         """
         modes = []
-        # Convert list of Fock dimensions to an array if it's not already an array
         fock_dimensions = self.convert_to_array(fock_dimensions).astype(int)
         for idx, fock_dim in zip(mode_indices, fock_dimensions):
             mode_index = idx # or some function of this index
@@ -1000,7 +989,6 @@ def calculate_lamb_dicke_parameters_full(self, wavevector: Vector) -> Matrix:
         num_modes = 3 * self.num_ions
         total_ld_params = np.zeros((self.num_ions, num_modes), dtype=complex)
 
-        # TODO: Replace with dot product for efficiency 
         for i in range(3):
             total_ld_params += wavevector[i] * mode_pf[i, :, :]
 
@@ -1101,4 +1089,3 @@ def print_chain_summary(self):
         print(f"\nAxial mode frequencies: {axial_freqs/(2*np.pi)/1e6} MHz")
         print(f"Radial X mode frequencies: {radial_x_freqs/(2*np.pi)/1e6} MHz")
         print(f"Radial Y mode frequencies: {radial_y_freqs/(2*np.pi)/1e6} MHz")
-