Compute co-occurences in pandas dataframe for column values grouped by another column values

Question

Question

I am using Pandas on Python 3.7.7. I would like to compute the mutual information between categorical values of a variable x grouped by another variable's values y. My data looks like the following table:

+-----+-----+
|  x  |  y  |
+-----+-----+
| x_1 | y_1 |
| x_2 | y_1 |
| x_3 | y_1 |
| x_1 | y_2 |
| x_2 | y_2 |
| x_4 | y_3 |
| x_6 | y_3 |
| x_9 | y_3 |
| x_1 | y_4 |
| ... | ... |
+-----+-----+

I would like have a data structure (a pandas MultiIndex series/dataframe or a numpy matrix or anything suitable) which stores the number of co-occurrences of (x_i, x_j) pair given a certain y_k value. In fact, this would be great, for example, to easily compute PMI:

+-----+-----+--------+-------+
| x_i | x_j |  cooc  |  pmi  |
+-----+-----+--------+-------+
| x_1 | x_2 |        |       |
| x_1 | x_3 |        |       |
| x_1 | x_4 |        |       |
| x_1 | x_5 |        |       |
| ... | ... |   ...  |  ...  |
+-----+-----+--------+-------+

Is there any suitable-memory efficient way?

SIDE NOTE: I am using quite big data (40k distinct x values and 8k distinct y values, with a total of 300k (x,y) entries so a memory friendly and optimized approach would be great (maybe relying on third part libraries as Dask)

Update

Non-optimized solution

I came up with a solution using pd.crosstab. I provide here a small example:

import numpy as np
import pandas as pd
df = pd.DataFrame(np.random.randint(0,100,size=(100, 2)), columns=list('xy'))
"""
df:
  +-----+-----+
  |  x  |  y  |
  +-----+-----+
  | 4   | 99  |
  | 1   | 39  |
  | 39  | 56  |
  | ..  | ..  |
  | 59  | 20  |
  | 82  | 57  |
  +-----+-----+
 100 rows × 2 columns
"""
# Compute cross tabulation:
crosstab = pd.crosstab(df["x"], df["y"])
"""
crosstab:
  +------+-----+-----+-----+-----+
  |  y   |  0  |  2  |  3  | ... |
  |  x   +-----+-----+-----+-----+
  |  1   |  0  |  0  |  0  | ... |
  |  2   |  0  |  0  |  0  | ... |
  | ...  | ... | ... | ... | ... |
  +------+-----+-----+-----+-----+
 62 rows × 69 columns
"""
# Initialize a pandas MultiIndex Series storing PMI values
import itertools
x_pairs = list(itertools.combinations(crosstab.index, 2))
pmi = pd.Series(0, index = pd.MultiIndex.from_tuples(x_pairs))
"""
pmi:
  +-------------+-----+
  |    index    | val |
  +------+------|     |
  |  x_i |  x_j |     |
  +------+------+-----+
  |  1   |  2   |  0  |
  |      |  4   |  0  |
  |  ... |  ... | ... |
  |  95  |  98  |  0  |
  |      |  99  |  0  |
  |  96  |  98  |  0  |
  +------+------+-----+
 Length: 1891, dtype: int64
"""

Then, the loop I use to fill the Series is structured as follows:

for x1, x2 in x_pairs:
    pmi.loc[x1, x2] = crosstab.loc[[x1, x2]].min().sum() / (crosstab.loc[x1].sum() * crosstab.loc[x2].sum())

This is not an optional solution, poorly performing even with small use cases.

Answer 1

Optimized solution

In the end, I managed to compute cross occurrences in a memory friendly way using scipy sparse matrices for the intermediate computations:

import pandas as pd
import numpy as np
from scipy.sparse import csr_matrix

def df_compute_cooccurrences(df: pd.DataFrame, column1: str, column2: str) -> pd.DataFrame:
   
    # pd.factorize encode the object as an enumerated type or categorical variable, returning:
    # - `codes` (ndarray): an integer ndarray that’s an indexer into `uniques`.
    # - `uniques` (ndarray, Index, or Categorical): the unique valid values
    # see more at https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.factorize.html

    i, rows = pd.factorize(df[column1])
    # i    -> array([     0,      0,      0, ..., 449054,      0,      1])
    # rows -> Index(['column1_label1', 'column1_label2', ...])

    j, cols = pd.factorize(df[column2])
    # j    -> array([    0,     1,     2, ..., 28544,    -1,    -1])
    # cols -> Float64Index([column2_label1, column2_label2, ...])

    ij, tups = pd.factorize(list(zip(i, j)))
    # ij   -> array([      0,       1,       2, ..., 2878026, 2878027, 2878028])
    # tups -> array([(0, 0), (0, 1), (0, 2), ..., (449054, 28544), (0, -1), (1, -1)]

    # Then we can finally compute the crosstabulation matrix
    crosstab = csr_matrix((np.bincount(ij), tuple(zip(*tups))))
    # If we convert directly this into a Dataframe with
    # pd.DataFrame.sparse.from_spmatrix(crosstab, rows, cols)
    # we have the same result as using 
    # https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.crosstab.html
    # but we obtained it in a memory-friendly way (allowing big data processing)

    # In order to obtain the co-occurrences matrix for column 1, 
    # we have to multiply the crosstab matrix for its transposed 
    coocc = crosstab.dot(crosstab.transpose())
    
    # Then we can finally return the co-occurence matrix in in a DataFrame form
    return pd.DataFrame.sparse.from_spmatrix(coocc, rows, rows)

Here a small example is provided:

import pandas as pd
import numpy as np
from scipy.sparse import csr_matrix

def df_compute_cooccurrences(df: pd.DataFrame, column1: str, column2: str) -> pd.DataFrame:
    i, rows = pd.factorize(df[column1])
    j, cols = pd.factorize(df[column2])
    ij, tups = pd.factorize(list(zip(i, j)))
    crosstab = csr_matrix((np.bincount(ij), tuple(zip(*tups))))
    coocc = crosstab.dot(crosstab.transpose())
    return pd.DataFrame.sparse.from_spmatrix(coocc, rows, rows)

df = pd.DataFrame(zip([1,1,1,2,2,3,4],["a","a","a","a","a","b","b"]), columns=list('xy'))
"""
df:
  +-----+-----+
  ¦  x  ¦  y  ¦
  +-----+-----+
  |  1  |  a  |
  |  1  |  a  |
  |  1  |  a  |
  |  2  |  a  |
  |  2  |  a  |
  |  3  |  b  |
  |  4  |  b  |
  +-----+-----+
"""
cooc_df = df_compute_cooccurrences(df, "x", "y")
"""
cooc_df:
    +---+---+---+---+
    ¦ 1 | 2 | 3 | 4 |
+---+---+---+---+---+
¦ 1 ¦ 9 | 6 | 0 | 0 |
¦ 2 ¦ 6 | 4 | 0 | 0 |
¦ 3 ¦ 0 | 0 | 1 | 1 |
¦ 4 ¦ 0 | 0 | 1 | 1 |
+---+---+---+---+---+
"""
cooc_df2 = df_compute_cooccurrences(df, "y", "x")
"""
cooc_df2:
    +----+----+
    ¦  a ¦  b ¦
+---+----+----+
¦ a ¦ 13 |  0 |
¦ b ¦  0 |  2 |
+---+----+----+
"""