How to create empty rows for missing time periods, quickly

Viewed 42

Take the following data structure:

  id  t  x  y     0
0  0  2  1 -1  0.37
1  0  3  2 -2 -0.31
2  0  4  0  0 -0.48
3  1  3  1 -1 -0.12
4  1  4  0  0 -0.08
5  2  1  0  0  0.56
6  2  3  1 -1 -0.18

There are N observations of each "id", observed n times each across several integer-valued time periods. I want to expand the data to add rows for each possible time for each id:

        x    y     0
id t                
0  1  NaN  NaN   NaN
   2  1.0 -1.0  0.37
   3  2.0 -2.0 -0.31
   4  0.0  0.0 -0.48
1  1  NaN  NaN   NaN
   2  NaN  NaN   NaN
   3  1.0 -1.0 -0.12
   4  0.0  0.0 -0.08
2  1  0.0  0.0  0.56
   2  NaN  NaN   NaN
   3  1.0 -1.0 -0.18
   4  NaN  NaN   NaN

So far I am using pandas, but I would be open to using numpy or something else. I just need a fast way to identify which rows need to be added, and add them in order.

Code for replication:

import numpy as np
import pandas as pd
from itertools import product
import time


def expand_grid(grid: dict) -> pd.DataFrame:
    return pd.DataFrame([row for row in product(*grid.values())],
                        columns=grid.keys())

def makedf(N, n, k):
    np.random.seed(4)
    df = []
    for i in range(N):
        l = np.random.choice(n)+1
        id = [str(i) for j in range(l)]
        t = np.random.choice(n, size = l, replace=False)
        x = list(range(l))
        y = [-1*j for j in x]
        res = pd.DataFrame(dict(id = id, t = t, x = x, y = y))
        df.append(res)
    df = pd.concat(df).sort_values(['id', 't']).reset_index(drop=True)
    xx = np.round(np.stack([np.random.normal(size=len(df)) for i in range(k)]),2)
    df = pd.concat([df, pd.DataFrame(xx.T, columns=[str(i) for i in range(k)])], axis=1)
    return df


df = makedf(N=3, n=5, k=1)

Here's a more realistically-sized example along with my (slow) pandas join-based implementation:

df = makedf(N=1000, n=1000, k=700)
idvar = 'id'
tvar = 't'
dd = df.copy()
st = time.time()
full_idx = expand_grid({idvar: sorted(list(dd[idvar].unique())),
                        tvar: sorted(list(dd[tvar].unique()))})
A = time.time()
print(f"Expanded grid {A-st}")
idx = full_idx.set_index([idvar, tvar])
dd = dd.set_index([idvar, tvar])
B = time.time()
print(f"create index {B-A}")
expanded_df = idx.join(dd)
C = time.time()
print(f"join {C-B}")

Expanded grid 0.34296512603759766
create index 0.891047477722168
join 3.028926134109497
1 Answers

This should work:

i = pd.MultiIndex.from_product([df['id'].unique(),df['t'].unique()])
df.set_index(['id','t']).reindex(i).sort_index()
Related