diff --git a/giant/data/loader.py b/giant/data/loader.py index e84e492..7897b2e 100644 --- a/giant/data/loader.py +++ b/giant/data/loader.py @@ -1,13 +1,16 @@ from dataclasses import dataclass, field from pathlib import Path -from typing import Iterator +from typing import TYPE_CHECKING, Any, Iterator, Mapping import numpy as np -import pandas as pd +import polars as pl import pyarrow.parquet as pq from giant.constants import K_MAX +if TYPE_CHECKING: + from giant.data.scan import ValueStat + # A manifest is a plain text file listing one parquet path per line, used to # name a curated subset of files (e.g. a train/holdout pool) without copying # or symlinking the underlying parquet files. Lines are resolved relative to @@ -77,88 +80,75 @@ def find_parquet_files(path: str | Path) -> list[Path]: return [p] -def _pad_list_col(series: pd.Series, K: int, fill: float = 0.0) -> np.ndarray: - """Pad / truncate a list-valued Series to fixed width K → (N, K) float32.""" - out = np.full((len(series), K), fill, dtype=np.float32) - for i, lst in enumerate(series): - if lst is not None and len(lst) > 0: - n = min(len(lst), K) - out[i, :n] = lst[:n] - return out +def _pad_list_column(df: pl.DataFrame, col: str, k: int, fill, dtype: type[pl.DataType] | pl.DataType) -> np.ndarray: + """Pad / truncate a list-valued column to fixed width `k` → (N, k) numpy array. - -def _pad_list_col_int(series: pd.Series, K: int, fill: int = 0) -> np.ndarray: - """Pad / truncate a list-valued integer Series to fixed width K → (N, K) int64.""" - out = np.full((len(series), K), fill, dtype=np.int64) - for i, lst in enumerate(series): - if lst is not None and len(lst) > 0: - n = min(len(lst), K) - out[i, :n] = lst[:n] - return out - - -def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndarray: - """Pad three list-valued direction columns → (N, K, 3) float32. - - Padding direction defaults to (0,0,1) (forward) so it is a valid unit vector. + Concatenating `k` fill values before truncating to `k` guarantees every + row ends up with exactly `k` non-null elements regardless of how short + (including empty) or long the original list was, so `list.to_array(k)` + (a fixed-size-array dtype) converts to a plain 2D numpy array with a + single vectorized expression — no per-row Python loop. """ - N = len(dx) - out = np.zeros((N, K, 3), dtype=np.float32) - out[:, :, 2] = 1.0 - for i in range(N): - lx, ly, lz = dx.iloc[i], dy.iloc[i], dz.iloc[i] - if lx is not None and len(lx) > 0: - n = min(len(lx), K) - out[i, :n, 0] = lx[:n] - out[i, :n, 1] = ly[:n] - out[i, :n, 2] = lz[:n] - return out + fill_tail = pl.lit([fill] * k, dtype=pl.List(dtype)) + out = df.select(pl.col(col).cast(pl.List(dtype)).list.concat(fill_tail).list.head(k).list.to_array(k).alias("_p")) + return out["_p"].to_numpy() -def _df_to_dict(df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]: +def _pad_dir_col(df: pl.DataFrame, dx: str, dy: str, dz: str, k: int) -> np.ndarray: + """Pad three list-valued direction columns → (N, k, 3) float32. + + Padding direction defaults to (0, 0, 1) (forward) so it is a valid unit vector. + """ + px = _pad_list_column(df, dx, k, 0.0, pl.Float64) + py = _pad_list_column(df, dy, k, 0.0, pl.Float64) + pz = _pad_list_column(df, dz, k, 1.0, pl.Float64) + return np.stack([px, py, pz], axis=-1).astype(np.float32) + + +def _df_to_dict(df: pl.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]: has_sec_lists = "sec_E_list" in df.columns d: dict[str, np.ndarray] = { "event_id": _offset_event_id(df["event_id"].to_numpy(), offset), - "pdg": df["pdg"].to_numpy(dtype=np.int32), - "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32), - "pre_E": df["pre_E"].to_numpy(dtype=np.float32), - "pre_dir": df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32), - "material": df["material"].to_numpy(dtype=object), - "layer_id": df["layer_id"].to_numpy(dtype=np.int32), - "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32), - "e_sec": df["e_sec"].to_numpy(dtype=np.float32), + "pdg": df["pdg"].to_numpy().astype(np.int32), + "pre_pos": df.select(["pre_x", "pre_y", "pre_z"]).to_numpy().astype(np.float32), + "pre_E": df["pre_E"].to_numpy().astype(np.float32), + "pre_dir": df.select(["pre_dx", "pre_dy", "pre_dz"]).to_numpy().astype(np.float32), + "material": df["material"].to_numpy().astype(object), + "layer_id": df["layer_id"].to_numpy().astype(np.int32), + "n_sec": df["child_track_ids"].list.len().to_numpy().astype(np.int32), + "e_sec": df["e_sec"].to_numpy().astype(np.float32), # The physics process that ended the step (e.g. "compt", "phot", # "eBrem") — a post-step outcome, so it's a router/classifier # supervision label only, never conditioning (see build_process_map* # / ProcessRouter). Guarded like has_sec_lists: older parquet # conversions predating this column still load fine. "process": ( - df["process"].to_numpy(dtype=object) if "process" in df.columns else np.full(len(df), "", dtype=object) + df["process"].to_numpy().astype(object) if "process" in df.columns else np.full(len(df), "", dtype=object) ), - "step_length": df["step_length"].to_numpy(dtype=np.float32), - "post_E": df["post_E"].to_numpy(dtype=np.float32), - "delta_e": (df["pre_E"] - df["post_E"]).to_numpy(dtype=np.float32), - "edep": df["edep"].to_numpy(dtype=np.float32), - "post_dir": df[["post_dx", "post_dy", "post_dz"]].to_numpy(dtype=np.float32), - "post_pos": df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32), + "step_length": df["step_length"].to_numpy().astype(np.float32), + "post_E": df["post_E"].to_numpy().astype(np.float32), + "delta_e": (df["pre_E"] - df["post_E"]).to_numpy().astype(np.float32), + "edep": df["edep"].to_numpy().astype(np.float32), + "post_dir": df.select(["post_dx", "post_dy", "post_dz"]).to_numpy().astype(np.float32), + "post_pos": df.select(["post_x", "post_y", "post_z"]).to_numpy().astype(np.float32), } if has_sec_lists: - d["sec_E_list"] = _pad_list_col(df["sec_E_list"], k_max) - d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], k_max) - d["sec_dir_list"] = _pad_dir_col(df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], k_max) + d["sec_E_list"] = _pad_list_column(df, "sec_E_list", k_max, 0.0, pl.Float64).astype(np.float32) + d["sec_pdg_list"] = _pad_list_column(df, "sec_pdg_list", k_max, 0, pl.Int64).astype(np.int64) + d["sec_dir_list"] = _pad_dir_col(df, "sec_dx_list", "sec_dy_list", "sec_dz_list", k_max) return d def load_steps(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]: - return _df_to_dict(pd.read_parquet(path), offset=offset, k_max=k_max) + return _df_to_dict(pl.read_parquet(path), offset=offset, k_max=k_max) def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray: """Read only the event_id column — cheap scan for split assignment.""" - ids = pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy() + ids = pl.read_parquet(path, columns=["event_id"])["event_id"].to_numpy() return _offset_event_id(ids, offset) @@ -170,7 +160,7 @@ def iter_file_chunks(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> I module constant for callers that don't care (e.g. Stage-1-only reads).""" pf = pq.ParquetFile(path) for i in range(pf.num_row_groups): - yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset, k_max=k_max) + yield _df_to_dict(pl.DataFrame(pf.read_row_group(i)), offset=offset, k_max=k_max) _COND_COLS = [ @@ -190,17 +180,17 @@ _COND_COLS = [ ] -def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]: +def _cond_df_to_dict(df: pl.DataFrame, offset: int = 0) -> dict[str, np.ndarray]: return { "event_id": _offset_event_id(df["event_id"].to_numpy(), offset), - "pdg": df["pdg"].to_numpy(dtype=np.int32), - "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32), - "pre_E": df["pre_E"].to_numpy(dtype=np.float32), - "pre_dir": df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32), - "material": df["material"].to_numpy(dtype=object), - "layer_id": df["layer_id"].to_numpy(dtype=np.int32), - "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32), - "e_sec": df["e_sec"].to_numpy(dtype=np.float32), + "pdg": df["pdg"].to_numpy().astype(np.int32), + "pre_pos": df.select(["pre_x", "pre_y", "pre_z"]).to_numpy().astype(np.float32), + "pre_E": df["pre_E"].to_numpy().astype(np.float32), + "pre_dir": df.select(["pre_dx", "pre_dy", "pre_dz"]).to_numpy().astype(np.float32), + "material": df["material"].to_numpy().astype(object), + "layer_id": df["layer_id"].to_numpy().astype(np.int32), + "n_sec": df["child_track_ids"].list.len().to_numpy().astype(np.int32), + "e_sec": df["e_sec"].to_numpy().astype(np.float32), } @@ -208,7 +198,7 @@ def iter_cond_chunks(path: str | Path, offset: int = 0) -> Iterator[dict[str, np """Yield conditioning-only row-groups (no post-step columns read from disk).""" pf = pq.ParquetFile(path) for i in range(pf.num_row_groups): - yield _cond_df_to_dict(pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset) + yield _cond_df_to_dict(pl.DataFrame(pf.read_row_group(i, columns=_COND_COLS)), offset=offset) def build_index_maps( @@ -225,42 +215,28 @@ def build_index_maps( def build_index_maps_from_files( files: list[Path], ) -> tuple[dict[int, int], dict[str, int]]: - """Scan only pdg and material columns across all files (2-column read).""" - pdg_vals: set[int] = set() - mat_vals: set[str] = set() - for path in files: - df = pd.read_parquet(path, columns=["pdg", "material"]) - pdg_vals.update(int(v) for v in df["pdg"].unique()) - mat_vals.update(str(v) for v in df["material"].unique()) + """Scan only pdg and material columns across all files (fused single-pass scan).""" + from giant.data.scan import ScanRequest, scan_metadata + + result = scan_metadata(files, ScanRequest(pdg=True, material=True)) + assert result.pdg is not None and result.material is not None return ( - {v: i for i, v in enumerate(sorted(pdg_vals))}, - {v: i for i, v in enumerate(sorted(mat_vals))}, + {v: i for i, v in enumerate(sorted(result.pdg))}, + {v: i for i, v in enumerate(sorted(result.material))}, ) -def _accumulate_value_counts(counts: dict, series: pd.Series, cast) -> None: - for name, count in series.value_counts().items(): - name = cast(name) - counts[name] = counts.get(name, 0) + int(count) - - -def _rank_by_frequency_from_files(files: list[Path], column: str, cast) -> dict: - """Scan `column` across `files` and return `{cast(value): total_count}`, - accumulated in file order (see `fingerprint_files`'s docstring on why - scan order — not a normalized/sorted order — is preserved: it drives - tie-breaking in the frequency ranking below).""" - counts: dict = {} - for path in files: - df = pd.read_parquet(path, columns=[column]) - _accumulate_value_counts(counts, df[column], cast) - return counts - - -def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict, dict]: +def _topn_plus_other_map(counts: "Mapping[Any, ValueStat]", n_classes: int) -> tuple[dict, dict, dict]: """Frequency-capped value->index map: the `n_classes - 1` most frequent keys get their own index; every rarer key is bucketed into a shared "other" index (`n_classes - 1`). + `counts` maps each key to something with `.count` and `.first_seen` + attributes (`giant.data.scan.ValueStat`) — ties in `.count` are broken by + `.first_seen` (whichever value was scanned first: file order, then row + order within a file — see `giant.data.scan`'s module docstring). This is + an explicit, documented contract, not an accident of iteration order. + Returns `(class_map, other_members, class_counts)` — `other_members` is `{key: count}` for every key bucketed into "other" (the empirical within-bucket distribution, for `other_policy = "sample"` at rollout); @@ -270,15 +246,15 @@ def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict, dict (gitea #44) needs and that would otherwise be dropped once `counts` is collapsed into `class_map`. """ - ranked = sorted(counts, key=lambda k: counts[k], reverse=True) + ranked = sorted(counts, key=lambda k: (-counts[k].count, counts[k].first_seen)) keep = ranked[: max(n_classes - 1, 0)] class_map = {k: i for i, k in enumerate(keep)} - class_counts = {i: counts[k] for i, k in enumerate(keep)} + class_counts = {i: counts[k].count for i, k in enumerate(keep)} other_idx = n_classes - 1 other_members: dict = {} for k in ranked[len(keep) :]: class_map[k] = other_idx - other_members[k] = counts[k] + other_members[k] = counts[k].count if other_members: class_counts[other_idx] = sum(other_members.values()) return class_map, other_members, class_counts @@ -294,8 +270,11 @@ def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str, mirrors how `build_features` clamps the n_sec label to K_MAX for the fixed-width n_sec_head classifier. """ - counts = _rank_by_frequency_from_files(files, "process", str) - class_map, _, _ = _topn_plus_other_map(counts, n_experts) + from giant.data.scan import ScanRequest, scan_metadata + + result = scan_metadata(files, ScanRequest(process=True)) + assert result.process is not None + class_map, _, _ = _topn_plus_other_map(result.process, n_experts) return class_map @@ -327,8 +306,13 @@ def build_topn_map_from_files(files: list[Path], column: str, n_classes: int, ca later for `other_policy = "sample"` at rollout — computed now since it's free during this same scan. """ - counts = _rank_by_frequency_from_files(files, column, cast) - class_map, other_members, class_counts = _topn_plus_other_map(counts, n_classes) + from giant.data.scan import ScanRequest, scan_metadata + + if column != "material": + raise ValueError(f"build_topn_map_from_files only supports column='material', got {column!r}") + result = scan_metadata(files, ScanRequest(material=True)) + assert result.material is not None + class_map, other_members, class_counts = _topn_plus_other_map(result.material, n_classes) return TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts) @@ -349,16 +333,9 @@ def build_pdg_topn_map_from_files(files: list[Path], n_classes: int) -> TopNMap: join (see `_df_to_dict`'s `has_sec_lists` guard) — silently skipped for those, same convention as elsewhere in this module. """ - counts: dict = {} - for path in files: - columns = ["pdg"] - has_sec = "sec_pdg_list" in pq.ParquetFile(path).schema_arrow.names - if has_sec: - columns.append("sec_pdg_list") - df = pd.read_parquet(path, columns=columns) - _accumulate_value_counts(counts, df["pdg"], int) - if has_sec: - exploded = df["sec_pdg_list"].explode().dropna() - _accumulate_value_counts(counts, exploded, int) - class_map, other_members, class_counts = _topn_plus_other_map(counts, n_classes) + from giant.data.scan import ScanRequest, scan_metadata + + result = scan_metadata(files, ScanRequest(pooled_pdg=True)) + assert result.pooled_pdg is not None + class_map, other_members, class_counts = _topn_plus_other_map(result.pooled_pdg, n_classes) return TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts) diff --git a/giant/data/scan.py b/giant/data/scan.py new file mode 100644 index 0000000..b26bd74 --- /dev/null +++ b/giant/data/scan.py @@ -0,0 +1,171 @@ +"""Fused metadata scan over one or more parquet files. + +`giant.pipeline.run_setup_stage` needs several distinct frequency summaries +before training can start — the event-id → row-count index (for the train/val +split), the pdg/material vocabularies, an optional physics-process count, and +a pooled pdg count (primary + secondary species, for onehot conditioning). +Each of those used to be its own full `pd.read_parquet(path, columns=[...])` +per file (`giant.data.loader`'s old `_rank_by_frequency_from_files` / +`build_index_maps_from_files` / `build_pdg_topn_map_from_files`) — up to five +separate reads of the same file. `scan_metadata` answers all of them in one +`pl.collect_all` per file instead, sharing the file open/decompress cost. + +Every requested count comes back keyed by value, as a `ValueStat(count, +first_seen)`. `first_seen` is the value's row ordinal — file order (as given +in `files`), then row order within a file — via `row_index_name` on the +per-file lazy scan plus a running row offset across files. This is what +`giant.data.loader._topn_plus_other_map`'s frequency-ranking tie-break keys +on: among equally-frequent values, whichever was scanned first wins its own +class slot. That is an explicit, documented contract (this module is where +it's implemented), not an accident of iteration order. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import polars as pl + +from giant.data.loader import _offset_event_id, event_id_offset + + +@dataclass(frozen=True) +class ScanRequest: + """Which aggregations to compute. Every field defaults off so a caller + only pays for what it actually needs.""" + + event_index: bool = False + pdg: bool = False + material: bool = False + process: bool = False + pooled_pdg: bool = False + """pdg ∪ exploded sec_pdg_list — both roles a PDG code plays (primary + species and secondary species), pooled into one count per code. See + `giant.data.loader.build_pdg_topn_map_from_files`'s docstring for why.""" + + +@dataclass(frozen=True) +class ValueStat: + count: int + first_seen: int + + +@dataclass +class MetadataScan: + event_index: tuple[np.ndarray, np.ndarray] | None = None + """(unique_ids, counts), ids ascending — matches + `setup_cache.compute_event_index_from_files`'s return shape.""" + pdg: dict[int, ValueStat] | None = None + material: dict[str, ValueStat] | None = None + process: dict[str, ValueStat] | None = None + pooled_pdg: dict[int, ValueStat] | None = None + + +def _group_lazy(path: Path, column: str) -> pl.LazyFrame: + return ( + pl.scan_parquet(path, row_index_name="__row") + .select(column, "__row") + .group_by(column) + .agg(pl.len().alias("__count"), pl.col("__row").min().alias("__first_row")) + ) + + +def _pooled_pdg_lazy(path: Path, has_sec_pdg_list: bool) -> pl.LazyFrame: + lf = pl.scan_parquet(path, row_index_name="__row") + parts = [lf.select(pl.col("pdg").alias("__val"), "__row")] + if has_sec_pdg_list: + parts.append(lf.select(pl.col("sec_pdg_list").alias("__val"), "__row").explode("__val").drop_nulls("__val")) + combined = pl.concat(parts) + return combined.group_by("__val").agg(pl.len().alias("__count"), pl.col("__row").min().alias("__first_row")) + + +def _merge_counts(acc: dict, df: pl.DataFrame, column: str, row_offset: int, cast) -> None: + for key, count, first_row in zip( + df[column].to_list(), df["__count"].to_list(), df["__first_row"].to_list(), strict=True + ): + key = cast(key) + first_seen = row_offset + int(first_row) + if key in acc: + prev_count, prev_first = acc[key] + acc[key] = (prev_count + int(count), min(prev_first, first_seen)) + else: + acc[key] = (int(count), first_seen) + + +def scan_metadata(files: list[Path], request: ScanRequest) -> MetadataScan: + """Scan `files` once (one `pl.collect_all` per file) and return every + aggregation `request` asks for. Files with zero rows contribute nothing + but still advance nothing (no row_offset change, nothing to merge).""" + event_id_parts: list[tuple[np.ndarray, np.ndarray]] = [] + pdg_acc: dict[int, tuple[int, int]] = {} + material_acc: dict[str, tuple[int, int]] = {} + process_acc: dict[str, tuple[int, int]] = {} + pooled_pdg_acc: dict[int, tuple[int, int]] = {} + + row_offset = 0 + for file_idx, path in enumerate(files): + keys: list[str] = [] + lazies: list[pl.LazyFrame] = [] + + if request.event_index: + keys.append("event_id") + lazies.append(_group_lazy(path, "event_id")) + if request.pdg: + keys.append("pdg") + lazies.append(_group_lazy(path, "pdg")) + if request.material: + keys.append("material") + lazies.append(_group_lazy(path, "material")) + if request.process: + keys.append("process") + lazies.append(_group_lazy(path, "process")) + if request.pooled_pdg: + has_sec = "sec_pdg_list" in pl.scan_parquet(path).collect_schema().names() + keys.append("pooled_pdg") + lazies.append(_pooled_pdg_lazy(path, has_sec)) + + keys.append("__n") + lazies.append(pl.scan_parquet(path).select(pl.len().alias("__n"))) + + results = dict(zip(keys, pl.collect_all(lazies, engine="streaming"), strict=True)) + n_rows = int(results["__n"].item()) if len(results["__n"]) else 0 + + if request.event_index: + df = results["event_id"] + ids = _offset_event_id(df["event_id"].to_numpy(), event_id_offset(file_idx)) + counts = df["__count"].to_numpy().astype(np.int64) + if ids.size: + event_id_parts.append((ids, counts)) + if request.pdg: + _merge_counts(pdg_acc, results["pdg"], "pdg", row_offset, int) + if request.material: + _merge_counts(material_acc, results["material"], "material", row_offset, str) + if request.process: + _merge_counts(process_acc, results["process"], "process", row_offset, str) + if request.pooled_pdg: + _merge_counts(pooled_pdg_acc, results["pooled_pdg"], "__val", row_offset, int) + + row_offset += n_rows + + event_index = None + if request.event_index: + if event_id_parts: + all_ids = np.concatenate([p[0] for p in event_id_parts]) + all_counts = np.concatenate([p[1] for p in event_id_parts]) + order = np.argsort(all_ids, kind="stable") + event_index = (all_ids[order], all_counts[order]) + else: + event_index = (np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64)) + + def _to_stats(acc: dict) -> dict: + return {k: ValueStat(*v) for k, v in acc.items()} + + return MetadataScan( + event_index=event_index, + pdg=_to_stats(pdg_acc) if request.pdg else None, + material=_to_stats(material_acc) if request.material else None, + process=_to_stats(process_acc) if request.process else None, + pooled_pdg=_to_stats(pooled_pdg_acc) if request.pooled_pdg else None, + ) diff --git a/giant/data/setup_cache.py b/giant/data/setup_cache.py index 1800b27..6147392 100644 --- a/giant/data/setup_cache.py +++ b/giant/data/setup_cache.py @@ -23,7 +23,7 @@ import numpy as np from giant import config from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM -from giant.data.loader import TopNMap, event_id_offset, load_event_ids +from giant.data.loader import TopNMap from giant.data.transforms import Normalizer, sorted_membership # Bump manually on a change to the data-encoding semantics (e.g. a future @@ -349,12 +349,21 @@ def save( def compute_event_index_from_files(files: list[Path]) -> tuple[np.ndarray, np.ndarray]: - """Unique event ids + per-event row (step) counts, across all `files`.""" + """Unique event ids + per-event row (step) counts, across all `files`. + + Computed via a streaming per-file `group_by("event_id")` (see + `giant.data.scan.scan_metadata`) rather than concatenating every row's + raw event_id across every file before `np.unique` — the latter's peak + memory is 8 bytes x total row count; this is bounded by the (much + smaller) unique event count instead. + """ + from giant.data.scan import ScanRequest, scan_metadata + if not files: return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64) - all_ids = np.concatenate([load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)]) - unique_ids, counts = np.unique(all_ids, return_counts=True) - return unique_ids, counts + result = scan_metadata(files, ScanRequest(event_index=True)) + assert result.event_index is not None + return result.event_index def n_train_steps_for_split(unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray) -> int: diff --git a/giant/geometry.py b/giant/geometry.py index 3d137e8..824ddf0 100644 --- a/giant/geometry.py +++ b/giant/geometry.py @@ -28,7 +28,7 @@ from pathlib import Path from typing import Any, Iterable import numpy as np -import pandas as pd +import polars as pl import pyarrow.parquet as pq _INSTALL_HINT = "the geometry oracle needs scikit-learn — install it with `uv sync --extra cpu --extra geometry`" @@ -301,14 +301,18 @@ def _fit_slab_lookup( edges = np.linspace(z_min, z_max, n_bins + 1) bin_idx = np.clip(np.searchsorted(edges, z, side="right") - 1, 0, n_bins - 1) + # pandas' groupby(...).size() sorts group keys ascending by default, so + # a tie in `n` for the same bin (equal counts split between two + # material/layer_id combos) resolves to the lexicographically-first + # combo — matched here by sorting on the keys first, then a + # maintain_order-stable sort on `n` so ties keep that key order. counts = ( - pd.DataFrame({"bin": bin_idx, "material": mat, "layer_id": lay}) - .groupby(["bin", "material", "layer_id"]) - .size() - .to_frame("n") - .reset_index() - .sort_values("n", ascending=False) - .drop_duplicates("bin") + pl.DataFrame({"bin": bin_idx, "material": mat, "layer_id": lay}) + .group_by(["bin", "material", "layer_id"]) + .agg(pl.len().alias("n")) + .sort(["bin", "material", "layer_id"]) + .sort("n", descending=True, maintain_order=True) + .unique(subset="bin", keep="first", maintain_order=True) ) bin_material = np.full(n_bins, "", dtype=object) diff --git a/giant/pipeline.py b/giant/pipeline.py index 79b3402..f9c4618 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -15,14 +15,12 @@ from giant.constants import ( from giant.data import setup_cache from giant.data.loader import ( TopNMap, + _topn_plus_other_map, event_id_offset, find_parquet_files, iter_file_chunks, - build_index_maps_from_files, - build_pdg_topn_map_from_files, - build_process_map_from_files, - build_topn_map_from_files, ) +from giant.data.scan import MetadataScan, ScanRequest, scan_metadata from giant.data.transforms import ( Normalizer, build_features, @@ -127,12 +125,71 @@ def run_setup_stage( loaded = setup_cache.load(data, files, echo=echo) cache = loaded if loaded is not None else setup_cache.SetupCache.empty(files) - if cache is not None and cache.event_index is not None: + # Every section below first asks the cache; whatever's missing is + # collected into one ScanRequest and answered by a single fused scan + # (giant.data.scan.scan_metadata), instead of a separate full pass per + # section (event index, vocab, process counts, pdg/material top-N counts + # used to each re-open and re-read every file on their own). + particle_cfg = cfg["conditioning"]["particle"] + material_cfg = cfg["conditioning"]["material"] + particle_type_cfg = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type")) + particle_type_target = particle_type_cfg.target + + # A process map is needed if either stage's router reads the physics + # process label (type="process"). Only one map is built even if both + # stages want one — see the module-level note in giant/cli.py's + # _router_total_experts for why composed-router n_experts isn't a plain + # int; process routers are never composed in practice, so this doesn't + # need that generality. + process_router_cfg = next( + (r for r in (stage1_router, stage2_router) if r.get("enabled") and r.get("type") == "process"), + None, + ) + process_n_experts = process_router_cfg["n_experts"] if process_router_cfg is not None else None + + need_pdg_onehot = particle_cfg["type"] == "onehot" + need_sec_type_onehot = particle_type_target == "onehot" + sec_type_n_classes = ( + resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"]) if need_sec_type_onehot else None + ) + need_material_onehot = material_cfg["type"] == "onehot" + material_n_classes = material_cfg["emb_dim"] if need_material_onehot else None + + def _topn_cached(axis: str, n_classes: int) -> TopNMap | None: + return cache.topn_maps.get(setup_cache.topn_key(axis, n_classes)) if cache is not None else None + + need_event_index = cache is None or cache.event_index is None + need_vocab = cache is None or cache.vocab is None + need_process = process_n_experts is not None and (cache is None or cache.proc_maps.get(process_n_experts) is None) + # The PDG axis is used independently by conditioning.particle.type="onehot" + # (cond_cat's onehot feature) and stage2_model.particle_type.target="onehot" + # (secondary-species decode) — their class counts can now differ (gitea + # #29: stage2_model.particle_type.n_classes, 0 = inherit + # conditioning.particle.emb_dim), but both are built from the same + # pooled pdg-count scan, so a cache miss on either one asks for it. + need_pdg_pooled = (need_pdg_onehot and _topn_cached("pdg", particle_cfg["emb_dim"]) is None) or ( + need_sec_type_onehot and sec_type_n_classes is not None and _topn_cached("pdg", sec_type_n_classes) is None + ) + need_material_topn = ( + need_material_onehot and material_n_classes is not None and _topn_cached("material", material_n_classes) is None + ) + + request = ScanRequest( + event_index=need_event_index, + pdg=need_vocab, + material=need_vocab or need_material_topn, + process=need_process, + pooled_pdg=need_pdg_pooled, + ) + scan = scan_metadata(files, request) if request != ScanRequest() else MetadataScan() + + if not need_event_index: unique_ids, counts = cache.event_index echo(f"event index: cache hit ({len(unique_ids):,} unique events)") else: echo("scanning event IDs …") - unique_ids, counts = setup_cache.compute_event_index_from_files(files) + assert scan.event_index is not None + unique_ids, counts = scan.event_index if cache is not None: cache.event_index = (unique_ids, counts) @@ -141,55 +198,34 @@ def run_setup_stage( n_train_steps = setup_cache.n_train_steps_for_split(unique_ids, counts, events_arr) echo(f" {int(counts.sum()):,} steps | {len(train_events)} train events | {len(val_events)} val events") - if cache is not None and cache.vocab is not None: + if not need_vocab: pdg_map, mat_map = cache.vocab echo(f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, {len(mat_map)} materials)") else: echo("building vocabulary maps …") - pdg_map, mat_map = build_index_maps_from_files(files) + assert scan.pdg is not None and scan.material is not None + pdg_map = {v: i for i, v in enumerate(sorted(scan.pdg))} + mat_map = {v: i for i, v in enumerate(sorted(scan.material))} echo(f" {len(pdg_map)} PDG codes | {len(mat_map)} materials") if cache is not None: cache.vocab = (pdg_map, mat_map) - # A process map is needed if either stage's router reads the physics - # process label (type="process"). Only one map is built even if both - # stages want one — see the module-level note in giant/cli.py's - # _router_total_experts for why composed-router n_experts isn't a plain - # int; process routers are never composed in practice, so this doesn't - # need that generality. proc_map: dict[str, int] | None = None - process_router_cfg = next( - (r for r in (stage1_router, stage2_router) if r.get("enabled") and r.get("type") == "process"), - None, - ) if process_router_cfg is not None: - n_experts = process_router_cfg["n_experts"] - cached_proc_map = cache.proc_maps.get(n_experts) if cache is not None else None - if cached_proc_map is not None: + assert process_n_experts is not None + if not need_process: + assert cache is not None + cached_proc_map = cache.proc_maps.get(process_n_experts) + assert cached_proc_map is not None proc_map = cached_proc_map - echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {n_experts} experts)") + echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {process_n_experts} experts)") else: echo("building process vocabulary …") - proc_map = build_process_map_from_files(files, n_experts=n_experts) - echo(f" {len(proc_map)} process labels mapped to {n_experts} experts") + assert scan.process is not None + proc_map, _, _ = _topn_plus_other_map(scan.process, process_n_experts) + echo(f" {len(proc_map)} process labels mapped to {process_n_experts} experts") if cache is not None: - cache.proc_maps[n_experts] = proc_map - - # Top-N-plus-other maps for onehot conditioning/type axes. - # The PDG axis is used independently by conditioning.particle.type="onehot" - # (cond_cat's onehot feature) and stage2_model.particle_type.target="onehot" - # (secondary-species decode) — their class counts can now differ (gitea - # #29: stage2_model.particle_type.n_classes, 0 = inherit - # conditioning.particle.emb_dim), so each is resolved and built - # independently via _pdg_topn below. cache.topn_maps is keyed by - # (axis, n_classes) (setup_cache.topn_key), so when the two resolve to - # the same N the second call is a cache hit against the first — no extra - # scan in the common case where they still match. The material axis is - # independent of both. - particle_cfg = cfg["conditioning"]["particle"] - material_cfg = cfg["conditioning"]["material"] - particle_type_cfg = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type")) - particle_type_target = particle_type_cfg.target + cache.proc_maps[process_n_experts] = proc_map def _pdg_topn(n_classes: int) -> TopNMap: cache_key = setup_cache.topn_key("pdg", n_classes) @@ -198,33 +234,36 @@ def run_setup_stage( echo(f"pdg top-N map: cache hit ({len(cached.class_map)} codes, {n_classes} classes)") return cached echo("building pdg top-N map …") - topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes) + assert scan.pooled_pdg is not None + class_map, other_members, class_counts = _topn_plus_other_map(scan.pooled_pdg, n_classes) + topn_map = TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts) echo(f" {len(topn_map.class_map)} pdg codes mapped to {n_classes} classes") if cache is not None: cache.topn_maps[cache_key] = topn_map return topn_map - pdg_topn_map: TopNMap | None = None - if particle_cfg["type"] == "onehot": - pdg_topn_map = _pdg_topn(particle_cfg["emb_dim"]) - + pdg_topn_map: TopNMap | None = _pdg_topn(particle_cfg["emb_dim"]) if need_pdg_onehot else None sec_type_topn_map: TopNMap | None = None - if particle_type_target == "onehot": - sec_type_n_classes = resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"]) + if need_sec_type_onehot: + assert sec_type_n_classes is not None sec_type_topn_map = _pdg_topn(sec_type_n_classes) mat_topn_map: TopNMap | None = None - if material_cfg["type"] == "onehot": - n_classes = material_cfg["emb_dim"] - cache_key = setup_cache.topn_key("material", n_classes) + if need_material_onehot: + assert material_n_classes is not None + cache_key = setup_cache.topn_key("material", material_n_classes) cached = cache.topn_maps.get(cache_key) if cache is not None else None if cached is not None: mat_topn_map = cached - echo(f"material top-N map: cache hit ({len(mat_topn_map.class_map)} materials, {n_classes} classes)") + echo( + f"material top-N map: cache hit ({len(mat_topn_map.class_map)} materials, {material_n_classes} classes)" + ) else: echo("building material top-N map …") - mat_topn_map = build_topn_map_from_files(files, "material", n_classes=n_classes, cast=str) - echo(f" {len(mat_topn_map.class_map)} materials mapped to {n_classes} classes") + assert scan.material is not None + class_map, other_members, class_counts = _topn_plus_other_map(scan.material, material_n_classes) + mat_topn_map = TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts) + echo(f" {len(mat_topn_map.class_map)} materials mapped to {material_n_classes} classes") if cache is not None: cache.topn_maps[cache_key] = mat_topn_map @@ -456,17 +495,30 @@ def run_train_job( ) pin = device.type == "cuda" + # DataLoader worker subprocesses default to fork() on Linux, but by the + # time they're created this process has already run polars queries + # (run_setup_stage's fused metadata scan, above) — polars' native + # (rayon) thread pool doesn't survive a fork: a worker that inherits it + # mid-fork deadlocks the instant it touches polars itself, which + # StreamingStepsDataset's iter_file_chunks now does on every row group. + # "spawn" starts each worker as a fresh interpreter with no inherited + # thread-pool state, avoiding that hazard entirely. Only matters when + # workers actually exist — num_workers=0 runs the dataset in-process and + # never forks. + mp_context = "spawn" if num_workers > 0 else None train_loader = DataLoader( train_ds, batch_size=None, num_workers=num_workers, pin_memory=pin, + multiprocessing_context=mp_context, ) val_loader = DataLoader( val_ds, batch_size=None, num_workers=num_workers, pin_memory=pin, + multiprocessing_context=mp_context, ) model_config = { diff --git a/giant/tools/profile_setup_scan.py b/giant/tools/profile_setup_scan.py new file mode 100644 index 0000000..57c8b62 --- /dev/null +++ b/giant/tools/profile_setup_scan.py @@ -0,0 +1,162 @@ +"""Benchmark `giant.pipeline.run_setup_stage`'s cold-cache scan against synthetic data. + +Generates a schema-complete synthetic steps parquet (matching +`tests/test_pipeline.py`'s `_make_synthetic_steps`, but built with vectorized +numpy instead of a per-row Python loop so it scales to millions of rows) at a +few row counts, times `run_setup_stage` with `cache_setup=False` (so every +call is a genuine cold scan, never served from the sidecar), and prints a +before/after-style table. Run this on `master` before a change and again +after to see what a step actually bought — see the "speed up dwarf +warm-cache" plan for the pass-by-pass breakdown this benchmark is meant to +attribute (giant/data/loader.py, giant/data/scan.py, giant/pipeline.py). + +Usage: ``uv run python giant/tools/profile_setup_scan.py`` +""" + +from __future__ import annotations + +import time +from pathlib import Path +from tempfile import TemporaryDirectory + +import numpy as np +import polars as pl + +from giant import config as gconfig +from giant.pipeline import run_setup_stage + +ROW_COUNTS = [20_000, 100_000, 500_000, 2_000_000] + +_MATERIALS = ["G4_AIR", "G4_Fe"] +_PDGS = [11, 22] +_PROCESSES = ["eIoni", "phot", "compt"] + + +def _unit_vectors(n: int, rng: np.random.Generator) -> np.ndarray: + v = rng.normal(size=(n, 3)) + return v / np.linalg.norm(v, axis=1, keepdims=True) + + +def _ragged_lists(k: np.ndarray, rng: np.random.Generator, lo: float, hi: float) -> list[list[float]]: + total = int(k.sum()) + flat = rng.uniform(lo, hi, size=total) + idx = np.cumsum(k)[:-1] + return [arr.tolist() for arr in np.split(flat, idx)] + + +def _make_synthetic_steps(n: int, seed: int = 0) -> pl.DataFrame: + """Vectorized equivalent of tests/test_pipeline.py's `_make_synthetic_steps`. + + event_id is assigned so each event gets 2-3 steps (matching that + fixture's structure), and pdg/material/process cycle deterministically + by row index rather than being drawn at random, same as the original. + """ + rng = np.random.default_rng(seed) + n_events = max(n // 3, 1) + + pre_E = rng.uniform(50.0, 500.0, size=n) + n_sec = rng.integers(0, 3, size=n) + frac_dep = rng.uniform(0.05, 0.3, size=n) + frac_sec = np.where(n_sec > 0, rng.uniform(0.05, 0.2, size=n), 0.0) + frac_post = 1.0 - frac_dep - frac_sec + edep = pre_E * frac_dep + e_sec = pre_E * frac_sec + post_E = pre_E * frac_post + pre_pos = rng.uniform(-10, 10, size=(n, 3)) + step_length = rng.uniform(0.1, 5.0, size=n) + pre_dir = np.zeros((n, 3)) + pre_dir[:, 2] = 1.0 + post_dir = _unit_vectors(n, rng) + post_pos = pre_pos + step_length[:, None] * pre_dir + + row_idx = np.arange(n) + event_id = row_idx % n_events + + sec_E = _ragged_lists(n_sec, rng, 0.1, 1.0) # placeholder magnitude, rescaled below + sec_dx = _ragged_lists(n_sec, rng, -1.0, 1.0) + sec_dy = _ragged_lists(n_sec, rng, -1.0, 1.0) + sec_dz = _ragged_lists(n_sec, rng, -1.0, 1.0) + total_sec = int(n_sec.sum()) + flat_pdg = [_PDGS[(row_idx[i] + j) % 2] for i in range(n) for j in range(n_sec[i])] + idx = np.cumsum(n_sec)[:-1] + sec_pdg = ( + [list(x) for x in np.split(np.array(flat_pdg, dtype=np.int64), idx)] if total_sec else [[] for _ in range(n)] + ) + # Rescale each row's secondary energies to sum to that row's e_sec (a + # Dirichlet split, like the original fixture) rather than the raw + # uniform placeholder. + sec_E_scaled = [] + for i in range(n): + vals = np.array(sec_E[i]) + if vals.size: + sec_E_scaled.append((vals / vals.sum() * e_sec[i]).tolist()) + else: + sec_E_scaled.append([]) + + return pl.DataFrame( + { + "event_id": event_id, + "pdg": np.array(_PDGS)[row_idx % 2], + "pre_x": pre_pos[:, 0], + "pre_y": pre_pos[:, 1], + "pre_z": pre_pos[:, 2], + "pre_E": pre_E, + "pre_dx": pre_dir[:, 0], + "pre_dy": pre_dir[:, 1], + "pre_dz": pre_dir[:, 2], + "material": np.array(_MATERIALS)[row_idx % 2], + "layer_id": row_idx % 5, + "child_track_ids": [list(range(int(k))) for k in n_sec], + "e_sec": e_sec, + "process": np.array(_PROCESSES)[row_idx % 3], + "step_length": step_length, + "post_E": post_E, + "edep": edep, + "post_dx": post_dir[:, 0], + "post_dy": post_dir[:, 1], + "post_dz": post_dir[:, 2], + "post_x": post_pos[:, 0], + "post_y": post_pos[:, 1], + "post_z": post_pos[:, 2], + "sec_E_list": sec_E_scaled, + "sec_pdg_list": sec_pdg, + "sec_dx_list": sec_dx, + "sec_dy_list": sec_dy, + "sec_dz_list": sec_dz, + } + ) + + +def _time_setup_stage(data: Path) -> float: + cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {}) + gconfig.validate_config(cfg) + t0 = time.perf_counter() + run_setup_stage( + data, + val_fraction=cfg["train"]["val_fraction"], + seed=cfg["train"]["seed"], + cfg=cfg, + cache_setup=False, + echo=lambda *a, **k: None, + ) + return time.perf_counter() - t0 + + +def main() -> None: + with TemporaryDirectory(prefix="giant-setup-scan-profile-") as tmp: + tmp_path = Path(tmp) + print(f"{'n_rows':>10s} {'time (s)':>10s} {'rows/s':>12s}") + for n in ROW_COUNTS: + path = tmp_path / f"steps_{n}.parquet" + _make_synthetic_steps(n).write_parquet(path) + # warm the OS page cache so the timed pass measures compute, not + # the one-time cold read of a freshly-written file. + pl.scan_parquet(path).select(pl.len()).collect() + + dt = _time_setup_stage(path) + print(f"{n:>10,d} {dt:>10.3f} {n / dt:>12,.0f}") + path.unlink() + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index be34148..29d6f84 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ readme = "README.md" requires-python = ">=3.12" dependencies = [ "numpy>=1.26,<3", - "pandas>=2.2,<4", + "polars>=1.0,<2", "pyarrow>=16,<25", "tqdm>=4.60,<5", "typer>=0.12,<1", @@ -28,6 +28,10 @@ dev = [ "ty>=0.0.50,<0.1", "bump-my-version>=1.2,<2", "git-cliff>=2,<3", + # Only used by test fixtures (writing small parquet files) — not a + # runtime dependency of giant itself since the pandas -> polars + # data-loading rewrite. + "pandas>=2.2,<4", "giant[convert,analysis,geometry,wandb]", ] geometry = [ diff --git a/tests/test_loader.py b/tests/test_loader.py index b7b0ffd..abaf06b 100644 --- a/tests/test_loader.py +++ b/tests/test_loader.py @@ -108,12 +108,12 @@ def test_build_process_map_from_files_spans_multiple_files(tmp_path): def test_build_process_map_from_files_tie_breaking_pins_first_seen_order(tmp_path): """When two processes end up with equal total counts, ranking falls back - to whichever was accumulated first (`sorted(..., reverse=True)` is stable, - and `counts` is built in file/row-scan order) — this is implementation- - defined, not a documented contract, so pin it explicitly: a future - rewrite (e.g. a polars-based single-scan) that ties differently would - silently reshuffle which processes get their own expert slot across a - retrain, and this test is what should catch that.""" + to whichever was scanned first — file order, then row order within a + file (`giant.data.scan`'s `first_seen` ordinal, ranked by + `giant.data.loader._topn_plus_other_map`'s `(-count, first_seen)` key). + This is an explicit, documented contract (not an accident of iteration + order), pinned here so a future change to the ranking can't silently + reshuffle which processes get their own expert slot across a retrain.""" path = tmp_path / "a.parquet" pd.DataFrame({"process": ["compt", "phot", "compt", "phot"]}).to_parquet(path) @@ -233,6 +233,23 @@ def test_build_pdg_topn_map_from_files_pools_primary_and_secondary_pdg(tmp_path) assert m.class_counts == {0: 11, 1: 5} +def test_build_pdg_topn_map_from_files_pooled_tie_breaks_by_row_position(tmp_path): + """Pooled pdg counting merges the primary `pdg` column and the exploded + `sec_pdg_list` column via one `group_by` over both (see + `giant.data.scan._pooled_pdg_lazy`), keyed by row position regardless of + which role (primary or secondary) a code was seen in — not "all + primaries before all secondaries" the way a two-pass accumulation would. + 11 (primary, row 0), 33 (primary, row 1) and 22 (secondary, row 1) all + end up with count 1; 11's strictly earlier row wins the tie over both, + whatever order 33/22 (tied with each other, same row) land in.""" + path = tmp_path / "a.parquet" + pd.DataFrame({"pdg": [11, 33], "sec_pdg_list": [[], [22]]}).to_parquet(path) + + m = build_pdg_topn_map_from_files([path], n_classes=4) + + assert m.class_map[11] == 0 + + def test_build_pdg_topn_map_from_files_missing_sec_pdg_list_column(tmp_path): """Files predating the parent->child join have no sec_pdg_list column — must not raise, just count the primary pdg column alone.""" diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index ed2e4ce..c5470a4 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -142,7 +142,7 @@ def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch): def _forbidden(*a, **k): raise AssertionError("should be served from cache, not recomputed") - monkeypatch.setattr("giant.pipeline.build_index_maps_from_files", _forbidden) + monkeypatch.setattr("giant.pipeline.scan_metadata", _forbidden) monkeypatch.setattr("giant.pipeline.iter_file_chunks", _forbidden) echo2 = _run(data, tmp_path / "out2") @@ -283,7 +283,7 @@ def test_run_train_job_new_val_fraction_is_partial_hit(tmp_path, data, monkeypat def _forbidden(*a, **k): raise AssertionError("vocab should be served from cache") - monkeypatch.setattr("giant.pipeline.build_index_maps_from_files", _forbidden) + monkeypatch.setattr("giant.pipeline.scan_metadata", _forbidden) echo2 = _run(data, tmp_path / "out2", cfg=_tiny_cfg(val_fraction=0.3)) joined = "\n".join(echo2) diff --git a/uv.lock b/uv.lock index 0f29335..784a00d 100644 --- a/uv.lock +++ b/uv.lock @@ -679,8 +679,8 @@ version = "0.3.16" source = { editable = "." } dependencies = [ { name = "numpy" }, - { name = "pandas" }, { name = "particle" }, + { name = "polars" }, { name = "pyarrow" }, { name = "pyyaml" }, { name = "tqdm" }, @@ -712,6 +712,7 @@ dev = [ { name = "git-cliff" }, { name = "ipykernel" }, { name = "matplotlib" }, + { name = "pandas" }, { name = "plotstyle" }, { name = "polars" }, { name = "pytest" }, @@ -738,9 +739,10 @@ requires-dist = [ { name = "ipykernel", marker = "extra == 'analysis'", specifier = ">=7.3.0" }, { name = "matplotlib", marker = "extra == 'analysis'", specifier = ">=3.8,<4" }, { name = "numpy", specifier = ">=1.26,<3" }, - { name = "pandas", specifier = ">=2.2,<4" }, + { name = "pandas", marker = "extra == 'dev'", specifier = ">=2.2,<4" }, { name = "particle", specifier = ">=1.0,<2" }, { name = "plotstyle", marker = "extra == 'analysis'", specifier = ">=1.0.0", index = "https://git.larsbogner.de/api/packages/lars/pypi/simple/" }, + { name = "polars", specifier = ">=1.0,<2" }, { name = "polars", marker = "extra == 'analysis'", specifier = ">=1.0,<2" }, { name = "polars", marker = "extra == 'convert'", specifier = ">=1.0,<2" }, { name = "pyarrow", specifier = ">=16,<25" },