Smoothed Particle Hydrodynamics in Rust

Question

So I'm simultaneously learning both Rust and Smoothed Particle Hydrodynamics.

I've been using this video from AMD as a reference and the provided smoothing kernels.

At this stage, I've not optimised the simulation with voxelisation, nor have I added hardware acceleration, so it is very slow. I've also not added a visualisation step to check if it actually works.

At this point, it compiles, and it runs. So what I'm interested in is whether this code looks reasonably understandable from a purely rust code point of view. Although if you feel like reviewing the actual SPH implementation, I wouldn't say no...

//use std::collections::HashMap;
use std::f32::consts::PI;
use cgmath::{InnerSpace, Vector3, Vector4, Zero};
use rand::Rng;

const NUM_PARTICLES: u32 = 1000; // 1.2 Litres
const BUFFER_LEN: usize = NUM_PARTICLES as usize;
const BOX_DIMENSIONS_M: Vector3<f32> = Vector3::new(0.30, 0.10, 0.10);
const SMOOTHING_RADIUS: f32 = NUM_PARTICLES as f32 / (0.3 * 0.1 * 0.1 * 60.0);
const TIME_STEP_SECONDS: f32 = 0.01;
const GRAVITY: Vector4<f32> = Vector4::new(0.0, -9.8, 0.0, 0.0);
const SIM_LENGTH_SECONDS: f32 = 10.0;
const PARTICLE_MASS_KG: f32 = 0.001; //1mL of water
const FLUID_CONST: f32 = 2.2; // Possibly needs a x10^9
const VISCOUS_CONST: f32 = 0.0001;
const RHO_ZERO: f32 = 1000.0;

//TODO - Voxelise

fn main() {
    let mut sim_data: SimulationData = setup();
    let mut current_time: f32 = 0.0;
    while current_time < SIM_LENGTH_SECONDS {
        sim_step(&mut sim_data);
        current_time += TIME_STEP_SECONDS;
        println!("Time: {}s", current_time);
    }
}

fn sim_step(sim_data: &mut SimulationData) {

    let densities = sim_data.simulation_space.positions.iter().enumerate()
        .map(|(i, _)| calculate_density_at_point(&sim_data.simulation_space, i))
        .collect::<Vec<f32>>();

    let pressures = sim_data.simulation_space.positions.iter().enumerate()
        .map(|(i, _)| calculate_pressure_at_point(densities[i]))
        .collect::<Vec<f32>>();

    let pressure_grad_terms = sim_data.simulation_space.positions.iter().enumerate()
        .map(|(i, _)| calculate_pressure_grad_term_at_point(&sim_data.simulation_space, &pressures, &densities, i))
        .collect::<Vec<Vector4<f32>>>();

    let viscosity_terms = sim_data.simulation_space.positions.iter().enumerate()
        .map(|(i, _)| calculate_viscosity_term_at_point(&sim_data.simulation_space, &densities, i))
        .collect::<Vec<Vector4<f32>>>();

    let accelerations = sim_data.simulation_space.positions.iter().enumerate()
        .map(|(i, _)| calculate_acceleration_at_point(&pressure_grad_terms, &viscosity_terms, i))
        .collect::<Vec<Vector4<f32>>>();

    let new_velocities = sim_data.simulation_space.velocities.iter().enumerate()
        .map(|(i, v)| v + accelerations[i] * TIME_STEP_SECONDS)
        .collect::<Vec<Vector4<f32>>>();

    let new_positions = sim_data.simulation_space.positions.iter().enumerate()
        .map(|(i, p)| p + new_velocities[i] * TIME_STEP_SECONDS)
        .collect::<Vec<Vector4<f32>>>();

    //TODO - Check for collisions with walls and floor

    // This will make rustaceans cry
    sim_data.simulation_space.velocities = new_velocities.try_into().unwrap();
    sim_data.simulation_space.positions = new_positions.try_into().unwrap();
}

fn calculate_acceleration_at_point(pressure_terms: &Vec<Vector4<f32>>, viscosity_terms: &Vec<Vector4<f32>>, i: usize) -> Vector4<f32> {
    let acceleration: Vector4<f32> = GRAVITY + pressure_terms[i] + viscosity_terms[i];
    return acceleration;
}

