How to use rayon to update a personal struct containing an Array in Rust

Question

CONTEXT

General overview

(Here is the github page with the minimal example of my problem, and the page of my whole project) I'm very new to rust and I'm trying to simulate the behavior of a fluid in Rust. This is easy: computing large arrays with some functions for each timestep.

I'd like to parallelize the computations done each timestep using rayon. But the compiler doesn't want me to access a mutable struct containing an Array that I want to modify, even if I'm sure that each modification will be on different places in the array: which assure me that it's safe. (I think?).

So my question is: should I use unsafe rust in here? If so, how? And: is it possible to make Rust understand that it's safe, or to do it properly without unsafe rust? I saw this post which kind of resembled my problem, but couldn't find a way to use the solution for my problem.

I also tried to put unsafe {...} keywords everywhere, but the compilator still complains...

You may only need to read "Structs" subsection to understand the problem, but I will also put a "Function" subsection, in case it can be important. If you think it might not be necessary, you can skip to "Main function" subsection.

Structs

Here are my structs: I'd like to keep them that way, as they would give me (I think) more flexibility with setters/getters, but I'm open to change the way it's implemented right now.

#[derive(Debug, PartialEq)]
struct vec2D {pub x: f64, pub y: f64}

#[derive(Debug, PartialEq)]
 struct ScalarField2D {
     s: Array2<f64>,
}    

#[derive(Debug, PartialEq)]
 struct VectorField2D {
     x: ScalarField2D,
     y: ScalarField2D
}

impl ScalarField2D {

     // also a constructor new() but not shown for simplicity

     fn get_pos(&self, x: usize, y: usize) -> f64
    {return self.s[[y,x]];}

    fn set_pos(&mut self, x: usize, y: usize, f: f64)
    {self.s[[y,x]] = f;}
}

impl VectorField2D {

     // also a constructor new() but not shown for simplicity

     fn get_pos(&self, x: usize, y: usize) -> vec2D
    {let vec_at_pos = vec2D {
        x: self.x.get_pos(x, y),
        y: self.y.get_pos(x, y)};
     return vec_at_pos;}

    fn set_pos(&mut self, x: usize, y: usize, vec: &vec2D)
    {self.x.set_pos(x, y, vec.x);
     self.y.set_pos(x, y, vec.y);}

}

Function

Here is my function: which takes a ScalarField2D struct, and computes a vector called the "gradient" at a particular position of the ScalarField2D array, and then returning this vector as a "vec2D" struct.

// computes the gradient of a scalar field at a given position
fn grad_scalar(a: &ScalarField2D,
               x: i32, y: i32,
               x_max: i32, y_max: i32) -> vec2D
{
    let ip = ((x+1) % x_max) as usize;
    // i-1 with Periodic Boundaries
    let im = ((x - 1 + x_max) % x_max) as usize;
    // j+1 with Periodic Boundaries
    let jp = ((y+1) % y_max) as usize;
    // j-1 with Periodic Boundaries
    let jm = ((y - 1 + y_max) % y_max) as usize;
    let (i, j) = (x as usize, y as usize);

    let grad = vec2D {
        x: (a.get_pos(ip, j) - a.get_pos(im, j))/(2.),
        y: (a.get_pos(i, jp) - a.get_pos(i, jm))/(2.)};
        
    return grad;
}

Main function

Here is my problem: I try to iterate over all the possible x and y using (0..x_max).into_par_iter() (or y_max for y), compute the gradient associated with each position, and then set the value to the ScalarField2D struct using the set_pos method... Here is the main function, and the import commands, and I will show you the error message I get in the next subsection

use ndarray::prelude::*;
use rayon::prelude::*;

fn main() {

    let (x_max, y_max) = (2usize, 50usize);
    let (x_maxi32, y_maxi32) = (x_max as i32, y_max as i32);

    let mut GD_grad_rho = VectorField2D::new(x_max, y_max);
    let mut GD_rho = ScalarField2D::new(x_max, y_max);    

    let x_iterator = (0..x_max).into_par_iter();
    x_iterator.map(|xi| {
        let y_iterator = (0..y_max).into_par_iter();
        y_iterator.map(|yi| {

            // unsafe here?
            GD_grad_rho
                .set_pos(xi, yi,
                         &grad_scalar(&GD_rho,
                                      xi as i32, yi as i32,
                                      x_maxi32, y_maxi32));
            
        });});
}

Error message

Here is the compilation error I get

