efficient symmertic computation in spark

Question

A common construct I've seen in algorithms that contain symmetry is

for (int i = 0; i < n ; i++) {
    for (int j = i+1; j < n ; j++) {
        [compute x]
        objects[i][j] += x;
        objects[j][i] -= x;
    }
}

that (while still having O(n^2) complexity) reduces the amount of calculations needed by exploiting the symmetry. Could you tell me what's the way to introduce such an optimization in a pyspark code?

For example, I've written code that computes the force per unit of mass acting on each particle in a system per the formula (where r is the position):

         N    m_j*(r_i - r_j)
F = -G * Σ   -----------------
        i!=j   |r_i - r_j|^3

In it I first do a cross product of my dataframe with itself to get every pairwise interaction, and then aggregate them all by id to get the total force acting on each particle:

def calc_F(df_clust, G=1):

    # cartesian product of the dataframe with itself
    renameCols = [f"`{col}` as `{col}_other`" for col in df_clust.columns]
    df_cart = df_clust.crossJoin(df_clust.selectExpr(renameCols))
    df_clust_cartesian = df_cart.filter("id != id_other")

    df_F_cartesian = df_clust_cartesian.selectExpr("id", "id_other", "m_other",
                                                   "`x` - `x_other` as `diff(x)`",
                                                   "`y` - `y_other` as `diff(y)`",
                                                   "`z` - `z_other` as `diff(z)`"
                                                   )
    df_F_cartesian = df_F_cartesian.selectExpr("id", "id_other",
                                               "`diff(x)` * `m_other` as `num(x)`",
                                               "`diff(y)` * `m_other` as `num(y)`",
                                               "`diff(z)` * `m_other` as `num(z)`",
                                               "sqrt(`diff(x)` * `diff(x)` + `diff(y)`"
                                               "* `diff(y)` + `diff(z)` * `diff(z)`) as `denom`",
                                               )
    df_F_cartesian = df_F_cartesian.selectExpr("id", "id_other",
                                               "`num(x)` / pow(`denom`, 3) as `Fx`",
                                               "`num(y)` / pow(`denom`, 3) as `Fy`",
                                               "`num(z)` / pow(`denom`, 3) as `Fz`",
                                               )
    # squish back to inital particles
    sumCols = ["Fx", "Fy", "Fz"]
    df_agg = df_F_cartesian.groupBy("id").sum(*sumCols)
    renameCols = [f"`sum({col})` as `{col}`" for col in sumCols]
    df_F = df_agg.selectExpr("id", *renameCols)

    df_F = df_F.selectExpr("id",
                           f"`Fx` * {-G} as Fx",
                           f"`Fy` * {-G} as Fy",
                           f"`Fz` * {-G} as Fz")

    return df_F

But I know that the force between two particles is symmetrical - F_ij = -F_ji (I am assuming all masses to be equal) - so here I am computing a number of forces twice, instead of reusing them. So in this particular case I'd want to turn df_clust_cartesian = df_cart.filter("id != id_other") into df_clust_cartesian = df_cart.filter("id < id_other") for example, and somehow reuse these forces when computing the total force in the second part of the function. (Of course ideally I want to learn to do it in general)

Example input for this case would be

a = sc.parallelize([
    [0.48593906,-0.52435857,-0.53198230,0.46153894,-0.33775792E-01,-0.32276499,0.15625001E-04,1],
    [-0.65960690E-01,0.80844238E-01,-0.27603051,-0.57578009,1.1078150,-0.29340765,0.15625001E-04,2],
    [-0.34809157E-01,0.76795481E-01,-0.39087987,-0.55399138,-0.17386098,0.59250806E-01,0.15625001E-04,3]
])                                                           

from pyspark.sql.types import * 
 
clust_input = StructType([ 
    StructField('x',  DoubleType(), False), 
    StructField('y',  DoubleType(), False), 
    StructField('z',  DoubleType(), False), 
    StructField('vx', DoubleType(), False), 
    StructField('vy', DoubleType(), False), 
    StructField('vz', DoubleType(), False), 
    StructField('m',  DoubleType(), False), 
    StructField('id', IntegerType(), False) 
])    

df_clust = a.toDF(schema=clust_input)

user_8893686 · Answer

Basically you want to only compute your formula when id < other_id, and use that result to generate by symmetry all the elements for which id > other_id.

You would simply need to modify your filter by this

df_clust_cartesian = df_cart.filter("id < id_other")

Then, once you have the dataframe df_F_cartesian you have one line per (id, id_other) pair. You can use that line to generate the line corresponding to(id_other, id) and add a minus sign to Fx, Fy and Fz.

This could be done by adding the following step right before the aggregation:

from pyspark.sql import functions as F

sumCols = ["Fx", "Fy", "Fz"]
oppositeSums = [ (-F.col(c)).alias(c) for c in sumCols]
df_F_cartesian = df_F_cartesian.select(F.explode(F.array(
    F.struct(F.col("id"), *sumCols),
    F.struct(F.col("id_other").alias("id"), *oppositeSums)
)).alias("s")).select("s.*")

efficient symmertic computation in spark

1 Answers