fn calculate_viscosity_term_at_point(sim_space: &SimulationSpace, densities: &Vec<f32>, i: usize) -> Vector4<f32> {
    let viscosity_term: Vector4<f32> = sim_space.positions.iter()
        .filter(|&x| is_in_interaction_radius_and_not_self(x, sim_space.positions[i]))
        .enumerate()// exclude self
        .fold(Vector4::zero(), |acc: Vector4<f32>, (j, _)| {
            acc + (VISCOUS_CONST / densities[j]) * PARTICLE_MASS_KG * (sim_space.velocities[j] - sim_space.velocities[i]) / densities[j] * laplacian_smooth(sim_space.positions[i], sim_space.positions[j])
        });
    return viscosity_term;
}

fn calculate_pressure_at_point(density: f32) -> f32 {
    let pressure_at_point: f32 = FLUID_CONST * (density - RHO_ZERO);
    return pressure_at_point;
}

fn calculate_pressure_grad_term_at_point(sim_space: &SimulationSpace, pressures: &Vec<f32>, densities: &Vec<f32>, i: usize) -> Vector4<f32> {
    let pressure_grad_term: Vector4<f32> = sim_space.positions.iter()
        .filter(|&x| is_in_interaction_radius_and_not_self(x, sim_space.positions[i]))
        .enumerate()// exclude self
        .fold(Vector4::zero(), |acc: Vector4<f32>, (j, _)| {
            acc + PARTICLE_MASS_KG * (pressures[i] / (densities[i].powi(2)) + pressures[j] / (densities[j].powi(2))) * grad_smooth(sim_space.positions[i], sim_space.positions[j])
        });
    return pressure_grad_term;
}

fn calculate_density_at_point(sim_space: &SimulationSpace, i: usize) -> f32 {
    let density = sim_space.positions.iter()
        .filter(|&x| is_in_interaction_radius_and_not_self(x, sim_space.positions[i]))
        .enumerate()// exclude self
        .fold(0.0, |acc: f32, (j, _)| acc + PARTICLE_MASS_KG * smooth(sim_space.positions[i], sim_space.positions[j]));

    return density;
}

fn is_in_interaction_radius_and_not_self(current: &Vector4<f32>, other: Vector4<f32>) -> bool {
    return *current != other && (current - other).magnitude() < SMOOTHING_RADIUS;
}

// There are some terms reused between smooth, grad_smooth and laplacian_smooth, this could be optimised
// Not to mention, some of these terms are constants...
fn smooth(current_position: Vector4<f32>, other_position: Vector4<f32>) -> f32
{
    return (315.0/(64.0*PI*(SMOOTHING_RADIUS.powi(9)))) * (SMOOTHING_RADIUS.powi(2) - (current_position - other_position).magnitude2()).powi(3);
}

fn grad_smooth(current_position: Vector4<f32>, other_position: Vector4<f32>) -> Vector4<f32>
{

    return (-45.0/(PI*(SMOOTHING_RADIUS.powi(6)))) * (SMOOTHING_RADIUS - (current_position - other_position).magnitude2()).powi(2) * ((current_position - other_position) / (current_position - other_position).magnitude2());
}

fn laplacian_smooth(current_position: Vector4<f32>, other_position: Vector4<f32>) -> f32
{
    return (45.0/(PI*(SMOOTHING_RADIUS.powi(6)))) * (SMOOTHING_RADIUS - (current_position - other_position).magnitude());
}

fn setup() -> SimulationData {
    let sim_space = SimulationSpace{
        positions: get_initial_positions(),
        velocities: get_initial_velocities(),
        accelerations: get_initial_accelerations()
    };

    let sim_data = SimulationData {
        simulation_space: sim_space,
        //voxel_pixel_map: HashMap::new()
    };

    return sim_data;
}

fn get_initial_positions() -> [Vector4<f32>; BUFFER_LEN] {
    let mut rng = rand::thread_rng();
    let mut positions: [Vector4<f32>; BUFFER_LEN] = [Vector4::zero(); BUFFER_LEN];
    for i in 0..BUFFER_LEN {
        positions[i] = Vector4::new(
            rng.gen_range(0.0..0.1),
            rng.gen_range(0.0..0.01),
            rng.gen_range(4.9..5.0),
            0.0
        );
    }
    return positions;
}

