What Is The Recommended Way To Compute A Weighted Sum Of Selected Columns Of A Pandas Dataframe?

For example, I would like to compute the weighted sum of columns 'a' and 'c' for the below matrix, with weights defined in the dictionary w. df = pd.DataFrame({'a': [1,2,3],

Solution 1:

You could use a Series as in your first example, just use reindex afterwards:

import pandas as pd

df = pd.DataFrame({'a': [1,2,3],
                   'b': [10,20,30],
                   'c': [100,200,300],
                   'd': [1000,2000,3000]})

w = {'a': 1000., 'c': 10.}
print(df.dot(pd.Series(w).reindex(df.columns, fill_value=0)))

Output

02000.014000.026000.0dtype:float64

Solution 2:

Here's an option without having to create a pd.Series:

(df.loc[:,w.keys()] * list(w.values())).sum(axis=1)
02000.014000.026000.0

Solution 3:

I stumbled on my own question again and benchmarked the available answers.

Observation: It is worthwhile to fill the incomplete weight vector with zeros first, rather than first capturing a view on the columns and then having the resulting subframe dot-multiplied.

---

import pandas as pd
import numpy as np

defbenchmark(n_rows, n_cols, n_ws):
    print("n_rows:%d, n_cols:%d, n_ws:%d" % (n_rows, n_cols, n_ws))
    df = pd.DataFrame(np.random.randn(n_rows, n_cols), 
                      columns=range(n_cols))
    w = dict(zip(np.random.choice(np.arange(n_cols), n_ws), 
                 np.random.randn(n_ws)))
    w0 = pd.Series(w).reindex(df.columns, fill_value=0).values

    # Method 0 (aligned vector w0, reference!)deffun0(df, w0): return df.values.dot(w0)
    # Method 1 (reindex)deffun1(df, w): return df.dot(pd.Series(w).reindex(df.columns, fill_value=0))
    # Method 2 (column view)deffun2(df, w): return (df.loc[:,w.keys()] * list(w.values())).sum(axis=1)
    # Method 3 (column view, faster)deffun3(df, w): return df.loc[:, w].dot(pd.Series(w))
    # Method 4 (column view, numpy)deffun4(df, w): return df[list(w.keys())].values.dot(list(w.values()))

    # Assert equivalence
    np.testing.assert_array_almost_equal(fun0(df,w0), fun1(df,w), decimal=10)
    np.testing.assert_array_almost_equal(fun0(df,w0), fun2(df,w), decimal=10)
    np.testing.assert_array_almost_equal(fun0(df,w0), fun3(df,w), decimal=10)
    np.testing.assert_array_almost_equal(fun0(df,w0), fun4(df,w), decimal=10)

    print("fun0:", end=" ")
    %timeit fun0(df, w0)
    print("fun1:", end=" ")
    %timeit fun1(df, w)
    print("fun2:", end=" ")
    %timeit fun2(df, w)
    print("fun3:", end=" ")
    %timeit fun3(df, w)
    print("fun4:", end=" ")
    %timeit fun4(df, w)

benchmark(n_rows = 200000, n_cols = 11, n_ws = 3)
benchmark(n_rows = 200000, n_cols = 11, n_ws = 9)
benchmark(n_rows = 200000, n_cols = 31, n_ws = 5)

The output (fun0() is the reference using the zero-filled vector w0):

n_rows:200000, n_cols:11, n_ws:3
fun1: 1.98 ms ± 86.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
fun2: 9.66 ms ± 32.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
fun3: 2.68 ms ± 90.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
fun4: 2.2 ms ± 45.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

n_rows:200000, n_cols:11, n_ws:9
fun1: 1.85 ms ± 28.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
fun2: 11.7 ms ± 54.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
fun3: 3.7 ms ± 84.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
fun4: 3.17 ms ± 29.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

n_rows:200000, n_cols:31, n_ws:5
fun1: 3.08 ms ± 42.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
fun2: 13.1 ms ± 260 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
fun3: 5.48 ms ± 57 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
fun4: 4.98 ms ± 49.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

I've tested with pandas 1.2.3, numpy 1.20.1 and Python 3.9.0. on a MacBookPro (Late 2015). (Similar results hold for older Python versions).

Solution 4:

Using numpydot with values

df[list(w.keys())].values.dot(list(w.values()))
array([2000., 4000., 6000.])

---

Fixed your error

df.mul(pd.Series(w),1).sum(axis=1)02000.014000.026000.0dtype:float64