Jan Kieseler
92 lines
2.6 KiB
Python
92 lines
2.6 KiB
Python
import numpy as np
|
|
import os
|
|
import minipandas as mpd # the teaching version of pandas :)
|
|
|
|
print("=== Creating test MiniFrames ===")
|
|
|
|
# Mix of scalar + 1D + 2D columns
|
|
mf1 = mpd.MiniFrame({
|
|
"run": 42, # scalar → broadcast
|
|
"event": np.arange(5), # 1D
|
|
"energy": np.random.rand(5, 3), # 2D fixed-size
|
|
})
|
|
|
|
mf2 = mpd.MiniFrame({
|
|
"run": 43,
|
|
"event": np.arange(5, 10),
|
|
"energy": np.random.rand(5, 3),
|
|
})
|
|
|
|
print(mf1)
|
|
print("\nScalar columns:", mf1.scalar_cols)
|
|
print("Vector columns:", mf1.vector_cols)
|
|
print("Vector length(energy):", mf1.vector_length("energy"))
|
|
|
|
print("\n=== Slicing and access ===")
|
|
print("First 3 rows:")
|
|
print(mf1[:3])
|
|
print("\nEnergy column sample:")
|
|
print(mf1["energy"][:2])
|
|
|
|
print("\n=== Concatenating like pandas.concat ===")
|
|
mf_all = mpd.concat([mf1, mf2])
|
|
print(mf_all)
|
|
print("Shape:", mf_all.shape)
|
|
|
|
print("\n=== Converting to pandas DataFrame ===")
|
|
df = mf_all.to_pandas()
|
|
print(df.head())
|
|
|
|
print("\n=== Saving and reloading with pickle ===")
|
|
path = "minipandas_test.pkl"
|
|
mf_all.to_pickle(path)
|
|
mf_loaded = mpd.MiniFrame.read_pickle(path)
|
|
print("Reloaded MiniFrame:")
|
|
print(mf_loaded)
|
|
os.remove(path)
|
|
|
|
print("\n=== Describe ===")
|
|
desc = mf_all.describe()
|
|
print(desc)
|
|
|
|
print("\nAll basic minipandas features appear to work correctly.")
|
|
|
|
|
|
|
|
# --- Create a simple MiniFrame ---
|
|
N, M = 3, 4
|
|
mf = mpd.MiniFrame({
|
|
"run": 1, # scalar -> broadcast
|
|
"event": np.arange(N), # 1D scalar column
|
|
"energy": np.arange(N * M).reshape(N, M), # 2D vector column
|
|
})
|
|
|
|
# --- Case 1: indiv_cols=True (default) ---
|
|
df_wide = mf.to_pandas(indiv_cols=True)
|
|
print("=== Expanded columns ===")
|
|
print(df_wide)
|
|
|
|
# Expected shape: N rows, 1 scalar column ('event'), 1 broadcasted column ('run'), and M expanded columns
|
|
assert list(df_wide.columns) == ["run", "event"] + [f"energy_{i}" for i in range(M)]
|
|
assert df_wide.shape == (N, 2 + M)
|
|
assert np.allclose(df_wide["energy_0"], [0, 4, 8])
|
|
|
|
# --- Case 2: indiv_cols=False (legacy nested format) ---
|
|
df_nested = mf.to_pandas(indiv_cols=False)
|
|
print("\n=== Nested legacy format ===")
|
|
print(df_nested)
|
|
|
|
# Expected columns: run, event, energy
|
|
assert list(df_nested.columns) == ["run", "event", "energy"]
|
|
assert isinstance(df_nested.loc[0, "energy"], list)
|
|
assert df_nested.loc[0, "energy"] == [0, 1, 2, 3]
|
|
|
|
# --- Round-trip compatibility check ---
|
|
mf2 = mpd.MiniFrame({
|
|
"run": df_nested["run"].to_numpy(),
|
|
"event": df_nested["event"].to_numpy(),
|
|
"energy": np.stack(df_nested["energy"].to_numpy()),
|
|
})
|
|
|
|
assert np.allclose(mf["energy"], mf2["energy"])
|
|
print("\n✅ All checks passed successfully!") |