Define augmented phase space symplectic integrator

.gitignore

@@ -0,0 +1 @@
+/target

src/augmented.rs

@@ -0,0 +1,128 @@
+use crate::hamiltonian::*;
+
+use std::ops::{Add, Sub, Mul, Div, AddAssign};
+
+#[derive(Clone, Copy, Debug)]
+pub struct Augmented<T> {
+	omega: f64,
+	pub a: State<T>,
+	pub b: State<T>,
+}
+
+impl<T: Copy> Augmented<T> {
+	pub fn from(state: State<T>, omega: f64) -> Self {
+		Self {
+			omega,
+			a: state,
+			b: state,
+		}
+	}
+}
+
+impl<T> Augmented<T>
+where T: Copy + Add<Output = T> + Mul<f64, Output = T>
+{
+	pub fn average(&self) -> State<T> {
+		(self.a + self.b) * 0.5
+	}
+}
+
+impl<T> Augmented<T>
+where
+	T: Copy + Add<Output = T> + Sub<Output = T> + Mul<f64, Output = T> + Div<f64, Output = T> + AddAssign,
+	State<T>: Hamiltonian
+{
+	pub fn evolve(&mut self, h: f64) {
+		let cbrt2 = 2.0_f64.powf(1.0 / 3.0);
+		let z0 = - cbrt2 / (2.0 - cbrt2);
+		let z1 = 1.0 / (2.0 - cbrt2);
+
+		self.strang(z1 * h);
+		self.strang(z0 * h);
+		self.strang(z1 * h);
+	}
+
+	fn strang(&mut self, h: f64) {
+		self.evolve_a(h / 2.0);
+		self.evolve_b(h / 2.0);
+		self.evolve_c(h);
+		self.evolve_b(h / 2.0);
+		self.evolve_a(h / 2.0);
+	}
+
+	fn evolve_a(&mut self, h: f64) {
+		let dots = State { q: self.a.q, p: self.b.p }.dots();
+		self.a.p += dots.p * h;
+		self.b.q += dots.q * h;
+	}
+
+	fn evolve_b(&mut self, h: f64) {
+		let dots = State { q: self.b.q, p: self.a.p }.dots();
+		self.a.q += dots.q * h;
+		self.b.p += dots.p * h;
+	}
+
+	fn evolve_c(&mut self, h: f64) {
+		let alpha = 2.0 * self.omega * h;
+		let c = alpha.cos();
+		let s = alpha.sin();
+
+		let diff = self.a - self.b;
+		let r = State {
+			q: diff.q * c + diff.p * s,
+			p: diff.p * c - diff.q * s,
+		};
+
+		let sum = self.a + self.b;
+		self.a = (sum + r) / 2.0;
+		self.b = (sum - r) / 2.0;
+	}
+}
+
+impl<T: Add<Output = T>> Add for Augmented<T> {
+	type Output = Self;
+
+	fn add(self, rhs: Self) -> Self {
+		Self {
+			omega: self.omega,
+			a: self.a + rhs.a,
+			b: self.b + rhs.b,
+		}
+	}
+}
+
+impl<T: Sub<Output = T>> Sub for Augmented<T> {
+	type Output = Self;
+
+	fn sub(self, rhs: Self) -> Self {
+		Self {
+			omega: self.omega,
+			a: self.a - rhs.a,
+			b: self.b - rhs.b,
+		}
+	}
+}
+
+impl<T: Mul<f64, Output = T>> Mul<f64> for Augmented<T> {
+	type Output = Self;
+
+	fn mul(self, rhs: f64) -> Self {
+		Self {
+			omega: self.omega,
+			a: self.a * rhs,
+			b: self.b * rhs,
+		}
+	}
+}
+
+impl<T: Div<f64, Output = T>> Div<f64> for Augmented<T> {
+	type Output = Self;
+
+	fn div(self, rhs: f64) -> Self {
+		Self {
+			omega: self.omega,
+			a: self.a / rhs,
+			b: self.b / rhs,
+		}
+	}
+}

src/lib.rs

@@ -1,3 +1,5 @@
 pub mod hamiltonian;
+pub mod augmented;
 
 pub use hamiltonian::*;
+pub use augmented::*;