How to apply a function to elements of a Polars Series in Rust?

Question

I am trying to convert some code using an ndarray to a polars DataFrame.

I have a function that samples from multiple distributions. I create an ndarray with number_of_distributions x number_of_samples shape. Then I iterate over each column, and for each column I look up the corresponding mean and standard deviation from the arguments provided and sample n times from that distribution, updating the ndarray.

fn multi_rnorm(n: usize, means: Vec<f64>, sds: Vec<f64>) -> Array2<f64> {

    let mut preds: Array2<f64> = Array2::zeros((means.len(), n));

    preds.axis_iter_mut(ndarray::Axis(0)).into_par_iter().enumerate().for_each(|(i, mut row)| {

        let mut rng = rand::thread_rng();
        (0..n).into_iter().for_each(|j| {
            let normal = Normal::new(means[i], sds[i]).unwrap();
            row[j as usize] = normal.sample(&mut rng);
        })
    });
    preds
}

I am trying to do it with a DataFrame, but cannot figure out how to change the values of the Series/columns. Series cannot be indexed, and I cannot find any setter() methods. The code below does not compile as series cannot be indexed.

fn multi_rnorm(n: usize, means: Vec<f64>, sds: Vec<f64>) -> Result<DataFrame, PolarsError> {

    let mut df: DataFrame = DataFrame::new(c![Series::new(&f!("s{i}"), vec![0.0; 99_128]), for i in 0..n])?;

    df.iter().enumerate().for_each(|(i, mut col)| {

        let mut rng = rand::thread_rng();
        col.iter().map(|j| {
            let normal = Normal::new(means[i], sds[i]).unwrap();
            col[j as usize] = normal.sample(&mut rng);
        });
    });
    Ok(df)
}

Does anybody know how to mutate a specific value of a Polars Series in Rust?

Try creating each series by collecting an iterator (https://docs.rs/polars/latest/polars/series/struct.Series.html#creation), then creating the data frame (https://docs.rs/polars/latest/polars/frame/struct.DataFrame.html#method.new) from the already-initialized series. — Solomon Ucko, Feb 14 '23 at 20:01
@SolomonUcko So instead of creating the data frame and iterating over it to sample, you mean create some collectable, add the samples to it, then turn them into series and create the data frame from them? — Jage, Feb 14 '23 at 20:15

How to apply a function to elements of a Polars Series in Rust?

0 Answers0