use std::{ collections::HashMap, f64::consts::{ PI, TAU, }, hash::Hash, fmt::Debug, }; use rand::{ prelude as rnd, Rng, distributions::Distribution, }; use statrs::distribution::Exp; use thiserror::Error; use crate::{ newton::{ ThreeVector, PhaseSpace, }, phys::hbar, trap::Trap, }; #[derive(Error, Debug)] pub enum AtomError { #[error("reached dark state")] DarkState, #[error("simultaneous excitation and decay pathways are not supported")] ExciteAndDecay, #[error("trap missing for state {0}")] TrapUndefined(String), } pub type AtomResult<T> = Result<T, AtomError>; pub trait RadiationPattern: Copy + Clone + PartialEq { /// Generate angles (theta, phi) where phi is the azimuthal angle. fn sample_angles_rng<R>(&self, rng: &mut R) -> (f64, f64) where R: Rng + ?Sized; fn sample_angles(&self) -> (f64, f64) { let mut rng = rnd::thread_rng(); return self.sample_angles_rng(&mut rng); } fn sample_momentum_rng<R>(&self, k: f64, rng: &mut R) -> ThreeVector where R: Rng + ?Sized { let (th, ph): (f64, f64) = self.sample_angles_rng(rng); return hbar * k * ThreeVector( ph.cos() * th.sin(), ph.sin() * th.sin(), th.cos(), ); } fn sample_momentum(&self, k: f64) -> ThreeVector { let mut rng = rnd::thread_rng(); return self.sample_momentum_rng(k, &mut rng); } fn sample_momentum_kick_rng<R>(&self, k: f64, rng: &mut R) -> ThreeVector where R: Rng + ?Sized { let (th, ph): (f64, f64) = self.sample_angles_rng(rng); return -hbar * k * ThreeVector( ph.cos() * th.sin(), ph.sin() * th.sin(), th.cos(), ); } fn sample_momentum_kick(&self, k: f64) -> ThreeVector { let mut rng = rnd::thread_rng(); return self.sample_momentum_kick_rng(k, &mut rng); } } #[derive(Copy, Clone, Debug, Default, PartialEq, Eq)] pub struct RadUniform { } impl RadUniform { pub fn new() -> Self { Self { } } } impl RadiationPattern for RadUniform { fn sample_angles_rng<R>(&self, rng: &mut R) -> (f64, f64) where R: Rng + ?Sized { return ( rng.gen::<f64>() * PI, rng.gen::<f64>() * TAU, ); } } #[derive(Copy, Clone, Debug, PartialEq)] pub struct RadDipole { pi: f64, sigma: f64, } impl RadDipole { /// Give the relative proportion of radiation scattered into pi-polarized /// (oscillating dipole) and sigma-polarized (rotating dipole) modes. The /// oscillation axis is fixed to Z. Proportions will be automatically /// normalized. pub fn new(pi: f64, sigma: f64) -> Self { let pi_norm: f64 = pi / (pi + sigma); let sigma_norm: f64 = sigma / (pi + sigma); return Self { pi: pi_norm, sigma: sigma_norm }; } /// Special case for which the distribution is uniform. pub fn uniform() -> Self { return Self { pi: 0.25, sigma: 0.75 }; } pub fn pdf_theta(&self, theta: f64) -> f64 { return self.pi * 2.0 / PI * theta.sin().powi(2) + self.sigma * 2.0 / 3.0 / PI * (1.0 + theta.cos().powi(2)); } pub fn cdf_theta(&self, theta: f64) -> f64 { return theta / PI + (self.sigma / (3.0 * TAU) - self.pi / TAU) * (2.0 * theta).sin(); } pub fn cdf_inv_theta(&self, r: f64) -> f64 { // cdf = r is transcendental, so have to use newton-raphson let mut th: f64 = r * PI; let mut dth: f64; for _ in 0..1000 { dth = (self.cdf_theta(th) - r) / self.pdf_theta(th); th -= dth; if dth.abs() < 1e-6 { return th; } th = th.min(PI).max(0.0); } return th; } } impl RadiationPattern for RadDipole { fn sample_angles_rng<R>(&self, rng: &mut R) -> (f64, f64) where R: Rng + ?Sized { return ( self.cdf_inv_theta(rng.gen::<f64>()), rng.gen::<f64>() * TAU, ); } } pub trait State: Copy + Clone + Debug + PartialEq + Eq + Hash { } #[derive(Copy, Clone, Debug, PartialEq, Eq)] pub enum TransitionKind { Exciting = 0, Decaying = 1, } /// Assume that: /// - no state is simultaneously able to decay and be excited /// - Incoming light saturates all relevant transitions #[derive(Copy, Clone, Debug, PartialEq)] pub enum Transition<S, R> where S: State, R: RadiationPattern, { Exciting { ground: S, excited: S, prob: f64, /// m^-1 wavelength: f64, laser_dir: ThreeVector, }, Decaying { ground: S, excited: S, prob: f64, /// m^-1 wavelength: f64, /// Hz linewidth: f64, radiation: R, }, } impl<S, R> Transition<S, R> where S: State, R: RadiationPattern, { /// Probability will be normalized when loaded into a StateMap. pub fn new_exciting( ground: S, excited: S, prob: f64, wavelength: f64, laser_dir: ThreeVector, ) -> Self { return Self::Exciting { ground, excited, prob: prob.abs(), wavelength: wavelength.abs(), laser_dir: laser_dir.normalized(), }; } /// Probability will be normalized when loaded into a StateMap. pub fn new_decaying( ground: S, excited: S, prob: f64, wavelength: f64, linewidth: f64, radiation: R, ) -> Self { return Self::Decaying { ground, excited, prob: prob.abs(), wavelength: wavelength.abs(), linewidth: linewidth.abs(), radiation, }; } pub fn kind(&self) -> TransitionKind { return match self { Self::Exciting { .. } => TransitionKind::Exciting, Self::Decaying { .. } => TransitionKind::Decaying, }; } pub fn is_exciting(&self) -> bool { return matches!(self, Self::Exciting { .. }); } pub fn is_decaying(&self) -> bool { return matches!(self, Self::Decaying { .. }); } pub fn ground_state(&self) -> S { return match self { Self::Exciting { ground, .. } => *ground, Self::Decaying { ground, .. } => *ground, }; } pub fn get_ground_state(&self) -> &S { return match self { Self::Exciting { ground, .. } => ground, Self::Decaying { ground, .. } => ground, }; } pub fn excited_state(&self) -> S { return match self { Self::Exciting { ground: _, excited, .. } => *excited, Self::Decaying { ground: _, excited, .. } => *excited, }; } pub fn get_excited_state(&self) -> &S { return match self { Self::Exciting { ground: _, excited, .. } => excited, Self::Decaying { ground: _, excited, .. } => excited, }; } pub fn start_state(&self) -> S { return match self { Self::Exciting { ground, .. } => *ground, Self::Decaying { ground: _, excited, .. } => *excited, }; } pub fn get_start_state(&self) -> &S { return match self { Self::Exciting { ground, .. } => ground, Self::Decaying { ground: _, excited, .. } => excited, }; } pub fn end_state(&self) -> S { return match self { Self::Exciting { ground: _, excited, .. } => *excited, Self::Decaying { ground, .. } => *ground, }; } pub fn get_end_state(&self) -> &S { return match self { Self::Exciting { ground: _, excited, .. } => excited, Self::Decaying { ground, .. } => ground, }; } pub fn probability(&self) -> f64 { return match self { Self::Exciting { ground: _, excited: _, prob, .. } => *prob, Self::Decaying { ground: _, excited: _, prob, .. } => *prob, }; } pub fn with_probability(self, prob: f64) -> Self { return match self { Self::Exciting { ground, excited, prob: _, wavelength, laser_dir, } => Self::Exciting { ground, excited, prob, wavelength, laser_dir, }, Self::Decaying { ground, excited, prob: _, wavelength, linewidth, radiation, } => Self::Decaying { ground, excited, prob, wavelength, linewidth, radiation, }, }; } pub fn with_prob_normalization(self, N: f64) -> Self { return match self { Self::Exciting { ground, excited, prob, wavelength, laser_dir, } => Self::Exciting { ground, excited, prob: prob / N, wavelength, laser_dir, }, Self::Decaying { ground, excited, prob, wavelength, linewidth, radiation, } => Self::Decaying { ground, excited, prob: prob / N, wavelength, linewidth, radiation, }, }; } pub fn wavelength(&self) -> f64 { return match self { Self::Exciting { ground: _, excited: _, prob: _, wavelength, .. } => *wavelength, Self::Decaying { ground: _, excited: _, prob: _, wavelength, .. } => *wavelength, }; } pub fn absorption(&self) -> Option<Absorption> { return match self { Self::Exciting { ground: _, excited: _, prob: _, wavelength, laser_dir, .. } => Some( Absorption { wavelength: *wavelength, laser_dir: *laser_dir, } ), Self::Decaying { .. } => None, }; } pub fn linewidth(&self) -> Option<f64> { return match self { Self::Exciting { .. } => None, Self::Decaying { ground: _, excited: _, prob: _, wavelength: _, linewidth, .. } => Some(*linewidth), }; } pub fn radiation_pattern(&self) -> Option<&R> { return match self { Self::Exciting { .. } => None, Self::Decaying { ground: _, excited: _, prob: _, wavelength: _, linewidth: _, radiation, .. } => Some(radiation), }; } pub fn radiation(&self) -> Option<Radiation<R>> { return match self { Self::Exciting { .. } => None, Self::Decaying { ground: _, excited: _, prob: _, wavelength, linewidth, radiation, .. } => Some( Radiation { wavelength: *wavelength, linewidth: *linewidth, pattern: *radiation, lifetime_dist: Exp::new(linewidth / 2.0).unwrap(), } ), }; } pub fn photon_interaction(&self) -> PhotonInteraction<R> { return match self { Self::Exciting { .. } => PhotonInteraction::Absorption(self.absorption().unwrap()), Self::Decaying { .. } => PhotonInteraction::Radiation(self.radiation().unwrap()), }; } pub fn starts_with(&self, state: &S) -> bool { return match self { Self::Exciting { ground, .. } => ground == state, Self::Decaying { ground: _, excited, .. } => excited == state, }; } pub fn exciting_starts_with(&self, state: &S) -> Option<bool> { return match self { Self::Exciting { ground, .. } => Some(ground == state), _ => None, }; } pub fn decaying_starts_with(&self, state: &S) -> Option<bool> { return match self { Self::Decaying { ground: _, excited, .. } => Some(excited == state), _ => None, }; } pub fn same_starts_with(&self, other: &Self) -> bool { return match (self, other) { ( Self::Exciting { ground: g0, .. }, Self::Exciting { ground: g1, .. }, ) => g0 == g1, ( Self::Decaying { ground: _, excited: e0, .. }, Self::Decaying { ground: _, excited: e1, .. }, ) => e0 == e1, _ => false, }; } pub fn ends_with(&self, state: &S) -> bool { return match self { Self::Exciting { ground: _, excited, .. } => excited == state, Self::Decaying { ground, .. } => ground == state, }; } pub fn exciting_ends_with(&self, state: &S) -> Option<bool> { return match self { Self::Exciting { ground: _, excited, .. } => Some(excited == state), _ => None, }; } pub fn decaying_ends_with(&self, state: &S) -> Option<bool> { return match self { Self::Decaying { ground, .. } => Some(ground == state), _ => None, }; } pub fn same_ends_with(&self, other: &Self) -> bool { return match (self, other) { ( Self::Exciting { ground: _, excited: e0, .. }, Self::Exciting { ground: _, excited: e1, .. }, ) => e0 == e1, ( Self::Decaying { ground: g0, .. }, Self::Decaying { ground: g1, .. }, ) => g0 == g1, _ => false, }; } } #[derive(Clone, Debug)] pub struct StateMap<S, T, R> where S: State, T: Trap, R: RadiationPattern, { transitions: Vec<Transition<S, R>>, traps: HashMap<S, T>, } impl<S, T, R> StateMap<S, T, R> where S: State, T: Trap, R: RadiationPattern, { /// Verifies that all transition probabilities are properly normalized. /// Duplicates are all counted. pub fn new<I, V>(transitions: I, traps: V) -> AtomResult<Self> where I: IntoIterator<Item = Transition<S, R>>, V: IntoIterator<Item = (S, T)>, { let mut transition_list: Vec<Transition<S, R>> = transitions.into_iter().collect(); let state_traps: HashMap<S, T> = traps.into_iter().collect(); let mut verified: Vec<Transition<S, R>> = Vec::new(); let mut N: f64; let mut take: Vec<usize> = Vec::new(); let mut transition_kind: TransitionKind; while let Some(transition) = transition_list.first() { transition_kind = transition.kind(); if !transition_list.iter().all(|t| t.kind() == transition_kind) { return Err(AtomError::ExciteAndDecay); } transition_list.iter() .map(|t| { if !state_traps.contains_key(t.get_ground_state()) { Err(AtomError::TrapUndefined( format!("{:?}", t.get_ground_state()) )) } else if !state_traps.contains_key(t.get_excited_state()) { Err(AtomError::TrapUndefined( format!("{:?}", t.get_excited_state()) )) } else { Ok(()) } }) .collect::<AtomResult<Vec<()>>>()?; N = transition_list.iter().enumerate() .filter_map(|(k, t)| { if transition.same_starts_with(t) { take.push(k); Some(t.probability()) } else { None } }) .sum(); take.drain(..) .for_each(|k| { verified.push( transition_list.swap_remove(k) .with_prob_normalization(N) ) }); } return Ok(Self { transitions: verified, traps: state_traps }); } pub fn get_trap(&self, state: &S) -> Option<&T> { self.traps.get(state) } /// Returns Err if `init_state` has nowhere to go. pub fn next_state_checked_rng<G>( &self, current_state: &S, rng: &mut G ) -> AtomResult<(S, PhotonInteraction<R>)> where G: Rng + ?Sized { let r: f64 = rng.gen(); let mut acc: f64 = 0.0; let transitions = self.transitions.iter() .filter_map(|t| { t.starts_with(current_state) .then_some( (t.probability(), t.end_state(), t.photon_interaction()) ) }); for (prob, state, photon) in transitions { acc += prob; if r < acc { return Ok((state, photon)); } } return Err(AtomError::DarkState); } pub fn next_state_checked(&self, current_state: &S) -> AtomResult<(S, PhotonInteraction<R>)> { let mut rng = rnd::thread_rng(); return self.next_state_checked_rng(current_state, &mut rng); } pub fn next_state_rng<G>(&self, current_state: &S, rng: &mut G) -> (S, Option<PhotonInteraction<R>>) where G: Rng + ?Sized { return self.next_state_checked_rng(current_state, rng) .map(|(state, photon)| (state, Some(photon))) .unwrap_or((*current_state, None)); } pub fn next_state(&self, current_state: &S) -> (S, Option<PhotonInteraction<R>>) { return self.next_state_checked(current_state) .map(|(state, photon)| (state, Some(photon))) .unwrap_or((*current_state, None)); } } #[derive(Copy, Clone, Debug)] pub struct Absorption { pub wavelength: f64, pub laser_dir: ThreeVector, } impl Absorption { pub fn momentum_kick(&self) -> ThreeVector { return -TAU / self.wavelength * self.laser_dir.normalized(); } } #[derive(Copy, Clone, Debug)] pub struct Radiation<R> where R: RadiationPattern { pub wavelength: f64, pub linewidth: f64, pub pattern: R, lifetime_dist: Exp, } impl<R> Radiation<R> where R: RadiationPattern { pub fn momentum_kick_rng<G>(&self, rng: &mut G) -> ThreeVector where G: Rng + ?Sized { return self.pattern .sample_momentum_kick_rng(TAU / self.wavelength, rng); } /// Assume the transition is saturated. pub fn lifetime(&self) -> f64 { return 2.0 / self.linewidth.abs(); } } #[derive(Copy, Clone, Debug)] pub enum PhotonInteraction<R> where R: RadiationPattern { Absorption(Absorption), Radiation(Radiation<R>), } impl<R> PhotonInteraction<R> where R: RadiationPattern { pub fn momentum_kick_rng<G>(&self, rng: &mut G) -> ThreeVector where G: Rng + ?Sized { return match self { Self::Absorption(a) => a.momentum_kick(), Self::Radiation(r) => r.momentum_kick_rng(rng), }; } /// If `self` is a radiation, return a sampled atomic lifetime with the /// mean lifetime. pub fn momentum_kick_lifetime<G>(&self, rng: &mut G) -> (ThreeVector, Option<(f64, f64)>) where G: Rng + ?Sized { return match self { Self::Absorption(a) => (a.momentum_kick(), None), Self::Radiation(r) => ( r.momentum_kick_rng(rng), Some( (r.lifetime_dist.sample(rng), r.lifetime()) ), ), }; } } #[derive(Clone, Debug)] pub struct Atom<S, T, R> where S: State, T: Trap, R: RadiationPattern, { pub state: S, state_map: StateMap<S, T, R>, pub mass: f64, pub temperature: f64, } impl<S, T, R> Atom<S, T, R> where S: State, T: Trap, R: RadiationPattern, { /// Verifies that `state`'s trap is defined in `state_map`. pub fn new( state: S, state_map: StateMap<S, T, R>, mass: f64, temperature: f64, ) -> AtomResult<Self> { state_map.get_trap(&state) .ok_or_else(|| AtomError::TrapUndefined(format!("{:?}", state)))?; return Ok(Self { state, state_map, mass, temperature, }); } pub fn state(&self) -> S { self.state } pub fn get_trap_for(&self, state: &S) -> &T { return self.state_map.get_trap(state).unwrap(); } pub fn get_trap(&self) -> &T { return self.state_map.get_trap(&self.state).unwrap(); } /// Returns an `Iterator` for the Markov Chain with an initial /// position/momentum. pub fn state_iter( init: S, state_map: StateMap<S, T, R>, mass: f64, temperature: f64, ) -> AtomResult<(AtomIter<S, T, R>, PhaseSpace)> { let mut rng = rnd::thread_rng(); let init_phasespace: PhaseSpace = state_map.get_trap(&init).unwrap() .sample_phasespace_rng(mass, temperature, &mut rng); let atom = Self::new(init, state_map, mass, temperature)?; return Ok((AtomIter { atom, rng }, init_phasespace)); } /// Clones `self` into an `Iterator` Markov Chain with an initial /// position/momentum. pub fn to_state_iter(&self) -> (AtomIter<S, T, R>, PhaseSpace) { return Self::state_iter( self.state, self.state_map.clone(), self.mass, self.temperature, ).unwrap(); } /// Converts `self` to an `Iterator` Markov Chain with an initial /// position/momentum. pub fn into_state_iter(self) -> (AtomIter<S, T, R>, PhaseSpace) { return Self::state_iter( self.state, self.state_map, self.mass, self.temperature, ).unwrap(); } pub fn sample_position<G>(&self, rng: &mut G) -> ThreeVector where G: Rng + ?Sized { return self.state_map.get_trap(&self.state).unwrap() .sample_position_rng(self.mass, self.temperature, rng); } pub fn sample_velocity<G>(&self, rng: &mut G) -> ThreeVector where G: Rng + ?Sized { return self.state_map.get_trap(&self.state).unwrap() .sample_velocity_rng(self.mass, self.temperature, rng); } pub fn sample_momentum<G>(&self, rng: &mut G) -> ThreeVector where G: Rng + ?Sized { return self.state_map.get_trap(&self.state).unwrap() .sample_momentum_rng(self.mass, self.temperature, rng); } pub fn sample_phasespace<G>(&self, rng: &mut G) -> PhaseSpace where G: Rng + ?Sized { return self.state_map.get_trap(&self.state).unwrap() .sample_phasespace_rng(self.mass, self.temperature, rng); } pub fn next_state_checked<G>(&mut self, rng: &mut G) -> AtomResult<(S, PhotonInteraction<R>)> where G: Rng + ?Sized { let (state, photon): (S, PhotonInteraction<R>) = self.state_map.next_state_checked_rng(&self.state, rng)?; self.state = state; return Ok((state, photon)); } pub fn next_state<G>(&mut self, rng: &mut G) -> (S, Option<PhotonInteraction<R>>) where G: Rng + ?Sized { let (state, maybe_photon): (S, Option<PhotonInteraction<R>>) = self.state_map.next_state_rng(&self.state, rng); self.state = state; return (state, maybe_photon); } } #[derive(Clone)] pub struct AtomIter<S, T, R> where S: State, T: Trap, R: RadiationPattern, { atom: Atom<S, T, R>, rng: rnd::ThreadRng, } impl<S, T, R> AtomIter<S, T, R> where S: State, T: Trap, R: RadiationPattern, { pub fn new(atom: Atom<S, T, R>, rng: rnd::ThreadRng) -> Self { return Self { atom, rng }; } pub fn get_trap(&self) -> &T { self.atom.get_trap() } pub fn dump(self) -> Atom<S, T, R> { self.atom } } impl<S, T, R> Iterator for AtomIter<S, T, R> where S: State, T: Trap, R: RadiationPattern, { type Item = (S, PhotonInteraction<R>); fn next(&mut self) -> Option<Self::Item> { return match self.atom.next_state(&mut self.rng) { (s, Some(p)) => Some((s, p)), (_, None) => None, }; } }