use std::path::PathBuf;
use std::sync::Arc;
use std::sync::mpsc::{Receiver, Sender};
use rsplot::egui;
use tomoxide::io::{DatasetReader, RowBandReader};
use tomoxide::{
Algorithm, Angles, BackendKind, Center, Engine, FilterName, Geometry, PhaseMethod, ReconParams,
StripeMethod,
};
pub struct DatasetMeta {
pub path: PathBuf,
pub nproj: usize,
pub nz: usize,
pub nx: usize,
pub nflat: usize,
pub ndark: usize,
pub theta: Vec<f32>,
pub data_range: (f32, f32),
}
#[derive(Clone)]
pub struct PreviewSpec {
pub slice: usize,
pub algorithm: Algorithm,
pub center: Option<f32>,
pub filter: FilterName,
pub num_iter: usize,
pub reg_par: Vec<f32>,
pub ext_pad: bool,
pub stripe: StripeMethod,
pub phase: PhaseMethod,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CenterMethod {
Vo,
Entropy,
Pc,
Sift,
}
impl CenterMethod {
pub fn label(self) -> &'static str {
match self {
CenterMethod::Vo => "vo",
CenterMethod::Entropy => "entropy",
CenterMethod::Pc => "pc",
CenterMethod::Sift => "sift",
}
}
}
pub enum Job {
OpenDataset(PathBuf),
ReadSinogram { row: usize },
Preview { generation: u64, spec: PreviewSpec },
FindCenter {
method: CenterMethod,
row: usize,
init: Option<f32>,
},
CenterSweep { row: usize, range: (f32, f32, f32) },
LambdaSweep {
spec: PreviewSpec,
lambdas: Vec<f32>,
},
Shutdown,
}
pub enum Event {
BackendReady(String),
DatasetOpened(Arc<DatasetMeta>),
Sinogram {
row: usize,
nproj: usize,
nx: usize,
data: Vec<f32>,
},
CenterFound {
method: CenterMethod,
center: f32,
millis: u128,
},
CenterSweep {
centers: Vec<f32>,
ny: usize,
nx: usize,
frames: Vec<f32>,
millis: u128,
},
LambdaSweep {
lambdas: Vec<f32>,
ny: usize,
nx: usize,
frames: Vec<f32>,
residual: Vec<f64>,
roughness: Vec<f64>,
millis: u128,
},
Preview {
generation: u64,
slice: usize,
ny: usize,
nx: usize,
data: Vec<f32>,
millis: u128,
},
JobFailed { what: String, error: String },
}
pub struct Worker {
pub jobs: Sender<Job>,
pub events: Receiver<Event>,
thread: Option<std::thread::JoinHandle<()>>,
}
impl Worker {
pub fn spawn(ctx: egui::Context) -> Self {
let (job_tx, job_rx) = std::sync::mpsc::channel();
let (event_tx, event_rx) = std::sync::mpsc::channel();
let thread = std::thread::Builder::new()
.name("tomoxide-worker".into())
.spawn(move || worker_main(job_rx, event_tx, ctx))
.expect("spawning the worker thread");
Worker {
jobs: job_tx,
events: event_rx,
thread: Some(thread),
}
}
}
impl Drop for Worker {
fn drop(&mut self) {
let _ = self.jobs.send(Job::Shutdown);
if let Some(t) = self.thread.take() {
let _ = t.join();
}
}
}
fn worker_main(jobs: Receiver<Job>, events: Sender<Event>, ctx: egui::Context) {
let send = |event: Event| {
let _ = events.send(event);
ctx.request_repaint();
};
let engine = match Engine::new(BackendKind::Auto) {
Ok(engine) => {
send(Event::BackendReady(engine.name().to_string()));
engine
}
Err(e) => {
send(Event::JobFailed {
what: "engine init".into(),
error: e.to_string(),
});
return;
}
};
let mut current: Option<PathBuf> = None;
while let Ok(job) = jobs.recv() {
match job {
Job::Shutdown => break,
Job::OpenDataset(path) => match probe(&path) {
Ok(meta) => {
current = Some(path);
send(Event::DatasetOpened(Arc::new(meta)));
}
Err(e) => send(Event::JobFailed {
what: format!("open {}", path.display()),
error: e.to_string(),
}),
},
Job::ReadSinogram { row } => {
let Some(path) = ¤t else {
continue;
};
match read_sinogram(path, row) {
Ok((nproj, nx, data)) => send(Event::Sinogram {
row,
nproj,
nx,
data,
}),
Err(e) => send(Event::JobFailed {
what: format!("sinogram row {row}"),
error: e.to_string(),
}),
}
}
Job::FindCenter { method, row, init } => {
let Some(path) = ¤t else {
continue;
};
let t0 = std::time::Instant::now();
match run_find_center(&engine, path, method, row, init) {
Ok(center) => send(Event::CenterFound {
method,
center,
millis: t0.elapsed().as_millis(),
}),
Err(e) => send(Event::JobFailed {
what: format!("center ({})", method.label()),
error: e.to_string(),
}),
}
}
Job::CenterSweep { row, range } => {
let Some(path) = ¤t else {
continue;
};
let t0 = std::time::Instant::now();
match run_center_sweep(&engine, path, row, range) {
Ok((centers, ny, nx, frames)) => send(Event::CenterSweep {
centers,
ny,
nx,
frames,
millis: t0.elapsed().as_millis(),
}),
Err(e) => send(Event::JobFailed {
what: "center sweep".into(),
error: e.to_string(),
}),
}
}
Job::LambdaSweep { spec, lambdas } => {
let Some(path) = ¤t else {
continue;
};
let t0 = std::time::Instant::now();
match run_lambda_sweep(&engine, path, &spec, &lambdas) {
Ok((ny, nx, frames, residual, roughness)) => send(Event::LambdaSweep {
lambdas,
ny,
nx,
frames,
residual,
roughness,
millis: t0.elapsed().as_millis(),
}),
Err(e) => send(Event::JobFailed {
what: "lambda sweep".into(),
error: e.to_string(),
}),
}
}
Job::Preview { generation, spec } => {
let Some(path) = ¤t else {
continue;
};
let t0 = std::time::Instant::now();
match run_preview(&engine, path, &spec) {
Ok((ny, nx, data)) => send(Event::Preview {
generation,
slice: spec.slice,
ny,
nx,
data,
millis: t0.elapsed().as_millis(),
}),
Err(e) => send(Event::JobFailed {
what: format!("preview slice {}", spec.slice),
error: e.to_string(),
}),
}
}
}
}
}
fn run_preview(
engine: &Engine,
path: &std::path::Path,
spec: &PreviewSpec,
) -> tomoxide::Result<(usize, usize, Vec<f32>)> {
let backend = engine.backend();
let (one, geom) = prep_slice(engine, path, spec)?;
let params = recon_params(spec, one.array.dim().2, spec.reg_par.clone());
let vol = tomoxide::recon::recon(&one, &geom, spec.algorithm, ¶ms, backend)?;
let (_nz1, ny, nxo) = vol.dims();
Ok((ny, nxo, vol.array.iter().copied().collect()))
}
fn prep_slice(
engine: &Engine,
path: &std::path::Path,
spec: &PreviewSpec,
) -> tomoxide::Result<(tomoxide::Tomo, Geometry)> {
let backend = engine.backend();
let mut probe = tomoxide::io::open_dxchange(&path.to_string_lossy())?;
let (_nproj, nz, nx, _nflat, _ndark) = probe.read_sizes()?;
drop(probe);
let slice = spec.slice.min(nz.saturating_sub(1));
let m = tomoxide::prep::phase::margin_rows(&spec.phase);
let z0 = slice.saturating_sub(m);
let z1 = (slice + m + 1).min(nz);
let inner = tomoxide::io::open_dxchange(&path.to_string_lossy())?;
let mut ds = RowBandReader::new(inner, z0, z1)?.read_all()?;
tomoxide::prep::normalize_dataset(&mut ds, backend)?;
tomoxide::prep::retrieve_phase(&mut ds.data, spec.phase, backend)?;
let sino = ds.data.to_layout(tomoxide::Layout::Sinogram);
let row = slice - z0;
let mut one = tomoxide::Tomo::new(
sino.array
.slice(ndarray::s![row..row + 1, .., ..])
.to_owned(),
tomoxide::Layout::Sinogram,
);
tomoxide::prep::remove_stripe(&mut one, spec.stripe)?;
let mut geom = Geometry::parallel(Angles(ds.theta), nx, 1, 1.0);
if let Some(c) = spec.center {
geom.center = Center::Scalar(c);
}
Ok((one, geom))
}
fn recon_params(spec: &PreviewSpec, nx: usize, reg_par: Vec<f32>) -> ReconParams {
ReconParams {
num_gridx: Some(nx),
filter_name: spec.filter,
num_iter: spec.num_iter,
reg_par,
ext_pad: spec.ext_pad,
..Default::default()
}
}
type LambdaSweepOut = (usize, usize, Vec<f32>, Vec<f64>, Vec<f64>);
fn run_lambda_sweep(
engine: &Engine,
path: &std::path::Path,
spec: &PreviewSpec,
lambdas: &[f32],
) -> tomoxide::Result<LambdaSweepOut> {
let backend = engine.backend();
let (one, geom) = prep_slice(engine, path, spec)?;
let nx = one.array.dim().2;
let mut frames = Vec::new();
let mut residual = Vec::with_capacity(lambdas.len());
let mut roughness = Vec::with_capacity(lambdas.len());
let (mut ny, mut nxo) = (0usize, 0usize);
for &lam in lambdas {
let mut reg_par = spec.reg_par.clone();
match reg_par.first_mut() {
Some(first) => *first = lam,
None => reg_par.push(lam),
}
let params = recon_params(spec, nx, reg_par);
let vol = tomoxide::recon::recon(&one, &geom, spec.algorithm, ¶ms, backend)?;
let (_nz1, y, x) = vol.dims();
(ny, nxo) = (y, x);
let sino = tomoxide::sim::project(&vol, &geom, backend)?;
residual.push(l2_diff(&sino.array, &one.array));
roughness.push(tv_seminorm(&vol.array));
frames.extend(vol.array.iter().copied());
}
Ok((ny, nxo, frames, residual, roughness))
}
fn l2_diff(a: &ndarray::Array3<f32>, b: &ndarray::Array3<f32>) -> f64 {
a.iter()
.zip(b.iter())
.map(|(&x, &y)| {
let d = x as f64 - y as f64;
d * d
})
.sum::<f64>()
.sqrt()
}
fn tv_seminorm(vol: &ndarray::Array3<f32>) -> f64 {
let (_nz, ny, nx) = vol.dim();
let mut sum = 0.0f64;
for y in 0..ny {
for x in 0..nx {
let v = vol[[0, y, x]] as f64;
let dx = if x + 1 < nx {
vol[[0, y, x + 1]] as f64 - v
} else {
0.0
};
let dy = if y + 1 < ny {
vol[[0, y + 1, x]] as f64 - v
} else {
0.0
};
sum += (dx * dx + dy * dy).sqrt();
}
}
sum
}
fn run_center_sweep(
engine: &Engine,
path: &std::path::Path,
row: usize,
range: (f32, f32, f32),
) -> tomoxide::Result<(Vec<f32>, usize, usize, Vec<f32>)> {
let backend = engine.backend();
let mut reader = tomoxide::io::open_dxchange(&path.to_string_lossy())?;
let (_nproj, nz, _nx, _nflat, _ndark) = reader.read_sizes()?;
let row = row.min(nz.saturating_sub(1));
let mut ds = reader.read_chunk(row, row + 1)?;
tomoxide::prep::normalize_dataset(&mut ds, backend)?;
let (centers, vols) = tomoxide::recon::center::write_center(
&ds.data,
&ds.theta,
backend,
Some(range),
None,
true,
1.0,
)?;
let (_n, ny, nx) = vols.dim();
Ok((centers, ny, nx, vols.iter().copied().collect()))
}
fn run_find_center(
engine: &Engine,
path: &std::path::Path,
method: CenterMethod,
row: usize,
init: Option<f32>,
) -> tomoxide::Result<f32> {
let backend = engine.backend();
let mut reader = tomoxide::io::open_dxchange(&path.to_string_lossy())?;
match method {
CenterMethod::Vo | CenterMethod::Entropy => {
let (_nproj, nz, _nx, _nflat, _ndark) = reader.read_sizes()?;
let row = row.min(nz.saturating_sub(1));
let mut ds = reader.read_chunk(row, row + 1)?;
tomoxide::prep::normalize_dataset(&mut ds, backend)?;
match method {
CenterMethod::Vo => tomoxide::recon::center::find_center_vo(
&ds.data, backend, None, -50.0, 50.0, 6.0, 0.25, 0.5, 20,
),
_ => tomoxide::recon::center::find_center(
&ds.data, &ds.theta, backend, None, init, 0.5,
),
}
}
CenterMethod::Pc | CenterMethod::Sift => {
let mut ds = reader.read_all()?;
tomoxide::prep::normalize_dataset(&mut ds, backend)?;
let proj = ds.data.to_layout(tomoxide::Layout::Projection);
let nproj = proj.array.dim().0;
if nproj < 2 {
return Err(tomoxide::Error::InvalidParam(
"center pc/sift needs at least two projections".into(),
));
}
let theta0 = ds.theta[0];
let i180 = ds
.theta
.iter()
.enumerate()
.min_by(|(_, a), (_, b)| {
let da = ((**a - theta0).abs() - std::f32::consts::PI).abs();
let db = ((**b - theta0).abs() - std::f32::consts::PI).abs();
da.total_cmp(&db)
})
.map(|(i, _)| i)
.unwrap_or(nproj - 1);
let proj0 = proj.array.index_axis(ndarray::Axis(0), 0).to_owned();
let proj180 = proj.array.index_axis(ndarray::Axis(0), i180).to_owned();
match method {
CenterMethod::Pc => {
tomoxide::recon::center::find_center_pc(&proj0, &proj180, backend, 0.25, init)
}
_ => tomoxide::recon::center::find_center_sift(&proj0, &proj180, 0.5),
}
}
}
}
fn probe(path: &std::path::Path) -> tomoxide::Result<DatasetMeta> {
let mut reader = tomoxide::io::open_dxchange(&path.to_string_lossy())?;
let (nproj, nz, nx, nflat, ndark) = reader.read_sizes()?;
let theta = reader.read_theta()?;
let (_ny, _nx, frame0) =
tomoxide::io::read_h5_frame(&path.to_string_lossy(), tomoxide::io::dxchange::DATA, 0)?;
let mut lo = f32::INFINITY;
let mut hi = f32::NEG_INFINITY;
for &v in &frame0 {
if v.is_finite() {
lo = lo.min(v);
hi = hi.max(v);
}
}
if !(lo.is_finite() && hi.is_finite() && lo < hi) {
(lo, hi) = (0.0, 1.0);
}
Ok(DatasetMeta {
path: path.to_path_buf(),
nproj,
nz,
nx,
nflat,
ndark,
theta,
data_range: (lo, hi),
})
}
fn read_sinogram(path: &std::path::Path, row: usize) -> tomoxide::Result<(usize, usize, Vec<f32>)> {
let mut reader = tomoxide::io::open_dxchange(&path.to_string_lossy())?;
let ds = reader.read_chunk(row, row + 1)?;
let (nproj, _rows, nx) = ds.data.array.dim();
let flat: Vec<f32> = ds.data.array.iter().copied().collect();
Ok((nproj, nx, flat))
}
#[cfg(test)]
mod tests {
use super::*;
fn fixture() -> PathBuf {
PathBuf::from(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../tomoxide/tests/fixtures/streaming_dxchange.h5"
))
}
#[test]
fn preview_reconstructs_one_slice() {
let engine = Engine::new(BackendKind::Cpu).unwrap();
let spec = PreviewSpec {
slice: 3,
algorithm: Algorithm::Fbp,
center: None,
filter: FilterName::Parzen,
num_iter: 1,
reg_par: Vec::new(),
ext_pad: false,
stripe: StripeMethod::None,
phase: PhaseMethod::None,
};
let (ny, nx, data) = run_preview(&engine, &fixture(), &spec).unwrap();
assert_eq!(data.len(), ny * nx);
assert!(data.iter().any(|&v| v != 0.0), "all-zero reconstruction");
}
#[test]
fn preview_phase_banded_runs() {
let engine = Engine::new(BackendKind::Cpu).unwrap();
let spec = PreviewSpec {
slice: 3,
algorithm: Algorithm::Fbp,
center: None,
filter: FilterName::Parzen,
num_iter: 1,
reg_par: Vec::new(),
ext_pad: false,
stripe: StripeMethod::None,
phase: PhaseMethod::Paganin {
pixel_size: 1e-4,
dist: 50.0,
energy: 30.0,
alpha: 1e-3,
},
};
let (ny, nx, data) = run_preview(&engine, &fixture(), &spec).unwrap();
assert_eq!(data.len(), ny * nx);
assert!(data.iter().all(|v| v.is_finite()), "non-finite preview");
assert!(data.iter().any(|&v| v != 0.0), "all-zero reconstruction");
}
#[cfg(feature = "cuda")]
#[test]
fn preview_iterative_cuda_is_finite_and_nonzero() {
if tomoxide::CudaBackend::new().is_err() {
eprintln!("skipping: no usable CUDA device");
return;
}
let engine = Engine::new(BackendKind::Cuda).unwrap();
if engine.name() != "cuda" {
eprintln!("skipping: engine resolved to {}", engine.name());
return;
}
let spec = PreviewSpec {
slice: 3,
algorithm: Algorithm::Sirt,
center: None,
filter: FilterName::Parzen,
num_iter: 10,
reg_par: Vec::new(),
ext_pad: true,
stripe: StripeMethod::None,
phase: PhaseMethod::None,
};
let (ny, nx, data) = run_preview(&engine, &fixture(), &spec).unwrap();
assert_eq!(data.len(), ny * nx);
assert!(
data.iter().all(|v| v.is_finite()),
"non-finite values in iterative preview"
);
assert!(data.iter().any(|&v| v != 0.0), "all-zero reconstruction");
}
#[test]
fn preview_clamps_slice_and_applies_stripe() {
let engine = Engine::new(BackendKind::Cpu).unwrap();
let spec = PreviewSpec {
slice: usize::MAX,
algorithm: Algorithm::Fbp,
center: Some(16.0),
filter: FilterName::Shepp,
num_iter: 1,
reg_par: Vec::new(),
ext_pad: false,
stripe: StripeMethod::Sf { size: 3 },
phase: PhaseMethod::None,
};
let (ny, nx, data) = run_preview(&engine, &fixture(), &spec).unwrap();
assert_eq!(data.len(), ny * nx);
assert!(data.iter().any(|&v| v != 0.0), "all-zero reconstruction");
}
#[cfg(feature = "sift-center")]
#[test]
fn find_center_sift_links_and_runs() {
let engine = Engine::new(BackendKind::Cpu).unwrap();
let meta = probe(&fixture()).unwrap();
match run_find_center(&engine, &fixture(), CenterMethod::Sift, meta.nz / 2, None) {
Ok(c) => assert!(c.is_finite(), "sift center returned non-finite {c}"),
Err(e) => eprintln!("sift center errored on the synthetic fixture (acceptable): {e}"),
}
}
#[test]
fn find_center_vo_and_entropy_run() {
let engine = Engine::new(BackendKind::Cpu).unwrap();
let meta = probe(&fixture()).unwrap();
let mid = meta.nx as f32 / 2.0;
for method in [CenterMethod::Vo, CenterMethod::Entropy] {
let c = run_find_center(&engine, &fixture(), method, meta.nz / 2, None).unwrap();
assert!(
(c - mid).abs() < meta.nx as f32 / 4.0,
"{}: center {c} implausibly far from midline {mid}",
method.label()
);
}
}
#[test]
fn center_sweep_shapes_match_candidates() {
let engine = Engine::new(BackendKind::Cpu).unwrap();
let meta = probe(&fixture()).unwrap();
let mid = meta.nx as f32 / 2.0;
let range = (mid - 1.0, mid + 1.25, 0.5);
let (centers, ny, nx, frames) =
run_center_sweep(&engine, &fixture(), meta.nz / 2, range).unwrap();
assert_eq!(centers.len(), 5);
assert_eq!((ny, nx), (meta.nx, meta.nx));
assert_eq!(frames.len(), centers.len() * ny * nx);
assert!(centers.windows(2).all(|w| w[1] > w[0]));
assert!(frames.iter().all(|v| v.is_finite()));
}
#[test]
fn lambda_sweep_scores_lcurve() {
let engine = Engine::new(BackendKind::Cpu).unwrap();
let meta = probe(&fixture()).unwrap();
let spec = PreviewSpec {
slice: meta.nz / 2,
algorithm: Algorithm::Tv,
center: None,
filter: FilterName::Parzen,
num_iter: 20,
reg_par: vec![0.001],
ext_pad: false,
stripe: StripeMethod::None,
phase: PhaseMethod::None,
};
let lambdas = [0.001_f32, 0.01, 0.1];
let (ny, nx, frames, residual, roughness) =
run_lambda_sweep(&engine, &fixture(), &spec, &lambdas).unwrap();
assert_eq!((ny, nx), (meta.nx, meta.nx));
assert_eq!(frames.len(), lambdas.len() * ny * nx);
assert_eq!(residual.len(), lambdas.len());
assert_eq!(roughness.len(), lambdas.len());
assert!(residual.iter().all(|v| v.is_finite() && *v >= 0.0));
assert!(roughness.iter().all(|v| v.is_finite() && *v >= 0.0));
assert!(
roughness.windows(2).all(|w| w[1] <= w[0] * 1.001),
"roughness not monotone in λ: {roughness:?}"
);
let size = ny * nx;
assert!(
frames[..size] != frames[size..2 * size],
"λ sweep produced identical frames"
);
}
#[cfg(feature = "cuda")]
#[test]
fn lambda_sweep_cuda_is_finite_and_varies() {
if tomoxide::CudaBackend::new().is_err() {
eprintln!("skipping: no usable CUDA device");
return;
}
let engine = Engine::new(BackendKind::Cuda).unwrap();
if engine.name() != "cuda" {
eprintln!("skipping: engine resolved to {}", engine.name());
return;
}
let meta = probe(&fixture()).unwrap();
let spec = PreviewSpec {
slice: meta.nz / 2,
algorithm: Algorithm::Tv,
center: None,
filter: FilterName::Parzen,
num_iter: 20,
reg_par: vec![0.001],
ext_pad: true,
stripe: StripeMethod::None,
phase: PhaseMethod::None,
};
let lambdas = [0.001_f32, 0.01, 0.1];
let (_ny, _nx, _frames, residual, roughness) =
run_lambda_sweep(&engine, &fixture(), &spec, &lambdas).unwrap();
assert!(residual.iter().all(|v| v.is_finite() && *v > 0.0));
assert!(roughness.iter().all(|v| v.is_finite()));
let (rmin, rmax) = residual
.iter()
.fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), &v| {
(lo.min(v), hi.max(v))
});
assert!(rmax > rmin, "residual L-curve axis is flat: {residual:?}");
}
#[test]
fn find_center_pc_runs() {
let engine = Engine::new(BackendKind::Cpu).unwrap();
let meta = probe(&fixture()).unwrap();
let c = run_find_center(&engine, &fixture(), CenterMethod::Pc, 0, None).unwrap();
assert!(
c > 0.0 && c < meta.nx as f32,
"pc center {c} outside the detector"
);
}
#[test]
fn probe_reads_sizes_and_theta() {
let meta = probe(&fixture()).unwrap();
assert!(meta.nproj > 0 && meta.nz > 0 && meta.nx > 0);
assert_eq!(meta.theta.len(), meta.nproj);
}
}