use crate::{settings::RansacSettings, types::DataMatrix};
pub trait Estimator {
type Model: Clone;
fn sample_size(&self) -> usize;
fn non_minimal_sample_size(&self) -> usize {
self.sample_size()
}
fn is_valid_sample(&self, data: &DataMatrix, sample: &[usize]) -> bool;
fn estimate_model(&self, data: &DataMatrix, sample: &[usize]) -> Vec<Self::Model>;
fn estimate_model_nonminimal(
&self,
data: &DataMatrix,
sample: &[usize],
_weights: Option<&[f64]>,
) -> Vec<Self::Model> {
self.estimate_model(data, sample)
}
fn is_valid_model(
&self,
model: &Self::Model,
data: &DataMatrix,
sample: &[usize],
threshold: f64,
) -> bool;
}
pub trait Sampler {
fn sample(&mut self, data: &DataMatrix, sample_size: usize, out_indices: &mut [usize]) -> bool;
fn update(&mut self, sample: &[usize], sample_size: usize, iteration: usize, score_hint: f64);
}
pub trait Scoring<M> {
type Score: Clone + PartialOrd;
fn threshold(&self) -> f64;
fn score(&self, data: &DataMatrix, model: &M, inliers_out: &mut Vec<usize>) -> Self::Score;
}
pub trait LocalOptimizer<M, S: Clone> {
fn run(
&mut self,
data: &DataMatrix,
inliers: &[usize],
model: &M,
best_score: &S,
) -> (M, S, Vec<usize>);
}
pub struct NoopLocalOptimizer;
impl<M: Clone, S: Clone> LocalOptimizer<M, S> for NoopLocalOptimizer {
fn run(
&mut self,
_data: &DataMatrix,
inliers: &[usize],
model: &M,
best_score: &S,
) -> (M, S, Vec<usize>) {
(model.clone(), best_score.clone(), inliers.to_vec())
}
}
pub struct LeastSquaresOptimizer<E>
where
E: Estimator,
{
estimator: E,
use_inliers: bool,
}
pub struct IteratedLeastSquaresOptimizer<E>
where
E: Estimator,
{
estimator: E,
max_iterations: usize,
use_inliers: bool,
}
impl<E> LeastSquaresOptimizer<E>
where
E: Estimator,
{
pub fn new(estimator: E) -> Self {
Self {
estimator,
use_inliers: true,
}
}
pub fn set_use_inliers(&mut self, use_inliers: bool) {
self.use_inliers = use_inliers;
}
}
impl<E, S> LocalOptimizer<E::Model, S> for LeastSquaresOptimizer<E>
where
E: Estimator,
E::Model: Clone,
S: Clone,
{
fn run(
&mut self,
data: &DataMatrix,
inliers: &[usize],
model: &E::Model,
best_score: &S,
) -> (E::Model, S, Vec<usize>) {
if !self.use_inliers || inliers.len() < self.estimator.non_minimal_sample_size() {
return (model.clone(), best_score.clone(), inliers.to_vec());
}
let refined_models = self
.estimator
.estimate_model_nonminimal(data, inliers, None);
if refined_models.is_empty() {
(model.clone(), best_score.clone(), inliers.to_vec())
} else {
(
refined_models[0].clone(),
best_score.clone(),
inliers.to_vec(),
)
}
}
}
impl<E> IteratedLeastSquaresOptimizer<E>
where
E: Estimator,
{
pub fn new(estimator: E) -> Self {
Self {
estimator,
max_iterations: 5,
use_inliers: true,
}
}
pub fn set_max_iterations(&mut self, max_iterations: usize) {
self.max_iterations = max_iterations;
}
pub fn set_use_inliers(&mut self, use_inliers: bool) {
self.use_inliers = use_inliers;
}
}
impl<E, S> LocalOptimizer<E::Model, S> for IteratedLeastSquaresOptimizer<E>
where
E: Estimator,
E::Model: Clone,
S: Clone + PartialOrd,
{
fn run(
&mut self,
data: &DataMatrix,
inliers: &[usize],
model: &E::Model,
best_score: &S,
) -> (E::Model, S, Vec<usize>) {
if !self.use_inliers || inliers.len() < self.estimator.non_minimal_sample_size() {
return (model.clone(), best_score.clone(), inliers.to_vec());
}
let mut current_model = model.clone();
let current_inliers = inliers.to_vec();
for _iteration in 0..self.max_iterations {
let refined_models =
self.estimator
.estimate_model_nonminimal(data, ¤t_inliers, None);
if refined_models.is_empty() {
break;
}
current_model = refined_models[0].clone();
}
(current_model, best_score.clone(), current_inliers)
}
}
pub struct NestedRansacOptimizer<E>
where
E: Estimator,
{
estimator: E,
max_iterations: usize,
sample_size_multiplier: usize,
}
impl<E> NestedRansacOptimizer<E>
where
E: Estimator,
{
pub fn new(estimator: E) -> Self {
Self {
estimator,
max_iterations: 50,
sample_size_multiplier: 7,
}
}
pub fn set_max_iterations(&mut self, max_iterations: usize) {
self.max_iterations = max_iterations;
}
pub fn set_sample_size_multiplier(&mut self, multiplier: usize) {
self.sample_size_multiplier = multiplier;
}
}
impl<E, S> LocalOptimizer<E::Model, S> for NestedRansacOptimizer<E>
where
E: Estimator,
E::Model: Clone,
S: Clone + PartialOrd,
{
fn run(
&mut self,
data: &DataMatrix,
inliers: &[usize],
model: &E::Model,
best_score: &S,
) -> (E::Model, S, Vec<usize>) {
use crate::samplers::UniformRandomSampler;
if inliers.len() < self.estimator.sample_size() {
return (model.clone(), best_score.clone(), inliers.to_vec());
}
let mut current_model = model.clone();
let current_inliers = inliers.to_vec();
let current_score = best_score.clone();
let non_minimal_sample_size = self.sample_size_multiplier * self.estimator.sample_size();
let mut sampler = UniformRandomSampler::new();
for _iteration in 0..self.max_iterations {
let mut current_sample_size = current_inliers.len().saturating_sub(1);
if current_sample_size >= non_minimal_sample_size {
current_sample_size = non_minimal_sample_size;
}
if current_sample_size < self.estimator.sample_size() {
break;
}
let mut sample = vec![0usize; current_sample_size];
if current_sample_size == current_inliers.len() {
sample.copy_from_slice(¤t_inliers[..current_sample_size]);
} else {
let mut temp_indices = vec![0usize; current_sample_size];
let dummy_data = DataMatrix::zeros(current_inliers.len(), data.ncols());
if !sampler.sample(&dummy_data, current_sample_size, &mut temp_indices) {
continue;
}
for (i, &temp_idx) in temp_indices.iter().enumerate() {
sample[i] = current_inliers[temp_idx];
}
}
let refined_models = self
.estimator
.estimate_model_nonminimal(data, &sample, None);
if refined_models.is_empty() {
continue;
}
current_model = refined_models[0].clone();
}
(current_model, current_score, current_inliers)
}
}
pub struct IRLSOptimizer<E, F>
where
E: Estimator,
F: Fn(&DataMatrix, &E::Model, usize) -> f64,
{
estimator: E,
residual_fn: F,
threshold: f64,
max_iterations: usize,
use_inliers: bool,
}
impl<E, F> IRLSOptimizer<E, F>
where
E: Estimator,
F: Fn(&DataMatrix, &E::Model, usize) -> f64,
{
pub fn new(estimator: E, residual_fn: F, threshold: f64) -> Self {
Self {
estimator,
residual_fn,
threshold,
max_iterations: 100,
use_inliers: true,
}
}
pub fn set_max_iterations(&mut self, max_iterations: usize) {
self.max_iterations = max_iterations;
}
pub fn set_use_inliers(&mut self, use_inliers: bool) {
self.use_inliers = use_inliers;
}
fn compute_weights(&self, data: &DataMatrix, model: &E::Model, indices: &[usize]) -> Vec<f64> {
let thresh_sq = self.threshold * self.threshold;
let mut weights = Vec::with_capacity(indices.len());
for &idx in indices {
let r = (self.residual_fn)(data, model, idx);
let r_sq = r * r;
if r_sq < thresh_sq {
weights.push(1.0 - r_sq / thresh_sq);
} else {
weights.push(0.0);
}
}
weights
}
}
impl<E, F, S> LocalOptimizer<E::Model, S> for IRLSOptimizer<E, F>
where
E: Estimator,
E::Model: Clone,
F: Fn(&DataMatrix, &E::Model, usize) -> f64,
S: Clone,
{
fn run(
&mut self,
data: &DataMatrix,
inliers: &[usize],
model: &E::Model,
best_score: &S,
) -> (E::Model, S, Vec<usize>) {
if !self.use_inliers || inliers.len() < self.estimator.sample_size() {
return (model.clone(), best_score.clone(), inliers.to_vec());
}
let mut current_model = model.clone();
for _iteration in 0..self.max_iterations {
let weights = self.compute_weights(data, ¤t_model, inliers);
let refined_models =
self.estimator
.estimate_model_nonminimal(data, inliers, Some(&weights));
if refined_models.is_empty() {
break;
}
current_model = refined_models[0].clone();
}
(current_model, best_score.clone(), inliers.to_vec())
}
}
pub struct CrossValidationOptimizer<E, F>
where
E: Estimator,
F: Fn(&DataMatrix, &E::Model, usize) -> f64,
{
estimator: E,
residual_fn: F,
threshold: f64,
repetitions: usize,
sample_size_multiplier: f64,
use_inliers: bool,
}
impl<E, F> CrossValidationOptimizer<E, F>
where
E: Estimator,
F: Fn(&DataMatrix, &E::Model, usize) -> f64,
{
pub fn new(estimator: E, residual_fn: F, threshold: f64) -> Self {
Self {
estimator,
residual_fn,
threshold,
repetitions: 100,
sample_size_multiplier: 0.5,
use_inliers: false,
}
}
pub fn with_repetitions(mut self, repetitions: usize) -> Self {
self.repetitions = repetitions;
self
}
pub fn with_sample_size_multiplier(mut self, multiplier: f64) -> Self {
self.sample_size_multiplier = multiplier;
self
}
pub fn set_use_inliers(&mut self, use_inliers: bool) {
self.use_inliers = use_inliers;
}
}
impl<E, F, S> LocalOptimizer<E::Model, S> for CrossValidationOptimizer<E, F>
where
E: Estimator,
E::Model: Clone,
F: Fn(&DataMatrix, &E::Model, usize) -> f64,
S: Clone,
{
fn run(
&mut self,
data: &DataMatrix,
inliers: &[usize],
model: &E::Model,
best_score: &S,
) -> (E::Model, S, Vec<usize>) {
let inlier_count = inliers.len();
let minimal_sample_size = self.estimator.sample_size();
if inlier_count < minimal_sample_size {
return (model.clone(), best_score.clone(), inliers.to_vec());
}
let sample_size = (self.sample_size_multiplier * inlier_count as f64)
.max(minimal_sample_size as f64) as usize;
let mut accumulated_scores = vec![0.0; inlier_count];
let mut rng = rand::rng();
for _rep in 0..self.repetitions {
let mut bootstrap_sample = Vec::with_capacity(sample_size);
for _ in 0..sample_size {
let idx = rand::Rng::random_range(&mut rng, 0..inlier_count);
bootstrap_sample.push(inliers[idx]);
}
let bootstrap_models = self.estimator.estimate_model(data, &bootstrap_sample);
if bootstrap_models.is_empty() {
continue;
}
let bootstrap_model = &bootstrap_models[0];
for (i, &inlier_idx) in inliers.iter().enumerate() {
let error = (self.residual_fn)(data, bootstrap_model, inlier_idx);
let score = (1.0 - error / self.threshold).max(0.0);
accumulated_scores[i] += score;
}
}
let weights: Vec<f64> = accumulated_scores
.iter()
.map(|&score| score / self.repetitions as f64)
.collect();
let refined_models =
self.estimator
.estimate_model_nonminimal(data, inliers, Some(&weights));
if refined_models.is_empty() {
return (model.clone(), best_score.clone(), inliers.to_vec());
}
(
refined_models[0].clone(),
best_score.clone(),
inliers.to_vec(),
)
}
}
pub trait TerminationCriterion<S> {
fn check(
&mut self,
data: &DataMatrix,
best_score: &S,
sample_size: usize,
max_iterations: &mut usize,
) -> bool;
}
pub struct RansacTerminationCriterion {
pub confidence: f64,
}
impl crate::core::TerminationCriterion<crate::scoring::Score> for RansacTerminationCriterion {
fn check(
&mut self,
data: &DataMatrix,
best_score: &crate::scoring::Score,
sample_size: usize,
max_iterations: &mut usize,
) -> bool {
let n = data.nrows() as f64;
if n <= 0.0 {
return false;
}
let inlier_ratio = (best_score.inlier_count as f64 / n).clamp(0.0, 1.0);
if inlier_ratio <= 0.0 || inlier_ratio >= 1.0 {
return false;
}
let p_good_sample = inlier_ratio.powi(sample_size as i32);
if p_good_sample <= 0.0 || p_good_sample >= 1.0 {
return false;
}
let log_one_minus_conf = (1.0 - self.confidence).ln();
let log_one_minus_p = (1.0 - p_good_sample).ln();
if !log_one_minus_conf.is_finite() || !log_one_minus_p.is_finite() {
return false;
}
let required = (log_one_minus_conf / log_one_minus_p).ceil().max(1.0) as usize;
if required < *max_iterations {
*max_iterations = required;
}
false
}
}
pub trait InlierSelector<M> {
fn select(&mut self, data: &DataMatrix, model: &M) -> Vec<usize>;
}
pub struct NoopInlierSelector;
impl<M> InlierSelector<M> for NoopInlierSelector {
fn select(&mut self, _data: &DataMatrix, _model: &M) -> Vec<usize> {
Vec::new()
}
}
pub struct SpacePartitioningInlierSelector<F>
where
F: Fn(&DataMatrix, usize) -> f64,
{
coord_fn: F,
cell_size: f64,
min_cells: usize,
}
impl<F> SpacePartitioningInlierSelector<F>
where
F: Fn(&DataMatrix, usize) -> f64,
{
pub fn new(coord_fn: F, cell_size: f64) -> Self {
Self {
coord_fn,
cell_size,
min_cells: 1,
}
}
pub fn set_min_cells(&mut self, min_cells: usize) {
self.min_cells = min_cells;
}
}
impl<M, F> InlierSelector<M> for SpacePartitioningInlierSelector<F>
where
F: Fn(&DataMatrix, usize) -> f64,
{
fn select(&mut self, data: &DataMatrix, _model: &M) -> Vec<usize> {
let n = data.nrows();
if n == 0 {
return Vec::new();
}
let mut cells: std::collections::HashMap<i32, Vec<usize>> =
std::collections::HashMap::new();
for i in 0..n {
let coord = (self.coord_fn)(data, i);
let cell_idx = (coord / self.cell_size).floor() as i32;
cells.entry(cell_idx).or_default().push(i);
}
let mut candidates = Vec::new();
for (_, indices) in cells.iter() {
if indices.len() >= self.min_cells {
candidates.extend(indices.iter().cloned());
}
}
if candidates.is_empty() {
(0..n).collect()
} else {
candidates
}
}
}
#[derive(Debug)]
pub struct SuperRansac<E, Sa, Sc, LO, T, IS>
where
E: Estimator,
Sa: Sampler,
Sc: Scoring<E::Model>,
Sc::Score: Clone + PartialOrd,
LO: LocalOptimizer<E::Model, Sc::Score>,
T: TerminationCriterion<Sc::Score>,
IS: InlierSelector<E::Model>,
{
pub settings: RansacSettings,
pub estimator: E,
pub sampler: Sa,
pub scoring: Sc,
pub local_optimizer: Option<LO>,
pub final_optimizer: Option<LO>,
pub termination: T,
pub inlier_selector: Option<IS>,
pub best_model: Option<E::Model>,
pub best_inliers: Vec<usize>,
pub best_score: Option<Sc::Score>,
pub iteration: usize,
}
impl<E, Sa, Sc, LO, T, IS> SuperRansac<E, Sa, Sc, LO, T, IS>
where
E: Estimator,
Sa: Sampler,
Sc: Scoring<E::Model>,
Sc::Score: Clone + PartialOrd,
LO: LocalOptimizer<E::Model, Sc::Score>,
T: TerminationCriterion<Sc::Score>,
IS: InlierSelector<E::Model>,
{
#[allow(clippy::too_many_arguments)]
pub fn new(
settings: RansacSettings,
estimator: E,
sampler: Sa,
scoring: Sc,
local_optimizer: Option<LO>,
final_optimizer: Option<LO>,
termination: T,
inlier_selector: Option<IS>,
) -> Self {
Self {
settings,
estimator,
sampler,
scoring,
local_optimizer,
final_optimizer,
termination,
inlier_selector,
best_model: None,
best_inliers: Vec::new(),
best_score: None,
iteration: 0,
}
}
pub fn run(&mut self, data: &DataMatrix) {
let sample_size = self.estimator.sample_size();
let mut sample = vec![0usize; sample_size];
let mut tmp_inliers = Vec::new();
let mut max_iterations = self.settings.max_iterations;
let min_iterations = self.settings.min_iterations;
self.best_inliers.clear();
self.best_model = None;
self.best_score = None;
self.iteration = 0;
let threshold = self.scoring.threshold();
while self.iteration < max_iterations || self.iteration < min_iterations {
let mut have_model = false;
let mut models: Vec<E::Model> = Vec::new();
for _ in 0..100 {
if !self.sampler.sample(data, sample_size, &mut sample[..]) {
self.sampler
.update(&sample, sample_size, self.iteration, 0.0);
continue;
}
if !self.estimator.is_valid_sample(data, &sample) {
self.sampler
.update(&sample, sample_size, self.iteration, 0.0);
continue;
}
models = self.estimator.estimate_model(data, &sample);
if models.is_empty() {
self.sampler
.update(&sample, sample_size, self.iteration, 0.0);
continue;
}
have_model = true;
break;
}
if !have_model {
self.iteration += 1;
continue;
}
let mut iteration_improved_best = false;
for model in models.iter() {
if !self
.estimator
.is_valid_model(model, data, &sample, threshold)
{
continue;
}
tmp_inliers.clear();
let _selected_inliers = if let Some(selector) = &mut self.inlier_selector {
selector.select(data, model)
} else {
Vec::new()
};
let score = self.scoring.score(data, model, &mut tmp_inliers);
let better = match &self.best_score {
None => true,
Some(best) => score > *best,
};
if better {
self.best_score = Some(score.clone());
self.best_model = Some(model.clone());
self.best_inliers.clear();
self.best_inliers.extend_from_slice(&tmp_inliers);
iteration_improved_best = true;
}
}
if iteration_improved_best {
if let (Some(lo), Some(best_model), Some(best_score)) = (
&mut self.local_optimizer,
&self.best_model,
&self.best_score,
) {
let (refined_model, refined_score, refined_inliers) =
lo.run(data, &self.best_inliers, best_model, best_score);
if refined_score > *best_score {
self.best_model = Some(refined_model);
self.best_score = Some(refined_score);
self.best_inliers = refined_inliers;
}
}
if let (Some(best_score), Some(_best_model)) = (&self.best_score, &self.best_model)
{
let should_terminate =
self.termination
.check(data, best_score, sample_size, &mut max_iterations);
if should_terminate {
break;
}
}
}
self.sampler
.update(&sample, sample_size, self.iteration, 0.0);
self.iteration += 1;
}
if let (Some(final_opt), Some(best_model), Some(best_score)) = (
&mut self.final_optimizer,
&self.best_model,
&self.best_score,
) && self.best_inliers.len() > sample_size
{
let (refined_model, refined_score, refined_inliers) =
final_opt.run(data, &self.best_inliers, best_model, best_score);
if refined_score > *best_score {
self.best_model = Some(refined_model);
self.best_score = Some(refined_score);
self.best_inliers = refined_inliers;
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::settings::RansacSettings;
use crate::types::DataMatrix;
#[derive(Clone, Debug)]
struct MockModel;
#[derive(Clone, Debug, PartialEq, PartialOrd)]
struct MockScore(f64);
struct MockEstimator;
impl Estimator for MockEstimator {
type Model = MockModel;
fn sample_size(&self) -> usize {
2
}
fn is_valid_sample(&self, _data: &DataMatrix, _sample: &[usize]) -> bool {
true
}
fn estimate_model(&self, _data: &DataMatrix, _sample: &[usize]) -> Vec<Self::Model> {
vec![MockModel]
}
fn is_valid_model(
&self,
_model: &Self::Model,
_data: &DataMatrix,
_sample: &[usize],
_threshold: f64,
) -> bool {
true
}
}
struct MockSampler {
pub sample_calls: usize,
pub update_calls: usize,
}
impl Sampler for MockSampler {
fn sample(
&mut self,
data: &DataMatrix,
sample_size: usize,
out_indices: &mut [usize],
) -> bool {
self.sample_calls += 1;
for (i, v) in out_indices.iter_mut().enumerate().take(sample_size) {
*v = i.min(data.nrows().saturating_sub(1));
}
true
}
fn update(
&mut self,
_sample: &[usize],
_sample_size: usize,
_iteration: usize,
_score_hint: f64,
) {
self.update_calls += 1;
}
}
struct MockScoring;
impl Scoring<MockModel> for MockScoring {
type Score = MockScore;
fn threshold(&self) -> f64 {
1.0
}
fn score(
&self,
_data: &DataMatrix,
_model: &MockModel,
inliers_out: &mut Vec<usize>,
) -> Self::Score {
inliers_out.clear();
inliers_out.push(0);
MockScore(1.0)
}
}
struct MockLocalOptimizer;
impl LocalOptimizer<MockModel, MockScore> for MockLocalOptimizer {
fn run(
&mut self,
_data: &DataMatrix,
inliers: &[usize],
model: &MockModel,
best_score: &MockScore,
) -> (MockModel, MockScore, Vec<usize>) {
(model.clone(), best_score.clone(), inliers.to_vec())
}
}
struct MockTerminationCriterion {
pub called: usize,
}
impl TerminationCriterion<MockScore> for MockTerminationCriterion {
fn check(
&mut self,
_data: &DataMatrix,
_best_score: &MockScore,
_sample_size: usize,
max_iterations: &mut usize,
) -> bool {
self.called += 1;
*max_iterations = (*max_iterations).min(1);
true
}
}
struct MockInlierSelector;
impl InlierSelector<MockModel> for MockInlierSelector {
fn select(&mut self, _data: &DataMatrix, _model: &MockModel) -> Vec<usize> {
vec![0]
}
}
#[test]
fn superransac_runs_and_updates_best_model() {
let data = DataMatrix::zeros(4, 2);
let settings = RansacSettings::default();
let estimator = MockEstimator;
let sampler = MockSampler {
sample_calls: 0,
update_calls: 0,
};
let scoring = MockScoring;
let local_optimizer = Some(MockLocalOptimizer);
let final_optimizer = None;
let termination = MockTerminationCriterion { called: 0 };
let inlier_selector = Some(MockInlierSelector);
let mut pipeline = SuperRansac::new(
settings,
estimator,
sampler,
scoring,
local_optimizer,
final_optimizer,
termination,
inlier_selector,
);
pipeline.run(&data);
assert!(
pipeline.best_model.is_some(),
"Best model should be set after run"
);
assert!(
pipeline.best_score.is_some(),
"Best score should be set after run"
);
assert!(
!pipeline.best_inliers.is_empty(),
"Best inliers should not be empty"
);
}
}