giant/tests/test_loader.py

import numpy as np
import pandas as pd
import pytest

from giant.data.loader import (
    EVENT_ID_FILE_STRIDE,
    build_index_maps,
    build_index_maps_from_files,
    build_process_map_from_files,
    event_id_offset,
    find_parquet_files,
    iter_cond_chunks,
    iter_file_chunks,
    load_event_ids,
    load_steps,
)


def _touch(path):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.touch()
    return path


def test_find_parquet_files_single_file(tmp_path):
    f = _touch(tmp_path / "shard-000.parquet")
    assert find_parquet_files(f) == [f]


def test_find_parquet_files_directory_glob(tmp_path):
    a = _touch(tmp_path / "shard-000.parquet")
    b = _touch(tmp_path / "shard-001.parquet")
    _touch(tmp_path / "not_a_parquet.root")
    assert find_parquet_files(tmp_path) == sorted([a, b])


def test_find_parquet_files_empty_directory_raises(tmp_path):
    with pytest.raises(FileNotFoundError):
        find_parquet_files(tmp_path)


def test_manifest_resolves_relative_to_its_own_directory(tmp_path):
    target = _touch(tmp_path / "processed" / "pbwo4" / "shard-000.parquet")
    manifest_dir = tmp_path / "pools" / "pbwo4"
    manifest_dir.mkdir(parents=True)
    manifest = manifest_dir / "full.manifest"
    manifest.write_text("../../processed/pbwo4/shard-000.parquet\n")

    assert find_parquet_files(manifest) == [target.resolve()]


def test_manifest_skips_blank_lines_and_comments(tmp_path):
    target = _touch(tmp_path / "shard-000.parquet")
    manifest = tmp_path / "full.manifest"
    manifest.write_text("\n# a comment\nshard-000.parquet\n\n")

    assert find_parquet_files(manifest) == [target.resolve()]


def test_manifest_missing_file_raises(tmp_path):
    manifest = tmp_path / "full.manifest"
    manifest.write_text("does_not_exist.parquet\n")

    with pytest.raises(FileNotFoundError):
        find_parquet_files(manifest)


def test_manifest_with_no_entries_raises(tmp_path):
    manifest = tmp_path / "full.manifest"
    manifest.write_text("# only comments\n")

    with pytest.raises(FileNotFoundError):
        find_parquet_files(manifest)


def test_build_process_map_from_files_keeps_most_frequent(tmp_path):
    """process counts: eIoni=5, phot=3, compt=2, Rayl=1 — with n_experts=3, only
    the top 2 (eIoni, phot) get their own index; compt/Rayl share the "other"
    (last) index."""
    process = ["eIoni"] * 5 + ["phot"] * 3 + ["compt"] * 2 + ["Rayl"] * 1
    path = tmp_path / "shard-000.parquet"
    pd.DataFrame({"process": process}).to_parquet(path)

    proc_map = build_process_map_from_files([path], n_experts=3)

    assert proc_map["eIoni"] == 0
    assert proc_map["phot"] == 1
    assert proc_map["compt"] == 2
    assert proc_map["Rayl"] == 2
    assert set(proc_map.values()) <= {0, 1, 2}


