Rust标准惯性权重全局最优粒子群算法性能优化求助

问题背景 #

作为Rust新手，我实现了针对经典Sphere基准函数的标准惯性权重全局最优粒子群算法，但运行速度比等效Python版本慢3倍。由于维度和种群规模无法在编译期确定，我使用了向量；已知热循环中存在clone操作但不知替代方案，在不采用并行的前提下陷入性能提升瓶颈，恳请提供代码提速建议。

原始代码 #

use std::fmt::Display;

use rand::Rng;

struct ObjectiveFunctionStruct {
    name: String,
    function: fn(&Vec<f64>) -> f64,
    lower_bound: Vec<f64>,
    upper_bound: Vec<f64>,
}

struct Particle {
    position: Vec<f64>,
    velocity: Vec<f64>,
    personal_best_position: Vec<f64>,
    personal_best_fitness: f64,
}

impl Display for Particle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "(Position: {:?}, Velocity: {:?}, Personal Best Position: {:?}, Personal Best Fitness: {})", self.position, self.velocity, self.personal_best_position, self.personal_best_fitness)
    }
}

struct Swarm {
    w: f64,
    c1: f64,
    c2: f64,
    particles: Vec<Particle>,
    global_best_position: Vec<f64>,
    global_best_fitness: f64,
    objective_function_struct: ObjectiveFunctionStruct,
    rng_thread: rand::rngs::ThreadRng,
}

impl Swarm {
    fn new(w: f64, c1: f64, c2: f64, swarm_size: usize, objective_function_struct: ObjectiveFunctionStruct) -> Swarm {
        let dimension: usize = objective_function_struct.lower_bound.len();
        let mut particles: Vec<Particle> = Vec::new();
        let mut rng_thread: rand::rngs::ThreadRng = rand::thread_rng();
        let mut global_best_position: Vec<f64> = vec![0.0; dimension];
        let mut global_best_fitness: f64 = std::f64::MAX;
        for _ in 0..swarm_size {
            let mut particle: Particle = Particle {
                position: (0..dimension).map(|i| rng_thread.gen_range(objective_function_struct.lower_bound[i]..objective_function_struct.upper_bound[i])).collect(),
                velocity: vec![0.0; dimension],
                personal_best_position: vec![0.0; dimension],
                personal_best_fitness: std::f64::MAX,
            };
            particle.personal_best_position = particle.position.clone();
            particle.personal_best_fitness = (objective_function_struct.function)(&particle.position);
            if particle.personal_best_fitness < global_best_fitness {
                global_best_fitness = particle.personal_best_fitness;
                global_best_position = particle.personal_best_position.clone();
            }
            particles.push(particle);
        }
        Swarm {
            w,
            c1,
            c2,
            particles,
            global_best_position,
            global_best_fitness,
            objective_function_struct,
            rng_thread,
        }
    }

    fn update_particles(&mut self) {
        for particle in &mut self.particles {
            let dimension: usize = particle.position.len();
            for i in 0..dimension {
                particle.velocity[i] = self.w * particle.velocity[i] + self.c1 * self.rng_thread.gen_range(0.0..1.0) * (particle.personal_best_position[i] - particle.position[i]) + self.c2 * self.rng_thread.gen_range(0.0..1.0) * (self.global_best_position[i] - particle.position[i]);
                particle.position[i] += particle.velocity[i];
            }
            let fitness: f64 = (self.objective_function_struct.function)(&particle.position);
            if fitness < self.global_best_fitness {
                particle.personal_best_fitness = fitness;
                particle.personal_best_position = particle.position.clone();
                self.global_best_fitness = fitness;
                self.global_best_position = particle.position.clone();
            } else if fitness < particle.personal_best_fitness {
                particle.personal_best_fitness = fitness;
                particle.personal_best_position = particle.position.clone();
            }
        }
    }


    fn run(&mut self, iterations: usize) {
        for _ in 0..iterations {
            self.update_particles();
        }
    }

    fn print(&self) {
        println!("Global Best Position: {:?}", self.global_best_position);
        println!("Global Best Fitness: {}", self.global_best_fitness);
    }
}

fn sphere_function(x: &Vec<f64>) -> f64 {
    x.iter().map(|a: &f64| a.powi(2)).sum()
}

fn main() {
    use std::time::Instant;
    let now = Instant::now();
    let dim = 100;
    let objective_function_struct: ObjectiveFunctionStruct = ObjectiveFunctionStruct {
        name: "Sphere Function".to_string(),
        function: sphere_function,
        lower_bound: vec![-5.12; dim],
        upper_bound: vec![5.12; dim],
    };
    let mut swarm: Swarm = Swarm::new(0.729, 1.49445, 1.49445, 1000, objective_function_struct);
    swarm.run(10000);
    swarm.print();
    let elapsed = now.elapsed();
    println!("Elapsed: {} ms", elapsed.as_millis());
}

具体优化建议 #

1. 替换clone为内存原地复制，避免不必要的内存分配 #