error[E0596]: cannot borrow `GD_grad_rho` as mutable, as it is a captured variable in a `Fn` closure
   --> src/main.rs:104:13
    |
104 | /             GD_grad_rho
105 | |                 .set_pos(xi, yi,
106 | |                          &grad_scalar(&GD_rho,
107 | |                                       xi as i32, yi as i32,
108 | |                                       x_maxi32, y_maxi32));
    | |__________________________________________________________^ cannot borrow as mutable

error[E0596]: cannot borrow `GD_grad_rho` as mutable, as it is a captured variable in a `Fn` closure
   --> src/main.rs:101:24
    |
101 |         y_iterator.map(|yi| {
    |                        ^^^^ cannot borrow as mutable
...
104 |             GD_grad_rho
    |             ----------- mutable borrow occurs due to use of `GD_grad_rho` in closure

For more information about this error, try `rustc --explain E0596`.
error: could not compile `minimal_example_para` due to 2 previous errors

If you want the complete thing, I created a github repo with everything in it.

Tests after susitsm answer

So I did something like this:

use ndarray::prelude::*;
use rayon::prelude::*;

fn grad_scalar(a: &Array2<f64>,
          i: usize, j: usize) -> (f64, f64)
{
    let array_shape = a.shape();
    let imax = array_shape[0];
    let jmax = array_shape[1];

    // i-1 with Periodic Boundaries
    let ip = ((i + 1) % imax);
    // i-1 with Periodic Boundaries
    let im = ((i + imax) - 1) % imax;
    // j+1 with Periodic Boundaries
    let jp = ((j + 1) % jmax);
    // j-1 with Periodic Boundaries
    let jm = ((j + jmax) - 1) % jmax;

    let grad_i = (a[[ip, j]] - a[[im, j]])/2.;
    let grad_j = (a[[i, jp]] - a[[i, jm]])/2.;

    return (grad_i, grad_j);
}

fn main() {

    let a = Array::<f64, Ix2>::ones((dim, dim));

    let index_list: Vec<(_, _)> = (0..a.len())
        .into_par_iter()
        .map(|i| (i / a.dim().0, i % a.dim().1))
        .collect();

    let (r1, r2): (Vec<_>, Vec<_>) = (0..a.len())
        .into_par_iter()
        .map(|i| (index_list[i].0, index_list[i].1))
        .map(|(i, j)| grad_scalar(&a, i, j))
        .collect();

    let grad_row = Array2::from_shape_vec(a.dim(), r1).unwrap();
    let grad_col = Array2::from_shape_vec(a.dim(), r2).unwrap();

}

Which gives me the result I want, even though I wanted to access the values through my Structs. Which is not exactly what I want but we're getting closer But I wonder about the efficiency for more arrays, I'll post a separate question for that

Answer 1

You can do something like this:

use ndarray::Array2;
use rayon::prelude::*;

fn main() {
    let a: Vec<u64> = (0..20000).collect();
    let a = Array2::from_shape_vec((100, 200), a).unwrap();

    let stuff = |arr, i, j| (i + j, i * j);
    let (x, y): (Vec<_>, Vec<_>) = (0..a.len())
        .into_par_iter()
        .map(|i| (i / a.dim().0, i % a.dim().1))
        .map(|(i, j)| stuff(&a, i, j))
        .collect();
    let grad_x = Array2::from_shape_vec(a.dim(), x);
    let grad_y = Array2::from_shape_vec(a.dim(), y);
    let grad_vector_field = VectorField2D {
        x: ScalarField2D { s: grad_x },
        y: ScalarField2D { s: grad_y },
    };
}

Or implement the FromParallelIterator<vec2D>

impl FromParallelIterator<vec2D> for VectorField2D {
    fn from_par_iter<I>(par_iter: I) -> Self
        where I: IntoParallelIterator<Item = T>
    {
        let (x, y): (Vec<_>, Vec<_>) = par_iter
            .into_par_iter()
            .map(|vec_2D| {
                let vec2D { x, y } = vec_2D;
                (x, y)
            })
            .collect();
        Self {
            x: ScalarField2D { s: Array2::from_shape_vec(a.dim(), x) },
            y: ScalarField2D { s: Array2::from_shape_vec(a.dim(), y) },
        }
    }
}

This will enable using collect for your type when using parallel iterators

let vector_field: VectorField2D = (0..a.len())
        .into_par_iter()
        .map(|i| (index_list[i].0, index_list[i].1))
        .map(|(i, j)| grad_scalar(&a, i, j))
        .collect();