giant/tests/test_geometry.py

"""Tests for the geometry oracle (position -> material/layer_id + escape)."""

from pathlib import Path
from unittest.mock import patch

import numpy as np
import pandas as pd
import pytest

from giant import geometry as g

# scikit-learn is an optional extra; skip the whole module if it's missing.
pytest.importorskip("sklearn")


def _steps_frame(n=5, with_post=True):
    rng = np.random.default_rng(0)
    data = {
        "pre_x": rng.uniform(-10, 10, n),
        "pre_y": rng.uniform(-10, 10, n),
        "pre_z": rng.uniform(-10, 10, n),
        "material": ["G4_AIR"] * n,
        "layer_id": np.arange(n, dtype=np.int64),
    }
    if with_post:
        data["post_x"] = rng.uniform(-10, 10, n)
        data["post_y"] = rng.uniform(-10, 10, n)
        data["post_z"] = rng.uniform(-10, 10, n)
    return pd.DataFrame(data)


def test_iter_point_batches_missing_columns_raises(tmp_path):
    path = tmp_path / "steps.parquet"
    pd.DataFrame({"pre_x": [0.0]}).to_parquet(path)
    with pytest.raises(ValueError, match="missing columns"):
        next(g._iter_point_batches(path))


def test_iter_point_batches_without_post_columns_yields_pre_only(tmp_path):
    path = tmp_path / "steps.parquet"
    df = _steps_frame(n=5, with_post=False)
    df.to_parquet(path)

    (pos, mat, lay) = next(g._iter_point_batches(path))

    assert pos.shape == (5, 3)
    np.testing.assert_allclose(
        pos, df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32)
    )
    assert list(mat) == ["G4_AIR"] * 5
    np.testing.assert_array_equal(lay, np.arange(5))


def test_iter_point_batches_with_post_columns_doubles_and_concatenates_points(
    tmp_path,
):
    path = tmp_path / "steps.parquet"
    df = _steps_frame(n=5, with_post=True)
    df.to_parquet(path)

    (pos, mat, lay) = next(g._iter_point_batches(path))

    # Every step contributes both its pre_pos and post_pos, sharing the
    # step's material/layer_id label — so batches double in length.
    assert pos.shape == (10, 3)
    np.testing.assert_allclose(
        pos[:5], df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32)
    )
    np.testing.assert_allclose(
        pos[5:], df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32)
    )
    assert list(mat) == ["G4_AIR"] * 10
    np.testing.assert_array_equal(lay, np.concatenate([np.arange(5), np.arange(5)]))


def test_iter_point_batches_respects_batch_size(tmp_path):
    path = tmp_path / "steps.parquet"
    df = _steps_frame(n=10, with_post=False)
    df.to_parquet(path, row_group_size=10)

    batches = list(g._iter_point_batches(path, batch_size=4))

    assert [len(pos) for pos, _, _ in batches] == [4, 4, 2]


def _box_batch(n, rng):
    """A labelled point cloud: inside a 100mm box -> PbWO4/0, else AIR/-1."""
    pos = rng.uniform(-200, 200, (n, 3)).astype(np.float32)
    inside = (np.abs(pos) < 100).all(axis=1)
    mat = np.where(inside, "G4_PbWO4", "G4_AIR").astype(object)
    lay = np.where(inside, 0, -1).astype(np.int64)
    return pos, mat, lay


def _build(subsample=30000, method="knn", escape_factor=5.0):
    rng = np.random.default_rng(0)
    batches = [_box_batch(20000, rng) for _ in range(3)]
    with patch.object(g, "_iter_point_batches", lambda p: iter(batches)):
        return g.build_geometry_oracle(
            [Path("x")],
            method=method,
            subsample=subsample,
            escape_factor=escape_factor,
        )


def test_classes_discovered():
    orc = _build()
    assert set(orc.classes) == {("G4_PbWO4", 0), ("G4_AIR", -1)}


def test_query_labels_inside_and_outside():
    orc = _build()
    pos = np.array([[0.0, 0.0, 0.0], [150.0, 150.0, 150.0]])
    material, layer_id, _ = orc.query(pos)
    assert material[0] == "G4_PbWO4" and layer_id[0] == 0
    assert material[1] == "G4_AIR" and layer_id[1] == -1


def test_escape_flag_fires_far_from_data():
    orc = _build()
    pos = np.array([[0.0, 0.0, 0.0], [1e5, 0.0, 0.0]])
    _, _, escaped = orc.query(pos)
    assert not escaped[0]
    assert escaped[1]


def test_query_empty():
    orc = _build()
    material, layer_id, escaped = orc.query(np.empty((0, 3)))
    assert len(material) == len(layer_id) == len(escaped) == 0


