Use nested iter_map() with ndarrays

I'm following along with some of the suggestions here.

The one concerning this question involves transforming this:

// There's an outer loop that calls this a lot
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;
}

Into this:

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
            }
        })
    })
}

Great, I love it and I've changed all of my vecs into ndarray:Array1s which will be a godsend for performance.

Now, I'm trying to replicate this for another one of my functions which is a bit more complicated:

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;
}

For this one, I need to pass in the simulation space as well as the pressures and densities that have been calculated in previous steps. So again, I'm trying to get rid of the outer loop, and use a map and a fold to do what I want.

This is called from the main loop like this:

let densities = calculate_densities(&sim_data.simulation_space);
let pressures = calculate_pressures(&densities);
let pressure_grad_terms = calculate_pressure_grad_terms(&sim_data.simulation_space, &pressures, &densities);

Unfortunately, I'm struggling with the transformation on this one... As I'm having to reference the density and pressure values of the current particle and the one I'm comparing it against in every loop. Which means I've had to add indexed_iter() into my method chain. It looks like:

fn calculate_pressure_grad_terms(simulation_space: &SimulationSpace, densities: &Array1<f64>, pressures: &Array1<f64>) -> Array1<Vector4<f64>> {
    simulation_space.positions.indexed_iter().map(|(current_index,&current)| {
        simulation_space.positions.indexed_iter().fold(Vector4::zero(), | acc,(other_index, &other) | {
            if is_in_interaction_radius_and_not_self(current, other) {
                acc + PARTICLE_MASS_KG * (pressures[current_index] / (densities[current_index].powi(2)) + pressures[other_index] / (densities[other_index].powi(2)) * grad_smooth(current, other))
            } else {
                acc
            }
        })
    })
}

Now, this is giving me a mismatched type error that I'm really not sure how to address:

    = note: expected struct `ArrayBase<OwnedRepr<Vector4<f64>>, Dim<[usize; 1]>>`
               found struct `Map<IndexedIter<'_, Vector4<f64>, Dim<[usize; 1]>>, [closure@src/main.rs:144:51: 144:77]>`

So I'm thinking this whole process might be easier if I refactor the data to have an array of a struct holding all of these properties and I just directly operate on that, which would remove the need for any of this indexing, something like:

struct Particle {
    position: Vector3<f64>,
    velocity: Vector3<f64>,
    acceleration: Vector3<f64>,
    mass: f64,
    pressure_at_location: f64,
    density_at_location: f64
}

But I'm somewhat reluctant to do that just yet because, 1) the tutorial had separate arrays and I think that might be a requirement for OpenCL (I've not got that far yet), and 2) I want to understand what this problem is not pivot to avoid it.

So what's gone wrong here?

Solution

To operate over a bunch of arrays, you'll want to use Zip.

Something like this, (Code typed out, not even a little tested)

Zip::from(&simulation_space.positions)
    .and(densities)
    .and(pressures)
    .map_collect(|current_position, current_density, current_pressure| 
       Zip::from(&simulation_space.positions)
           .and(pressures)
           .and(densities)
           .fold(Vector4::zero(), | acc, other_position, other_pressure, other_density | {
            if is_in_interaction_radius_and_not_self(current_position, other_position) {
                acc + PARTICLE_MASS_KG * (current_pressure / (current_density.powi(2)) + other_pressure / (other_density.powi(2)) * grad_smooth(current_position, other_index))
            } else {
                acc
            }
        })

Your indexed_iter code didn't work because indexed_iter returns an iterator. As @cafce25 noted, you need to collect() to convert the iterator back into an array.