Jan Kieseler
256 lines
8.4 KiB
Python
256 lines
8.4 KiB
Python
# 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
|
|
|
|
|