Using binomal distribution with Spark to calculate an expected value and a variance ("number of times getting a six when throwing a dice three times") - apache-spark

Question

I'm training myself to resolve classical statistics exercices with Spark and it's MLib module when needed, in order to face any situation later.
Spark is dedicated to calculation on matrices, but let's say that I have a side calculation to do in my program, at a time, and that I would like to resolve it with Spark, without adding others APIs.
The calculation is simple : learn the expected value and the variance
of having a six on a dice, when you throw it three times.
Currently, I resolve it by the help of the Apache maths API and a bit of Spark :
/** Among imports : */
import org.apache.commons.math3.distribution.*;
/**
* E8.5 : On lance trois dés. Probabilité d'un six obtenu ?
* Suit une loi binomiale B(3, 1/6)
* Par calcul : F(0) = 125/216, F(1) = 125/216 + 75/216, F(2) = (125 + 75 + 15)/216,
* F(3) = (125 + 75 + 15 + 1)/216 = 1
*/
#Test
#DisplayName("On lance trois dés. Probabilité d'un six obtenu ?. Caractéristiques de la loi B(3, 1/6)")
public void troisDésDonnentUnSix() {
int n = 3; // Nombre de lancés de dés.
double v[] = {0, 1, 2, 3}; // Valeurs possibles de la variable aléatoire.
double f[] = new double[n+1]; // Fréquences cumulées.
double p[] = new double[n+1]; // Probabilités.
// Calculer les probabilités et les fréquences.
BinomialDistribution loiBinomale = new BinomialDistribution(n, 1.0/6.0);
for(int i=0; i <= n; i++) {
p[i] = loiBinomale.probability(i);
f[i] = loiBinomale.cumulativeProbability(i);
}
LOGGER.info("P(x) = {}", p);
LOGGER.info("F(x) = {}", f);
Dataset<Row> ds = fromDouble(v, p);
LOGGER.info("E(X) = {}, V(X) = {}", esperance(ds), variance(ds));
}
where fromDouble(v, p) method creates a Dataset from a list of random variable values (column x) and their associated frequencies (column Px) :
/**
* Renvoyer un Dataset depuis une série de valeurs entières à probabilités.
* #param valeurs Valeurs.
* #param probabilites Probabilités.
* #return Dataset avec une colonne x, entière, contenant les valeurs<br>
* et une colonne Px, décimale, contenant les probabilités.
*/
protected Dataset<Row> fromDouble(double[] valeurs, double[] probabilites) {
StructType schema = new StructType()
.add("x", DoubleType, false)
.add("Px", DoubleType, false);
List<Row> rows = new ArrayList<>();
for(int index=0; index < valeurs.length; index ++) {
rows.add(RowFactory.create(valeurs[index], probabilites[index]));
}
return this.session.createDataFrame(rows, schema);
}
And esperance (= expected value) and variance methods called are doing these calculations :
/**
* Calculer l'espérance sur un Dataset avec valeurs et probabilités.
* #param ds Dataset avec colonnes : <br>
* x : valeur<br>
* Px : fréquence<br>
* #return espérance.
*/
protected double esperance(Dataset<Row> ds) {
return ds.agg(sum(col("x").multiply(col("Px")))).first().getDouble(0);
}
/**
* Calculer la variance sur un Dataset avec valeurs et probabilités.
* #param ds Dataset avec colonnes : <br>
* x : valeur<br>
* Px : fréquence<br>
* #return espérance.
*/
protected double variance(Dataset<Row> ds) {
Column variation = col("x").minus(esperance(ds));
Column variationCarre = variation.multiply(variation);
Column termeCalculVariance = col("Px").multiply(variationCarre);
return ds.agg(sum(termeCalculVariance)).first().getDouble(0);
}
LOGGER output :
P(x) = [0.5787037037037037, 0.34722222222222215, 0.06944444444444445, 0.0046296296296296285]
F(x) = [0.5787037037037035, 0.9259259259259259, 0.9953703703703703, 1.0]
E(X) = 0.49999999999999994, V(X) = 0.41666666666666663
It works (? I have caculatd by hand "E(X) = 53/108", and it finds 54/108 = 0.5..., I might be wrong), but it's not perfect.
Is there a more elegant way to solve this problem using Spark (and Spark-MLib, if needed) ?