Skip to main content

qsm_core/bet/
evolution.rs

1//! BET surface evolution algorithm
2//!
3//! Based on: Smith, S.M. (2002) "Fast robust automated brain extraction"
4//! Human Brain Mapping, 17(3):143-155
5//!
6//! Aligned with FSL-BET2 implementation.
7
8use crate::Grid;
9use super::icosphere::create_icosphere;
10use super::mesh::{build_neighbor_matrix, compute_vertex_normals, compute_mean_edge_length_mm, self_intersection_heuristic};
11use std::collections::VecDeque;
12
13/// Brain parameters struct (like FSL's bet_parameters)
14struct BetParameters {
15    t2: f64,       // 2nd percentile (robust min)
16    t98: f64,      // 98th percentile (robust max)
17    t: f64,        // threshold = t2 + 0.1*(t98-t2)
18    tm: f64,       // median within-brain intensity (critical for proper surface evolution)
19    cog: [f64; 3], // center of gravity in voxel coordinates
20    cog_mm: [f64; 3], // center of gravity in mm (for z-gradient)
21    radius: f64,   // estimated brain radius in mm
22}
23
24/// Estimate brain parameters from the image (matches FSL-BET2's adjust_initial_mesh)
25fn estimate_brain_parameters(
26    data: &[f64],
27    nx: usize, ny: usize, nz: usize,
28    voxel_size: &[f64; 3],
29) -> BetParameters {
30    // Collect non-zero values
31    let nonzero: Vec<f64> = data.iter().copied().filter(|&v| v > 0.0).collect();
32
33    if nonzero.is_empty() {
34        let cog = [(nx as f64) / 2.0, (ny as f64) / 2.0, (nz as f64) / 2.0];
35        let cog_mm = [cog[0] * voxel_size[0], cog[1] * voxel_size[1], cog[2] * voxel_size[2]];
36        return BetParameters { t2: 0.0, t98: 1.0, t: 0.1, tm: 0.5, cog, cog_mm, radius: 50.0 };
37    }
38
39    // Sort for percentiles
40    let mut sorted = nonzero.clone();
41    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
42
43    let t2 = percentile(&sorted, 2.0);
44    let t98 = percentile(&sorted, 98.0);
45    let t = t2 + 0.1 * (t98 - t2);
46
47    // Find center of gravity (weighted by intensity, like FSL)
48    let mut sum_x = 0.0;
49    let mut sum_y = 0.0;
50    let mut sum_z = 0.0;
51    let mut sum_weight = 0.0;
52    let mut n_voxels = 0usize;
53
54    // Use Fortran order: index = x + y*nx + z*nx*ny
55    for k in 0..nz {
56        for j in 0..ny {
57            for i in 0..nx {
58                let idx = i + j * nx + k * nx * ny;
59                let val = data[idx];
60                if val > t {
61                    // FSL: weight = min(c, t98 - t2) where c = val - t2
62                    let c = (val - t2).min(t98 - t2);
63                    sum_x += (i as f64) * c;
64                    sum_y += (j as f64) * c;
65                    sum_z += (k as f64) * c;
66                    sum_weight += c;
67                    n_voxels += 1;
68                }
69            }
70        }
71    }
72
73    let cog = if sum_weight > 0.0 {
74        [sum_x / sum_weight, sum_y / sum_weight, sum_z / sum_weight]
75    } else {
76        [(nx as f64) / 2.0, (ny as f64) / 2.0, (nz as f64) / 2.0]
77    };
78
79    // Estimate brain radius
80    let voxel_volume = voxel_size[0] * voxel_size[1] * voxel_size[2];
81    let brain_volume = (n_voxels as f64) * voxel_volume;
82    let radius = (3.0 * brain_volume / (4.0 * std::f64::consts::PI)).powf(1.0 / 3.0);
83
84    // Compute tm: median intensity within a sphere centered at COG with radius = brain radius
85    // This is critical for proper intensity-based surface evolution (FSL bet2.cpp lines 385-403)
86    let cog_mm = [cog[0] * voxel_size[0], cog[1] * voxel_size[1], cog[2] * voxel_size[2]];
87    let radius_sq = radius * radius;
88
89    let mut within_brain_values: Vec<f64> = Vec::new();
90    for k in 0..nz {
91        for j in 0..ny {
92            for i in 0..nx {
93                let idx = i + j * nx + k * nx * ny;
94                let val = data[idx];
95                // Only consider voxels with intensity between t2 and t98
96                if val > t2 && val < t98 {
97                    // Check if within sphere of radius centered at COG
98                    let px = (i as f64) * voxel_size[0];
99                    let py = (j as f64) * voxel_size[1];
100                    let pz = (k as f64) * voxel_size[2];
101                    let dx = px - cog_mm[0];
102                    let dy = py - cog_mm[1];
103                    let dz = pz - cog_mm[2];
104                    let dist_sq = dx * dx + dy * dy + dz * dz;
105                    if dist_sq < radius_sq {
106                        within_brain_values.push(val);
107                    }
108                }
109            }
110        }
111    }
112
113    // Compute median (tm)
114    let tm = if within_brain_values.is_empty() {
115        (t2 + t98) / 2.0 // fallback
116    } else {
117        within_brain_values.sort_by(|a, b| a.partial_cmp(b).unwrap());
118        let mid = within_brain_values.len() / 2;
119        within_brain_values[mid]
120    };
121
122    BetParameters { t2, t98, t, tm, cog, cog_mm, radius }
123}
124
125/// Compute percentile of sorted array
126fn percentile(sorted: &[f64], p: f64) -> f64 {
127    if sorted.is_empty() {
128        return 0.0;
129    }
130    let idx = (p / 100.0 * (sorted.len() - 1) as f64).round() as usize;
131    sorted[idx.min(sorted.len() - 1)]
132}
133
134/// Trilinear interpolation
135fn sample_intensity(data: &[f64], nx: usize, ny: usize, nz: usize, x: f64, y: f64, z: f64) -> f64 {
136    let x = x.max(0.0).min((nx - 1) as f64);
137    let y = y.max(0.0).min((ny - 1) as f64);
138    let z = z.max(0.0).min((nz - 1) as f64);
139
140    let x0 = x.floor() as usize;
141    let y0 = y.floor() as usize;
142    let z0 = z.floor() as usize;
143    let x1 = (x0 + 1).min(nx - 1);
144    let y1 = (y0 + 1).min(ny - 1);
145    let z1 = (z0 + 1).min(nz - 1);
146
147    let xd = x - x0 as f64;
148    let yd = y - y0 as f64;
149    let zd = z - z0 as f64;
150
151    // Fortran order: x + y*nx + z*nx*ny
152    let idx = |i: usize, j: usize, k: usize| i + j * nx + k * nx * ny;
153
154    let c000 = data[idx(x0, y0, z0)];
155    let c001 = data[idx(x0, y0, z1)];
156    let c010 = data[idx(x0, y1, z0)];
157    let c011 = data[idx(x0, y1, z1)];
158    let c100 = data[idx(x1, y0, z0)];
159    let c101 = data[idx(x1, y0, z1)];
160    let c110 = data[idx(x1, y1, z0)];
161    let c111 = data[idx(x1, y1, z1)];
162
163    let c00 = c000 * (1.0 - xd) + c100 * xd;
164    let c01 = c001 * (1.0 - xd) + c101 * xd;
165    let c10 = c010 * (1.0 - xd) + c110 * xd;
166    let c11 = c011 * (1.0 - xd) + c111 * xd;
167
168    let c0 = c00 * (1.0 - yd) + c10 * yd;
169    let c1 = c01 * (1.0 - yd) + c11 * yd;
170
171    c0 * (1.0 - zd) + c1 * zd
172}
173
174/// Sample min/max intensities along inward normal (matches FSL-BET2's step_of_computation)
175///
176/// FSL samples from 1mm to d1 (7mm) for Imin, and up to d2 (3mm) for Imax.
177/// Initial values: Imin = tm, Imax = t
178/// Final clamps: Imin >= t2, Imax <= tm
179///
180/// point_mm: vertex position in mm coordinates
181/// normal: unit normal vector (outward pointing)
182///
183/// Returns (i_min, i_max, success) where success=false means sampling failed
184/// and the caller should use f3=0 (like FSL does).
185fn sample_intensities_fsl(
186    data: &[f64],
187    nx: usize, ny: usize, nz: usize,
188    point_mm: &[f64; 3],
189    normal: &[f64; 3],
190    voxel_size: &[f64; 3],
191    t2: f64,
192    t: f64,
193    tm: f64,
194) -> (f64, f64, bool) {
195    let d1 = 7.0; // max search distance for Imin (mm)
196    let d2 = 3.0; // max search distance for Imax (mm)
197    let dscale = voxel_size[0].min(voxel_size[1]).min(voxel_size[2]).min(1.0);
198
199    // Initialize like FSL does
200    let mut i_min = tm;
201    let mut i_max = t;
202
203    // Convert mm to voxel helper
204    let mm_to_voxel = |p_mm: &[f64; 3]| -> [f64; 3] {
205        [p_mm[0] / voxel_size[0], p_mm[1] / voxel_size[1], p_mm[2] / voxel_size[2]]
206    };
207
208    // Check if voxel position is in bounds
209    let in_bounds = |v: &[f64; 3]| -> bool {
210        v[0] >= 0.0 && v[0] < (nx - 1) as f64 &&
211        v[1] >= 0.0 && v[1] < (ny - 1) as f64 &&
212        v[2] >= 0.0 && v[2] < (nz - 1) as f64
213    };
214
215    // Starting position in mm (1mm inward along normal)
216    let mut p_mm = [
217        point_mm[0] - normal[0],
218        point_mm[1] - normal[1],
219        point_mm[2] - normal[2],
220    ];
221    let mut p_vox = mm_to_voxel(&p_mm);
222
223    // Check if starting point is in bounds (FSL: first bounds check)
224    if !in_bounds(&p_vox) {
225        // FSL: if first bounds check fails, f3 = 0 (no intensity force)
226        return (i_min, i_max, false);
227    }
228
229    let im = sample_intensity(data, nx, ny, nz, p_vox[0], p_vox[1], p_vox[2]);
230    i_min = i_min.min(im);
231    i_max = i_max.max(im);
232
233    // Check far point at d1-1 (FSL: second bounds check)
234    let p_far_mm = [
235        point_mm[0] - (d1 - 1.0) * normal[0],
236        point_mm[1] - (d1 - 1.0) * normal[1],
237        point_mm[2] - (d1 - 1.0) * normal[2],
238    ];
239    let p_far_vox = mm_to_voxel(&p_far_mm);
240
241    if !in_bounds(&p_far_vox) {
242        // FSL: if second bounds check fails, f3 = 0 (no intensity force)
243        return (i_min, i_max, false);
244    }
245
246    let im = sample_intensity(data, nx, ny, nz, p_far_vox[0], p_far_vox[1], p_far_vox[2]);
247    i_min = i_min.min(im);
248
249    // Sample from 2mm to d1 (stepping by dscale mm)
250    let mut gi = 2.0;
251    while gi < d1 {
252        p_mm[0] -= normal[0] * dscale;
253        p_mm[1] -= normal[1] * dscale;
254        p_mm[2] -= normal[2] * dscale;
255        p_vox = mm_to_voxel(&p_mm);
256
257        if in_bounds(&p_vox) {
258            let im = sample_intensity(data, nx, ny, nz, p_vox[0], p_vox[1], p_vox[2]);
259            i_min = i_min.min(im);
260
261            // Only update Imax for samples within d2
262            if gi < d2 {
263                i_max = i_max.max(im);
264            }
265        }
266
267        gi += dscale;
268    }
269
270    // Clamp like FSL does (this is critical for sinus exclusion)
271    i_min = i_min.max(t2);    // Imin can't go below noise floor
272    i_max = i_max.min(tm);    // Imax can't go above median brain intensity
273
274    (i_min, i_max, true)
275}
276
277/// Convert surface mesh to binary mask using flood fill
278///
279/// vertices_mm: vertex positions in mm coordinates
280/// voxel_size: voxel dimensions in mm
281fn surface_to_mask(
282    vertices_mm: &[[f64; 3]],
283    faces: &[[usize; 3]],
284    nx: usize, ny: usize, nz: usize,
285    voxel_size: &[f64; 3],
286) -> Vec<u8> {
287    // Convert mm vertices to voxel coordinates
288    let vertices: Vec<[f64; 3]> = vertices_mm
289        .iter()
290        .map(|v| [
291            v[0] / voxel_size[0],
292            v[1] / voxel_size[1],
293            v[2] / voxel_size[2],
294        ])
295        .collect();
296
297    let mininc = 0.5;
298
299    // Start with all 1s (outside)
300    let mut grid: Vec<u8> = vec![1; nx * ny * nz];
301    // Fortran order: x + y*nx + z*nx*ny
302    let idx = |i: usize, j: usize, k: usize| i + j * nx + k * nx * ny;
303
304    // Draw mesh surface as 0s
305    for &[i0, i1, i2] in faces {
306        let v0 = vertices[i0];
307        let v1 = vertices[i1];
308        let v2 = vertices[i2];
309
310        // Edge from v1 to v0
311        let edge = [v0[0] - v1[0], v0[1] - v1[1], v0[2] - v1[2]];
312        let edge_len = (edge[0].powi(2) + edge[1].powi(2) + edge[2].powi(2)).sqrt();
313
314        if edge_len < 0.001 {
315            continue;
316        }
317
318        let edge_dir = [edge[0] / edge_len, edge[1] / edge_len, edge[2] / edge_len];
319        let n_edge_steps = (edge_len / mininc).ceil() as usize + 1;
320
321        for j in 0..n_edge_steps {
322            let d = (j as f64) * mininc;
323            let p_edge = if d > edge_len {
324                v0
325            } else {
326                [v1[0] + d * edge_dir[0], v1[1] + d * edge_dir[1], v1[2] + d * edge_dir[2]]
327            };
328
329            // Draw segment from p_edge to v2
330            let seg = [v2[0] - p_edge[0], v2[1] - p_edge[1], v2[2] - p_edge[2]];
331            let seg_len = (seg[0].powi(2) + seg[1].powi(2) + seg[2].powi(2)).sqrt();
332
333            if seg_len < 0.001 {
334                let ix = p_edge[0].round() as isize;
335                let iy = p_edge[1].round() as isize;
336                let iz = p_edge[2].round() as isize;
337                if ix >= 0 && ix < nx as isize && iy >= 0 && iy < ny as isize && iz >= 0 && iz < nz as isize {
338                    grid[idx(ix as usize, iy as usize, iz as usize)] = 0;
339                }
340                continue;
341            }
342
343            let seg_dir = [seg[0] / seg_len, seg[1] / seg_len, seg[2] / seg_len];
344            let n_seg_steps = (seg_len / mininc).ceil() as usize + 1;
345
346            for k in 0..n_seg_steps {
347                let sd = (k as f64) * mininc;
348                let p = if sd > seg_len {
349                    v2
350                } else {
351                    [p_edge[0] + sd * seg_dir[0], p_edge[1] + sd * seg_dir[1], p_edge[2] + sd * seg_dir[2]]
352                };
353
354                let ix = p[0].round() as isize;
355                let iy = p[1].round() as isize;
356                let iz = p[2].round() as isize;
357
358                if ix >= 0 && ix < nx as isize && iy >= 0 && iy < ny as isize && iz >= 0 && iz < nz as isize {
359                    grid[idx(ix as usize, iy as usize, iz as usize)] = 0;
360                }
361            }
362        }
363    }
364
365    // Flood fill from center of mesh
366    let mut center = [0.0, 0.0, 0.0];
367    for v in &vertices {
368        center[0] += v[0];
369        center[1] += v[1];
370        center[2] += v[2];
371    }
372    center[0] /= vertices.len() as f64;
373    center[1] /= vertices.len() as f64;
374    center[2] /= vertices.len() as f64;
375
376    let mut cx = center[0].round() as isize;
377    let mut cy = center[1].round() as isize;
378    let mut cz = center[2].round() as isize;
379
380    cx = cx.max(0).min(nx as isize - 1);
381    cy = cy.max(0).min(ny as isize - 1);
382    cz = cz.max(0).min(nz as isize - 1);
383
384    // If center is on surface, find nearby interior point
385    if grid[idx(cx as usize, cy as usize, cz as usize)] == 0 {
386        'search: for dx in -5..=5 {
387            for dy in -5..=5 {
388                for dz in -5..=5 {
389                    let nx_ = cx + dx;
390                    let ny_ = cy + dy;
391                    let nz_ = cz + dz;
392                    if nx_ >= 0 && nx_ < nx as isize && ny_ >= 0 && ny_ < ny as isize && nz_ >= 0 && nz_ < nz as isize {
393                        if grid[idx(nx_ as usize, ny_ as usize, nz_ as usize)] == 1 {
394                            cx = nx_;
395                            cy = ny_;
396                            cz = nz_;
397                            break 'search;
398                        }
399                    }
400                }
401            }
402        }
403    }
404
405    // BFS flood fill
406    let mut queue: VecDeque<(usize, usize, usize)> = VecDeque::new();
407    let cx = cx as usize;
408    let cy = cy as usize;
409    let cz = cz as usize;
410    grid[idx(cx, cy, cz)] = 0;
411    queue.push_back((cx, cy, cz));
412
413    let neighbors: [(isize, isize, isize); 6] = [
414        (-1, 0, 0), (1, 0, 0), (0, -1, 0), (0, 1, 0), (0, 0, -1), (0, 0, 1)
415    ];
416
417    while let Some((x, y, z)) = queue.pop_front() {
418        for &(dx, dy, dz) in &neighbors {
419            let nx_ = x as isize + dx;
420            let ny_ = y as isize + dy;
421            let nz_ = z as isize + dz;
422
423            if nx_ >= 0 && nx_ < nx as isize && ny_ >= 0 && ny_ < ny as isize && nz_ >= 0 && nz_ < nz as isize {
424                let ni = idx(nx_ as usize, ny_ as usize, nz_ as usize);
425                if grid[ni] == 1 {
426                    grid[ni] = 0;
427                    queue.push_back((nx_ as usize, ny_ as usize, nz_ as usize));
428                }
429            }
430        }
431    }
432
433    // Invert: 0 = brain (inside + surface), we want 1 = brain
434    for v in grid.iter_mut() {
435        *v = if *v == 0 { 1 } else { 0 };
436    }
437
438    // Fill holes using simple morphological closing
439    fill_holes(&mut grid, nx, ny, nz);
440
441    grid
442}
443
444/// Simple hole filling
445fn fill_holes(mask: &mut [u8], nx: usize, ny: usize, nz: usize) {
446    // Fortran order: x + y*nx + z*nx*ny
447    let idx = |i: usize, j: usize, k: usize| i + j * nx + k * nx * ny;
448
449    // Flood fill from corners to find exterior
450    let mut exterior: Vec<bool> = vec![false; nx * ny * nz];
451    let mut queue: VecDeque<(usize, usize, usize)> = VecDeque::new();
452
453    // Start from all boundary voxels that are 0
454    for i in 0..nx {
455        for j in 0..ny {
456            for k in 0..nz {
457                if i == 0 || i == nx - 1 || j == 0 || j == ny - 1 || k == 0 || k == nz - 1 {
458                    if mask[idx(i, j, k)] == 0 {
459                        exterior[idx(i, j, k)] = true;
460                        queue.push_back((i, j, k));
461                    }
462                }
463            }
464        }
465    }
466
467    let neighbors: [(isize, isize, isize); 6] = [
468        (-1, 0, 0), (1, 0, 0), (0, -1, 0), (0, 1, 0), (0, 0, -1), (0, 0, 1)
469    ];
470
471    while let Some((x, y, z)) = queue.pop_front() {
472        for &(dx, dy, dz) in &neighbors {
473            let nx_ = x as isize + dx;
474            let ny_ = y as isize + dy;
475            let nz_ = z as isize + dz;
476
477            if nx_ >= 0 && nx_ < nx as isize && ny_ >= 0 && ny_ < ny as isize && nz_ >= 0 && nz_ < nz as isize {
478                let ni = idx(nx_ as usize, ny_ as usize, nz_ as usize);
479                if mask[ni] == 0 && !exterior[ni] {
480                    exterior[ni] = true;
481                    queue.push_back((nx_ as usize, ny_ as usize, nz_ as usize));
482                }
483            }
484        }
485    }
486
487    // Any voxel that is 0 but not exterior is interior -> fill it
488    for i in 0..mask.len() {
489        if mask[i] == 0 && !exterior[i] {
490            mask[i] = 1;
491        }
492    }
493}
494
495/// Self-intersection threshold (matches FSL-BET2)
496const SELF_INTERSECTION_THRESHOLD: f64 = 4000.0;
497
498/// Maximum number of recovery passes (matches FSL-BET2)
499const MAX_PASSES: usize = 10;
500
501/// Run a single pass of surface evolution
502///
503/// Returns the final vertices after evolution
504fn evolution_pass(
505    data: &[f64],
506    nx: usize, ny: usize, nz: usize,
507    voxel_size: &[f64; 3],
508    bp: &BetParameters,
509    vertices: &mut Vec<[f64; 3]>,
510    faces: &[[usize; 3]],
511    neighbor_matrix: &[Vec<usize>],
512    neighbor_counts: &[usize],
513    bt: f64,
514    smoothness_factor: f64,
515    gradient_threshold: f64,
516    iterations: usize,
517    pass: usize,
518    progress_callback: &mut Option<&mut dyn FnMut(usize, usize)>,
519) {
520    let n_vertices = vertices.len();
521
522    // BET parameters (from FSL) - adjusted by smoothness_factor
523    let rmin = 3.33 * smoothness_factor;
524    let rmax = 10.0 * smoothness_factor;
525    let e = (1.0 / rmin + 1.0 / rmax) / 2.0;
526    let f = 6.0 / (1.0 / rmin - 1.0 / rmax);
527    let normal_max_update_fraction = 0.5;
528    let lambda_fit = 0.1;
529
530    // Initial mean edge length
531    // Vertices are in mm, so use the mm version
532    let mut l = compute_mean_edge_length_mm(vertices, faces);
533
534    // Smoothing increase factor for recovery passes (FSL: 10^(pass+1))
535    let base_increase = if pass > 0 { 10.0_f64.powi((pass + 1) as i32) } else { 1.0 };
536
537    // Report progress at start if first pass
538    let progress_interval = (iterations / 20).max(1);
539
540    // Debug counter for sampling failures
541    let mut sample_fail_count: usize = 0;
542
543    for iteration in 0..iterations {
544        // Report progress periodically
545        if let Some(ref mut cb) = progress_callback {
546            if iteration % progress_interval == 0 {
547                cb(iteration, iterations);
548            }
549        }
550
551        // Compute increase factor with tapering in later iterations (FSL: after 75%)
552        let incfactor = if pass > 0 && iteration > (0.75 * iterations as f64) as usize {
553            let t = iteration as f64 / iterations as f64;
554            4.0 * (1.0 - t) * (base_increase - 1.0) + 1.0
555        } else {
556            base_increase
557        };
558
559        // Compute vertex normals
560        let normals = compute_vertex_normals(vertices, faces);
561
562        // Compute updates for each vertex
563        let mut updates: Vec<[f64; 3]> = vec![[0.0, 0.0, 0.0]; n_vertices];
564
565        for i in 0..n_vertices {
566            let v = vertices[i];
567            let n = normals[i];
568
569            // Compute mean neighbor position
570            let mut mean_neighbor = [0.0, 0.0, 0.0];
571            let count = neighbor_counts[i];
572            for j in 0..count {
573                let ni = neighbor_matrix[i][j];
574                mean_neighbor[0] += vertices[ni][0];
575                mean_neighbor[1] += vertices[ni][1];
576                mean_neighbor[2] += vertices[ni][2];
577            }
578            if count > 0 {
579                mean_neighbor[0] /= count as f64;
580                mean_neighbor[1] /= count as f64;
581                mean_neighbor[2] /= count as f64;
582            }
583
584            // Vector from vertex to mean neighbor
585            let dv = [mean_neighbor[0] - v[0], mean_neighbor[1] - v[1], mean_neighbor[2] - v[2]];
586
587            // Dot product with normal
588            let dv_dot_n = dv[0] * n[0] + dv[1] * n[1] + dv[2] * n[2];
589
590            // Normal component
591            let sn = [dv_dot_n * n[0], dv_dot_n * n[1], dv_dot_n * n[2]];
592
593            // Tangential component
594            let st = [dv[0] - sn[0], dv[1] - sn[1], dv[2] - sn[2]];
595
596            // Force 1: Tangential (vertex spacing)
597            let u1 = [st[0] * 0.5, st[1] * 0.5, st[2] * 0.5];
598
599            // Force 2: Normal (smoothness)
600            let sn_mag = dv_dot_n.abs();
601            let rinv = (2.0 * sn_mag) / (l * l);
602            let mut f2 = (1.0 + (f * (rinv - e)).tanh()) * 0.5;
603
604            // In recovery passes, increase smoothing for outward-pointing updates (FSL behavior)
605            if pass > 0 && dv_dot_n > 0.0 {
606                f2 *= incfactor;
607                f2 = f2.min(1.0);
608            }
609
610            let u2 = [f2 * sn[0], f2 * sn[1], f2 * sn[2]];
611
612            // Force 3: Intensity-based (using FSL-style sampling with tm)
613            let (i_min, i_max, sample_ok) = sample_intensities_fsl(
614                data, nx, ny, nz, &v, &n, voxel_size, bp.t2, bp.t, bp.tm
615            );
616
617            // FSL: if sampling fails (out of bounds), f3 = 0 (no intensity force)
618            let u3 = if sample_ok {
619                // Apply z-gradient to local threshold (FSL's -g option)
620                let local_bt = if gradient_threshold.abs() > 1e-10 {
621                    // Vertex is already in mm coordinates
622                    let z_offset = (v[2] - bp.cog_mm[2]) / bp.radius;
623                    (bt + gradient_threshold * z_offset).clamp(0.0, 1.0)
624                } else {
625                    bt
626                };
627
628                // Compute local threshold and force (matches FSL exactly)
629                let t_l = (i_max - bp.t2) * local_bt + bp.t2;
630                let f3 = if i_max - bp.t2 > 0.0 {
631                    2.0 * (i_min - t_l) / (i_max - bp.t2)
632                } else {
633                    2.0 * (i_min - t_l)
634                };
635                let f3 = f3 * normal_max_update_fraction * lambda_fit * l;
636
637                [f3 * n[0], f3 * n[1], f3 * n[2]]
638            } else {
639                // Sampling failed - use f3 = 0 like FSL does
640                sample_fail_count += 1;
641                [0.0, 0.0, 0.0]
642            };
643
644            // Combined update
645            updates[i] = [u1[0] + u2[0] + u3[0], u1[1] + u2[1] + u3[1], u1[2] + u2[2] + u3[2]];
646        }
647
648        // Apply updates
649        for i in 0..n_vertices {
650            vertices[i][0] += updates[i][0];
651            vertices[i][1] += updates[i][1];
652            vertices[i][2] += updates[i][2];
653        }
654
655        // Update edge length periodically
656        if iteration % 100 == 0 {
657            l = compute_mean_edge_length_mm(vertices, faces);
658        }
659    }
660
661    // Debug: report sampling failures
662    if sample_fail_count > 0 {
663        let total_samples = iterations * n_vertices;
664        let fail_pct = 100.0 * sample_fail_count as f64 / total_samples as f64;
665        eprintln!("[BET] Sampling fallback used: {} / {} ({:.1}%)",
666                  sample_fail_count, total_samples, fail_pct);
667    }
668}
669
670/// BET algorithm parameters
671#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
672#[derive(Clone, Debug)]
673pub struct BetParams {
674    /// Fractional intensity threshold (0.0-1.0, smaller = larger brain)
675    pub fractional_intensity: f64,
676    /// Surface smoothness factor
677    pub smoothness: f64,
678    /// Gradient threshold (-1 to 1)
679    pub gradient_threshold: f64,
680    /// Number of iterations
681    pub iterations: usize,
682    /// Icosphere subdivision level
683    pub subdivisions: usize,
684}
685
686impl Default for BetParams {
687    fn default() -> Self {
688        Self {
689            fractional_intensity: 0.5,
690            smoothness: 1.0,
691            gradient_threshold: 0.0,
692            iterations: 1000,
693            subdivisions: 4,
694        }
695    }
696}
697
698/// Run BET brain extraction
699///
700/// # Arguments
701/// * `data` - 3D magnitude image data (nx * ny * nz, Fortran order)
702/// * `grid` - Volume grid (dimensions and voxel sizes)
703/// * `params` - BET parameters (intensity threshold, smoothness, gradient, iterations, subdivisions)
704/// * `progress` - Progress callback: `(iteration, total_iterations)`
705///
706/// # Returns
707/// Binary mask (1 = brain, 0 = background)
708pub fn run_bet<F>(
709    data: &[f64],
710    grid: &Grid,
711    params: &BetParams,
712    mut progress: F,
713) -> Vec<u8>
714where
715    F: FnMut(usize, usize),
716{
717    let fractional_intensity = params.fractional_intensity;
718    let smoothness_factor = params.smoothness;
719    let gradient_threshold = params.gradient_threshold;
720    let iterations = params.iterations;
721    let subdivisions = params.subdivisions;
722    let (nx, ny, nz) = grid.dims;
723    let voxel_size = [grid.vsx(), grid.vsy(), grid.vsz()];
724
725    // Step 1: Estimate brain parameters
726    progress(0, iterations);
727    let bp = estimate_brain_parameters(data, nx, ny, nz, &voxel_size);
728
729    // Step 2: Create icosphere
730    let (unit_vertices, faces) = create_icosphere(subdivisions);
731    let n_vertices = unit_vertices.len();
732
733    // Scale and position sphere in mm coordinates (start at 50% of estimated radius)
734    // Like FSL, we work entirely in mm - voxel conversion only happens at sampling/masking
735    let initial_radius_mm = bp.radius * 0.5;
736
737    let initial_vertices: Vec<[f64; 3]> = unit_vertices
738        .iter()
739        .map(|v| [
740            v[0] * initial_radius_mm + bp.cog_mm[0],
741            v[1] * initial_radius_mm + bp.cog_mm[1],
742            v[2] * initial_radius_mm + bp.cog_mm[2],
743        ])
744        .collect();
745
746    // Build neighbor structure
747    let (neighbor_matrix, neighbor_counts) = build_neighbor_matrix(n_vertices, &faces, 6);
748
749    // FSL power transform: bt = pow(f, 0.275)
750    // This raises 0.5 -> 0.826, which makes the surface more aggressive in expanding
751    let bt = fractional_intensity.powf(0.275);
752
753    // Multi-pass evolution with self-intersection recovery (like FSL-BET2)
754    let mut vertices = initial_vertices.clone();
755    let mut pass = 0;
756
757    loop {
758        // Run evolution pass
759        let mut cb: Option<&mut dyn FnMut(usize, usize)> = Some(&mut progress);
760        evolution_pass(
761            data, nx, ny, nz, &voxel_size, &bp,
762            &mut vertices, &faces,
763            &neighbor_matrix, &neighbor_counts,
764            bt, smoothness_factor, gradient_threshold,
765            iterations, pass,
766            &mut cb,
767        );
768
769        // Check for self-intersection
770        let si_score = self_intersection_heuristic(&vertices, &initial_vertices, &faces, &voxel_size);
771        let has_self_intersection = si_score > SELF_INTERSECTION_THRESHOLD;
772
773        if has_self_intersection {
774            eprintln!("[BET] Self-intersection detected (score={:.0}, threshold={}), pass {}",
775                      si_score, SELF_INTERSECTION_THRESHOLD, pass + 1);
776        }
777
778        // Exit if no self-intersection or max passes reached
779        if !has_self_intersection || pass >= MAX_PASSES {
780            if pass > 0 {
781                eprintln!("[BET] Completed after {} recovery pass(es)", pass);
782            }
783            break;
784        }
785
786        // Reset to original mesh and try again with higher smoothing
787        vertices = initial_vertices.clone();
788        pass += 1;
789    }
790
791    // Final progress update
792    progress(iterations, iterations);
793
794    // Step 4: Convert surface to binary mask
795    surface_to_mask(&vertices, &faces, nx, ny, nz, &voxel_size)
796}
797
798#[cfg(test)]
799mod tests {
800    use super::*;
801
802    /// Helper: create a 3D sphere volume (Fortran order) centered in the grid.
803    /// Returns (data, nx, ny, nz) where the sphere has intensity `intensity`
804    /// and background is 0. Sphere radius in voxels is `radius`.
805    fn make_sphere_volume(n: usize, radius: f64, intensity: f64) -> (Vec<f64>, usize, usize, usize) {
806        let center = (n as f64) / 2.0;
807        let mut data = vec![0.0; n * n * n];
808        for k in 0..n {
809            for j in 0..n {
810                for i in 0..n {
811                    let di = (i as f64) - center;
812                    let dj = (j as f64) - center;
813                    let dk = (k as f64) - center;
814                    let dist = (di * di + dj * dj + dk * dk).sqrt();
815                    if dist <= radius {
816                        // Smoothly varying intensity: bright at center, dimmer at edge
817                        data[i + j * n + k * n * n] = intensity * (1.0 - 0.5 * dist / radius);
818                    }
819                }
820            }
821        }
822        (data, n, n, n)
823    }
824
825    #[test]
826    fn test_estimate_brain_parameters() {
827        let nx = 10;
828        let ny = 10;
829        let nz = 10;
830        let mut data = vec![0.0; nx * ny * nz];
831
832        // Create a sphere with varying intensity (like a real brain)
833        // Fortran order: index = i + j*nx + k*nx*ny
834        for k in 0..nz {
835            for j in 0..ny {
836                for i in 0..nx {
837                    let di = (i as f64) - 5.0;
838                    let dj = (j as f64) - 5.0;
839                    let dk = (k as f64) - 5.0;
840                    let dist = (di*di + dj*dj + dk*dk).sqrt();
841                    if dist <= 4.0 {
842                        // Intensity varies from 50 to 150 based on distance from center
843                        data[i + j * nx + k * nx * ny] = 150.0 - dist * 25.0;
844                    }
845                }
846            }
847        }
848
849        let bp = estimate_brain_parameters(&data, nx, ny, nz, &[1.0, 1.0, 1.0]);
850
851        assert!(bp.t2 >= 0.0);
852        assert!(bp.t98 >= bp.t2); // Allow equal for edge cases
853        assert!((bp.cog[0] - 5.0).abs() < 1.0);
854        assert!((bp.cog[1] - 5.0).abs() < 1.0);
855        assert!((bp.cog[2] - 5.0).abs() < 1.0);
856        assert!(bp.radius > 0.0);
857        assert!(bp.tm > bp.t2 && bp.tm < bp.t98); // tm should be between t2 and t98
858        // Check cog_mm is correctly computed
859        assert!((bp.cog_mm[0] - bp.cog[0]).abs() < 1e-10);
860    }
861
862    #[test]
863    fn test_estimate_brain_parameters_empty_data() {
864        // All zeros should trigger the nonzero.is_empty() branch
865        let nx = 4;
866        let ny = 4;
867        let nz = 4;
868        let data = vec![0.0; nx * ny * nz];
869        let bp = estimate_brain_parameters(&data, nx, ny, nz, &[1.0, 1.0, 1.0]);
870
871        assert!((bp.t2 - 0.0).abs() < 1e-10);
872        assert!((bp.t98 - 1.0).abs() < 1e-10);
873        assert!((bp.t - 0.1).abs() < 1e-10);
874        assert!((bp.tm - 0.5).abs() < 1e-10);
875        assert!((bp.cog[0] - 2.0).abs() < 1e-10);
876        assert!((bp.cog[1] - 2.0).abs() < 1e-10);
877        assert!((bp.cog[2] - 2.0).abs() < 1e-10);
878        assert!((bp.radius - 50.0).abs() < 1e-10);
879    }
880
881    #[test]
882    fn test_estimate_brain_parameters_anisotropic_voxels() {
883        let (data, nx, ny, nz) = make_sphere_volume(12, 4.0, 200.0);
884        let voxel_size = [2.0, 2.0, 2.0];
885        let bp = estimate_brain_parameters(&data, nx, ny, nz, &voxel_size);
886
887        // cog_mm should be cog * voxel_size
888        assert!((bp.cog_mm[0] - bp.cog[0] * voxel_size[0]).abs() < 1e-10);
889        assert!((bp.cog_mm[1] - bp.cog[1] * voxel_size[1]).abs() < 1e-10);
890        assert!((bp.cog_mm[2] - bp.cog[2] * voxel_size[2]).abs() < 1e-10);
891        assert!(bp.radius > 0.0);
892        assert!(bp.t98 > bp.t2);
893    }
894
895    #[test]
896    fn test_sample_intensity() {
897        let data = vec![
898            0.0, 1.0, 2.0, 3.0,
899            4.0, 5.0, 6.0, 7.0,
900        ];
901        let val = sample_intensity(&data, 2, 2, 2, 0.5, 0.5, 0.5);
902        // Trilinear interpolation of cube corners
903        assert!((val - 3.5).abs() < 0.01);
904    }
905
906    #[test]
907    fn test_sample_intensity_at_corners() {
908        // 2x2x2 cube: values 0..7
909        let data = vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0];
910        // Sampling at integer corners should return exact values
911        assert!((sample_intensity(&data, 2, 2, 2, 0.0, 0.0, 0.0) - 0.0).abs() < 1e-10);
912        assert!((sample_intensity(&data, 2, 2, 2, 1.0, 0.0, 0.0) - 1.0).abs() < 1e-10);
913        assert!((sample_intensity(&data, 2, 2, 2, 0.0, 1.0, 0.0) - 2.0).abs() < 1e-10);
914        assert!((sample_intensity(&data, 2, 2, 2, 1.0, 1.0, 0.0) - 3.0).abs() < 1e-10);
915        assert!((sample_intensity(&data, 2, 2, 2, 0.0, 0.0, 1.0) - 4.0).abs() < 1e-10);
916    }
917
918    #[test]
919    fn test_sample_intensity_clamping() {
920        // Out-of-bounds coordinates should be clamped
921        let data = vec![10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0];
922        // Negative coordinates should clamp to 0
923        let val = sample_intensity(&data, 2, 2, 2, -1.0, -1.0, -1.0);
924        assert!((val - 10.0).abs() < 1e-10, "Expected 10.0, got {}", val);
925        // Beyond max should clamp to max
926        let val = sample_intensity(&data, 2, 2, 2, 5.0, 5.0, 5.0);
927        assert!((val - 80.0).abs() < 1e-10, "Expected 80.0, got {}", val);
928    }
929
930    #[test]
931    fn test_sample_intensity_larger_volume() {
932        // 4x4x4 volume with known pattern
933        let n = 4;
934        let mut data = vec![0.0; n * n * n];
935        for k in 0..n {
936            for j in 0..n {
937                for i in 0..n {
938                    data[i + j * n + k * n * n] = (i + j + k) as f64;
939                }
940            }
941        }
942        // At center (1.5, 1.5, 1.5), each corner = i+j+k
943        // Average of corners: (1+1+1=3, 2+1+1=4, 1+2+1=4, 2+2+1=5, 1+1+2=4, 2+1+2=5, 1+2+2=5, 2+2+2=6) / 8 = 4.5
944        let val = sample_intensity(&data, n, n, n, 1.5, 1.5, 1.5);
945        assert!((val - 4.5).abs() < 1e-10, "Expected 4.5, got {}", val);
946    }
947
948    #[test]
949    fn test_percentile() {
950        let sorted = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];
951        let p0 = percentile(&sorted, 0.0);
952        assert!((p0 - 1.0).abs() < 1e-10);
953
954        let p50 = percentile(&sorted, 50.0);
955        // Index = round(0.5 * 9) = round(4.5) = 5, value = 6.0
956        assert!((p50 - 5.0).abs() < 1.5, "p50={}", p50);
957
958        let p100 = percentile(&sorted, 100.0);
959        assert!((p100 - 10.0).abs() < 1e-10);
960
961        // Empty array
962        let empty: Vec<f64> = vec![];
963        assert!((percentile(&empty, 50.0) - 0.0).abs() < 1e-10);
964    }
965
966    #[test]
967    fn test_power_transform() {
968        // Verify power transform matches FSL
969        let f = 0.5_f64;
970        let bt = f.powf(0.275);
971        // 0.5^0.275 ~ 0.826
972        assert!((bt - 0.826).abs() < 0.01);
973    }
974
975    #[test]
976    fn test_sample_intensities_fsl_basic() {
977        // Create a 16x16x16 uniform volume
978        let n = 16;
979        let data = vec![100.0; n * n * n];
980        let voxel_size = [1.0, 1.0, 1.0];
981        // Point at the center, normal pointing inward
982        let point_mm = [8.0, 8.0, 8.0];
983        let normal = [0.0, 0.0, 1.0]; // outward along z
984
985        let t2 = 10.0;
986        let t = 20.0;
987        let tm = 100.0;
988
989        let (i_min, i_max, ok) = sample_intensities_fsl(
990            &data, n, n, n, &point_mm, &normal, &voxel_size, t2, t, tm,
991        );
992
993        assert!(ok, "Sampling should succeed for center point");
994        assert!(i_min.is_finite());
995        assert!(i_max.is_finite());
996        // Uniform volume: i_min should be clamped to >= t2
997        assert!(i_min >= t2, "i_min={} should be >= t2={}", i_min, t2);
998        // i_max should be clamped to <= tm
999        assert!(i_max <= tm, "i_max={} should be <= tm={}", i_max, tm);
1000    }
1001
1002    #[test]
1003    fn test_sample_intensities_fsl_out_of_bounds() {
1004        // Point near edge, normal pointing outward -- should fail bounds check
1005        let n = 8;
1006        let data = vec![100.0; n * n * n];
1007        let voxel_size = [1.0, 1.0, 1.0];
1008        // Point at the edge
1009        let point_mm = [0.5, 0.5, 0.5];
1010        let normal = [-1.0, 0.0, 0.0]; // pointing outward (further out of bounds)
1011
1012        let (_, _, _ok) = sample_intensities_fsl(
1013            &data, n, n, n, &point_mm, &normal, &voxel_size, 10.0, 20.0, 100.0,
1014        );
1015
1016        // Construct a case that goes out of bounds on the first inward step:
1017        // point - normal = (0.5 - 1.0, 0.5, 0.5) = (-0.5, 0.5, 0.5) -> voxel (-0.5) is OOB
1018        let point_mm2 = [0.5, 0.5, 0.5];
1019        let normal2 = [1.0, 0.0, 0.0]; // point - normal = (-0.5, 0.5, 0.5) -> voxel (-0.5) out of bounds
1020        let (_, _, ok2) = sample_intensities_fsl(
1021            &data, n, n, n, &point_mm2, &normal2, &voxel_size, 10.0, 20.0, 100.0,
1022        );
1023        assert!(!ok2, "Sampling should fail when initial step goes out of bounds");
1024    }
1025
1026    #[test]
1027    fn test_sample_intensities_fsl_varying_intensity() {
1028        // Create a volume with intensity gradient along z
1029        let n = 16;
1030        let mut data = vec![0.0; n * n * n];
1031        for k in 0..n {
1032            for j in 0..n {
1033                for i in 0..n {
1034                    data[i + j * n + k * n * n] = 50.0 + (k as f64) * 10.0;
1035                }
1036            }
1037        }
1038        let voxel_size = [1.0, 1.0, 1.0];
1039        let point_mm = [8.0, 8.0, 8.0];
1040        // Normal pointing in -z direction (inward sampling goes into higher z)
1041        let normal = [0.0, 0.0, -1.0];
1042
1043        let t2 = 50.0;
1044        let t = 55.0;
1045        let tm = 120.0;
1046
1047        let (i_min, i_max, ok) = sample_intensities_fsl(
1048            &data, n, n, n, &point_mm, &normal, &voxel_size, t2, t, tm,
1049        );
1050
1051        assert!(ok, "Sampling should succeed");
1052        assert!(i_min.is_finite() && i_max.is_finite());
1053        assert!(i_min >= t2);
1054        assert!(i_max <= tm);
1055    }
1056
1057    #[test]
1058    fn test_fill_holes_no_holes() {
1059        // A solid 3x3x3 cube with no holes should remain unchanged
1060        let nx = 5;
1061        let ny = 5;
1062        let nz = 5;
1063        let mut mask = vec![0u8; nx * ny * nz];
1064        let idx = |i: usize, j: usize, k: usize| i + j * nx + k * nx * ny;
1065
1066        // Place a solid 3x3x3 block in the center
1067        for k in 1..4 {
1068            for j in 1..4 {
1069                for i in 1..4 {
1070                    mask[idx(i, j, k)] = 1;
1071                }
1072            }
1073        }
1074
1075        let original = mask.clone();
1076        fill_holes(&mut mask, nx, ny, nz);
1077
1078        assert_eq!(mask, original, "Solid mask should not change");
1079    }
1080
1081    #[test]
1082    fn test_fill_holes_with_interior_hole() {
1083        // A 7x7x7 volume with a shell of 1s and a hole (0) inside
1084        let n = 7;
1085        let idx = |i: usize, j: usize, k: usize| i + j * n + k * n * n;
1086        let mut mask = vec![0u8; n * n * n];
1087
1088        // Create a hollow cube: 1 on shell from (1,1,1) to (5,5,5), 0 inside
1089        for k in 1..6 {
1090            for j in 1..6 {
1091                for i in 1..6 {
1092                    mask[idx(i, j, k)] = 1;
1093                }
1094            }
1095        }
1096        // Carve out the interior
1097        for k in 2..5 {
1098            for j in 2..5 {
1099                for i in 2..5 {
1100                    mask[idx(i, j, k)] = 0;
1101                }
1102            }
1103        }
1104
1105        // The interior (2..5)^3 should be 0 before fill_holes
1106        assert_eq!(mask[idx(3, 3, 3)], 0, "Center should be 0 before fill");
1107
1108        fill_holes(&mut mask, n, n, n);
1109
1110        // After fill_holes, interior holes should be filled
1111        assert_eq!(mask[idx(3, 3, 3)], 1, "Center hole should be filled");
1112        assert_eq!(mask[idx(2, 2, 2)], 1, "Interior corner should be filled");
1113
1114        // Exterior should still be 0
1115        assert_eq!(mask[idx(0, 0, 0)], 0, "Exterior should remain 0");
1116        assert_eq!(mask[idx(6, 6, 6)], 0, "Exterior should remain 0");
1117    }
1118
1119    #[test]
1120    fn test_fill_holes_exterior_connected() {
1121        // All zeros (no brain): nothing to fill
1122        let n = 5;
1123        let mut mask = vec![0u8; n * n * n];
1124        fill_holes(&mut mask, n, n, n);
1125        assert!(mask.iter().all(|&v| v == 0), "All-zero mask should stay all-zero");
1126    }
1127
1128    #[test]
1129    fn test_surface_to_mask_small_sphere() {
1130        // Create a small icosphere mesh centered in a 16x16x16 grid
1131        let (unit_verts, faces) = create_icosphere(2);
1132        let n = 16;
1133        let center = (n as f64) / 2.0;
1134        let radius = 5.0; // 5 voxel radius
1135        let voxel_size = [1.0, 1.0, 1.0];
1136
1137        // Place vertices in mm (= voxel coords since voxel_size is 1)
1138        let vertices_mm: Vec<[f64; 3]> = unit_verts
1139            .iter()
1140            .map(|v| [
1141                v[0] * radius + center,
1142                v[1] * radius + center,
1143                v[2] * radius + center,
1144            ])
1145            .collect();
1146
1147        let mask = surface_to_mask(&vertices_mm, &faces, n, n, n, &voxel_size);
1148
1149        assert_eq!(mask.len(), n * n * n);
1150
1151        // Center should be inside the mask
1152        let idx = |i: usize, j: usize, k: usize| i + j * n + k * n * n;
1153        assert_eq!(mask[idx(8, 8, 8)], 1, "Center should be brain");
1154
1155        // Corners should be outside the mask
1156        assert_eq!(mask[idx(0, 0, 0)], 0, "Corner should be background");
1157        assert_eq!(mask[idx(15, 15, 15)], 0, "Corner should be background");
1158
1159        // Count brain voxels -- should be a reasonable fraction of total
1160        let brain_count: usize = mask.iter().map(|&v| v as usize).sum();
1161        assert!(brain_count > 0, "Should have some brain voxels");
1162        // Sphere volume ~ 4/3 * pi * 5^3 = 524 voxels; total = 4096
1163        // Allow generous range
1164        assert!(brain_count > 100 && brain_count < 2000,
1165                "Brain voxels ({}) should be roughly sphere-like", brain_count);
1166    }
1167
1168    #[test]
1169    fn test_surface_to_mask_anisotropic_voxels() {
1170        let (unit_verts, faces) = create_icosphere(1);
1171        let n = 16;
1172        let voxel_size = [2.0, 2.0, 2.0];
1173        let center_mm = [(n as f64) * voxel_size[0] / 2.0,
1174                         (n as f64) * voxel_size[1] / 2.0,
1175                         (n as f64) * voxel_size[2] / 2.0];
1176        let radius_mm = 8.0;
1177
1178        let vertices_mm: Vec<[f64; 3]> = unit_verts
1179            .iter()
1180            .map(|v| [
1181                v[0] * radius_mm + center_mm[0],
1182                v[1] * radius_mm + center_mm[1],
1183                v[2] * radius_mm + center_mm[2],
1184            ])
1185            .collect();
1186
1187        let mask = surface_to_mask(&vertices_mm, &faces, n, n, n, &voxel_size);
1188        assert_eq!(mask.len(), n * n * n);
1189
1190        let brain_count: usize = mask.iter().map(|&v| v as usize).sum();
1191        assert!(brain_count > 0, "Should have brain voxels with anisotropic voxels");
1192    }
1193
1194    #[test]
1195    fn test_run_bet_small_synthetic_volume() {
1196        // Create a 16x16x16 volume with a bright sphere
1197        let (data, nx, ny, nz) = make_sphere_volume(16, 6.0, 200.0);
1198        let grid = Grid::new(nx, ny, nz, 1.0, 1.0, 1.0);
1199
1200        let mask = run_bet(
1201            &data,
1202            &grid,
1203            &BetParams { fractional_intensity: 0.5, smoothness: 1.0, gradient_threshold: 0.0, iterations: 50, subdivisions: 1 },
1204            |_, _| {},
1205        );
1206
1207        assert_eq!(mask.len(), nx * ny * nz);
1208
1209        // Should produce a valid binary mask
1210        assert!(mask.iter().all(|&v| v == 0 || v == 1), "Mask should be binary");
1211
1212        // Should have some brain voxels
1213        let brain_count: usize = mask.iter().map(|&v| v as usize).sum();
1214        assert!(brain_count > 0, "BET should extract some brain voxels");
1215
1216        // Center should be brain
1217        let idx = |i: usize, j: usize, k: usize| i + j * nx + k * nx * ny;
1218        assert_eq!(mask[idx(8, 8, 8)], 1, "Center of sphere should be brain");
1219    }
1220
1221    #[test]
1222    fn test_run_bet_with_gradient_threshold() {
1223        // Exercise the gradient threshold code path
1224        let (data, nx, ny, nz) = make_sphere_volume(16, 6.0, 200.0);
1225        let grid = Grid::new(nx, ny, nz, 1.0, 1.0, 1.0);
1226
1227        let mask = run_bet(
1228            &data,
1229            &grid,
1230            &BetParams { fractional_intensity: 0.5, smoothness: 1.0, gradient_threshold: 0.3, iterations: 50, subdivisions: 1 },
1231            |_, _| {},
1232        );
1233
1234        assert_eq!(mask.len(), nx * ny * nz);
1235        assert!(mask.iter().all(|&v| v == 0 || v == 1));
1236        let brain_count: usize = mask.iter().map(|&v| v as usize).sum();
1237        assert!(brain_count > 0, "BET with gradient should extract brain voxels");
1238    }
1239
1240    #[test]
1241    fn test_run_bet_with_progress() {
1242        // Test the progress callback variant
1243        let (data, nx, ny, nz) = make_sphere_volume(16, 6.0, 200.0);
1244        let grid = Grid::new(nx, ny, nz, 1.0, 1.0, 1.0);
1245
1246        let mut progress_calls = Vec::new();
1247        let mask = run_bet(
1248            &data,
1249            &grid,
1250            &BetParams { fractional_intensity: 0.5, smoothness: 1.0, gradient_threshold: 0.0, iterations: 50, subdivisions: 1 },
1251            |current, total| {
1252                progress_calls.push((current, total));
1253            },
1254        );
1255
1256        assert_eq!(mask.len(), nx * ny * nz);
1257        assert!(mask.iter().all(|&v| v == 0 || v == 1));
1258
1259        // Progress callback should have been called at least twice (start + end)
1260        assert!(progress_calls.len() >= 2,
1261                "Progress callback should be called at least twice, got {} calls",
1262                progress_calls.len());
1263
1264        // First call should be (0, iterations)
1265        assert_eq!(progress_calls[0].0, 0);
1266
1267        // Last call should be (iterations, iterations)
1268        let last = progress_calls.last().unwrap();
1269        assert_eq!(last.0, last.1, "Final progress should be complete");
1270    }
1271
1272    #[test]
1273    fn test_run_bet_different_fractional_intensities() {
1274        let (data, nx, ny, nz) = make_sphere_volume(16, 6.0, 200.0);
1275        let grid = Grid::new(nx, ny, nz, 1.0, 1.0, 1.0);
1276
1277        // Higher fractional_intensity = smaller brain
1278        let mask_small = run_bet(
1279            &data,
1280            &grid,
1281            &BetParams { fractional_intensity: 0.7, smoothness: 1.0, gradient_threshold: 0.0, iterations: 50, subdivisions: 1 },
1282            |_, _| {},
1283        );
1284        let mask_large = run_bet(
1285            &data,
1286            &grid,
1287            &BetParams { fractional_intensity: 0.3, smoothness: 1.0, gradient_threshold: 0.0, iterations: 50, subdivisions: 1 },
1288            |_, _| {},
1289        );
1290
1291        let count_small: usize = mask_small.iter().map(|&v| v as usize).sum();
1292        let count_large: usize = mask_large.iter().map(|&v| v as usize).sum();
1293
1294        // Both should produce valid masks
1295        assert!(count_small > 0);
1296        assert!(count_large > 0);
1297    }
1298
1299    #[test]
1300    fn test_run_bet_anisotropic_voxels() {
1301        let (data, nx, ny, nz) = make_sphere_volume(16, 6.0, 200.0);
1302        let grid = Grid::new(nx, ny, nz, 2.0, 2.0, 2.0);
1303
1304        let mask = run_bet(
1305            &data,
1306            &grid,
1307            &BetParams { fractional_intensity: 0.5, smoothness: 1.0, gradient_threshold: 0.0, iterations: 50, subdivisions: 1 },
1308            |_, _| {},
1309        );
1310
1311        assert_eq!(mask.len(), nx * ny * nz);
1312        assert!(mask.iter().all(|&v| v == 0 || v == 1));
1313    }
1314
1315    #[test]
1316    fn test_evolution_pass_basic() {
1317        // Directly test evolution_pass with a small icosphere
1318        let n = 16;
1319        let (data, nx, ny, nz) = make_sphere_volume(n, 6.0, 200.0);
1320        let voxel_size = [1.0, 1.0, 1.0];
1321        let bp = estimate_brain_parameters(&data, nx, ny, nz, &voxel_size);
1322
1323        let (unit_verts, faces) = create_icosphere(1);
1324        let n_vertices = unit_verts.len();
1325        let initial_radius_mm = bp.radius * 0.5;
1326
1327        let mut vertices: Vec<[f64; 3]> = unit_verts
1328            .iter()
1329            .map(|v| [
1330                v[0] * initial_radius_mm + bp.cog_mm[0],
1331                v[1] * initial_radius_mm + bp.cog_mm[1],
1332                v[2] * initial_radius_mm + bp.cog_mm[2],
1333            ])
1334            .collect();
1335
1336        let (neighbor_matrix, neighbor_counts) = build_neighbor_matrix(n_vertices, &faces, 6);
1337        let bt = 0.5_f64.powf(0.275);
1338
1339        let vertices_before = vertices.clone();
1340
1341        evolution_pass(
1342            &data, nx, ny, nz, &voxel_size, &bp,
1343            &mut vertices, &faces,
1344            &neighbor_matrix, &neighbor_counts,
1345            bt, 1.0, 0.0,
1346            10, // few iterations
1347            0,  // first pass
1348            &mut None,
1349        );
1350
1351        // Vertices should have moved
1352        let mut any_moved = false;
1353        for (before, after) in vertices_before.iter().zip(vertices.iter()) {
1354            let dist = ((after[0] - before[0]).powi(2) +
1355                       (after[1] - before[1]).powi(2) +
1356                       (after[2] - before[2]).powi(2)).sqrt();
1357            if dist > 1e-10 {
1358                any_moved = true;
1359            }
1360            // All coordinates should be finite
1361            assert!(after[0].is_finite() && after[1].is_finite() && after[2].is_finite());
1362        }
1363        assert!(any_moved, "Evolution should move at least some vertices");
1364    }
1365
1366    #[test]
1367    fn test_evolution_pass_recovery_pass() {
1368        // Test with pass > 0 to exercise recovery smoothing code paths
1369        let n = 16;
1370        let (data, nx, ny, nz) = make_sphere_volume(n, 6.0, 200.0);
1371        let voxel_size = [1.0, 1.0, 1.0];
1372        let bp = estimate_brain_parameters(&data, nx, ny, nz, &voxel_size);
1373
1374        let (unit_verts, faces) = create_icosphere(1);
1375        let n_vertices = unit_verts.len();
1376        let initial_radius_mm = bp.radius * 0.5;
1377
1378        let mut vertices: Vec<[f64; 3]> = unit_verts
1379            .iter()
1380            .map(|v| [
1381                v[0] * initial_radius_mm + bp.cog_mm[0],
1382                v[1] * initial_radius_mm + bp.cog_mm[1],
1383                v[2] * initial_radius_mm + bp.cog_mm[2],
1384            ])
1385            .collect();
1386
1387        let (neighbor_matrix, neighbor_counts) = build_neighbor_matrix(n_vertices, &faces, 6);
1388        let bt = 0.5_f64.powf(0.275);
1389
1390        // Run with pass=1 and enough iterations to hit the 75% tapering code
1391        evolution_pass(
1392            &data, nx, ny, nz, &voxel_size, &bp,
1393            &mut vertices, &faces,
1394            &neighbor_matrix, &neighbor_counts,
1395            bt, 1.0, 0.0,
1396            20,  // enough iterations to hit 75% mark
1397            1,   // recovery pass
1398            &mut None,
1399        );
1400
1401        // All coordinates should remain finite
1402        for v in &vertices {
1403            assert!(v[0].is_finite() && v[1].is_finite() && v[2].is_finite());
1404        }
1405    }
1406
1407    #[test]
1408    fn test_evolution_pass_with_gradient() {
1409        // Exercise the z-gradient threshold branch
1410        let n = 16;
1411        let (data, nx, ny, nz) = make_sphere_volume(n, 6.0, 200.0);
1412        let voxel_size = [1.0, 1.0, 1.0];
1413        let bp = estimate_brain_parameters(&data, nx, ny, nz, &voxel_size);
1414
1415        let (unit_verts, faces) = create_icosphere(1);
1416        let n_vertices = unit_verts.len();
1417        let initial_radius_mm = bp.radius * 0.5;
1418
1419        let mut vertices: Vec<[f64; 3]> = unit_verts
1420            .iter()
1421            .map(|v| [
1422                v[0] * initial_radius_mm + bp.cog_mm[0],
1423                v[1] * initial_radius_mm + bp.cog_mm[1],
1424                v[2] * initial_radius_mm + bp.cog_mm[2],
1425            ])
1426            .collect();
1427
1428        let (neighbor_matrix, neighbor_counts) = build_neighbor_matrix(n_vertices, &faces, 6);
1429        let bt = 0.5_f64.powf(0.275);
1430
1431        evolution_pass(
1432            &data, nx, ny, nz, &voxel_size, &bp,
1433            &mut vertices, &faces,
1434            &neighbor_matrix, &neighbor_counts,
1435            bt, 1.0, 0.5, // gradient_threshold = 0.5
1436            10, 0,
1437            &mut None,
1438        );
1439
1440        for v in &vertices {
1441            assert!(v[0].is_finite() && v[1].is_finite() && v[2].is_finite());
1442        }
1443    }
1444
1445    #[test]
1446    fn test_evolution_pass_with_progress_callback() {
1447        let n = 16;
1448        let (data, nx, ny, nz) = make_sphere_volume(n, 6.0, 200.0);
1449        let voxel_size = [1.0, 1.0, 1.0];
1450        let bp = estimate_brain_parameters(&data, nx, ny, nz, &voxel_size);
1451
1452        let (unit_verts, faces) = create_icosphere(1);
1453        let n_vertices = unit_verts.len();
1454        let initial_radius_mm = bp.radius * 0.5;
1455
1456        let mut vertices: Vec<[f64; 3]> = unit_verts
1457            .iter()
1458            .map(|v| [
1459                v[0] * initial_radius_mm + bp.cog_mm[0],
1460                v[1] * initial_radius_mm + bp.cog_mm[1],
1461                v[2] * initial_radius_mm + bp.cog_mm[2],
1462            ])
1463            .collect();
1464
1465        let (neighbor_matrix, neighbor_counts) = build_neighbor_matrix(n_vertices, &faces, 6);
1466        let bt = 0.5_f64.powf(0.275);
1467
1468        let mut calls = 0usize;
1469        let mut callback = |_iter: usize, _total: usize| {
1470            calls += 1;
1471        };
1472        let mut cb: Option<&mut dyn FnMut(usize, usize)> = Some(&mut callback);
1473
1474        evolution_pass(
1475            &data, nx, ny, nz, &voxel_size, &bp,
1476            &mut vertices, &faces,
1477            &neighbor_matrix, &neighbor_counts,
1478            bt, 1.0, 0.0,
1479            20, 0,
1480            &mut cb,
1481        );
1482
1483        assert!(calls > 0, "Progress callback should have been called");
1484    }
1485}