def test_query_bad_shape_raises():
    orc = _build()
    with pytest.raises(ValueError):
        orc.query(np.zeros((4, 2)))


def test_save_load_roundtrip(tmp_path):
    orc = _build()
    p = tmp_path / "oracle.pkl"
    orc.save(p)
    loaded = g.GeometryOracle.load(p)

    pos = np.array([[0.0, 0.0, 0.0], [150.0, 150.0, 150.0], [1e5, 0.0, 0.0]])
    m0, l0, e0 = orc.query(pos)
    m1, l1, e1 = loaded.query(pos)
    assert (m0 == m1).all() and (l0 == l1).all() and (e0 == e1).all()
    assert loaded.escape_threshold == orc.escape_threshold
    assert loaded.classes == orc.classes


def test_svm_method_has_escape_tree():
    orc = _build(subsample=4000, method="svm")
    # SVM cannot answer NN-distance, so a reference tree backs the escape test.
    assert orc._ref_tree is not None
    _, _, escaped = orc.query(np.array([[1e5, 0.0, 0.0]]))
    assert escaped[0]


def _layer_batch(n, rng):
    """Two 100mm slabs along z (with an air gap between/around them), bounded
    to a 100x100mm transverse footprint — miniCaloSim's actual layer-stack
    shape."""
    z = rng.uniform(-20.0, 220.0, n).astype(np.float32)
    x = rng.uniform(-50.0, 50.0, n).astype(np.float32)
    y = rng.uniform(-50.0, 50.0, n).astype(np.float32)
    material = np.full(n, "G4_AIR", dtype=object)
    layer_id = np.full(n, -1, dtype=np.int64)
    in_l0 = (z >= 0.0) & (z < 100.0)
    in_l1 = (z >= 110.0) & (z < 210.0)
    material[in_l0] = "G4_PbWO4"
    layer_id[in_l0] = 0
    material[in_l1] = "G4_W"
    layer_id[in_l1] = 1
    pos = np.stack([x, y, z], axis=1)
    return pos, material, layer_id


def _build_slab(subsample=60000, n_bins=500, escape_factor=5.0):
    rng = np.random.default_rng(0)
    batches = [_layer_batch(20000, rng) for _ in range(3)]
    with patch.object(g, "_iter_point_batches", lambda p: iter(batches)):
        return g.build_geometry_oracle(
            [Path("x")],
            method="slab",
            subsample=subsample,
            escape_factor=escape_factor,
            depth_axis=2,
            n_bins=n_bins,
        )


def test_slab_is_default_method():
    rng = np.random.default_rng(0)
    batches = [_layer_batch(20000, rng)]
    with patch.object(g, "_iter_point_batches", lambda p: iter(batches)):
        orc = g.build_geometry_oracle([Path("x")], subsample=20000)
    assert orc.metadata["method"] == "slab"
    assert orc._slab is not None


def test_slab_classes_discovered():
    orc = _build_slab()
    assert set(orc.classes) == {("G4_PbWO4", 0), ("G4_W", 1), ("G4_AIR", -1)}


def test_slab_query_labels_by_depth():
    orc = _build_slab()
    pos = np.array(
        [[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0]]
    )  # layer 0, gap, layer 1
    material, layer_id, escaped = orc.query(pos)
    assert list(material) == ["G4_PbWO4", "G4_AIR", "G4_W"]
    assert list(layer_id) == [0, -1, 1]
    assert not escaped.any()


def test_slab_escape_beyond_depth_range():
    orc = _build_slab()
    pos = np.array([[0.0, 0.0, 50.0], [0.0, 0.0, 1e5]])
    _, _, escaped = orc.query(pos)
    assert not escaped[0]
    assert escaped[1]


def test_slab_escape_beyond_transverse_radius():
    orc = _build_slab()
    pos = np.array([[0.0, 0.0, 50.0], [1e5, 1e5, 50.0]])
    _, _, escaped = orc.query(pos)
    assert not escaped[0]
    assert escaped[1]


def test_slab_save_load_roundtrip(tmp_path):
    orc = _build_slab()
    p = tmp_path / "slab_oracle.pkl"
    orc.save(p)
    loaded = g.GeometryOracle.load(p)

    pos = np.array(
        [[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0], [0.0, 0.0, 1e5]]
    )
    m0, l0, e0 = orc.query(pos)
    m1, l1, e1 = loaded.query(pos)
    assert (m0 == m1).all() and (l0 == l1).all() and (e0 == e1).all()
    assert loaded.escape_threshold == orc.escape_threshold
    assert loaded._slab is not None