added slim pandas wrapper
This commit is contained in:
Jan Kieseler
@@ -0,0 +1,255 @@
|
||||
# minipandas.py
|
||||
|
||||
from __future__ import annotations
|
||||
import numpy as np
|
||||
from typing import Dict, List, Iterable, Mapping, Optional, Union
|
||||
|
||||
ArrayLike = np.ndarray
|
||||
|
||||
class MiniFrame:
|
||||
"""
|
||||
Minimal NumPy-backed dataframe-like container.
|
||||
Holds {column_name: np.ndarray}; all columns share length on axis 0.
|
||||
Supports:
|
||||
- len(mf), mf.shape, mf.keys()
|
||||
- column access: mf["col"] -> ndarray
|
||||
- row slicing/indexing: mf[:N], mf[idx], mf[mask] -> MiniFrame (views where possible)
|
||||
- head(), copy(), to_pandas()
|
||||
"""
|
||||
def __init__(self, data: Mapping[str, np.ndarray], *, length: int | None = None):
|
||||
if not data:
|
||||
raise ValueError("MiniFrame: empty data dict.")
|
||||
|
||||
norm: dict[str, np.ndarray] = {}
|
||||
col_lengths: set[int] = set()
|
||||
|
||||
# First pass: determine N if not provided
|
||||
N = length
|
||||
if N is None:
|
||||
for v in data.values():
|
||||
if np.isscalar(v):
|
||||
continue
|
||||
arr = np.asarray(v)
|
||||
if arr.ndim >= 1:
|
||||
N = len(arr)
|
||||
break
|
||||
if N is None:
|
||||
raise ValueError("MiniFrame: cannot infer length from only scalars; pass length=...")
|
||||
|
||||
# Second pass: normalize and validate
|
||||
for k, v in data.items():
|
||||
if np.isscalar(v):
|
||||
arr = np.full(N, v) # broadcast scalar -> (N,)
|
||||
else:
|
||||
arr = np.asarray(v)
|
||||
if arr.ndim == 0:
|
||||
arr = np.full(N, arr.item())
|
||||
elif arr.ndim not in (1, 2):
|
||||
raise ValueError(f"MiniFrame: column '{k}' has unsupported ndim={arr.ndim} (only 1D or 2D).")
|
||||
if len(arr) != N:
|
||||
raise ValueError(f"MiniFrame: column '{k}' length {len(arr)} != expected {N}.")
|
||||
|
||||
norm[k] = arr
|
||||
col_lengths.add(len(arr))
|
||||
|
||||
if len(col_lengths) != 1:
|
||||
raise ValueError(f"MiniFrame: inconsistent column lengths: {col_lengths}")
|
||||
|
||||
self._data = norm
|
||||
self._length = N
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self._length
|
||||
|
||||
@property
|
||||
def scalar_cols(self) -> list[str]:
|
||||
return [k for k, v in self._data.items() if v.ndim == 1]
|
||||
|
||||
@property
|
||||
def vector_cols(self) -> list[str]:
|
||||
return [k for k, v in self._data.items() if v.ndim == 2]
|
||||
|
||||
def vector_length(self, col: str) -> int | None:
|
||||
v = self._data[col]
|
||||
return v.shape[1] if v.ndim == 2 else None
|
||||
|
||||
@property
|
||||
def shape(self) -> tuple[int, int]:
|
||||
return (self._length, len(self._data))
|
||||
|
||||
def keys(self) -> List[str]:
|
||||
return list(self._data.keys())
|
||||
|
||||
def __contains__(self, key: str) -> bool:
|
||||
return key in self._data
|
||||
|
||||
def __repr__(self) -> str:
|
||||
cols = ", ".join(self._data.keys())
|
||||
return f"<MiniFrame n={self._length}, cols=[{cols}]>"
|
||||
|
||||
# --- access ---
|
||||
def __getitem__(self, key: Union[str, int, slice, np.ndarray, List[int]]):
|
||||
# column access
|
||||
if isinstance(key, str):
|
||||
return self._data[key]
|
||||
# row selection (returns a new MiniFrame of views where possible)
|
||||
idx = key
|
||||
view = {k: v[idx] for k, v in self._data.items()}
|
||||
return MiniFrame(view)
|
||||
|
||||
def head(self, n: int = 5) -> "MiniFrame":
|
||||
return self[: min(n, self._length)]
|
||||
|
||||
def copy(self) -> "MiniFrame":
|
||||
return MiniFrame({k: v.copy() for k, v in self._data.items()})
|
||||
|
||||
def to_pandas(self):
|
||||
import pandas as pd
|
||||
df = pd.DataFrame()
|
||||
for k, v in self._data.items():
|
||||
if v.ndim == 1:
|
||||
df[k] = v
|
||||
elif v.ndim == 2:
|
||||
# expand fixed-size vectors into wide columns
|
||||
for i in range(v.shape[1]):
|
||||
df[f"{k}_{i}"] = v[:, i]
|
||||
else:
|
||||
raise ValueError(f"Column '{k}' has ndim={v.ndim}, not supported.")
|
||||
return df
|
||||
|
||||
# convenience constructor
|
||||
@staticmethod
|
||||
def from_scalars_and_arrays(
|
||||
scalars: Mapping[str, ArrayLike],
|
||||
arrays: Mapping[str, ArrayLike],
|
||||
) -> "MiniFrame":
|
||||
return MiniFrame({**scalars, **arrays})
|
||||
|
||||
def __repr__(self) -> str:
|
||||
# keep a compact technical summary
|
||||
cols = ", ".join(self._data.keys())
|
||||
return f"<MiniFrame n={self._length}, cols=[{cols}]>"
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""Pretty print like a lightweight pandas.DataFrame."""
|
||||
# Number of rows/cols to display
|
||||
n_show = min(6, self._length)
|
||||
out_lines = []
|
||||
header = " | ".join(f"{k}" for k in self._data.keys())
|
||||
out_lines.append(header)
|
||||
out_lines.append("-" * len(header))
|
||||
|
||||
for i in range(n_show):
|
||||
row_vals = []
|
||||
for v in self._data.values():
|
||||
arr = v[i]
|
||||
# scalar column
|
||||
if v.ndim == 1:
|
||||
row_vals.append(f"{arr!r}")
|
||||
# vector column: show as short array summary
|
||||
elif v.ndim == 2:
|
||||
# show first 3 elements
|
||||
if v.shape[1] > 3:
|
||||
short = ", ".join(f"{x:.3g}" for x in arr[:3])
|
||||
row_vals.append(f"[{short}, …]")
|
||||
else:
|
||||
short = ", ".join(f"{x:.3g}" for x in arr)
|
||||
row_vals.append(f"[{short}]")
|
||||
out_lines.append(" | ".join(row_vals))
|
||||
|
||||
if self._length > n_show:
|
||||
out_lines.append(f"... ({self._length - n_show} more rows)")
|
||||
out_lines.append(f"[{self._length} rows x {len(self._data)} columns]")
|
||||
return "\n".join(out_lines)
|
||||
|
||||
def describe(self):
|
||||
stats = {}
|
||||
for k, v in self._data.items():
|
||||
if v.ndim == 1:
|
||||
stats[k] = {
|
||||
"mean": float(np.mean(v)),
|
||||
"std": float(np.std(v)),
|
||||
"min": float(np.min(v)),
|
||||
"max": float(np.max(v)),
|
||||
}
|
||||
elif v.ndim == 2:
|
||||
stats[k] = {
|
||||
"mean": float(np.mean(v)),
|
||||
"std": float(np.std(v)),
|
||||
"min": float(np.min(v)),
|
||||
"max": float(np.max(v)),
|
||||
"shape": v.shape[1],
|
||||
}
|
||||
import pandas as pd
|
||||
return pd.DataFrame(stats).T
|
||||
|
||||
def to_pickle(self, path: str, protocol: int = 4):
|
||||
"""
|
||||
Save the MiniFrame to a pickle file.
|
||||
Compatible in spirit with pandas.DataFrame.to_pickle().
|
||||
"""
|
||||
import pickle
|
||||
with open(path, "wb") as f:
|
||||
pickle.dump(self._data, f, protocol=protocol)
|
||||
|
||||
|
||||
@staticmethod
|
||||
def read_pickle(path: str) -> "MiniFrame":
|
||||
"""
|
||||
Load a MiniFrame from a pickle file created by to_pickle().
|
||||
"""
|
||||
import pickle
|
||||
with open(path, "rb") as f:
|
||||
data = pickle.load(f)
|
||||
return MiniFrame(data)
|
||||
|
||||
|
||||
# --- pandas-like top-level concat ---
|
||||
|
||||
def concat(
|
||||
objs: Iterable[MiniFrame],
|
||||
axis: int = 0,
|
||||
ignore_index: bool = False,
|
||||
copy: bool = False,
|
||||
) -> MiniFrame:
|
||||
"""
|
||||
Concatenate MiniFrames along rows (axis=0). Column sets and per-column shapes must match.
|
||||
Args match pandas.concat subset for familiarity; only axis=0 is supported.
|
||||
"""
|
||||
objs = list(objs)
|
||||
if not objs:
|
||||
raise ValueError("minipandas.concat: empty iterable.")
|
||||
|
||||
if axis != 0:
|
||||
raise NotImplementedError("minipandas.concat: only axis=0 is supported.")
|
||||
|
||||
if len(objs) == 1:
|
||||
return objs[0].copy() if copy else objs[0]
|
||||
|
||||
# check column names and per-column shapes (except first dimension)
|
||||
cols = objs[0].keys()
|
||||
for f in objs[1:]:
|
||||
if f.keys() != cols:
|
||||
raise ValueError("minipandas.concat: column mismatch among frames.")
|
||||
for c in cols:
|
||||
a0, a1 = objs[0]._data[c], f._data[c]
|
||||
if a0.ndim != a1.ndim or (a0.ndim == 2 and a0.shape[1] != a1.shape[1]):
|
||||
raise ValueError(f"minipandas.concat: shape mismatch in column '{c}'.")
|
||||
|
||||
# concatenate each column along axis 0
|
||||
new_data: Dict[str, ArrayLike] = {}
|
||||
for c in cols:
|
||||
stacks = [mf._data[c] for mf in objs]
|
||||
new_data[c] = np.concatenate(stacks, axis=0)
|
||||
|
||||
out = MiniFrame(new_data)
|
||||
|
||||
# ignore_index matches pandas semantics for RangeIndex; MiniFrame has no index,
|
||||
# so the flag is accepted for API similarity but is a no-op.
|
||||
_ = ignore_index # no-op
|
||||
|
||||
if copy:
|
||||
return out.copy()
|
||||
return out
|
||||
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
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.")
|
||||
Reference in New Issue
Block a user