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