diff --git a/bind/minipandas.py b/bind/minipandas.py new file mode 100644 index 0000000..22dc5cb --- /dev/null +++ b/bind/minipandas.py @@ -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"" + + # --- 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"" + + 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 + + diff --git a/bind/test_minipandas.py b/bind/test_minipandas.py new file mode 100644 index 0000000..95657fd --- /dev/null +++ b/bind/test_minipandas.py @@ -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.") \ No newline at end of file