原始代码中每次更新最优位置时调用Vec::clone()，会触发新内存分配和数据拷贝。由于personal_best_position和position的长度始终一致，可以用copy_from_slice直接在已有内存中复制数据，完全避免分配开销：

修改update_particles中的克隆逻辑：

// 替换原来的clone操作
particle.personal_best_position.copy_from_slice(&particle.position);
self.global_best_position.copy_from_slice(&particle.position);

初始化Particle时也可以直接用position克隆初始化，避免先分配空Vec再覆盖：

// 替换Swarm::new中的Particle初始化
let position = (0..dimension).map(|i| rng_thread.gen_range(objective_function_struct.lower_bound[i]..objective_function_struct.upper_bound[i])).collect();
let mut particle = Particle {
    position: position.clone(),
    velocity: vec![0.0; dimension],
    personal_best_position: position,
    personal_best_fitness: std::f64::MAX,
};

2. 批量生成随机数，减少函数调用开销 #

原始代码在内层循环中每次调用gen_range，频繁的函数调用会累积开销。可以批量生成足够的随机数，然后在循环中取用：

修改update_particles：

fn update_particles(&mut self) {
    let dimension = self.global_best_position.len();
    let total_rand = self.particles.len() * dimension * 2;
    // 批量生成所有需要的随机数
    let mut rands = vec![0.0; total_rand];
    self.rng_thread.fill(&mut rands[..]);
    
    let mut rand_idx = 0;
    for particle in &mut self.particles {
        for i in 0..dimension {
            let r1 = rands[rand_idx];
            let r2 = rands[rand_idx + 1];
            rand_idx += 2;
            
            particle.velocity[i] = self.w * particle.velocity[i]
                + self.c1 * r1 * (particle.personal_best_position[i] - particle.position[i])
                + self.c2 * r2 * (self.global_best_position[i] - particle.position[i]);
            particle.position[i] += particle.velocity[i];
        }
        // ... 后续fitness判断逻辑
    }
}

3. 优化目标函数的定义与调用 #

• 将函数参数从&Vec<f64>改为&[f64]，slice更轻量且通用，避免Vec的强制转换开销
• 替换powi(2)为直接乘法，a*a比通用幂运算更高效

修改Sphere函数：

fn sphere_function(x: &[f64]) -> f64 {
    x.iter().map(|&a| a * a).sum()
}

同时更新ObjectiveFunctionStruct的函数类型：

struct ObjectiveFunctionStruct {
    name: String,
    function: fn(&[f64]) -> f64,
    lower_bound: Vec<f64>,
    upper_bound: Vec<f64>,
}

4. 缓存固定维度值，避免重复计算 #

将维度值存储在Swarm结构体中，不用每次循环都从粒子位置获取：

struct Swarm {
    w: f64,
    c1: f64,
    c2: f64,
    particles: Vec<Particle>,
    global_best_position: Vec<f64>,
    global_best_fitness: f64,
    objective_function_struct: ObjectiveFunctionStruct,
    rng_thread: rand::rngs::ThreadRng,
    dimension: usize, // 新增：缓存维度值
}

// 在Swarm::new中初始化
let dimension = objective_function_struct.lower_bound.len();
// ...
Swarm {
    w,
    c1,
    c2,
    particles,
    global_best_position,
    global_best_fitness,
    objective_function_struct,
    rng_thread,
    dimension,
}

后续循环中直接使用self.dimension即可。

5. 使用静态分发的trait替代函数指针（可选） #

如果需要更灵活的目标函数且保持性能，可以定义trait实现静态分发：

trait ObjectiveFunction {
    fn evaluate(&self, x: &[f64]) -> f64;
    fn name(&self) -> &str;
    fn lower_bound(&self) -> &[f64];
    fn upper_bound(&self) -> &[f64];
}

struct SphereFunction {
    lower_bound: Vec<f64>,
    upper_bound: Vec<f64>,
}

impl ObjectiveFunction for SphereFunction {
    fn evaluate(&self, x: &[f64]) -> f64 {
        x.iter().map(|&a| a*a).sum()
    }
    
    fn name(&self) -> &str {
        "Sphere Function"
    }
    
    fn lower_bound(&self) -> &[f64] {
        &self.lower_bound
    }
    
    fn upper_bound(&self) -> &[f64] {
        &self.upper_bound
    }
}

// 修改Swarm为泛型结构体
struct Swarm<F: ObjectiveFunction> {
    w: f64,
    c1: f64,
    c2: f64,
    particles: Vec<Particle>,
    global_best_position: Vec<f64>,
    global_best_fitness: f64,
    objective_function: F,
    rng_thread: rand::rngs::ThreadRng,
    dimension: usize,
}

这种方式能在编译期确定函数调用，比函数指针的动态调用更高效。

6. 编译优化 #

在Cargo.toml中开启最高级别优化：

[profile.release]
opt-level = 3
lto = true
codegen-units = 1

这会让Rust编译器进行更深度的优化，比如函数内联、循环展开等，对数值计算代码的性能提升显著。

内容的提问来源于stack exchange，提问作者user9312688