rewrite atom/scatter/recapture for better ergonomics, start on scatter executable

Cargo.toml

-# [[bin]]
-# path = "src/scatter.rs"
-# name = "scatter"
+[[bin]]
+path = "src/scatter.rs"
+name = "scatter"
 
 [package]
 name = "pachinko"

lib/lib.rs

@@ -12,7 +12,6 @@ pub mod utils;
 pub mod newton;
 pub mod trap;
 pub mod atom;
-pub mod atom_new;
 pub mod recapture;
 pub mod scatter;
 

lib/newton.rs

@@ -31,6 +31,9 @@ pub type NewtonResult<T> = Result<T, NewtonError>;
 pub struct ThreeVector(pub f64, pub f64, pub f64);
 
 impl ThreeVector {
+    /// Return a new vector with zero magnitude.
+    pub fn zero() -> Self { Self(0.0, 0.0, 0.0) }
+
     /// Generate from spherical coordinates $`(r, \theta, \phi)`$ where $`\phi`$
     /// is the azimuthal angle.
     pub fn from_angles(r: f64, theta: f64, phi: f64) -> Self {

lib/recapture.rs

-use rand::Rng;
+use rand::{
+    prelude as rnd,
+    Rng,
+};
+use thiserror::Error;
 use crate::{
-    atom::{
+    atom_new::{
         Atom,
-        State,
         RadiationPattern,
+        State,
     },
     newton::{
-        ThreeVector,
-        PhaseSpace,
         Axis,
+        PhaseSpace,
+        ThreeVector,
     },
     phys::g as grav,
     trap::Trap,
 };
 
-pub fn recapture<S, T, R, G>(
+#[derive(Error, Debug)]
+pub enum RecaptureError {
+    #[error("trap missing for state {0}")]
+    TrapUndefined(String),
+}
+pub type RecaptureResult<T> = Result<T, RecaptureError>;
+
+/// Perform a series of `N` release-recapture trials for fixed time of flight in
+/// the trap for the current state of `atom`, returning the average recapture
+/// probability.
+pub fn recapture_rng<S, T, R, G>(
     atom: &Atom<S, T, R>,
     tof: f64,
     N: usize,
@@ -29,14 +43,80 @@ where
     let mut count: usize = 0;
     let mut q: PhaseSpace;
     for _ in 0..N {
-        q = atom.sample_phasespace(rng);
+        q = atom.sample_phasespace_rng(rng);
+        q.pos
+            += q.mom / atom.mass * tof
+            - 0.5 * ThreeVector::from_axis(-grav, Axis::Z) * tof.powi(2);
+        if atom.is_trapped(q) { count += 1; }
+    }
+    return count as f64 / N as f64;
+}
+
+/// Perform a series of `N` release-recapture trials for fixed time of flight in
+/// the trap for the current state of `atom`, returning the average recapture
+/// probability.
+pub fn recapture<S, T, R>(atom: &Atom<S, T, R>, tof: f64, N: usize) -> f64
+where
+    S: State,
+    T: Trap,
+    R: RadiationPattern,
+{
+    let mut rng = rnd::thread_rng();
+    return recapture_rng(atom, tof, N, &mut rng);
+}
+
+/// Perform a series of `N` release-recapture trials for fixed time of flight in
+/// the trap for a given atomic state, returning the average recapture
+/// probability.
+pub fn recapture_for_rng<S, T, R, G>(
+    atom: &Atom<S, T, R>,
+    state: &S,
+    tof: f64,
+    N: usize,
+    rng: &mut G,
+) -> RecaptureResult<f64>
+where
+    S: State,
+    T: Trap,
+    R: RadiationPattern,
+    G: Rng + ?Sized,
+{
+    let mut count: usize = 0;
+    let mut q: PhaseSpace;
+    for _ in 0..N {
+        q = atom.sample_phasespace_for_rng(state, rng)
+            .ok_or_else(
+                || RecaptureError::TrapUndefined(format!("{:?}", state))
+            )?;
         q.pos
             += q.mom / atom.mass * tof
             - 0.5 * ThreeVector::from_axis(-grav, Axis::Z) * tof.powi(2);
-        if atom.get_trap().is_trapped(atom.mass, q) {
+        if atom.is_trapped_for(state, q)
+            .ok_or_else(
+                || RecaptureError::TrapUndefined(format!("{:?}", state))
+            )?
+        {
             count += 1;
         }
     }
-    return count as f64 / N as f64;
+    return Ok(count as f64 / N as f64);
+}
+
+/// Perform a series of `N` release-recapture trials for fixed time of flight in
+/// the trap for a given atomic state, returning the average recapture
+/// probability.
+pub fn recapture_for<S, T, R>(
+    atom: &Atom<S, T, R>,
+    state: &S,
+    tof: f64,
+    N: usize
+) -> RecaptureResult<f64>
+where
+    S: State,
+    T: Trap,
+    R: RadiationPattern,
+{
+    let mut rng = rnd::thread_rng();
+    return recapture_for_rng(atom, state, tof, N, &mut rng);
 }
 

lib/scatter.rs

-use rand::Rng;
 use thiserror::Error;
 use crate::{
-    atom::{
+    atom_new::{
         Atom,
         AtomIter,
+        AtomIterStatic,
+        // PhotonInteraction,
         RadiationPattern,
         State,
+        // PhotonGen,
     },
     newton::{
         NewtonError,
+        ThreeVector,
         PhaseSpace,
         rka,
     },
@@ -19,187 +22,274 @@ use crate::{
 pub enum ScatterError {
     #[error("newton error {0}")]
     NewtonError(NewtonError),
+    #[error("StateHistory::new: components have unequal lengths")]
+    StateHistoryUnequalLength,
+    #[error("Trajectory::new: components have unequal lengths")]
+    TrajectoryUnequalLength,
 }
 pub type ScatterResult<T> = Result<T, ScatterError>;
 
-#[derive(Clone, Debug)]
-pub struct ScatterHistory<S>
-where S: State,
+/// Simple data struct to hold a time series of states.
+#[derive(Clone, Debug, Default)]
+pub struct StateHistory<S>
+where S: State
 {
-    pub state: Vec<S>,
-    pub time: Vec<f64>,
-    pub trajectory: Vec<PhaseSpace>,
+    t: Vec<f64>,
+    s: Vec<S>,
 }
 
-/// Calculate a single trajectory up to `t_final` *or* until the atom goes dark;
-/// i.e. the final time is not guaranteed to be equal to `t_final`.
-pub fn scatter_sim<S, T, R, G>(
-    atom: &Atom<S, T, R>,
-    t_final: f64,
-    epsilon: f64,
-    rng: &mut G,
-) -> ScatterResult<(ScatterHistory<S>, Atom<S, T, R>)>
-where
-    S: State,
-    T: Trap,
-    R: RadiationPattern,
-    G: Rng + ?Sized,
+impl<S> StateHistory<S>
+where S: State
 {
-    let m: f64 = atom.mass;
-    let mut state_history: Vec<S> = vec![atom.state];
-    let mut t_cur: f64 = 0.0;
-    let mut t: Vec<f64> = Vec::new();
-    let (mut atom_iter, mut q_cur): (AtomIter<S, T, R>, PhaseSpace)
-        = atom.to_state_iter();
-    let mut trajectory: Vec<PhaseSpace> = Vec::new();
-    // simulate until t_final or the atom goes dark
-    while let Some((state, photon)) = atom_iter.next() {
-        state_history.push(state);
-        match photon.momentum_kick_lifetime(rng) {
-            // absorb a photon
-            (kick, None) => {
-                q_cur.mom += kick;
-            },
-            // integrate position until photon emission
-            (kick, Some((lifetime, tau))) => {
-                let dt: f64
-                    = atom_iter.get_trap().period(m)
-                    .min(tau)
-                    .min(lifetime)
-                    .min(t_final - t_cur)
-                    / 10.0;
-                let t_int_end: f64 = (t_cur + lifetime).min(t_final);
-                let rhs
-                    = |_t: f64, q: PhaseSpace| -> PhaseSpace {
-                        PhaseSpace {
-                            pos: q.mom / m,
-                            mom: -atom_iter.get_trap().gradient(q.pos),
-                        }
-                    };
-                let (mut t_int, mut q_int): (Vec<f64>, Vec<PhaseSpace>)
-                    = rka((t_cur, t_int_end), q_cur, dt, rhs, epsilon)
-                    .map_err(ScatterError::NewtonError)?;
-                t_cur = t_int.pop().unwrap();
-                q_cur = q_int.pop().unwrap();
-                t.append(&mut t_int);
-                trajectory.append(&mut q_int);
-                q_cur.mom += kick;
-            },
+    /// Construct a new `StateHistory`. Fails if the two series are of unequal
+    /// lengths.
+    pub fn new(t: Vec<f64>, s: Vec<S>) -> ScatterResult<Self> {
+        return if t.len() == s.len() {
+            Ok(Self { t, s })
+        } else {
+            Err(ScatterError::StateHistoryUnequalLength)
+        };
+    }
+
+    pub fn len(&self) -> usize { self.t.len() }
+
+    pub fn is_empty(&self) -> bool { self.t.len() == 0 }
+
+    /// Push a new data point onto the end of the series.
+    pub fn push(&mut self, t: f64, s: S) {
+        self.t.push(t);
+        self.s.push(s);
+    }
+
+    /// Return an iterator over `(time, state)` pairs.
+    pub fn iter(&self) -> StateHistoryIter<S> {
+        return StateHistoryIter { state_history: self, idx: 0 };
+    }
+}
+
+impl<S> FromIterator<(f64, S)> for StateHistory<S>
+where S: State
+{
+    fn from_iter<I>(iter: I) -> Self
+    where I: IntoIterator<Item = (f64, S)>
+    {
+        let (t, s): (Vec<f64>, Vec<S>) = iter.into_iter().unzip();
+        return Self { t, s };
+    }
+}
+
+impl<S> IntoIterator for StateHistory<S>
+where S: State
+{
+    type Item = (f64, S);
+    type IntoIter
+        = std::iter::Zip<std::vec::IntoIter<f64>, std::vec::IntoIter<S>>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        return self.t.into_iter().zip(self.s.into_iter());
+    }
+}
+
+pub struct StateHistoryIter<'a, S>
+where S: State
+{
+    state_history: &'a StateHistory<S>,
+    idx: usize,
+}
+
+impl<'a, S> Iterator for StateHistoryIter<'a, S>
+where S: State
+{
+    type Item = (&'a f64, &'a S);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.idx < self.state_history.len() {
+            let ret
+                = Some((
+                    &self.state_history.t[self.idx],
+                    &self.state_history.s[self.idx],
+                ));
+            self.idx += 1;
+            return ret;
+        } else {
+            return None;
+        };
+    }
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct Trajectory
+{
+    t: Vec<f64>,
+    q: Vec<PhaseSpace>,
+}
+
+/// Simple data struct to hold a time series of `PhaseSpace` vectors.
+impl Trajectory {
+    /// Construct a new `Trajectory`. Fails if the two series are of unequal
+    /// lengths.
+    pub fn new(t: Vec<f64>, q: Vec<PhaseSpace>) -> ScatterResult<Self> {
+        return if t.len() == q.len() {
+            Ok(Self { t, q })
+        } else {
+            Err(ScatterError::TrajectoryUnequalLength)
+        };
+    }
+
+    /// Push a new data point onto the end of the series.
+    pub fn push(&mut self, t: f64, q: PhaseSpace) {
+        self.t.push(t);
+        self.q.push(q);
+    }
+
+    /// Move time series data into `self`.
+    pub fn append(&mut self, t: &mut Vec<f64>, q: &mut Vec<PhaseSpace>) {
+        self.t.append(t);
+        self.q.append(q);
+    }
+
+    pub fn len(&self) -> usize { self.t.len() }
+
+    pub fn is_empty(&self) -> bool { self.t.len() == 0 }
+
+    /// Return an iterator over `(time, phasespace)` pairs.
+    pub fn iter(&self) -> TrajectoryIter {
+        return TrajectoryIter { trajectory: self, idx: 0 };
+    }
+}
+
+impl IntoIterator for Trajectory {
+    type Item = (f64, PhaseSpace);
+    type IntoIter
+        = std::iter::Zip<
+            std::vec::IntoIter<f64>,
+            std::vec::IntoIter<PhaseSpace>,
+        >;
+
+    fn into_iter(self) -> Self::IntoIter {
+        return self.t.into_iter().zip(self.q.into_iter());
+    }
+}
+
+impl FromIterator<(f64, PhaseSpace)> for Trajectory {
+    fn from_iter<I>(iter: I) -> Self
+    where I: IntoIterator<Item = (f64, PhaseSpace)>
+    {
+        let (t, q): (Vec<f64>, Vec<PhaseSpace>) = iter.into_iter().unzip();
+        return Self { t, q };
+    }
+}
+
+pub struct TrajectoryIter<'a> {
+    trajectory: &'a Trajectory,
+    idx: usize,
+}
+
+impl<'a> Iterator for TrajectoryIter<'a> {
+    type Item = (&'a f64, &'a PhaseSpace);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.idx < self.trajectory.len() {
+            let ret
+                = Some((
+                    &self.trajectory.t[self.idx],
+                    &self.trajectory.q[self.idx],
+                ));
+            self.idx += 1;
+            return ret;
+        } else {
+            return None
         }
-        if t_cur >= t_final { break; }
     }
-    return Ok((
-        ScatterHistory {
-            state: state_history,
-            time: t,
-            trajectory,
-        },
-        atom_iter.dump(),
-    ));
 }
 
-/// Calculate a single trajectory, recording times and positions/momenta only
-/// when photons are emitted. Necessarily limited to when the atom goes dark.
-pub fn scatter_sim_emission<S, T, R, G>(
-    atom: &Atom<S, T, R>,
-    t_final: f64,
-    epsilon: f64,
-    rng: &mut G,
-) -> ScatterResult<(ScatterHistory<S>, Atom<S, T, R>)>
+/// Simple data class holding both a state history and a phase-space trajectory.
+#[derive(Clone, Debug, Default)]
+pub struct ScatterHistory<S>
+where S: State
+{
+    pub states: StateHistory<S>,
+    pub traj: Trajectory,
+}
+
+impl<S> ScatterHistory<S>
+where S: State
+{
+    /// Push a new data point onto the end of the state history.
+    pub fn push_state(&mut self, t: f64, s: S) { self.states.push(t, s); }
+
+    /// Push a new data point onto the end of the phase-space trajectory.
+    pub fn push_traj(&mut self, t: f64, q: PhaseSpace) { self.traj.push(t, q); }
+}
+
+/// Calculate a single trajectory up to `t_final` *or* until the atom goes dark;
+/// i.e. the final time in the trajectory is not guaranteed to be equal to
+/// `t_final`.
+pub fn scatter_sim<S, T, R>(atom: &Atom<S, T, R>, t_final: f64)
+    -> ScatterResult<(ScatterHistory<S>, Atom<S, T, R>)>
 where
     S: State,
     T: Trap,
     R: RadiationPattern,
-    G: Rng + ?Sized,
 {
     let m: f64 = atom.mass;
-    let mut state_history: Vec<S> = vec![atom.state];
     let mut t_cur: f64 = 0.0;
-    let mut t: Vec<f64> = vec![0.0];
-    let (mut atom_iter, mut q_cur): (AtomIter<S, T, R>, PhaseSpace)
-        = atom.to_state_iter();
-    let mut trajectory: Vec<PhaseSpace> = vec![q_cur];
-    while let Some((state, photon)) = atom_iter.next() {
-        state_history.push(state);
-        match photon.momentum_kick_lifetime(rng) {
-            (kick, None) => {
-                q_cur.mom += kick;
-            },
-            (kick, Some((lifetime, tau))) => {
-                let dt: f64
-                    = atom_iter.get_trap().period(m)
-                    .min(tau)
-                    .min(lifetime)
-                    .min(t_final - t_cur)
-                    / 10.0;
-                let t_int_end: f64 = (t_cur + lifetime).min(t_final);
-                let rhs
-                    = |_t: f64, q: PhaseSpace| -> PhaseSpace {
-                        PhaseSpace {
-                            pos: q.mom / m,
-                            mom: -atom_iter.get_trap().gradient(q.pos),
-                        }
-                    };
-                let (mut t_int, mut q_int): (Vec<f64>, Vec<PhaseSpace>)
-                    = rka((t_cur, t_int_end), q_cur, dt, rhs, epsilon)
-                    .map_err(ScatterError::NewtonError)?;
-                t_cur = t_int.pop().unwrap();
-                q_cur = q_int.pop().unwrap();
-                t.push(t_cur);
-                trajectory.push(q_cur);
-                q_cur.mom += kick;
-            },
-        }
+    let mut state_history = StateHistory::new(vec![t_cur], vec![atom.state])?;
+    let mut trajectory = Trajectory::new(Vec::new(), Vec::new())?;
+    let mut atom_iter: AtomIter<S, T, R> = atom.to_state_iter();
+    let (mut dt, mut t_int_end): (f64, f64);
+    let mut tq_int: (Vec<f64>, Vec<PhaseSpace>);
+    // simulate until t_final or the atom goes dark
+    while let Some((state, photon_int)) = atom_iter.next() {
+        dt
+            = atom_iter.get_trap().period(atom.mass)
+            .min(photon_int.time())
+            .min(photon_int.time_mean())
+            .min(t_final - t_cur)
+            / 10.0;
+        t_int_end = (t_cur + photon_int.time()).min(t_final);
+        let int_rhs
+            = |_t: f64, q: PhaseSpace| -> PhaseSpace {
+                PhaseSpace {
+                    pos: q.mom / m,
+                    mom: -atom_iter.get_trap().gradient(q.pos),
+                }
+            };
+        tq_int = rka((t_cur, t_int_end), atom_iter.q, dt, int_rhs, 1e-6)
+            .map_err(ScatterError::NewtonError)?;
+        t_cur = tq_int.0.pop().unwrap();
+        atom_iter.q = tq_int.1.pop().unwrap();
+        trajectory.append(&mut tq_int.0, &mut tq_int.1);
+        state_history.push(t_cur, state);
+        atom_iter.q.mom
+            += photon_int.momentum_kick().unwrap_or(ThreeVector::zero());
         if t_cur >= t_final { break; }
     }
     return Ok((
-        ScatterHistory {
-            state: state_history,
-            time: t,
-            trajectory,
-        },
-        atom_iter.dump(),
+        ScatterHistory { states: state_history, traj: trajectory },
+        atom_iter.dump().0,
     ));
 }
 
-/// Calculate a single trajectory up to `t_final`, continuing past the point at
-/// which the atom goes dark.
-pub fn scatter_sim_dark<S, T, R, G>(
-    atom: &Atom<S, T, R>,
-    t_final: f64,
-    epsilon: f64,
-    rng: &mut G,
-) -> ScatterResult<(ScatterHistory<S>, Atom<S, T, R>)>
+/// Calculate a single state history, disregarding position and momentum, up to
+/// `t_final` *or* until the atom goes dark; i.e. the final time in the
+/// trajectory is not guaranteed to be equal to `t_final`.
+pub fn scatter_sim_static<S, T, R>(atom: &Atom<S, T, R>, t_final: f64)
+    -> ScatterResult<(StateHistory<S>, Atom<S, T, R>)>
 where
     S: State,
     T: Trap,
     R: RadiationPattern,
-    G: Rng + ?Sized,
 {
-    let (mut scatter_history, atom_sim): (ScatterHistory<S>, Atom<S, T, R>)
-        = scatter_sim(atom, t_final, epsilon, rng)?;
-    if scatter_history.time.last().unwrap_or(&0.0) < &t_final {
-        let t_init: f64 = scatter_history.time.pop().unwrap_or(0.0);
-        let q_init: PhaseSpace
-            = scatter_history.trajectory.pop()
-            .unwrap_or_else(|| atom_sim.sample_phasespace(rng));
-        let dt: f64
-            = atom_sim.get_trap().period(atom_sim.mass)
-            .min(t_final - t_init)
-            / 10.0;
-        let rhs
-            = |_t: f64, q: PhaseSpace| -> PhaseSpace {
-                PhaseSpace {
-                    pos: q.mom / atom_sim.mass,
-                    mom: -atom_sim.get_trap().gradient(q.pos),
-                }
-            };
-        let (mut t_int, mut q_int): (Vec<f64>, Vec<PhaseSpace>)
-            = rka((t_init, t_final), q_init, dt, rhs, epsilon)
-                .map_err(ScatterError::NewtonError)?;
-        scatter_history.time.append(&mut t_int);
-        scatter_history.trajectory.append(&mut q_int);
+    let mut t_cur: f64 = 0.0;
+    let mut state_history = StateHistory::new(vec![t_cur], vec![atom.state])?;
+    let mut atom_iter: AtomIterStatic<S, T, R> = atom.to_state_iter_static();
+    // simulate until t_final or the atom goes dark
+    for (state, photon_int) in atom_iter.by_ref() {
+        t_cur += photon_int.time();
+        state_history.push(t_cur, state);
+        if t_cur >= t_final { break; }
     }
-    return Ok((scatter_history, atom_sim));
+    return Ok((state_history, atom_iter.dump()));
 }
 

src/scatter.rs

+#![allow(dead_code, non_snake_case, non_upper_case_globals)]
+#![allow(clippy::needless_return)]
+
 use anyhow::Result;
-use lib::newton::NewtonError;
+use pachinko::newton::NewtonError;
 
 fn main() -> Result<()> {
     return Err(NewtonError::RKAErrorBound)?;