fn get_initial_velocities() -> [Vector4<f32>; BUFFER_LEN] {
    return [ Vector4::zero(); BUFFER_LEN];
}

fn get_initial_accelerations() -> [Vector4<f32>; BUFFER_LEN] {
    return [ Vector4::zero(); BUFFER_LEN];
}

struct SimulationData {
    simulation_space: SimulationSpace,
    //voxel_pixel_map: HashMap<u16, Vec<u32>>
}

struct SimulationSpace {
    positions: [Vector4<f32>; BUFFER_LEN as usize],
    velocities: [Vector4<f32>; BUFFER_LEN as usize],
    accelerations: [Vector4<f32>; BUFFER_LEN as usize]
}

Answer 1

One big item: compile and run in --release mode. You mention it being slow, but it takes only a few second to run for me in release mode, mostly because of all the printing it does.

Typically, in Rust, we'd make the setup function a constructor method on SimulationData

impl SimulationData {
   fn new() -> SimulationData {
      ...
   }
}

We'd also typically make sim_step a method on SimulationData

impl SimulationData {
    fn step(&mut self) {
        ...
    }
}

This comment seem inscrutable and probably wrong

const NUM_PARTICLES: u32 = 1000; // 1.2 Litres

This constant could use a comment explaining something about the mysterious numbers being multiplied:

const SMOOTHING_RADIUS: f32 = NUM_PARTICLES as f32 / (0.3 * 0.1 * 0.1 * 60.0);

You've got a number of blocks that have this sort of pattern

let pressure_grad_term: Vector4<f32> = sim_space.positions.iter()
    .filter(|&x| is_in_interaction_radius_and_not_self(x, sim_space.positions[i]))
    .enumerate()// exclude self
    .fold(Vector4::zero(), |acc: Vector4<f32>, (j, _)| {
        acc + PARTICLE_MASS_KG * (pressures[i] / (densities[i].powi(2)) + pressures[j] / (densities[j].powi(2))) * grad_smooth(sim_space.positions[i], sim_space.positions[j])
    });

Firstly, it's wrong because you filter before you enumerate. This means that all of your indexes are off because the indexes are generated after the filter. If you want the indexes to match your arrays, you need to enumerate() before filtering items out.

Secondly, rust has a sum method you can use instead of the generic fold.

 let pressure_grad_term: Vector4<f32> = sim_space.positions.iter()
    .filter(|&x| is_in_interaction_radius_and_not_self(x, sim_space.positions[i]))
    .enumerate()// exclude self
    .map(|(j, _)| {
        PARTICLE_MASS_KG * (pressures[i] / (densities[i].powi(2)) + pressures[j] / (densities[j].powi(2))) * grad_smooth(sim_space.positions[i], sim_space.positions[j])
    })
    .sum()

However, for the type of coding you are doing here, you probably want to look at using the ndarray crate. It provides an Array type with useful functions that operate on arrays. In your case, this seems a better fit than converting to iterators and back. So, for example, here is a rewritten version of calculate_density_at_point

fn calculate_density_at_point(sim_space: &SimulationSpace, i: usize) -> f32 {
    sim_space.positions.fold(0.0, |acc, x| {
        if is_in_interaction_radius_and_not_self(x, sim_space.positions[i]) {
            acc + PARTICLE_MASS_KG * smooth(sim_space.positions[i], *x)
        } else {
            acc
        }
    })
}

You'll notice that its much simpler then the iterator version. We can go a step further and define a function that builds the entire density array.

fn calculate_density_at_point(sim_space: &SimulationSpace) -> Array1<f32> {
    sim_space.positions.map(|&position| {
        sim_space.positions.fold(0.0, |acc, x| {
            if is_in_interaction_radius_and_not_self(x, position) {
                acc + PARTICLE_MASS_KG * smooth(position, *x)
            } else {
                acc
            }
        })
    })
}

Going one step further, if you turn on the parallel features, you can easily parallelize the calculations:

fn calculate_density_at_point(sim_space: &SimulationSpace) -> Array1<f32> {
    ndarray::Zip::from(&sim_space.positions).par_map_collect(|&position| {
        sim_space.positions.fold(0.0, |acc, x| {
            if is_in_interaction_radius_and_not_self(x, position) {
                acc + PARTICLE_MASS_KG * smooth(position, *x)
            } else {
                acc
            }
        })
    })
}