From 75d1065db5432678f20f6ca9fbb158bc4e76fd4f Mon Sep 17 00:00:00 2001
From: Muthu Annamalai <muthiah.annamalai@tetramem.com>
Date: Tue, 17 Mar 2026 13:11:50 -0700
Subject: [PATCH 1/5] Fix build error with change in Progress Style Clippy fmt
 fix Add seeded RNG - not working fully; - add lambda function based objective
 function support in the docs

---
 Cargo.toml      |  3 ++-
 src/bin/e_lj.rs |  5 ++--
 src/bin/main.rs | 11 +++++++--
 src/lib.rs      | 39 ++++++++++++++++++++++++------
 src/model.rs    | 23 ++++++++++++++----
 src/pso.rs      | 64 ++++++++++++++++++++++++++++++-------------------
 tests/pso.rs    |  8 +++----
 7 files changed, 107 insertions(+), 46 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 109bbc2..b39cdd7 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -17,7 +17,8 @@ exclude = [
 
 [dependencies]
 rand = "^0.8.4"
-indicatif = "^0.16.2"
+rand_chacha = "^0.3"
+indicatif = "^0.18.4"
 rayon = "^1.5.1"
 
 # [profile.release]
diff --git a/src/bin/e_lj.rs b/src/bin/e_lj.rs
index 98f413b..9dd6947 100644
--- a/src/bin/e_lj.rs
+++ b/src/bin/e_lj.rs
@@ -1,7 +1,6 @@
 use pso_rs::*;
 use std::process;
-const N_PARTICLES:usize = 20;
-
+const N_PARTICLES: usize = 20;
 
 fn main() {
     let dimensions = vec![N_PARTICLES, 3];
@@ -30,7 +29,7 @@ fn main() {
     };
     use std::time::Instant;
     let before = Instant::now();
-    match pso_rs::run(config, e_lj, None) {
+    match pso_rs::run(config, e_lj, None, None) {
         Ok(pso) => {
             println!("Elapsed time: {:.2?}", before.elapsed());
             pso.write_f_to_file("./best_f_trajectory.txt")
diff --git a/src/bin/main.rs b/src/bin/main.rs
index 5d884af..43da428 100644
--- a/src/bin/main.rs
+++ b/src/bin/main.rs
@@ -1,6 +1,6 @@
 use pso_rs::*;
 
-const N_DIMENSIONS:usize = 3;
+const N_DIMENSIONS: usize = 3;
 
 fn main() {
     let config = Config {
@@ -12,7 +12,14 @@ fn main() {
     };
     use std::time::Instant;
     let before = Instant::now();
-    let pso = pso_rs::run(config, sum_squares, Some(|f_best| f_best < 1e-4)).unwrap();
+
+    let pso = pso_rs::run(
+        config,
+        sum_squares,
+        Some(|f_best| f_best < 1e-4),
+        Some(123456u64),
+    )
+    .unwrap();
     println!("Elapsed time: {:.2?}", before.elapsed());
     let model = pso.model;
     println!("Found minimum: {:#?} ", model.get_f_best());
diff --git a/src/lib.rs b/src/lib.rs
index 3ccbfc7..b16bfeb 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -50,7 +50,8 @@
 //! let pso = pso_rs::run(
 //!     config,
 //!     objective_function,
-//!     Some(terminate)
+//!     Some(terminate),
+//!     None,
 //! ).unwrap();
 //!     
 //! let model = pso.model;
@@ -84,7 +85,8 @@
 //!
 //! let mut pso = pso_rs::init(
 //!     config,
-//!     objective_function
+//!     objective_function,
+//!     None,
 //! ).unwrap();
 //!
 //! // run PSO with no termination condition
@@ -94,6 +96,27 @@
 //! println!("Found minimum: {:#?} ", model.get_f_best());
 //! println!("Minimizer: {:#?}", model.get_x_best());
 //! ```
+//! //example shows how to use lambda functions to model objective
+//! fn main() {
+//!    let config = Config {
+//!        dimensions: vec![N_DIMENSIONS],
+//!        population_size: 100,
+//!        bounds: vec![(-10.0, 10.0); N_DIMENSIONS],
+//!        t_max: 1e7 as usize,
+//!        ..Config::default()
+//!    };
+//!    use std::time::Instant;
+//!    let before = Instant::now();
+//!    let xsum_squares = |p: &Particle, _flat_dim: usize, dimensions: &Vec<usize>| -> f64 {
+//!         (0..dimensions[0]).map(|i| i as f64 * p[i].powf(2.0)).sum()
+//!     };
+//!     
+//!    let pso = pso_rs::run(config, xsum_squares, Some(|f_best| f_best < 1e-4),Some(123456u64)).unwrap();
+//!    println!("Elapsed time: {:.2?}", before.elapsed());
+//!    let model = pso.model;
+//!    println!("Found minimum: {:#?} ", model.get_f_best());
+//!    println!("Found minimizer: {:#?} ", model.get_x_best());
+//! }
 //!
 //! # Notes
 //!
@@ -151,6 +174,7 @@
 //! let pso = pso_rs::run(
 //!     config,
 //!     objective_function,
+//!     None,
 //!     None
 //! ).unwrap();
 //!
@@ -167,7 +191,6 @@ pub mod pso;
 
 pub use model::*;
 
-use model::Model;
 use pso::PSO;
 use std::error::Error;
 
@@ -180,9 +203,10 @@ pub fn run(
     config: Config,
     obj_f: fn(&Particle, usize, &Vec<usize>) -> f64,
     terminate_f: Option<fn(f64) -> bool>,
+    seed: Option<u64>,
 ) -> Result<PSO, Box<dyn Error>> {
     assert_config(&config)?;
-    let mut pso = init(config, obj_f).unwrap();
+    let mut pso = init(config, obj_f, seed).unwrap();
     let term_condition = match terminate_f {
         Some(terminate_f) => terminate_f,
         None => |_| false,
@@ -197,10 +221,11 @@ pub fn run(
 pub fn init(
     config: Config,
     obj_f: fn(&Particle, usize, &Vec<usize>) -> f64,
+    seed: Option<u64>,
 ) -> Result<PSO, &'static str> {
     assert_config(&config)?;
-    let model = Model::new(config, obj_f);
-    let pso = PSO::new(model);
+    let model = Model::new(config, obj_f, seed);
+    let pso = PSO::new(model, seed);
     Ok(pso)
 }
 
@@ -208,7 +233,7 @@ fn assert_config(config: &Config) -> Result<(), &'static str> {
     if config.c1 + config.c2 < 4.0 {
         return Err("c1 + c2 must be greater than 4");
     }
-    if config.dimensions.len() == 0 {
+    if config.dimensions.is_empty() {
         return Err("dimensions must be set");
     }
     if config.bounds.len() != config.dimensions[config.dimensions.len() - 1] {
diff --git a/src/model.rs b/src/model.rs
index 371dcb1..0121298 100644
--- a/src/model.rs
+++ b/src/model.rs
@@ -1,6 +1,8 @@
-use rand::{thread_rng, Rng};
+use rand::{thread_rng, Rng, SeedableRng};
+use rand_chacha::ChaCha8Rng;
 use rayon::prelude::*;
 use std::fmt;
+
 pub type Particle = Vec<f64>;
 pub type Population = Vec<Particle>;
 
@@ -14,6 +16,7 @@ pub struct Model {
     pub population_f_scores: Vec<f64>,
     pub x_best: Particle,
     pub f_best: f64,
+    pub seed: Option<u64>,
     obj_f: fn(&Particle, usize, &Vec<usize>) -> f64,
 }
 
@@ -22,9 +25,14 @@ impl Model {
     pub fn new(
         config: Config,
         obj_f: fn(p: &Particle, flat_dim: usize, dim: &Vec<usize>) -> f64,
+        seed: Option<u64>,
     ) -> Model {
         // init population
         let mut rng = thread_rng();
+        let mut seeded_rng = ChaCha8Rng::seed_from_u64(0);
+        if let Some(seedval) = seed {
+            seeded_rng = ChaCha8Rng::seed_from_u64(seedval);
+        }
         let mut flat_dim = 1;
         for d in config.dimensions.clone() {
             flat_dim *= d;
@@ -35,13 +43,19 @@ impl Model {
             let mut particle: Particle = vec![];
             for flat_i in 0..flat_dim {
                 let true_i = flat_i % config.dimensions[config.dimensions.len() - 1];
-                particle.push(rng.gen_range(config.bounds[true_i].0..config.bounds[true_i].1));
+                if let Some(_) = seed {
+                    particle.push(rng.gen_range(config.bounds[true_i].0..config.bounds[true_i].1));
+                } else {
+                    particle.push(
+                        seeded_rng.gen_range(config.bounds[true_i].0..config.bounds[true_i].1),
+                    );
+                }
             }
             population.push(particle);
         }
         let population_f_scores = vec![f64::INFINITY; config.population_size];
         let x_best = population[0].clone();
-        let f_best = population_f_scores[0].clone();
+        let f_best = population_f_scores[0];
         let mut model = Model {
             config,
             flat_dim,
@@ -49,7 +63,8 @@ impl Model {
             population_f_scores,
             x_best,
             f_best,
-            obj_f: obj_f,
+            seed,
+            obj_f,
         };
         model.get_f_values();
         model
diff --git a/src/pso.rs b/src/pso.rs
index e30302c..2d8bd56 100644
--- a/src/pso.rs
+++ b/src/pso.rs
@@ -1,6 +1,8 @@
 use crate::model::*;
 use indicatif::{ProgressBar, ProgressStyle};
-use rand::{thread_rng, Rng};
+use rand::rngs::ThreadRng;
+use rand::{thread_rng, Rng, SeedableRng};
+use rand_chacha::ChaCha8Rng;
 
 use std::error::Error;
 use std::fs::File;
@@ -19,11 +21,14 @@ pub struct PSO {
     pub best_f_values: Vec<f64>,
     pub best_f_trajectory: Vec<f64>,
     pub best_x_trajectory: Vec<Particle>,
+    pub seed: Option<u64>,
+    pub seeded_rng: ChaCha8Rng,
+    pub rng: ThreadRng,
 }
 
 impl PSO {
     /// Initialize Particle Swarm Optimization
-    pub fn new(model: Model) -> PSO {
+    pub fn new(model: Model, seed: Option<u64>) -> PSO {
         let phi = model.config.c1 + model.config.c2;
         let phi_squared = phi.powf(2.0);
         let tmp = phi_squared - (4.0 * phi);
@@ -34,11 +39,19 @@ impl PSO {
 
         // initialize
         let mut rng = thread_rng();
+        let mut seeded_rng = ChaCha8Rng::seed_from_u64(0);
+        if let Some(seedval) = seed {
+            seeded_rng = ChaCha8Rng::seed_from_u64(seedval);
+        }
         let mut velocities = vec![];
         for _ in 0..model.config.population_size {
             let mut tmp = vec![];
             for _ in 0..model.flat_dim {
-                tmp.push(rng.gen_range(v_max * -1.0..v_max * 1.0));
+                if let Some(_) = seed {
+                    tmp.push(seeded_rng.gen_range(-v_max..v_max));
+                } else {
+                    tmp.push(rng.gen_range(-v_max..v_max));
+                }
             }
             velocities.push(tmp);
         }
@@ -58,6 +71,9 @@ impl PSO {
             neigh_population,
             best_f_trajectory,
             best_x_trajectory,
+            seed,
+            seeded_rng,
+            rng,
         }
     }
 
@@ -70,13 +86,12 @@ impl PSO {
         let mut bar: Option<ProgressBar> = None;
         if self.model.config.progress_bar {
             bar = Some(ProgressBar::new(self.model.config.t_max as u64));
-            match bar {
-                Some(ref bar) => {
-                    bar.set_style(ProgressStyle::default_bar().template(
-                        "{msg} [{elapsed}] {bar:20.cyan/blue} {pos:>7}/{len:7} ETA: {eta}",
-                    ));
+            if let Some(ref bar) = bar {
+                if let Ok(value) = ProgressStyle::default_bar()
+                    .template("{msg} [{elapsed}] {bar:20.cyan/blue} {pos:>7}/{len:7} ETA: {eta}")
+                {
+                    bar.set_style(value);
                 }
-                None => {}
             }
         }
         let mut k = 0;
@@ -91,35 +106,34 @@ impl PSO {
 
             self.model.population = self.model.population.clone();
             k += pop_size;
-            match bar {
-                Some(ref bar) => {
-                    bar.inc(pop_size as u64);
-                    bar.set_message(format!("{:.6}", self.model.f_best));
-                }
-                None => {}
+            if let Some(ref bar) = bar {
+                bar.inc(pop_size as u64);
+                bar.set_message(format!("{:.6}", self.model.f_best));
             }
             if k > self.model.config.t_max || terminate(self.model.f_best) {
                 break;
             }
         }
-        match bar {
-            Some(ref bar) => {
-                bar.finish_and_clear();
-            }
-            None => {}
+        if let Some(ref bar) = bar {
+            bar.finish_and_clear();
         }
         k
     }
 
     /// Updates the velocity and position of each particle in the population
     fn update_velocity_and_pos(&mut self) {
-        let mut rng = thread_rng();
-
         for i in 0..self.model.config.population_size {
             let lbest = &self.neigh_population[self.local_best(i)];
             for j in 0..self.model.flat_dim {
-                let r1 = rng.gen_range(-1.0..1.0);
-                let r2 = rng.gen_range(-1.0..1.0);
+                let r1: f64;
+                let r2: f64;
+                if let Some(_) = self.seed {
+                    r1 = self.seeded_rng.gen_range(-1.0..1.0);
+                    r2 = self.seeded_rng.gen_range(-1.0..1.0);
+                } else {
+                    r1 = self.rng.gen_range(-1.0..1.0);
+                    r2 = self.rng.gen_range(-1.0..1.0);
+                }
                 let cog = self.model.config.c1
                     * r1
                     * (self.neigh_population[i][j] - self.model.population[i][j]);
@@ -173,7 +187,7 @@ impl PSO {
     fn local_best(&self, i: usize) -> usize {
         let best = PSO::argsort(&self.best_f_values);
         for b in best {
-            if self.neighborhoods[i].iter().any(|&n| n == b) {
+            if self.neighborhoods[i].contains(&b) {
                 return b;
             }
         }
diff --git a/tests/pso.rs b/tests/pso.rs
index e691457..8ddc8d8 100644
--- a/tests/pso.rs
+++ b/tests/pso.rs
@@ -15,7 +15,7 @@ fn it_runs_non_parallel() {
         parallelize: false,
         ..Config::default()
     };
-    let pso = pso_rs::run(config, rosenbrock, None).unwrap();
+    let pso = pso_rs::run(config, rosenbrock, None, None).unwrap();
 
     let mut model = pso.model;
 
@@ -46,7 +46,7 @@ fn it_computes_correct_minimum_rosenbrock_2d() {
         progress_bar: false,
         ..Config::default()
     };
-    let pso = pso_rs::run(config, rosenbrock, None).unwrap();
+    let pso = pso_rs::run(config, rosenbrock, None, None).unwrap();
 
     let mut model = pso.model;
 
@@ -79,7 +79,7 @@ fn it_computes_correct_minimum_rosenbrock_3d() {
         progress_bar: false,
         ..Config::default()
     };
-    let pso = pso_rs::run(config, rosenbrock, None).unwrap();
+    let pso = pso_rs::run(config, rosenbrock, None, None).unwrap();
 
     let mut model = pso.model;
 
@@ -139,7 +139,7 @@ fn it_computes_correct_minimum_e_lj() {
         ..Config::default()
     };
 
-    let pso = pso_rs::run(config, e_lj, Some(|_| true)).unwrap();
+    let pso = pso_rs::run(config, e_lj, Some(|_| true), None).unwrap();
 
     let mut model = pso.model;
 

From f117ae295cd4c02f8d9c49958500b6b7e11ba7ea Mon Sep 17 00:00:00 2001
From: Muthiah Annamalai <muthu@panmo.cloud>
Date: Wed, 8 Apr 2026 07:32:51 -0700
Subject: [PATCH 2/5] [pso]: run against parameters which can be boolean, f64
 or i32 (NumericKind) based heterogenity        -- Fix models to use
 polymorphic numeric type      : add timer information

---
 src/bin/e_lj.rs |   8 +--
 src/bin/main.rs |  17 +++++--
 src/lib.rs      |  37 ++++++++------
 src/model.rs    | 123 +++++++++++++++++++++++++++++++++++++++++++---
 src/pso.rs      |  54 +++++++++++---------
 tests/pso.rs    | 127 +++++++++++++++++++++++++++++++++---------------
 6 files changed, 274 insertions(+), 92 deletions(-)

diff --git a/src/bin/e_lj.rs b/src/bin/e_lj.rs
index 9dd6947..b832178 100644
--- a/src/bin/e_lj.rs
+++ b/src/bin/e_lj.rs
@@ -29,7 +29,7 @@ fn main() {
     };
     use std::time::Instant;
     let before = Instant::now();
-    match pso_rs::run(config, e_lj, None, None) {
+    match pso_rs::run(config, e_lj, None, None, None) {
         Ok(pso) => {
             println!("Elapsed time: {:.2?}", before.elapsed());
             pso.write_f_to_file("./best_f_trajectory.txt")
@@ -61,7 +61,7 @@ fn l2(x_i: Particle, x_j: Particle, particle_dim: usize) -> f64 {
     // calculated as the square root of the sum of the squared vector values
     let mut sum: f64 = 0.0;
     for i in 0..particle_dim {
-        sum += (x_i[i] - x_j[i]).powf(2.0);
+        sum += (x_i[i] - x_j[i]).value_f64().powf(2.0);
     }
     sum.sqrt()
 }
@@ -73,7 +73,7 @@ fn v_ij(x_i: Particle, x_j: Particle, particle_dim: usize) -> f64 {
 }
 
 /// Get potential energy of a cluster of particles
-fn e_lj(particle: &Particle, _flat_dim: usize, particle_dims: &Vec<usize>) -> f64 {
+fn e_lj(particle: &Particle, _flat_dim: usize, particle_dims: &[usize]) -> f64 {
     let mut sum = 0.0;
     for i in 0..particle_dims[0] - 1 {
         for j in (i + 1)..particle_dims[0] {
@@ -89,7 +89,7 @@ fn e_lj(particle: &Particle, _flat_dim: usize, particle_dims: &Vec<usize>) -> f6
     4.0 * sum
 }
 
-fn reshape(particle: &Particle, particle_dims: &Vec<usize>) -> Vec<Vec<f64>> {
+fn reshape(particle: &Particle, particle_dims: &[usize]) -> Vec<Vec<NumericKind>> {
     // reshape particle
     let mut reshaped_cluster = vec![];
     let mut i = 0;
diff --git a/src/bin/main.rs b/src/bin/main.rs
index 43da428..de4dbad 100644
--- a/src/bin/main.rs
+++ b/src/bin/main.rs
@@ -13,10 +13,17 @@ fn main() {
     use std::time::Instant;
     let before = Instant::now();
 
-    let pso = pso_rs::run(
+    let pso: pso::PSO = pso_rs::run(
         config,
         sum_squares,
-        Some(|f_best| f_best < 1e-4),
+        Some(|f| {
+            let mut r: Vec<NumericKind> = vec![];
+            for ff in f {
+                r.push(NumericKind::ValueF64(*ff));
+            }
+            r
+        }),
+        None,
         Some(123456u64),
     )
     .unwrap();
@@ -26,6 +33,8 @@ fn main() {
     println!("Found minimizer: {:#?} ", model.get_x_best());
 }
 
-fn sum_squares(p: &Particle, _flat_dim: usize, dimensions: &Vec<usize>) -> f64 {
-    (0..dimensions[0]).map(|i| i as f64 * p[i].powf(2.0)).sum()
+fn sum_squares(p: &Particle, _flat_dim: usize, dimensions: &[usize]) -> f64 {
+    (0..dimensions[0])
+        .map(|i| i as f64 * p[i].value_f64().powf(2.0))
+        .sum()
 }
diff --git a/src/lib.rs b/src/lib.rs
index b16bfeb..ec0872c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -23,11 +23,11 @@
 //! fn objective_function(
 //!     p: &Particle,
 //!     _flat_dim: usize,
-//!     dimensions: &Vec<usize>
+//!     dimensions: &[usize]
 //! ) -> f64 {
 //!     (0..dimensions[0] - 1).map(|i| {
-//!         100.0 * ((p[i+1]-p[i]).powf(2.0)).powf(2.0)
-//!             + (1.0-p[i]).powf(2.0)
+//!         100.0 * ((p[i+1]-p[i]).value_f64().powf(2.0)).powf(2.0)
+//!             + (1.0-p[i].value_f64()).powf(2.0)
 //!     }).sum()
 //! }
 //!
@@ -50,8 +50,9 @@
 //! let pso = pso_rs::run(
 //!     config,
 //!     objective_function,
-//!     Some(terminate),
 //!     None,
+//!     Some(terminate),
+//!     None
 //! ).unwrap();
 //!     
 //! let model = pso.model;
@@ -67,11 +68,11 @@
 //! fn objective_function(
 //!     p: &Particle,
 //!     _flat_dim: usize,
-//!     dimensions: &Vec<usize>
+//!     dimensions: &[usize]
 //! ) -> f64 {
 //!     (0..dimensions[0] - 1).map(|i| {
-//!         100.0 * ((p[i+1]-p[i]).powf(2.0)).powf(2.0)
-//!             + (1.0-p[i]).powf(2.0)
+//!         100.0 * ((p[i+1]-p[i]).value_f64().powf(2.0)).powf(2.0)
+//!             + (1.0-p[i].value_f64()).powf(2.0)
 //!     }).sum()
 //! }
 //!
@@ -87,6 +88,7 @@
 //!     config,
 //!     objective_function,
 //!     None,
+//!     None
 //! ).unwrap();
 //!
 //! // run PSO with no termination condition
@@ -107,11 +109,11 @@
 //!    };
 //!    use std::time::Instant;
 //!    let before = Instant::now();
-//!    let xsum_squares = |p: &Particle, _flat_dim: usize, dimensions: &Vec<usize>| -> f64 {
+//!    let xsum_squares = |p: &Particle, _flat_dim: usize, dimensions: &[usize]| -> f64 {
 //!         (0..dimensions[0]).map(|i| i as f64 * p[i].powf(2.0)).sum()
 //!     };
 //!     
-//!    let pso = pso_rs::run(config, xsum_squares, Some(|f_best| f_best < 1e-4),Some(123456u64)).unwrap();
+//!    let pso = pso_rs::run(config, xsum_squares, None, Some(|f_best| f_best < 1e-4),Some(123456u64)).unwrap();
 //!    println!("Elapsed time: {:.2?}", before.elapsed());
 //!    let model = pso.model;
 //!    println!("Found minimum: {:#?} ", model.get_f_best());
@@ -138,8 +140,8 @@
 //!
 //! fn reshape(
 //!     particle: &Particle,
-//!     particle_dims: &Vec<usize>
-//! ) -> Vec<Vec<f64>> {
+//!     particle_dims: &[usize]
+//! ) -> Vec<Vec<NumericKind>> {
 //!     let mut reshaped_cluster = vec![];
 //!     let mut i = 0;
 //!     for _ in 0..particle_dims[0] {
@@ -157,7 +159,7 @@
 //! fn objective_function(
 //!     p: &Particle,
 //!     _flat_dim: usize,
-//!     dimensions: &Vec<usize>
+//!     dimensions: &[usize]
 //! ) -> f64 {
 //!     let _reshaped_particle = reshape(p, dimensions);
 //!     /* Do stuff */
@@ -175,6 +177,7 @@
 //!     config,
 //!     objective_function,
 //!     None,
+//!     None,
 //!     None
 //! ).unwrap();
 //!
@@ -201,12 +204,13 @@ use std::error::Error;
 /// Panics if any particle coefficient becomes NaN (usually because of bad parameterization, e.g. c1 + c2 < 4)
 pub fn run(
     config: Config,
-    obj_f: fn(&Particle, usize, &Vec<usize>) -> f64,
+    obj_f: fn(&Particle, usize, &[usize]) -> f64,
+    cast_f: Option<fn(&Vec<f64>) -> Vec<NumericKind>>,
     terminate_f: Option<fn(f64) -> bool>,
     seed: Option<u64>,
 ) -> Result<PSO, Box<dyn Error>> {
     assert_config(&config)?;
-    let mut pso = init(config, obj_f, seed).unwrap();
+    let mut pso = init(config, obj_f, cast_f, seed).unwrap();
     let term_condition = match terminate_f {
         Some(terminate_f) => terminate_f,
         None => |_| false,
@@ -220,11 +224,12 @@ pub fn run(
 /// Useful for initializing an instance for running at a later time
 pub fn init(
     config: Config,
-    obj_f: fn(&Particle, usize, &Vec<usize>) -> f64,
+    obj_f: fn(&Particle, usize, &[usize]) -> f64,
+    cast_f: Option<fn(&Vec<f64>) -> Vec<NumericKind>>,
     seed: Option<u64>,
 ) -> Result<PSO, &'static str> {
     assert_config(&config)?;
-    let model = Model::new(config, obj_f, seed);
+    let model = Model::new(config, obj_f, cast_f, seed);
     let pso = PSO::new(model, seed);
     Ok(pso)
 }
diff --git a/src/model.rs b/src/model.rs
index 0121298..9bf672c 100644
--- a/src/model.rs
+++ b/src/model.rs
@@ -2,10 +2,105 @@ use rand::{thread_rng, Rng, SeedableRng};
 use rand_chacha::ChaCha8Rng;
 use rayon::prelude::*;
 use std::fmt;
+use std::ops::{Add, Mul, Sub};
 
-pub type Particle = Vec<f64>;
+#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
+pub enum NumericKind {
+    ValueF64(f64),
+    ValueBool(bool),
+    ValueI32(i32),
+}
+
+impl ToString for NumericKind {
+    fn to_string(&self) -> String {
+        match self {
+            NumericKind::ValueF64(v) => v.to_string(),
+            NumericKind::ValueBool(v) => v.to_string(),
+            NumericKind::ValueI32(v) => v.to_string(),
+        }
+    }
+}
+
+impl Mul for NumericKind {
+    type Output = Self;
+    fn mul(self, other: Self) -> Self {
+        match self {
+            NumericKind::ValueF64(v) => NumericKind::ValueF64 {
+                0: v * other.value_f64(),
+            },
+            NumericKind::ValueBool(v) => NumericKind::ValueBool {
+                0: v && other.value_bool(),
+            },
+            NumericKind::ValueI32(v) => NumericKind::ValueI32 {
+                0: v * other.value_i32(),
+            },
+        }
+    }
+}
+
+impl Add for NumericKind {
+    type Output = Self;
+    fn add(self, other: Self) -> Self {
+        match self {
+            NumericKind::ValueF64(v) => NumericKind::ValueF64 {
+                0: v + other.value_f64(),
+            },
+            NumericKind::ValueBool(v) => NumericKind::ValueBool {
+                0: v || other.value_bool(),
+            },
+            NumericKind::ValueI32(v) => NumericKind::ValueI32 {
+                0: v + other.value_i32(),
+            },
+        }
+    }
+}
+
+impl Sub for NumericKind {
+    type Output = Self;
+    fn sub(self, other: Self) -> Self {
+        match self {
+            NumericKind::ValueF64(v) => NumericKind::ValueF64 {
+                0: v - other.value_f64(),
+            },
+            NumericKind::ValueBool(v) => NumericKind::ValueBool {
+                0: v ^ other.value_bool(),
+            },
+            NumericKind::ValueI32(v) => NumericKind::ValueI32 {
+                0: v - other.value_i32(),
+            },
+        }
+    }
+}
+
+pub type Particle = Vec<NumericKind>;
 pub type Population = Vec<Particle>;
 
+impl NumericKind {
+    pub fn value_f64(&self) -> f64 {
+        match self {
+            NumericKind::ValueF64(v) => *v,
+            NumericKind::ValueBool(v) => *v as i32 as f64,
+            NumericKind::ValueI32(v) => *v as f64,
+        }
+    }
+
+    pub fn value_bool(&self) -> bool {
+        match self {
+            NumericKind::ValueF64(v) => *v as i32 != 0,
+            NumericKind::ValueBool(v) => *v,
+            NumericKind::ValueI32(v) => *v as i32 != 0,
+        }
+    }
+
+    pub fn value_i32(&self) -> i32 {
+        match self {
+            NumericKind::ValueF64(v) => *v as i32,
+            NumericKind::ValueBool(v) => *v as i32,
+            NumericKind::ValueI32(v) => *v,
+        }
+    }
+}
+
 /// Model struct
 ///
 /// It takes in a `Config` instance and `fn` pointer to an objective function and defines a `run` method for running Particle Swarm Optimization.
@@ -17,14 +112,16 @@ pub struct Model {
     pub x_best: Particle,
     pub f_best: f64,
     pub seed: Option<u64>,
-    obj_f: fn(&Particle, usize, &Vec<usize>) -> f64,
+    obj_f: fn(&Particle, usize, &[usize]) -> f64,
+    cast_f: Option<fn(&Vec<f64>) -> Vec<NumericKind>>, /* cast Particle -> Particle */
 }
 
 impl Model {
     /// Creates a new Model instance
     pub fn new(
         config: Config,
-        obj_f: fn(p: &Particle, flat_dim: usize, dim: &Vec<usize>) -> f64,
+        obj_f: fn(p: &Particle, flat_dim: usize, dim: &[usize]) -> f64,
+        cast_f: Option<fn(p: &Vec<f64>) -> Vec<NumericKind>>,
         seed: Option<u64>,
     ) -> Model {
         // init population
@@ -44,14 +141,25 @@ impl Model {
             for flat_i in 0..flat_dim {
                 let true_i = flat_i % config.dimensions[config.dimensions.len() - 1];
                 if let Some(_) = seed {
-                    particle.push(rng.gen_range(config.bounds[true_i].0..config.bounds[true_i].1));
+                    particle.push(NumericKind::ValueF64(
+                        rng.gen_range(config.bounds[true_i].0..config.bounds[true_i].1),
+                    ));
                 } else {
-                    particle.push(
+                    particle.push(NumericKind::ValueF64(
                         seeded_rng.gen_range(config.bounds[true_i].0..config.bounds[true_i].1),
-                    );
+                    ));
                 }
             }
-            population.push(particle);
+            population.push(match cast_f {
+                Some(caster) => {
+                    let v = particle
+                        .iter()
+                        .map(|f| (*f).value_f64())
+                        .collect::<Vec<f64>>();
+                    caster(&v)
+                }
+                _ => particle,
+            });
         }
         let population_f_scores = vec![f64::INFINITY; config.population_size];
         let x_best = population[0].clone();
@@ -65,6 +173,7 @@ impl Model {
             f_best,
             seed,
             obj_f,
+            cast_f,
         };
         model.get_f_values();
         model
diff --git a/src/pso.rs b/src/pso.rs
index 2d8bd56..a9e1c34 100644
--- a/src/pso.rs
+++ b/src/pso.rs
@@ -33,9 +33,9 @@ impl PSO {
         let phi_squared = phi.powf(2.0);
         let tmp = phi_squared - (4.0 * phi);
         let tmp = tmp.sqrt();
-        let chi = 2.0 / (2.0 - phi - tmp).abs();
+        let chi = ((2.0 / (2.0 - phi - tmp)) as f64).abs();
         let v_max = model.config.alpha * 5.0;
-        let neighborhoods = PSO::create_neighborhoods(&model);
+        let neighborhoods = Self::create_neighborhoods(&model);
 
         // initialize
         let mut rng = thread_rng();
@@ -48,18 +48,20 @@ impl PSO {
             let mut tmp = vec![];
             for _ in 0..model.flat_dim {
                 if let Some(_) = seed {
-                    tmp.push(seeded_rng.gen_range(-v_max..v_max));
+                    tmp.push(NumericKind::ValueF64(seeded_rng.gen_range(-v_max..v_max)));
                 } else {
-                    tmp.push(rng.gen_range(-v_max..v_max));
+                    tmp.push(NumericKind::ValueF64(rng.gen_range(-v_max..v_max)));
                 }
             }
             velocities.push(tmp);
         }
 
         let best_f_values = model.population_f_scores.clone();
-        let neigh_population = model.population.clone();
-        let best_f_trajectory = vec![model.f_best];
-        let best_x_trajectory = vec![model.x_best.clone()];
+        let neigh_population = (0..model.population_f_scores.len())
+            .map(|idx| model.population[idx].clone())
+            .collect();
+        let best_f_trajectory = vec![model.get_f_best()];
+        let best_x_trajectory = vec![model.get_x_best()];
 
         PSO {
             chi,
@@ -124,6 +126,7 @@ impl PSO {
     fn update_velocity_and_pos(&mut self) {
         for i in 0..self.model.config.population_size {
             let lbest = &self.neigh_population[self.local_best(i)];
+            #[allow(clippy::needless_range_loop)]
             for j in 0..self.model.flat_dim {
                 let r1: f64;
                 let r2: f64;
@@ -136,32 +139,37 @@ impl PSO {
                 }
                 let cog = self.model.config.c1
                     * r1
-                    * (self.neigh_population[i][j] - self.model.population[i][j]);
+                    * (self.neigh_population[i][j] - self.model.population[i][j]).value_f64();
 
-                let soc = self.model.config.c2 * r2 * (lbest[j] - self.model.population[i][j]);
-                let v = self.chi * (self.velocities[i][j] + cog + soc);
+                let soc = self.model.config.c2
+                    * r2
+                    * ((lbest[j] - self.model.population[i][j]).value_f64());
+                let v = self.chi * (self.velocities[i][j].value_f64() + cog + soc);
 
                 // check bounds
-                self.velocities[i][j] = if v.abs() > self.v_max {
+                self.velocities[i][j] = NumericKind::ValueF64(if v.abs() > self.v_max {
                     v.signum() * self.v_max
                 } else {
                     v
-                };
+                });
 
-                let x = self.model.population[i][j] + self.model.config.lr * self.velocities[i][j];
+                let x = NumericKind::ValueF64(
+                    self.model.population[i][j].value_f64()
+                        + self.model.config.lr * self.velocities[i][j].value_f64(),
+                );
 
                 let bound_index =
                     j % self.model.config.dimensions[self.model.config.dimensions.len() - 1];
                 let (lower_bound, upper_bound) = self.model.config.bounds[bound_index];
                 // check bounds
-                if x > upper_bound {
-                    self.model.population[i][j] = upper_bound;
-                } else if x < lower_bound {
-                    self.model.population[i][j] = lower_bound;
+                if x.value_f64() > upper_bound {
+                    self.model.population[i][j] = NumericKind::ValueF64(upper_bound);
+                } else if x.value_f64() < lower_bound {
+                    self.model.population[i][j] = NumericKind::ValueF64(lower_bound);
                 } else {
                     self.model.population[i][j] = x;
                 }
-                if x.is_nan() {
+                if x.value_f64().is_nan() {
                     panic!("A coefficient became NaN!");
                 }
             }
@@ -171,7 +179,7 @@ impl PSO {
     /// Updates the best found positions
     fn update_best_positions(&mut self) {
         for i in 0..self.best_f_values.len() {
-            let new = self.model.population_f_scores[i];
+            let new = self.model.get_f_values()[i];
             let old = self.best_f_values[i];
 
             if new < old {
@@ -179,8 +187,8 @@ impl PSO {
                 self.neigh_population[i] = self.model.population[i].clone();
             }
         }
-        self.best_f_trajectory.push(self.model.f_best);
-        self.best_x_trajectory.push(self.model.x_best.clone());
+        self.best_f_trajectory.push(self.model.get_f_best());
+        self.best_x_trajectory.push(self.model.get_x_best().clone());
     }
 
     /// Returns the neighborhood local best
@@ -230,7 +238,7 @@ impl PSO {
     }
 
     /// Returns the indices that would sort a vector
-    fn argsort(v: &Vec<f64>) -> Vec<usize> {
+    fn argsort(v: &[f64]) -> Vec<usize> {
         let mut idx = (0..v.len()).collect::<Vec<_>>();
         idx.sort_by(|&i, &j| v[i].partial_cmp(&v[j]).expect("NaN"));
         idx
@@ -259,7 +267,7 @@ impl PSO {
             .iter()
             .map(|x| {
                 x.iter()
-                    .map(|coef: &f64| coef.to_string())
+                    .map(|coef| coef.to_string())
                     .collect::<Vec<String>>()
                     .join(", ")
             })
diff --git a/tests/pso.rs b/tests/pso.rs
index 8ddc8d8..73219dd 100644
--- a/tests/pso.rs
+++ b/tests/pso.rs
@@ -1,10 +1,15 @@
+use core::f64;
+
 use pso_rs::*;
 
 #[test]
 fn it_runs_non_parallel() {
-    fn rosenbrock(p: &Particle, _flat_dim: usize, dimensions: &Vec<usize>) -> f64 {
+    fn rosenbrock(p: &Particle, _flat_dim: usize, dimensions: &[usize]) -> f64 {
         (0..dimensions[0] - 1)
-            .map(|i| 100.0 * ((p[i + 1] - p[i]).powf(2.0)).powf(2.0) + (1.0 - p[i]).powf(2.0))
+            .map(|i| {
+                100.0 * ((p[i + 1] - p[i]).value_f64().powf(2.0)).powf(2.0)
+                    + (1.0 - p[i].value_f64()).powf(2.0)
+            })
             .sum()
     }
 
@@ -15,18 +20,18 @@ fn it_runs_non_parallel() {
         parallelize: false,
         ..Config::default()
     };
-    let pso = pso_rs::run(config, rosenbrock, None, None).unwrap();
+    let pso = pso_rs::run(config, rosenbrock, None, None, None).unwrap();
 
     let mut model = pso.model;
 
-    model.population[0][0] = 2.0;
-    model.population[0][1] = -2.0;
+    model.population[0][0] = NumericKind::ValueF64(2.0);
+    model.population[0][1] = NumericKind::ValueF64(-2.0);
     model.get_f_values();
 
     assert_ne!(model.get_f_best(), 0.0);
 
-    model.population[0][0] = 1.0;
-    model.population[0][1] = 1.0;
+    model.population[0][0] = NumericKind::ValueF64(1.0);
+    model.population[0][1] = NumericKind::ValueF64(1.0);
     model.get_f_values();
 
     assert_eq!(model.get_f_best(), 0.0);
@@ -34,9 +39,12 @@ fn it_runs_non_parallel() {
 
 #[test]
 fn it_computes_correct_minimum_rosenbrock_2d() {
-    fn rosenbrock(p: &Particle, _flat_dim: usize, dimensions: &Vec<usize>) -> f64 {
+    fn rosenbrock(p: &Particle, _flat_dim: usize, dimensions: &[usize]) -> f64 {
         (0..dimensions[0] - 1)
-            .map(|i| 100.0 * ((p[i + 1] - p[i]).powf(2.0)).powf(2.0) + (1.0 - p[i]).powf(2.0))
+            .map(|i| {
+                100.0 * ((p[i + 1] - p[i]).value_f64().powf(2.0)).powf(2.0)
+                    + (1.0 - p[i].value_f64()).powf(2.0)
+            })
             .sum()
     }
 
@@ -46,18 +54,18 @@ fn it_computes_correct_minimum_rosenbrock_2d() {
         progress_bar: false,
         ..Config::default()
     };
-    let pso = pso_rs::run(config, rosenbrock, None, None).unwrap();
+    let pso = pso_rs::run(config, rosenbrock, None, None, None).unwrap();
 
     let mut model = pso.model;
 
-    model.population[0][0] = 2.0;
-    model.population[0][1] = -2.0;
+    model.population[0][0] = NumericKind::ValueF64(2.0);
+    model.population[0][1] = NumericKind::ValueF64(-2.0);
     model.get_f_values();
 
     assert_ne!(model.get_f_best(), 0.0);
 
-    model.population[0][0] = 1.0;
-    model.population[0][1] = 1.0;
+    model.population[0][0] = NumericKind::ValueF64(1.0);
+    model.population[0][1] = NumericKind::ValueF64(1.0);
     model.get_f_values();
 
     assert_eq!(model.get_f_best(), 0.0);
@@ -65,9 +73,12 @@ fn it_computes_correct_minimum_rosenbrock_2d() {
 
 #[test]
 fn it_computes_correct_minimum_rosenbrock_3d() {
-    fn rosenbrock(p: &Particle, _flat_dim: usize, dimensions: &Vec<usize>) -> f64 {
+    fn rosenbrock(p: &Particle, _flat_dim: usize, dimensions: &[usize]) -> f64 {
         (0..dimensions[0] - 1)
-            .map(|i| 100.0 * ((p[i + 1] - p[i]).powf(2.0)).powf(2.0) + (1.0 - p[i]).powf(2.0))
+            .map(|i| {
+                100.0 * ((p[i + 1] - p[i]).value_f64().powf(2.0)).powf(2.0)
+                    + (1.0 - p[i].value_f64()).powf(2.0)
+            })
             .sum()
     }
 
@@ -79,20 +90,20 @@ fn it_computes_correct_minimum_rosenbrock_3d() {
         progress_bar: false,
         ..Config::default()
     };
-    let pso = pso_rs::run(config, rosenbrock, None, None).unwrap();
+    let pso = pso_rs::run(config, rosenbrock, None, None, None).unwrap();
 
     let mut model = pso.model;
 
-    model.population[0][0] = 2.0;
-    model.population[0][1] = -2.0;
-    model.population[0][2] = -2.0;
+    model.population[0][0] = NumericKind::ValueF64(2.0);
+    model.population[0][1] = NumericKind::ValueF64(-2.0);
+    model.population[0][2] = NumericKind::ValueF64(-2.0);
     model.get_f_values();
 
     assert_ne!(model.get_f_best(), 0.0);
 
-    model.population[0][0] = 1.0;
-    model.population[0][1] = 1.0;
-    model.population[0][2] = 1.0;
+    model.population[0][0] = NumericKind::ValueF64(1.0);
+    model.population[0][1] = NumericKind::ValueF64(1.0);
+    model.population[0][2] = NumericKind::ValueF64(1.0);
     model.get_f_values();
 
     assert_eq!(model.get_f_best(), 0.0);
@@ -104,7 +115,7 @@ fn it_computes_correct_minimum_e_lj() {
     fn l2(x_i: Particle, x_j: Particle, particle_dim: usize) -> f64 {
         let mut sum: f64 = 0.0;
         for i in 0..particle_dim {
-            sum += (x_i[i] - x_j[i]).powf(2.0);
+            sum += (x_i[i] - x_j[i]).value_f64().powf(2.0);
         }
         sum.sqrt()
     }
@@ -116,7 +127,7 @@ fn it_computes_correct_minimum_e_lj() {
     }
 
     /// Get potential energy of a cluster of particles
-    fn e_lj(particle: &Particle, _flat_dim: usize, particle_dims: &Vec<usize>) -> f64 {
+    fn e_lj(particle: &Particle, _flat_dim: usize, particle_dims: &[usize]) -> f64 {
         let mut sum = 0.0;
         for i in 0..particle_dims[0] - 1 {
             for j in (i + 1)..particle_dims[0] {
@@ -139,22 +150,62 @@ fn it_computes_correct_minimum_e_lj() {
         ..Config::default()
     };
 
-    let pso = pso_rs::run(config, e_lj, Some(|_| true), None).unwrap();
+    let pso = pso_rs::run(config, e_lj, None, None, None).unwrap();
 
     let mut model = pso.model;
 
-    model.population[0][0] = -0.3616353090;
-    model.population[0][1] = 0.0439914505;
-    model.population[0][2] = 0.5828840628;
-    model.population[0][3] = 0.2505889242;
-    model.population[0][4] = 0.6193583398;
-    model.population[0][5] = -0.1614607010;
-    model.population[0][6] = -0.4082757926;
-    model.population[0][7] = -0.2212115329;
-    model.population[0][8] = -0.5067996704;
-    model.population[0][9] = 0.5193221773;
-    model.population[0][10] = -0.4421382574;
-    model.population[0][11] = 0.0853763087;
+    model.population[0][0] = NumericKind::ValueF64(-0.361635309);
+    model.population[0][1] = NumericKind::ValueF64(0.0439914505);
+    model.population[0][2] = NumericKind::ValueF64(0.5828840628);
+    model.population[0][3] = NumericKind::ValueF64(0.2505889242);
+    model.population[0][4] = NumericKind::ValueF64(0.6193583398);
+    model.population[0][5] = NumericKind::ValueF64(-0.161460701);
+    model.population[0][6] = NumericKind::ValueF64(-0.4082757926);
+    model.population[0][7] = NumericKind::ValueF64(-0.2212115329);
+    model.population[0][8] = NumericKind::ValueF64(-0.5067996704);
+    model.population[0][9] = NumericKind::ValueF64(0.5193221773);
+    model.population[0][10] = NumericKind::ValueF64(-0.4421382574);
+    model.population[0][11] = NumericKind::ValueF64(0.0853763087);
     model.get_f_values();
     assert!(model.get_f_best() < -5.9999999);
 }
+
+#[test]
+fn it_computes_boolean_sat_and_roots() {
+    fn bsat(p: &Particle, _flat_dim: usize, dimensions: &[usize]) -> f64 {
+        assert_eq!(dimensions[0], 3);
+        if !p[0].value_bool() {
+            f64::MAX
+        } else {
+            //(x-3)*(y-4)
+            ((p[1].value_f64() - 3.0) * (p[2].value_f64() - 4.0)).abs()
+        }
+    }
+
+    let config = Config {
+        dimensions: vec![3],
+        t_max: 1,
+        bounds: vec![(0.0, 1.0), (-5.0, 5.0), (-5.0, 5.0)],
+        population_size: 3,
+        progress_bar: true,
+        ..Config::default()
+    };
+    let pso = pso_rs::run(config, bsat, None, None, None).unwrap();
+
+    let mut model = pso.model;
+
+    model.population[0][0] = NumericKind::ValueBool(false);
+    model.population[0][1] = NumericKind::ValueF64(-2.0);
+    model.population[0][2] = NumericKind::ValueF64(-2.0);
+
+    model.population[1][0] = NumericKind::ValueBool(false);
+    model.population[1][1] = NumericKind::ValueF64(-2.0);
+    model.population[1][2] = NumericKind::ValueF64(2.0);
+
+    model.population[2][0] = NumericKind::ValueBool(true);
+    model.population[2][1] = NumericKind::ValueF64(2.0);
+    model.population[2][2] = NumericKind::ValueF64(-2.0);
+    model.get_f_values();
+
+    assert_ne!(model.get_f_best(), 0.0);
+}

From 3523e6227c1c3ebe815e6e5e90c69a424085382e Mon Sep 17 00:00:00 2001
From: Muthu Annamalai <muthiah.annamalai@tetramem.com>
Date: Sun, 19 Apr 2026 01:06:08 -0700
Subject: [PATCH 3/5] Add stylistic changes and fix lint / clippy issues

---
 src/lib.rs   |  6 +++---
 src/model.rs | 57 +++++++++++++++++++---------------------------------
 src/pso.rs   |  6 +++---
 3 files changed, 27 insertions(+), 42 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index ec0872c..d9814c3 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -205,12 +205,12 @@ use std::error::Error;
 pub fn run(
     config: Config,
     obj_f: fn(&Particle, usize, &[usize]) -> f64,
-    cast_f: Option<fn(&Vec<f64>) -> Vec<NumericKind>>,
+    cast_f: Option<CastFunctionT>,
     terminate_f: Option<fn(f64) -> bool>,
     seed: Option<u64>,
 ) -> Result<PSO, Box<dyn Error>> {
     assert_config(&config)?;
-    let mut pso = init(config, obj_f, cast_f, seed).unwrap();
+    let mut pso = init(config, obj_f, cast_f, seed)?;
     let term_condition = match terminate_f {
         Some(terminate_f) => terminate_f,
         None => |_| false,
@@ -225,7 +225,7 @@ pub fn run(
 pub fn init(
     config: Config,
     obj_f: fn(&Particle, usize, &[usize]) -> f64,
-    cast_f: Option<fn(&Vec<f64>) -> Vec<NumericKind>>,
+    cast_f: Option<CastFunctionT>,
     seed: Option<u64>,
 ) -> Result<PSO, &'static str> {
     assert_config(&config)?;
diff --git a/src/model.rs b/src/model.rs
index 9bf672c..fdd1894 100644
--- a/src/model.rs
+++ b/src/model.rs
@@ -11,13 +11,14 @@ pub enum NumericKind {
     ValueI32(i32),
 }
 
-impl ToString for NumericKind {
-    fn to_string(&self) -> String {
-        match self {
+impl std::fmt::Display for NumericKind {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let strval = match self {
             NumericKind::ValueF64(v) => v.to_string(),
             NumericKind::ValueBool(v) => v.to_string(),
             NumericKind::ValueI32(v) => v.to_string(),
-        }
+        };
+        write!(f, "{}", strval)
     }
 }
 
@@ -25,15 +26,9 @@ impl Mul for NumericKind {
     type Output = Self;
     fn mul(self, other: Self) -> Self {
         match self {
-            NumericKind::ValueF64(v) => NumericKind::ValueF64 {
-                0: v * other.value_f64(),
-            },
-            NumericKind::ValueBool(v) => NumericKind::ValueBool {
-                0: v && other.value_bool(),
-            },
-            NumericKind::ValueI32(v) => NumericKind::ValueI32 {
-                0: v * other.value_i32(),
-            },
+            NumericKind::ValueF64(v) => NumericKind::ValueF64(v * other.value_f64()),
+            NumericKind::ValueBool(v) => NumericKind::ValueBool(v && other.value_bool()),
+            NumericKind::ValueI32(v) => NumericKind::ValueI32(v * other.value_i32()),
         }
     }
 }
@@ -42,15 +37,9 @@ impl Add for NumericKind {
     type Output = Self;
     fn add(self, other: Self) -> Self {
         match self {
-            NumericKind::ValueF64(v) => NumericKind::ValueF64 {
-                0: v + other.value_f64(),
-            },
-            NumericKind::ValueBool(v) => NumericKind::ValueBool {
-                0: v || other.value_bool(),
-            },
-            NumericKind::ValueI32(v) => NumericKind::ValueI32 {
-                0: v + other.value_i32(),
-            },
+            NumericKind::ValueF64(v) => NumericKind::ValueF64(v + other.value_f64()),
+            NumericKind::ValueBool(v) => NumericKind::ValueBool(v || other.value_bool()),
+            NumericKind::ValueI32(v) => NumericKind::ValueI32(v + other.value_i32()),
         }
     }
 }
@@ -59,15 +48,9 @@ impl Sub for NumericKind {
     type Output = Self;
     fn sub(self, other: Self) -> Self {
         match self {
-            NumericKind::ValueF64(v) => NumericKind::ValueF64 {
-                0: v - other.value_f64(),
-            },
-            NumericKind::ValueBool(v) => NumericKind::ValueBool {
-                0: v ^ other.value_bool(),
-            },
-            NumericKind::ValueI32(v) => NumericKind::ValueI32 {
-                0: v - other.value_i32(),
-            },
+            NumericKind::ValueF64(v) => NumericKind::ValueF64(v - other.value_f64()),
+            NumericKind::ValueBool(v) => NumericKind::ValueBool(v ^ other.value_bool()),
+            NumericKind::ValueI32(v) => NumericKind::ValueI32(v - other.value_i32()),
         }
     }
 }
@@ -88,7 +71,7 @@ impl NumericKind {
         match self {
             NumericKind::ValueF64(v) => *v as i32 != 0,
             NumericKind::ValueBool(v) => *v,
-            NumericKind::ValueI32(v) => *v as i32 != 0,
+            NumericKind::ValueI32(v) => *v != 0,
         }
     }
 
@@ -101,6 +84,8 @@ impl NumericKind {
     }
 }
 
+pub type CastFunctionT = fn(p: &Vec<f64>) -> Vec<NumericKind>;
+
 /// Model struct
 ///
 /// It takes in a `Config` instance and `fn` pointer to an objective function and defines a `run` method for running Particle Swarm Optimization.
@@ -112,8 +97,8 @@ pub struct Model {
     pub x_best: Particle,
     pub f_best: f64,
     pub seed: Option<u64>,
-    obj_f: fn(&Particle, usize, &[usize]) -> f64,
-    cast_f: Option<fn(&Vec<f64>) -> Vec<NumericKind>>, /* cast Particle -> Particle */
+    pub obj_f: fn(&Particle, usize, &[usize]) -> f64,
+    pub cast_f: Option<CastFunctionT>, /* cast Particle -> Particle */
 }
 
 impl Model {
@@ -121,7 +106,7 @@ impl Model {
     pub fn new(
         config: Config,
         obj_f: fn(p: &Particle, flat_dim: usize, dim: &[usize]) -> f64,
-        cast_f: Option<fn(p: &Vec<f64>) -> Vec<NumericKind>>,
+        cast_f: Option<CastFunctionT>,
         seed: Option<u64>,
     ) -> Model {
         // init population
@@ -140,7 +125,7 @@ impl Model {
             let mut particle: Particle = vec![];
             for flat_i in 0..flat_dim {
                 let true_i = flat_i % config.dimensions[config.dimensions.len() - 1];
-                if let Some(_) = seed {
+                if seed.is_some() {
                     particle.push(NumericKind::ValueF64(
                         rng.gen_range(config.bounds[true_i].0..config.bounds[true_i].1),
                     ));
diff --git a/src/pso.rs b/src/pso.rs
index a9e1c34..67cb3e0 100644
--- a/src/pso.rs
+++ b/src/pso.rs
@@ -33,7 +33,7 @@ impl PSO {
         let phi_squared = phi.powf(2.0);
         let tmp = phi_squared - (4.0 * phi);
         let tmp = tmp.sqrt();
-        let chi = ((2.0 / (2.0 - phi - tmp)) as f64).abs();
+        let chi = 2.0 / (2.0 - phi - tmp);
         let v_max = model.config.alpha * 5.0;
         let neighborhoods = Self::create_neighborhoods(&model);
 
@@ -47,7 +47,7 @@ impl PSO {
         for _ in 0..model.config.population_size {
             let mut tmp = vec![];
             for _ in 0..model.flat_dim {
-                if let Some(_) = seed {
+                if seed.is_some() {
                     tmp.push(NumericKind::ValueF64(seeded_rng.gen_range(-v_max..v_max)));
                 } else {
                     tmp.push(NumericKind::ValueF64(rng.gen_range(-v_max..v_max)));
@@ -130,7 +130,7 @@ impl PSO {
             for j in 0..self.model.flat_dim {
                 let r1: f64;
                 let r2: f64;
-                if let Some(_) = self.seed {
+                if self.seed.is_some() {
                     r1 = self.seeded_rng.gen_range(-1.0..1.0);
                     r2 = self.seeded_rng.gen_range(-1.0..1.0);
                 } else {

From 9a0f960ff9719d072697177b12e73334ad134eae Mon Sep 17 00:00:00 2001
From: Muthu Annamalai <muthiah.annamalai@tetramem.com>
Date: Mon, 18 May 2026 22:00:08 -0700
Subject: [PATCH 4/5] [model evaluation]: reduce it during
 update_best_positions calculation - fix bug where eval fitness function is
 called 2x more; eliminate one source of clippy warning - add terminator
 function - we see the same results and times on this branch as the
 origin/main branch

---
 src/bin/e_lj.rs |  1 +
 src/bin/main.rs |  3 ++-
 src/pso.rs      | 35 +++++++++++++++++------------------
 3 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/src/bin/e_lj.rs b/src/bin/e_lj.rs
index b832178..de7ad15 100644
--- a/src/bin/e_lj.rs
+++ b/src/bin/e_lj.rs
@@ -25,6 +25,7 @@ fn main() {
         lr,
         bounds,
         t_max,
+        parallelize: true,
         ..Config::default()
     };
     use std::time::Instant;
diff --git a/src/bin/main.rs b/src/bin/main.rs
index de4dbad..6d27c46 100644
--- a/src/bin/main.rs
+++ b/src/bin/main.rs
@@ -8,6 +8,7 @@ fn main() {
         population_size: 100,
         bounds: vec![(-10.0, 10.0); N_DIMENSIONS],
         t_max: 1e7 as usize,
+        parallelize: true,
         ..Config::default()
     };
     use std::time::Instant;
@@ -23,7 +24,7 @@ fn main() {
             }
             r
         }),
-        None,
+        Some(|f_best| f_best < 1e-4),
         Some(123456u64),
     )
     .unwrap();
diff --git a/src/pso.rs b/src/pso.rs
index 67cb3e0..3f66883 100644
--- a/src/pso.rs
+++ b/src/pso.rs
@@ -33,7 +33,7 @@ impl PSO {
         let phi_squared = phi.powf(2.0);
         let tmp = phi_squared - (4.0 * phi);
         let tmp = tmp.sqrt();
-        let chi = 2.0 / (2.0 - phi - tmp);
+        let chi = 2.0 / (2.0 - phi - tmp).abs();
         let v_max = model.config.alpha * 5.0;
         let neighborhoods = Self::create_neighborhoods(&model);
 
@@ -103,8 +103,8 @@ impl PSO {
             self.update_velocity_and_pos();
 
             // Evaluate & update best
-            self.model.get_f_values();
-            self.update_best_positions();
+            let new_best_f_values = self.model.get_f_values();
+            self.update_best_positions(&new_best_f_values);
 
             self.model.population = self.model.population.clone();
             k += pop_size;
@@ -137,6 +137,7 @@ impl PSO {
                     r1 = self.rng.gen_range(-1.0..1.0);
                     r2 = self.rng.gen_range(-1.0..1.0);
                 }
+
                 let cog = self.model.config.c1
                     * r1
                     * (self.neigh_population[i][j] - self.model.population[i][j]).value_f64();
@@ -177,13 +178,11 @@ impl PSO {
     }
 
     /// Updates the best found positions
-    fn update_best_positions(&mut self) {
-        for i in 0..self.best_f_values.len() {
-            let new = self.model.get_f_values()[i];
-            let old = self.best_f_values[i];
-
-            if new < old {
-                self.best_f_values[i] = new;
+    fn update_best_positions(&mut self, new_best_f_values: &Vec<f64>) {
+        for (i, old) in self.best_f_values.iter_mut().enumerate() {
+            let new = new_best_f_values[i];
+            if new < *old {
+                *old = new;
                 self.neigh_population[i] = self.model.population[i].clone();
             }
         }
@@ -224,14 +223,14 @@ impl PSO {
                 }
             }
             NeighborhoodType::Gbest => {
-                neighborhoods = vec![];
-                for _ in 0..model.config.population_size {
-                    let mut tmp = vec![];
-                    for j in 0..model.config.population_size {
-                        tmp.push(j);
-                    }
-                    neighborhoods.push(tmp);
-                }
+                neighborhoods = (0..model.config.population_size)
+                    .into_iter()
+                    .map(|_| {
+                        (0..model.config.population_size)
+                            .into_iter()
+                            .collect::<Vec<usize>>()
+                    })
+                    .collect::<Vec<Vec<usize>>>();
             }
         }
         neighborhoods

From 0c95f1b9cc3d1c41616b124db80fc07de86635b4 Mon Sep 17 00:00:00 2001
From: Muthu Annamalai <muthiah.annamalai@tetramem.com>
Date: Thu, 21 May 2026 10:33:52 -0700
Subject: [PATCH 5/5] Add ObjectiveFunction trait so functors can be used with
 self contained state; update tests and examples

---
 src/bin/e_lj.rs |  28 +++++-----
 src/bin/main.rs |  13 +++--
 src/lib.rs      |  55 ++++++++++----------
 src/model.rs    |  30 +++++++++--
 src/pso.rs      |   9 +---
 tests/pso.rs    | 133 +++++++++++++++++++++++-------------------------
 6 files changed, 145 insertions(+), 123 deletions(-)

diff --git a/src/bin/e_lj.rs b/src/bin/e_lj.rs
index de7ad15..cb6f14d 100644
--- a/src/bin/e_lj.rs
+++ b/src/bin/e_lj.rs
@@ -30,6 +30,7 @@ fn main() {
     };
     use std::time::Instant;
     let before = Instant::now();
+    let e_lj = Box::new(Elj {});
     match pso_rs::run(config, e_lj, None, None, None) {
         Ok(pso) => {
             println!("Elapsed time: {:.2?}", before.elapsed());
@@ -74,20 +75,23 @@ fn v_ij(x_i: Particle, x_j: Particle, particle_dim: usize) -> f64 {
 }
 
 /// Get potential energy of a cluster of particles
-fn e_lj(particle: &Particle, _flat_dim: usize, particle_dims: &[usize]) -> f64 {
-    let mut sum = 0.0;
-    for i in 0..particle_dims[0] - 1 {
-        for j in (i + 1)..particle_dims[0] {
-            let true_i = i * particle_dims[1];
-            let true_j = j * particle_dims[1];
-            sum += v_ij(
-                particle[true_i..true_i + particle_dims[1]].to_vec(),
-                particle[true_j..true_j + particle_dims[1]].to_vec(),
-                particle_dims[1],
-            );
+struct Elj;
+impl ObjectiveFunction for Elj {
+    fn evaluate(&self, particle: &Particle, _flat_dim: usize, particle_dims: &[usize]) -> f64 {
+        let mut sum = 0.0;
+        for i in 0..particle_dims[0] - 1 {
+            for j in (i + 1)..particle_dims[0] {
+                let true_i = i * particle_dims[1];
+                let true_j = j * particle_dims[1];
+                sum += v_ij(
+                    particle[true_i..true_i + particle_dims[1]].to_vec(),
+                    particle[true_j..true_j + particle_dims[1]].to_vec(),
+                    particle_dims[1],
+                );
+            }
         }
+        4.0 * sum
     }
-    4.0 * sum
 }
 
 fn reshape(particle: &Particle, particle_dims: &[usize]) -> Vec<Vec<NumericKind>> {
diff --git a/src/bin/main.rs b/src/bin/main.rs
index 6d27c46..540c90a 100644
--- a/src/bin/main.rs
+++ b/src/bin/main.rs
@@ -13,7 +13,7 @@ fn main() {
     };
     use std::time::Instant;
     let before = Instant::now();
-
+    let sum_squares = Box::new(SumOfSquares {});
     let pso: pso::PSO = pso_rs::run(
         config,
         sum_squares,
@@ -34,8 +34,11 @@ fn main() {
     println!("Found minimizer: {:#?} ", model.get_x_best());
 }
 
-fn sum_squares(p: &Particle, _flat_dim: usize, dimensions: &[usize]) -> f64 {
-    (0..dimensions[0])
-        .map(|i| i as f64 * p[i].value_f64().powf(2.0))
-        .sum()
+struct SumOfSquares;
+impl ObjectiveFunction for SumOfSquares {
+    fn evaluate(&self, particle: &Particle, _flat_dim: usize, dimensions: &[usize]) -> f64 {
+        (0..dimensions[0])
+            .map(|i| i as f64 * particle[i].value_f64().powf(2.0))
+            .sum()
+    }
 }
diff --git a/src/lib.rs b/src/lib.rs
index d9814c3..96f94fd 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -20,17 +20,15 @@
 //! use pso_rs::*;
 //!
 //! // define objective function (d-dimensional Rosenbrock)
-//! fn objective_function(
-//!     p: &Particle,
-//!     _flat_dim: usize,
-//!     dimensions: &[usize]
-//! ) -> f64 {
-//!     (0..dimensions[0] - 1).map(|i| {
+//! struct ObjectiveObj;
+//! impl ObjectiveFunction for ObjectiveObj {
+//!     fn evaluate(&self, p: &Particle, _flat_dim: usize, dimensions: &[usize]) -> f64 {
+//!         (0..dimensions[0] - 1).map(|i| {
 //!         100.0 * ((p[i+1]-p[i]).value_f64().powf(2.0)).powf(2.0)
 //!             + (1.0-p[i].value_f64()).powf(2.0)
 //!     }).sum()
+//!     }
 //! }
-//!
 //! // define a termination condition (optional)
 //! fn terminate(f_best: f64) -> bool {
 //!     f_best - (0.0) < 1e-4
@@ -49,7 +47,7 @@
 //!
 //! let pso = pso_rs::run(
 //!     config,
-//!     objective_function,
+//!     Box::new(ObjectiveObj{}),
 //!     None,
 //!     Some(terminate),
 //!     None
@@ -65,18 +63,16 @@
 //! use pso_rs::*;
 //!
 //! // define objective function (d-dimensional Rosenbrock)
-//! fn objective_function(
-//!     p: &Particle,
-//!     _flat_dim: usize,
-//!     dimensions: &[usize]
-//! ) -> f64 {
+//! struct ObjectiveObj;
+//! impl ObjectiveFunction for ObjectiveObj {
+//!     fn evaluate(&self, p: &Particle, _flat_dim: usize, dimensions: &[usize]) -> f64 {
 //!     (0..dimensions[0] - 1).map(|i| {
 //!         100.0 * ((p[i+1]-p[i]).value_f64().powf(2.0)).powf(2.0)
 //!             + (1.0-p[i].value_f64()).powf(2.0)
 //!     }).sum()
+//!  }
 //! }
 //!
-//!
 //! let config = Config {
 //!     dimensions: vec![2],
 //!     bounds: vec![(-5.0, 10.0); 2],
@@ -86,7 +82,7 @@
 //!
 //! let mut pso = pso_rs::init(
 //!     config,
-//!     objective_function,
+//!     Box::new(ObjectiveObj{}),
 //!     None,
 //!     None
 //! ).unwrap();
@@ -98,6 +94,13 @@
 //! println!("Found minimum: {:#?} ", model.get_f_best());
 //! println!("Minimizer: {:#?}", model.get_x_best());
 //! ```
+//! struct XSumSquares;
+//! impl ObjectiveFunction for XSumSquares {
+//!    fn evaluate(&self, particle: &Particle, _flat_dim: usize, dimensions: &[usize]) -> f64 {
+//!         (0..dimensions[0]).map(|i| i as f64 * particle[i].powf(2.0)).sum()
+//!    };
+//! }
+//!
 //! //example shows how to use lambda functions to model objective
 //! fn main() {
 //!    let config = Config {
@@ -109,11 +112,8 @@
 //!    };
 //!    use std::time::Instant;
 //!    let before = Instant::now();
-//!    let xsum_squares = |p: &Particle, _flat_dim: usize, dimensions: &[usize]| -> f64 {
-//!         (0..dimensions[0]).map(|i| i as f64 * p[i].powf(2.0)).sum()
-//!     };
 //!     
-//!    let pso = pso_rs::run(config, xsum_squares, None, Some(|f_best| f_best < 1e-4),Some(123456u64)).unwrap();
+//!    let pso = pso_rs::run(config, Box::new(XSumSquares), None, Some(|f_best| f_best < 1e-4),Some(123456u64)).unwrap();
 //!    println!("Elapsed time: {:.2?}", before.elapsed());
 //!    let model = pso.model;
 //!    println!("Found minimum: {:#?} ", model.get_f_best());
@@ -156,16 +156,19 @@
 //! }
 //!
 //! // used in the objective function
-//! fn objective_function(
-//!     p: &Particle,
+//! struct ObjectiveObj;
+//! impl ObjectiveFunction for ObjectiveObj {
+//! fn evaluate(
+//!     &self,
+//!     particle: &Particle,
 //!     _flat_dim: usize,
 //!     dimensions: &[usize]
 //! ) -> f64 {
-//!     let _reshaped_particle = reshape(p, dimensions);
+//!     let _reshaped_particle = reshape(particle, dimensions);
 //!     /* Do stuff */
 //!     0.0
 //! }
-//!
+//! }
 //! let config = Config {
 //!     dimensions: vec![20, 3],
 //!     bounds: vec![(-2.5, 2.5); 3],
@@ -175,7 +178,7 @@
 //!
 //! let pso = pso_rs::run(
 //!     config,
-//!     objective_function,
+//!     Box::new(ObjectiveObj{}),
 //!     None,
 //!     None,
 //!     None
@@ -204,7 +207,7 @@ use std::error::Error;
 /// Panics if any particle coefficient becomes NaN (usually because of bad parameterization, e.g. c1 + c2 < 4)
 pub fn run(
     config: Config,
-    obj_f: fn(&Particle, usize, &[usize]) -> f64,
+    obj_f: Box<dyn ObjectiveFunction>,
     cast_f: Option<CastFunctionT>,
     terminate_f: Option<fn(f64) -> bool>,
     seed: Option<u64>,
@@ -224,7 +227,7 @@ pub fn run(
 /// Useful for initializing an instance for running at a later time
 pub fn init(
     config: Config,
-    obj_f: fn(&Particle, usize, &[usize]) -> f64,
+    obj_f: Box<dyn ObjectiveFunction>,
     cast_f: Option<CastFunctionT>,
     seed: Option<u64>,
 ) -> Result<PSO, &'static str> {
diff --git a/src/model.rs b/src/model.rs
index fdd1894..dc4c73e 100644
--- a/src/model.rs
+++ b/src/model.rs
@@ -11,6 +11,24 @@ pub enum NumericKind {
     ValueI32(i32),
 }
 
+impl From<f64> for NumericKind {
+    fn from(value: f64) -> Self {
+        NumericKind::ValueF64(value)
+    }
+}
+
+impl From<bool> for NumericKind {
+    fn from(value: bool) -> Self {
+        NumericKind::ValueBool(value)
+    }
+}
+
+impl From<i32> for NumericKind {
+    fn from(value: i32) -> Self {
+        NumericKind::ValueI32(value)
+    }
+}
+
 impl std::fmt::Display for NumericKind {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let strval = match self {
@@ -86,6 +104,10 @@ impl NumericKind {
 
 pub type CastFunctionT = fn(p: &Vec<f64>) -> Vec<NumericKind>;
 
+pub trait ObjectiveFunction: Send + Sync {
+    fn evaluate(&self, p: &Particle, flat_dim: usize, dimensions: &[usize]) -> f64;
+}
+
 /// Model struct
 ///
 /// It takes in a `Config` instance and `fn` pointer to an objective function and defines a `run` method for running Particle Swarm Optimization.
@@ -97,7 +119,7 @@ pub struct Model {
     pub x_best: Particle,
     pub f_best: f64,
     pub seed: Option<u64>,
-    pub obj_f: fn(&Particle, usize, &[usize]) -> f64,
+    pub obj_f: Box<dyn ObjectiveFunction>,
     pub cast_f: Option<CastFunctionT>, /* cast Particle -> Particle */
 }
 
@@ -105,7 +127,7 @@ impl Model {
     /// Creates a new Model instance
     pub fn new(
         config: Config,
-        obj_f: fn(p: &Particle, flat_dim: usize, dim: &[usize]) -> f64,
+        obj_f: Box<dyn ObjectiveFunction>,
         cast_f: Option<CastFunctionT>,
         seed: Option<u64>,
     ) -> Model {
@@ -175,7 +197,7 @@ impl Model {
             let iter = self.population.par_iter();
             self.population_f_scores = iter
                 .map(|particle| {
-                    (self.obj_f)(particle, self.flat_dim, &self.config.dimensions)
+                    (*self.obj_f).evaluate(particle, self.flat_dim, &self.config.dimensions)
                     // self.population_f_scores[i] = f_score;
                 })
                 .collect();
@@ -183,7 +205,7 @@ impl Model {
             let iter = self.population.iter();
             self.population_f_scores = iter
                 .map(|particle| {
-                    (self.obj_f)(particle, self.flat_dim, &self.config.dimensions)
+                    (*self.obj_f).evaluate(particle, self.flat_dim, &self.config.dimensions)
                     // self.population_f_scores[i] = f_score;
                 })
                 .collect();
diff --git a/src/pso.rs b/src/pso.rs
index 3f66883..2f2a03d 100644
--- a/src/pso.rs
+++ b/src/pso.rs
@@ -178,7 +178,7 @@ impl PSO {
     }
 
     /// Updates the best found positions
-    fn update_best_positions(&mut self, new_best_f_values: &Vec<f64>) {
+    fn update_best_positions(&mut self, new_best_f_values: &[f64]) {
         for (i, old) in self.best_f_values.iter_mut().enumerate() {
             let new = new_best_f_values[i];
             if new < *old {
@@ -224,12 +224,7 @@ impl PSO {
             }
             NeighborhoodType::Gbest => {
                 neighborhoods = (0..model.config.population_size)
-                    .into_iter()
-                    .map(|_| {
-                        (0..model.config.population_size)
-                            .into_iter()
-                            .collect::<Vec<usize>>()
-                    })
+                    .map(|_| (0..model.config.population_size).collect::<Vec<usize>>())
                     .collect::<Vec<Vec<usize>>>();
             }
         }
diff --git a/tests/pso.rs b/tests/pso.rs
index 73219dd..af9d4cd 100644
--- a/tests/pso.rs
+++ b/tests/pso.rs
@@ -2,17 +2,20 @@ use core::f64;
 
 use pso_rs::*;
 
-#[test]
-fn it_runs_non_parallel() {
-    fn rosenbrock(p: &Particle, _flat_dim: usize, dimensions: &[usize]) -> f64 {
+struct Rosenbrock;
+impl ObjectiveFunction for Rosenbrock {
+    fn evaluate(&self, particle: &Particle, _flat_dim: usize, dimensions: &[usize]) -> f64 {
         (0..dimensions[0] - 1)
             .map(|i| {
-                100.0 * ((p[i + 1] - p[i]).value_f64().powf(2.0)).powf(2.0)
-                    + (1.0 - p[i].value_f64()).powf(2.0)
+                100.0 * ((particle[i + 1] - particle[i]).value_f64().powf(2.0)).powf(2.0)
+                    + (1.0 - particle[i].value_f64()).powf(2.0)
             })
             .sum()
     }
+}
 
+#[test]
+fn it_runs_non_parallel() {
     let config = Config {
         t_max: 1,
         population_size: 1,
@@ -20,12 +23,13 @@ fn it_runs_non_parallel() {
         parallelize: false,
         ..Config::default()
     };
+    let rosenbrock = Box::new(Rosenbrock {});
     let pso = pso_rs::run(config, rosenbrock, None, None, None).unwrap();
 
     let mut model = pso.model;
 
-    model.population[0][0] = NumericKind::ValueF64(2.0);
-    model.population[0][1] = NumericKind::ValueF64(-2.0);
+    model.population[0][0] = 2.0.into();
+    model.population[0][1] = (-2.0).into();
     model.get_f_values();
 
     assert_ne!(model.get_f_best(), 0.0);
@@ -39,33 +43,24 @@ fn it_runs_non_parallel() {
 
 #[test]
 fn it_computes_correct_minimum_rosenbrock_2d() {
-    fn rosenbrock(p: &Particle, _flat_dim: usize, dimensions: &[usize]) -> f64 {
-        (0..dimensions[0] - 1)
-            .map(|i| {
-                100.0 * ((p[i + 1] - p[i]).value_f64().powf(2.0)).powf(2.0)
-                    + (1.0 - p[i].value_f64()).powf(2.0)
-            })
-            .sum()
-    }
-
     let config = Config {
         t_max: 1,
         population_size: 1,
         progress_bar: false,
         ..Config::default()
     };
-    let pso = pso_rs::run(config, rosenbrock, None, None, None).unwrap();
+    let pso = pso_rs::run(config, Box::new(Rosenbrock {}), None, None, None).unwrap();
 
     let mut model = pso.model;
 
-    model.population[0][0] = NumericKind::ValueF64(2.0);
-    model.population[0][1] = NumericKind::ValueF64(-2.0);
+    model.population[0][0] = (2.0).into();
+    model.population[0][1] = (-2.0).into();
     model.get_f_values();
 
     assert_ne!(model.get_f_best(), 0.0);
 
-    model.population[0][0] = NumericKind::ValueF64(1.0);
-    model.population[0][1] = NumericKind::ValueF64(1.0);
+    model.population[0][0] = (1.0).into();
+    model.population[0][1] = (1.0).into();
     model.get_f_values();
 
     assert_eq!(model.get_f_best(), 0.0);
@@ -73,15 +68,6 @@ fn it_computes_correct_minimum_rosenbrock_2d() {
 
 #[test]
 fn it_computes_correct_minimum_rosenbrock_3d() {
-    fn rosenbrock(p: &Particle, _flat_dim: usize, dimensions: &[usize]) -> f64 {
-        (0..dimensions[0] - 1)
-            .map(|i| {
-                100.0 * ((p[i + 1] - p[i]).value_f64().powf(2.0)).powf(2.0)
-                    + (1.0 - p[i].value_f64()).powf(2.0)
-            })
-            .sum()
-    }
-
     let config = Config {
         dimensions: vec![3],
         t_max: 1,
@@ -90,27 +76,28 @@ fn it_computes_correct_minimum_rosenbrock_3d() {
         progress_bar: false,
         ..Config::default()
     };
-    let pso = pso_rs::run(config, rosenbrock, None, None, None).unwrap();
+    let pso = pso_rs::run(config, Box::new(Rosenbrock {}), None, None, None).unwrap();
 
     let mut model = pso.model;
 
-    model.population[0][0] = NumericKind::ValueF64(2.0);
-    model.population[0][1] = NumericKind::ValueF64(-2.0);
-    model.population[0][2] = NumericKind::ValueF64(-2.0);
+    model.population[0][0] = (2.0).into();
+    model.population[0][1] = (-2.0).into();
+    model.population[0][2] = (-2.0).into();
     model.get_f_values();
 
     assert_ne!(model.get_f_best(), 0.0);
 
-    model.population[0][0] = NumericKind::ValueF64(1.0);
-    model.population[0][1] = NumericKind::ValueF64(1.0);
-    model.population[0][2] = NumericKind::ValueF64(1.0);
+    model.population[0][0] = (1.0).into();
+    model.population[0][1] = (1.0).into();
+    model.population[0][2] = (1.0).into();
     model.get_f_values();
 
     assert_eq!(model.get_f_best(), 0.0);
 }
 
-#[test]
-fn it_computes_correct_minimum_e_lj() {
+/// Get potential energy of a cluster of particles
+struct Elj;
+impl Elj {
     /// Get Euclidian distance of two particles
     fn l2(x_i: Particle, x_j: Particle, particle_dim: usize) -> f64 {
         let mut sum: f64 = 0.0;
@@ -122,18 +109,18 @@ fn it_computes_correct_minimum_e_lj() {
 
     /// Get potential energy of two particles
     fn v_ij(x_i: Particle, x_j: Particle, particle_dim: usize) -> f64 {
-        let denom: f64 = 1.0 / l2(x_i, x_j, particle_dim);
+        let denom: f64 = 1.0 / Elj::l2(x_i, x_j, particle_dim);
         denom.powf(12.0) - denom.powf(6.0)
     }
-
-    /// Get potential energy of a cluster of particles
-    fn e_lj(particle: &Particle, _flat_dim: usize, particle_dims: &[usize]) -> f64 {
+}
+impl ObjectiveFunction for Elj {
+    fn evaluate(&self, particle: &Particle, _flat_dim: usize, particle_dims: &[usize]) -> f64 {
         let mut sum = 0.0;
         for i in 0..particle_dims[0] - 1 {
             for j in (i + 1)..particle_dims[0] {
                 let true_i = i * particle_dims[1];
                 let true_j = j * particle_dims[1];
-                sum += v_ij(
+                sum += Elj::v_ij(
                     particle[true_i..true_i + particle_dims[1]].to_vec(),
                     particle[true_j..true_j + particle_dims[1]].to_vec(),
                     particle_dims[1],
@@ -142,6 +129,10 @@ fn it_computes_correct_minimum_e_lj() {
         }
         4.0 * sum
     }
+}
+
+#[test]
+fn it_computes_correct_minimum_e_lj() {
     let config = Config {
         dimensions: vec![4, 3],
         bounds: vec![(-2.5, 2.5); 3],
@@ -149,30 +140,30 @@ fn it_computes_correct_minimum_e_lj() {
         progress_bar: false,
         ..Config::default()
     };
-
+    let e_lj = Box::new(Elj);
     let pso = pso_rs::run(config, e_lj, None, None, None).unwrap();
 
     let mut model = pso.model;
 
-    model.population[0][0] = NumericKind::ValueF64(-0.361635309);
-    model.population[0][1] = NumericKind::ValueF64(0.0439914505);
-    model.population[0][2] = NumericKind::ValueF64(0.5828840628);
-    model.population[0][3] = NumericKind::ValueF64(0.2505889242);
-    model.population[0][4] = NumericKind::ValueF64(0.6193583398);
-    model.population[0][5] = NumericKind::ValueF64(-0.161460701);
-    model.population[0][6] = NumericKind::ValueF64(-0.4082757926);
-    model.population[0][7] = NumericKind::ValueF64(-0.2212115329);
-    model.population[0][8] = NumericKind::ValueF64(-0.5067996704);
-    model.population[0][9] = NumericKind::ValueF64(0.5193221773);
-    model.population[0][10] = NumericKind::ValueF64(-0.4421382574);
-    model.population[0][11] = NumericKind::ValueF64(0.0853763087);
+    model.population[0][0] = (-0.361635309f64).into();
+    model.population[0][1] = 0.0439914505f64.into();
+    model.population[0][2] = 0.5828840628f64.into();
+    model.population[0][3] = 0.2505889242f64.into();
+    model.population[0][4] = 0.6193583398f64.into();
+    model.population[0][5] = (-0.161460701f64).into();
+    model.population[0][6] = (-0.4082757926f64).into();
+    model.population[0][7] = (-0.2212115329f64).into();
+    model.population[0][8] = (-0.5067996704f64).into();
+    model.population[0][9] = 0.5193221773f64.into();
+    model.population[0][10] = (-0.4421382574f64).into();
+    model.population[0][11] = 0.0853763087f64.into();
     model.get_f_values();
     assert!(model.get_f_best() < -5.9999999);
 }
 
-#[test]
-fn it_computes_boolean_sat_and_roots() {
-    fn bsat(p: &Particle, _flat_dim: usize, dimensions: &[usize]) -> f64 {
+struct BSat;
+impl ObjectiveFunction for BSat {
+    fn evaluate(&self, p: &Particle, _flat_dim: usize, dimensions: &[usize]) -> f64 {
         assert_eq!(dimensions[0], 3);
         if !p[0].value_bool() {
             f64::MAX
@@ -181,7 +172,10 @@ fn it_computes_boolean_sat_and_roots() {
             ((p[1].value_f64() - 3.0) * (p[2].value_f64() - 4.0)).abs()
         }
     }
+}
 
+#[test]
+fn it_computes_boolean_sat_and_roots() {
     let config = Config {
         dimensions: vec![3],
         t_max: 1,
@@ -190,21 +184,22 @@ fn it_computes_boolean_sat_and_roots() {
         progress_bar: true,
         ..Config::default()
     };
+    let bsat = Box::new(BSat);
     let pso = pso_rs::run(config, bsat, None, None, None).unwrap();
 
     let mut model = pso.model;
 
-    model.population[0][0] = NumericKind::ValueBool(false);
-    model.population[0][1] = NumericKind::ValueF64(-2.0);
-    model.population[0][2] = NumericKind::ValueF64(-2.0);
+    model.population[0][0] = false.into();
+    model.population[0][1] = (-2.0).into();
+    model.population[0][2] = (-2.0).into();
 
-    model.population[1][0] = NumericKind::ValueBool(false);
-    model.population[1][1] = NumericKind::ValueF64(-2.0);
-    model.population[1][2] = NumericKind::ValueF64(2.0);
+    model.population[1][0] = false.into();
+    model.population[1][1] = (-2.0).into();
+    model.population[1][2] = (2.0).into();
 
-    model.population[2][0] = NumericKind::ValueBool(true);
-    model.population[2][1] = NumericKind::ValueF64(2.0);
-    model.population[2][2] = NumericKind::ValueF64(-2.0);
+    model.population[2][0] = true.into();
+    model.population[2][1] = (2.0).into();
+    model.population[2][2] = (-2.0).into();
     model.get_f_values();
 
     assert_ne!(model.get_f_best(), 0.0);