use core::f64::consts::TAU;
use std::path::{Path, PathBuf};
use super::Engine;
use super::fields::{FieldArrays, Fields};
use super::grid::Grid;
use super::operator::{Operator, StepContext};
use crate::axis::Axis;
use crate::io::{ComplexForm, Field, Layout, PvdCollection, VtrWriter};
use crate::post::Complex64;
use crate::post::dft::phase_cycles;
use crate::{Error, Result};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct DumpRegion {
lo: [usize; 3],
hi: [usize; 3],
stride: [usize; 3],
}
impl DumpRegion {
pub fn new(a: [usize; 3], b: [usize; 3]) -> Self {
Self {
lo: [0, 1, 2].map(|i| a[i].min(b[i])),
hi: [0, 1, 2].map(|i| a[i].max(b[i])),
stride: [1; 3],
}
}
pub fn whole(grid: &Grid) -> Self {
Self::new([0; 3], grid.nodes().map(|n| n - 1))
}
pub fn stride(mut self, stride: [usize; 3]) -> Self {
self.stride = stride.map(|s| s.max(1));
self
}
pub fn lo(&self) -> [usize; 3] {
self.lo
}
pub fn hi(&self) -> [usize; 3] {
self.hi
}
pub fn strides(&self) -> [usize; 3] {
self.stride
}
pub fn check_fits(&self, nodes: [usize; 3]) -> Result<()> {
for a in Axis::ALL {
let i = a.index();
if self.hi[i] >= nodes[i] {
return Err(Error::Invalid(format!(
"dump box {:?}..={:?} exceeds the grid ({} nodes along {a})",
self.lo, self.hi, nodes[i]
)));
}
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DumpQuantity {
E,
H,
}
impl DumpQuantity {
pub fn name(self) -> &'static str {
match self {
DumpQuantity::E => "E",
DumpQuantity::H => "H",
}
}
fn offset(self) -> f64 {
match self {
DumpQuantity::E => 0.0,
DumpQuantity::H => -0.5,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub enum DumpPosition {
#[default]
Nodes,
CellCenters,
}
type Taps = [(usize, f64); 2];
#[derive(Debug, Clone, Copy)]
enum Pos {
Node(usize),
Mid(usize),
}
fn taps(grid: &Grid, q: DumpQuantity, comp: Axis, axis: Axis, pos: Pos) -> Taps {
let n = grid.nodes()[axis.index()];
let d = |i: usize| grid.delta(axis, i);
let one = |i: usize, w: f64| [(i, w), (i, 0.0)];
let lin = |j: usize| -> [(usize, f64); 2] {
if j == 0 {
one(0, 1.0)
} else if j == n - 1 {
one(n - 2, 1.0)
} else {
let (dl, dh) = (d(j - 1), d(j));
[(j - 1, dh / (dl + dh)), (j, dl / (dl + dh))]
}
};
let mut t = match (q, comp == axis, pos) {
(DumpQuantity::E, true, Pos::Node(j)) => {
let [(a, wa), (b, wb)] = lin(j);
[(a, wa / d(a)), (b, wb / d(b))]
}
(DumpQuantity::E, true, Pos::Mid(j)) => one(j, 1.0 / d(j)),
(DumpQuantity::E, false, Pos::Node(j)) => one(j, 1.0),
(DumpQuantity::E, false, Pos::Mid(j)) => [(j, 0.5), (j + 1, 0.5)],
(DumpQuantity::H, true, Pos::Node(j)) => one(j, 1.0 / grid.dual_delta(axis, j)),
(DumpQuantity::H, true, Pos::Mid(j)) => [
(j, 0.5 / grid.dual_delta(axis, j)),
(j + 1, 0.5 / grid.dual_delta(axis, j + 1)),
],
(DumpQuantity::H, false, Pos::Node(j)) => lin(j),
(DumpQuantity::H, false, Pos::Mid(j)) => one(j, 1.0),
};
if grid.is_cylindrical() && comp == Axis::Y && axis == Axis::X {
let r = grid.lines(Axis::X);
let rho = |i: usize| match q {
DumpQuantity::E => r[i],
DumpQuantity::H => grid.dual_line(Axis::X, i.min(n - 2)),
};
let total: f64 = t.iter().map(|x| x.1).sum();
let valid: f64 = t.iter().filter(|x| rho(x.0) > 0.0).map(|x| x.1).sum();
for x in &mut t {
let p = rho(x.0);
x.1 = if p > 0.0 && valid != 0.0 {
x.1 * (total / valid) / p
} else {
0.0
};
}
}
t
}
#[derive(Debug, Clone)]
pub struct FieldSampler {
region: DumpRegion,
quantity: DumpQuantity,
position: DumpPosition,
coords: [Vec<f64>; 3],
taps: [[Vec<Taps>; 3]; 3],
}
impl FieldSampler {
pub fn new(
grid: &Grid,
region: DumpRegion,
quantity: DumpQuantity,
position: DumpPosition,
) -> Result<Self> {
region.check_fits(grid.nodes())?;
let pos = Axis::ALL.map(|a| {
let i = a.index();
let (lo, hi, s) = (region.lo[i], region.hi[i], region.stride[i]);
match position {
DumpPosition::CellCenters if hi > lo => {
(lo..hi).step_by(s).map(Pos::Mid).collect::<Vec<_>>()
}
_ => (lo..=hi).step_by(s).map(Pos::Node).collect(),
}
});
let coords = Axis::ALL.map(|a| {
pos[a.index()]
.iter()
.map(|&p| match p {
Pos::Node(j) => grid.lines(a)[j],
Pos::Mid(j) => grid.dual_line(a, j),
})
.collect()
});
let taps = Axis::ALL.map(|c| {
Axis::ALL.map(|a| {
pos[a.index()]
.iter()
.map(|&p| taps(grid, quantity, c, a, p))
.collect()
})
});
Ok(Self {
region,
quantity,
position,
coords,
taps,
})
}
pub fn region(&self) -> &DumpRegion {
&self.region
}
pub fn quantity(&self) -> DumpQuantity {
self.quantity
}
pub fn position(&self) -> DumpPosition {
self.position
}
pub fn dims(&self) -> [usize; 3] {
[0, 1, 2].map(|i| self.coords[i].len())
}
pub fn len(&self) -> usize {
self.dims().iter().product()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn coords(&self, axis: Axis) -> &[f64] {
&self.coords[axis.index()]
}
pub fn time(&self, step: u64, dt: f64) -> f64 {
(step as f64 + self.quantity.offset()) * dt
}
pub fn sample_with(&self, fields: &Fields<'_>, mut f: impl FnMut(usize, [f64; 3])) {
let arr: &FieldArrays<'_> = match self.quantity {
DumpQuantity::E => &fields.u,
DumpQuantity::H => &fields.i,
};
let [dx, dy, dz] = self.dims();
let mut idx = 0;
for p in 0..dx {
for q in 0..dy {
for r in 0..dz {
let mut v = [0.0; 3];
for c in Axis::ALL {
let t = &self.taps[c.index()];
let mut s = 0.0;
for &(i, wx) in &t[0][p] {
if wx == 0.0 {
continue;
}
for &(j, wy) in &t[1][q] {
if wy == 0.0 {
continue;
}
for &(k, wz) in &t[2][r] {
if wz != 0.0 {
s += wx * wy * wz * f64::from(arr.get(c, [i, j, k]));
}
}
}
}
v[c.index()] = s;
}
f(idx, v);
idx += 1;
}
}
}
}
pub fn sample(&self, fields: &Fields<'_>) -> Vec<[f32; 3]> {
let mut out = vec![[0.0f32; 3]; self.len()];
self.sample_with(fields, |i, v| out[i] = v.map(|x| x as f32));
out
}
pub fn snapshot(&self, engine: &Engine) -> Snapshot {
let fields = engine.shared.store.lock();
let step = engine.timestep();
Snapshot {
step,
time: self.time(step, engine.dt()),
data: self.sample(&fields),
}
}
fn vtr<'a>(&'a self, data: &'a [[f32; 3]], time: f64) -> VtrWriter<'a> {
VtrWriter::new(&self.coords[0], &self.coords[1], &self.coords[2])
.time(time)
.point_vector(self.quantity.name(), data, Layout::ZFastest)
}
pub fn write_vtr(&self, path: impl AsRef<Path>, snapshot: &Snapshot) -> Result<()> {
self.vtr(&snapshot.data, snapshot.time).write(path)
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct Snapshot {
pub step: u64,
pub time: f64,
pub data: Vec<[f32; 3]>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum Schedule {
Every(u64),
Steps(Vec<u64>),
Interval(f64),
}
#[derive(Debug, Clone, Copy, PartialEq)]
struct Window {
start: f64,
stop: f64,
}
impl Default for Window {
fn default() -> Self {
Self {
start: f64::NEG_INFINITY,
stop: 0.0,
}
}
}
impl Window {
fn contains(&self, t: f64, dt: f64) -> bool {
let eps = 1e-9 * dt;
t >= self.start - eps && (self.stop <= 0.0 || t <= self.stop + eps)
}
}
#[derive(Debug, Clone)]
struct Scheduler {
schedule: Schedule,
next: f64,
}
impl Scheduler {
fn new(mut schedule: Schedule) -> Self {
if let Schedule::Steps(v) = &mut schedule {
v.sort_unstable();
v.dedup();
}
Self {
schedule,
next: f64::NEG_INFINITY,
}
}
fn due(&mut self, m: u64, dt: f64) -> bool {
match &self.schedule {
Schedule::Every(k) => m.is_multiple_of((*k).max(1)),
Schedule::Steps(v) => v.binary_search(&m).is_ok(),
Schedule::Interval(t) => {
let t = *t;
if !(t > 0.0 && t.is_finite()) {
return true;
}
let now = m as f64 * dt;
if now < self.next - 1e-9 * dt {
return false;
}
self.next = ((now / t + 1e-9).floor() + 1.0) * t;
true
}
}
}
}
pub trait DumpSink: Send + 'static {
fn snapshot(&mut self, sampler: &FieldSampler, snapshot: Snapshot) -> Result<()>;
}
impl<F> DumpSink for F
where
F: FnMut(&FieldSampler, Snapshot) -> Result<()> + Send + 'static,
{
fn snapshot(&mut self, sampler: &FieldSampler, snapshot: Snapshot) -> Result<()> {
self(sampler, snapshot)
}
}
#[derive(Debug, Clone, Default)]
pub struct MemorySink {
snapshots: Vec<Snapshot>,
}
impl MemorySink {
pub fn new() -> Self {
Self::default()
}
pub fn snapshots(&self) -> &[Snapshot] {
&self.snapshots
}
pub fn take(&mut self) -> Vec<Snapshot> {
core::mem::take(&mut self.snapshots)
}
}
impl DumpSink for MemorySink {
fn snapshot(&mut self, _: &FieldSampler, snapshot: Snapshot) -> Result<()> {
self.snapshots.push(snapshot);
Ok(())
}
}
fn snapshot_file(basename: &str, step: u64) -> String {
format!("{basename}_{step:08}.vtr")
}
#[derive(Debug, Clone)]
pub struct VtrSeries {
dir: PathBuf,
basename: String,
pvd: PvdCollection,
}
impl VtrSeries {
pub fn new(dir: impl Into<PathBuf>, basename: impl Into<String>) -> Self {
Self {
dir: dir.into(),
basename: basename.into(),
pvd: PvdCollection::new(),
}
}
pub fn dir(&self) -> &Path {
&self.dir
}
pub fn pvd_path(&self) -> PathBuf {
self.dir.join(format!("{}.pvd", self.basename))
}
pub fn collection(&self) -> &PvdCollection {
&self.pvd
}
}
impl DumpSink for VtrSeries {
fn snapshot(&mut self, sampler: &FieldSampler, snapshot: Snapshot) -> Result<()> {
if self.pvd.entries().is_empty() {
std::fs::create_dir_all(&self.dir)?;
}
let file = snapshot_file(&self.basename, snapshot.step);
sampler.write_vtr(self.dir.join(&file), &snapshot)?;
self.pvd.push(snapshot.time, file);
self.pvd.write(self.pvd_path())
}
}
#[derive(Debug, Clone, Copy, Default)]
struct Cursor {
started: bool,
}
impl Cursor {
fn pre_h(&mut self, ctx: &StepContext<'_>) -> Option<u64> {
(!core::mem::replace(&mut self.started, true)).then_some(ctx.step)
}
fn post_e(&mut self, ctx: &StepContext<'_>) -> u64 {
self.started = true;
ctx.step + 1
}
}
#[derive(Debug)]
pub struct TimeDump<S = MemorySink> {
sampler: FieldSampler,
window: Window,
scheduler: Scheduler,
sink: S,
cursor: Cursor,
error: Option<Error>,
}
impl TimeDump<MemorySink> {
pub fn new(sampler: FieldSampler) -> Self {
Self::with_sink(sampler, MemorySink::new())
}
pub fn write_vtr(&self, dir: impl AsRef<Path>, basename: &str) -> Result<()> {
let mut series = VtrSeries::new(dir.as_ref(), basename);
for s in self.sink.snapshots() {
series.snapshot(&self.sampler, s.clone())?;
}
Ok(())
}
}
impl<S: DumpSink> TimeDump<S> {
pub fn with_sink(sampler: FieldSampler, sink: S) -> Self {
Self {
sampler,
window: Window::default(),
scheduler: Scheduler::new(Schedule::Every(1)),
sink,
cursor: Cursor::default(),
error: None,
}
}
pub fn schedule(mut self, schedule: Schedule) -> Self {
self.scheduler = Scheduler::new(schedule);
self
}
pub fn window(mut self, start: f64, stop: f64) -> Self {
self.window = Window { start, stop };
self
}
pub fn sampler(&self) -> &FieldSampler {
&self.sampler
}
pub fn sink(&self) -> &S {
&self.sink
}
pub fn sink_mut(&mut self) -> &mut S {
&mut self.sink
}
pub fn error(&self) -> Option<&Error> {
self.error.as_ref()
}
fn sample(&mut self, fields: &Fields<'_>, m: u64, dt: f64) {
if self.error.is_some()
|| !self.window.contains(self.sampler.time(m, dt), dt)
|| !self.scheduler.due(m, dt)
{
return;
}
let snap = Snapshot {
step: m,
time: self.sampler.time(m, dt),
data: self.sampler.sample(fields),
};
if let Err(e) = self.sink.snapshot(&self.sampler, snap) {
self.error = Some(e);
}
}
}
impl<S: DumpSink> Operator for TimeDump<S> {
fn pre_h(&mut self, fields: &mut Fields<'_>, ctx: &StepContext<'_>) {
if let Some(m) = self.cursor.pre_h(ctx) {
self.sample(fields, m, ctx.dt);
}
}
fn post_e(&mut self, fields: &mut Fields<'_>, ctx: &StepContext<'_>) {
let m = self.cursor.post_e(ctx);
self.sample(fields, m, ctx.dt);
}
}
#[derive(Debug, Clone)]
pub struct FreqDump {
sampler: FieldSampler,
freqs: Vec<f64>,
interval: u64,
window: Window,
acc: Vec<[Complex64; 3]>,
samples: u64,
dt: f64,
cursor: Cursor,
phasors: Vec<Complex64>,
}
impl FreqDump {
pub fn new(sampler: FieldSampler, freqs: impl Into<Vec<f64>>) -> Self {
let freqs = freqs.into();
let n = sampler.len() * freqs.len();
Self {
sampler,
phasors: vec![Complex64::ZERO; freqs.len()],
freqs,
interval: 1,
window: Window::default(),
acc: vec![[Complex64::ZERO; 3]; n],
samples: 0,
dt: 0.0,
cursor: Cursor::default(),
}
}
pub fn interval(mut self, k: u64) -> Self {
self.interval = k.max(1);
self
}
pub fn window(mut self, start: f64, stop: f64) -> Self {
self.window = Window { start, stop };
self
}
pub fn sampler(&self) -> &FieldSampler {
&self.sampler
}
pub fn freqs(&self) -> &[f64] {
&self.freqs
}
pub fn samples(&self) -> u64 {
self.samples
}
pub fn sample_interval(&self) -> f64 {
self.interval as f64 * self.dt
}
pub fn values(&self, k: usize) -> Vec<[Complex64; 3]> {
assert!(k < self.freqs.len(), "frequency index {k} out of range");
let n = self.sampler.len();
let s = self.sample_interval();
self.acc[k * n..(k + 1) * n]
.iter()
.map(|v| v.map(|c| c * s))
.collect()
}
pub(crate) fn add_imaginary(&mut self, imag: &FreqDump) {
debug_assert_eq!(self.acc.len(), imag.acc.len());
for (a, b) in self.acc.iter_mut().zip(&imag.acc) {
for c in 0..3 {
a[c] += Complex64::I * b[c];
}
}
}
pub fn reset(&mut self) {
self.acc.fill([Complex64::ZERO; 3]);
self.samples = 0;
}
pub fn write_vtr(&self, dir: impl AsRef<Path>, basename: &str) -> Result<()> {
let dir = dir.as_ref();
std::fs::create_dir_all(dir)?;
let s = &self.sampler;
let mut pvd = PvdCollection::new();
for (k, &f) in self.freqs.iter().enumerate() {
let v = self.values(k);
let re: Vec<[f64; 3]> = v.iter().map(|c| c.map(|z| z.re)).collect();
let im: Vec<[f64; 3]> = v.iter().map(|c| c.map(|z| z.im)).collect();
let fa = [f];
let file = format!("{basename}_f{k}.vtr");
VtrWriter::new(&s.coords[0], &s.coords[1], &s.coords[2])
.time(f)
.field("Frequency", &fa)
.point_complex(
s.quantity.name(),
Field::Vector(&re),
Field::Vector(&im),
Layout::ZFastest,
ComplexForm::ReIm,
)
.write(dir.join(&file))?;
pvd.push(f, file);
}
pvd.write(dir.join(format!("{basename}.pvd")))
}
fn sample(&mut self, fields: &Fields<'_>, m: u64, dt: f64) {
self.dt = dt;
if !m.is_multiple_of(self.interval) || !self.window.contains(self.sampler.time(m, dt), dt) {
return;
}
let t0 = self.sampler.quantity.offset() * dt;
for (p, &f) in self.phasors.iter_mut().zip(&self.freqs) {
*p = Complex64::cis(-TAU * phase_cycles(f, t0, dt, m as f64));
}
let n = self.sampler.len();
let (acc, phasors) = (&mut self.acc, &self.phasors);
self.sampler.sample_with(fields, |i, v| {
for (k, &p) in phasors.iter().enumerate() {
let a = &mut acc[k * n + i];
for (a, &x) in a.iter_mut().zip(&v) {
*a += p * x;
}
}
});
self.samples += 1;
}
}
impl Operator for FreqDump {
fn pre_h(&mut self, fields: &mut Fields<'_>, ctx: &StepContext<'_>) {
if let Some(m) = self.cursor.pre_h(ctx) {
self.sample(fields, m, ctx.dt);
}
}
fn post_e(&mut self, fields: &mut Fields<'_>, ctx: &StepContext<'_>) {
let m = self.cursor.post_e(ctx);
self.sample(fields, m, ctx.dt);
}
}