def test_build_process_map_from_files_spans_multiple_files(tmp_path):
    path_a = tmp_path / "a.parquet"
    path_b = tmp_path / "b.parquet"
    pd.DataFrame({"process": ["eIoni"] * 3 + ["phot"] * 1}).to_parquet(path_a)
    pd.DataFrame({"process": ["phot"] * 4 + ["compt"] * 1}).to_parquet(path_b)

    # phot: 1+4=5 total > eIoni: 3 > compt: 1
    proc_map = build_process_map_from_files([path_a, path_b], n_experts=3)

    assert proc_map["phot"] == 0
    assert proc_map["eIoni"] == 1
    assert proc_map["compt"] == 2


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."""
    path = tmp_path / "a.parquet"
    pd.DataFrame({"process": ["compt", "phot", "compt", "phot"]}).to_parquet(path)

    proc_map = build_process_map_from_files([path], n_experts=3)

    assert proc_map == {"compt": 0, "phot": 1}


def test_build_process_map_from_files_tie_breaking_favors_first_scanned_file(
    tmp_path,
):
    """Same total-count tie as above, but split across two files with equal
    per-file counts — the file listed first wins the tie."""
    path_a = tmp_path / "a.parquet"
    path_b = tmp_path / "b.parquet"
    pd.DataFrame({"process": ["zzz", "zzz"]}).to_parquet(path_a)
    pd.DataFrame({"process": ["aaa", "aaa"]}).to_parquet(path_b)

    forward = build_process_map_from_files([path_a, path_b], n_experts=3)
    backward = build_process_map_from_files([path_b, path_a], n_experts=3)

    assert forward == {"zzz": 0, "aaa": 1}
    assert backward == {"aaa": 0, "zzz": 1}


def test_build_process_map_from_files_fewer_processes_than_experts(tmp_path):
    """When there are fewer distinct processes than expert slots, every
    process gets its own index and the shared "other" bucket goes unused."""
    path = tmp_path / "a.parquet"
    pd.DataFrame({"process": ["eIoni", "phot"]}).to_parquet(path)

    proc_map = build_process_map_from_files([path], n_experts=5)

    assert proc_map == {"eIoni": 0, "phot": 1}
    assert 4 not in proc_map.values()  # the "other" slot (n_experts - 1) is unused


def test_build_process_map_from_files_n_experts_one_buckets_everything(tmp_path):
    """n_experts=1 leaves no room for a "most frequent" slot — every process
    (however frequent) is bucketed into the single shared index 0."""
    path = tmp_path / "a.parquet"
    pd.DataFrame({"process": ["eIoni"] * 10 + ["phot"] * 1}).to_parquet(path)

    proc_map = build_process_map_from_files([path], n_experts=1)

    assert proc_map == {"eIoni": 0, "phot": 0}


def test_build_process_map_from_files_three_files_partial_overlap(tmp_path):
    """Counts for a process appearing in only some of several files must sum
    correctly, not just match the two-file case already covered above."""
    path_a = tmp_path / "a.parquet"
    path_b = tmp_path / "b.parquet"
    path_c = tmp_path / "c.parquet"
    pd.DataFrame({"process": ["eIoni"] * 2}).to_parquet(path_a)
    pd.DataFrame({"process": ["phot"] * 3}).to_parquet(path_b)
    pd.DataFrame({"process": ["eIoni"] * 2 + ["compt"] * 1}).to_parquet(path_c)

    # eIoni: 2+2=4 > phot: 3 > compt: 1
    proc_map = build_process_map_from_files([path_a, path_b, path_c], n_experts=3)

    assert proc_map["eIoni"] == 0
    assert proc_map["phot"] == 1
    assert proc_map["compt"] == 2


# ── build_index_maps (in-memory) ────────────────────────────────────────────


def test_build_index_maps_sorts_numerically_not_lexicographically():
    """10-digit nuclear/ion PDG codes must sort numerically — a lexicographic
    sort would place "1000060120" before "22" since '1' < '2'."""
    data = {
        "pdg": np.array([22, 1000060120, 11], dtype=np.int64),
        "material": np.array(["G4_AIR", "PbWO4", "G4_Fe"], dtype=object),
    }
    pdg_map, mat_map = build_index_maps(data)
    assert list(pdg_map.keys()) == [11, 22, 1000060120]
    assert mat_map == {"G4_AIR": 0, "G4_Fe": 1, "PbWO4": 2}


def test_build_index_maps_dedups_repeated_values():
    data = {
        "pdg": np.array([11, 11, 22, 22, 22], dtype=np.int64),
        "material": np.array(["PbWO4"] * 5, dtype=object),
    }
    pdg_map, mat_map = build_index_maps(data)
    assert pdg_map == {11: 0, 22: 1}
    assert mat_map == {"PbWO4": 0}


def test_build_index_maps_handles_negative_pdg_codes():
    """Antiparticle codes (negative) must sort numerically, not by magnitude."""
    data = {
        "pdg": np.array([-13, 11, -11, 13], dtype=np.int64),
        "material": np.array(["X"] * 4, dtype=object),
    }
    pdg_map, _ = build_index_maps(data)
    assert list(pdg_map.keys()) == [-13, -11, 11, 13]


def test_build_index_maps_indices_are_dense_and_bijective():
    data = {
        "pdg": np.array([5, 1, 9, 1, 5], dtype=np.int64),
        "material": np.array(["a", "b", "c", "a", "b"], dtype=object),
    }
    pdg_map, mat_map = build_index_maps(data)
    assert sorted(pdg_map.values()) == list(range(len(pdg_map)))
    assert sorted(mat_map.values()) == list(range(len(mat_map)))


# ── build_index_maps_from_files ─────────────────────────────────────────────


def test_build_index_maps_from_files_unions_and_dedups_across_files(tmp_path):
    path_a = tmp_path / "a.parquet"
    path_b = tmp_path / "b.parquet"
    pd.DataFrame({"pdg": [11, 22], "material": ["G4_AIR", "PbWO4"]}).to_parquet(path_a)
    pd.DataFrame({"pdg": [22, 2112], "material": ["PbWO4", "G4_Fe"]}).to_parquet(path_b)

    pdg_map, mat_map = build_index_maps_from_files([path_a, path_b])

    assert pdg_map == {11: 0, 22: 1, 2112: 2}
    assert mat_map == {"G4_AIR": 0, "G4_Fe": 1, "PbWO4": 2}


def test_build_index_maps_from_files_ordering_independent_of_file_order(tmp_path):
    """Index assignment comes from the globally sorted union, not file-scan
    order — swapping which file is scanned first must not change the map,
    since the map is baked into a trained checkpoint's vocabulary."""
    path_a = tmp_path / "a.parquet"
    path_b = tmp_path / "b.parquet"
    pd.DataFrame({"pdg": [22], "material": ["PbWO4"]}).to_parquet(path_a)
    pd.DataFrame({"pdg": [11], "material": ["G4_AIR"]}).to_parquet(path_b)

    forward = build_index_maps_from_files([path_a, path_b])
    backward = build_index_maps_from_files([path_b, path_a])

    assert forward == backward
    assert forward == ({11: 0, 22: 1}, {"G4_AIR": 0, "PbWO4": 1})


def test_build_index_maps_from_files_numeric_sort_for_nuclear_codes(tmp_path):
    path = tmp_path / "a.parquet"
    pd.DataFrame({"pdg": [22, 1000060120, 11], "material": ["X", "X", "X"]}).to_parquet(
        path
    )

    pdg_map, _ = build_index_maps_from_files([path])
    assert list(pdg_map.keys()) == [11, 22, 1000060120]


def test_build_index_maps_from_files_single_file(tmp_path):
    path = tmp_path / "a.parquet"
    pd.DataFrame({"pdg": [11, 11, 22], "material": ["PbWO4"] * 3}).to_parquet(path)
    pdg_map, mat_map = build_index_maps_from_files([path])
    assert pdg_map == {11: 0, 22: 1}
    assert mat_map == {"PbWO4": 0}


def test_build_index_maps_from_files_matches_build_index_maps(tmp_path):
    """Sanity-pin: the file-scanning and in-memory variants must agree on the
    same data, since a future single-pass rewrite (pyarrow/polars) may
    replace one but not the other."""
    rng = np.random.default_rng(0)
    pdg = rng.choice([11, -11, 22, 2112, 1000060120], size=200)
    material = rng.choice(["G4_AIR", "PbWO4", "G4_Fe"], size=200)
    path = tmp_path / "a.parquet"
    pd.DataFrame({"pdg": pdg, "material": material}).to_parquet(path)

    from_files = build_index_maps_from_files([path])
    from_memory = build_index_maps({"pdg": pdg, "material": material})
    assert from_files == from_memory


# ── event_id_offset / per-file event_id offsetting ─────────────────────────


def _steps_df(event_ids):
    """Minimal schema-complete steps rows (no secondaries) for _df_to_dict."""
    return pd.DataFrame(
        [
            {
                "event_id": eid,
                "pdg": 11,
                "pre_x": 0.0,
                "pre_y": 0.0,
                "pre_z": 0.0,
                "pre_E": 100.0,
                "pre_dx": 0.0,
                "pre_dy": 0.0,
                "pre_dz": 1.0,
                "material": "G4_AIR",
                "layer_id": 0,
                "child_track_ids": [],
                "e_sec": 0.0,
                "step_length": 1.0,
                "post_E": 90.0,
                "edep": 10.0,
                "post_dx": 0.0,
                "post_dy": 0.0,
                "post_dz": 1.0,
                "post_x": 0.0,
                "post_y": 0.0,
                "post_z": 1.0,
            }
            for eid in event_ids
        ]
    )


def test_event_id_offset_scales_by_file_index():
    assert event_id_offset(0) == 0
    assert event_id_offset(1) == EVENT_ID_FILE_STRIDE
    assert event_id_offset(3) == 3 * EVENT_ID_FILE_STRIDE


def test_load_event_ids_default_offset_is_zero(tmp_path):
    path = tmp_path / "a.parquet"
    pd.DataFrame({"event_id": [5, 6, 7]}).to_parquet(path)
    np.testing.assert_array_equal(load_event_ids(path), [5, 6, 7])


def test_load_event_ids_applies_offset(tmp_path):
    path = tmp_path / "a.parquet"
    pd.DataFrame({"event_id": [0, 1, 2]}).to_parquet(path)
    offset = event_id_offset(1)
    np.testing.assert_array_equal(
        load_event_ids(path, offset=offset), [offset, offset + 1, offset + 2]
    )


def test_load_steps_applies_offset_to_event_id(tmp_path):
    path = tmp_path / "a.parquet"
    _steps_df([0, 1]).to_parquet(path)
    offset = event_id_offset(2)
    d = load_steps(path, offset=offset)
    np.testing.assert_array_equal(d["event_id"], [offset, offset + 1])


def test_iter_file_chunks_applies_offset(tmp_path):
    path = tmp_path / "a.parquet"
    _steps_df([0, 1, 2]).to_parquet(path)
    offset = event_id_offset(1)
    ids = np.concatenate([c["event_id"] for c in iter_file_chunks(path, offset=offset)])
    np.testing.assert_array_equal(sorted(ids), [offset, offset + 1, offset + 2])


def test_iter_cond_chunks_applies_offset(tmp_path):
    path = tmp_path / "a.parquet"
    _steps_df([0, 1]).to_parquet(path)
    offset = event_id_offset(5)
    ids = np.concatenate([c["event_id"] for c in iter_cond_chunks(path, offset=offset)])
    np.testing.assert_array_equal(sorted(ids), [offset, offset + 1])