celerpy/test/test_visualize.py at 5a60c0c9cb1ec26717b11002c40859cad17d016f · sethrj/celerpy · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
# Copyright 2024 UT-Battelle, LLC, and other Celeritas developers.
# See the top-level LICENSE file for details.
# SPDX-License-Identifier: Apache-2.0

from pathlib import Path

import numpy as np
import pytest
from numpy.testing import assert_array_equal

from celerpy import model, visualize
from celerpy.settings import settings

local_path = Path(__file__).parent
settings.prefix_path = local_path / "mock-prefix"


def test_CelerGeo():
    inp = local_path / "data" / "two-boxes.gdml"
    with visualize.CelerGeo.from_filename(inp) as cg:
        with pytest.raises(RuntimeError):
            cg.trace()
        (result, img) = cg.trace(
            model.ImageInput(upper_right=[1, 1, 0], vertical_pixels=4),
            geometry=model.GeometryEngine.orange,
        )
        assert isinstance(result, model.TraceOutput)
        assert img.shape == (4, 4)
        assert img.dtype == np.int32
        (result, img) = cg.trace(geometry=model.GeometryEngine.orange)
        assert isinstance(result, model.TraceOutput)
        assert img.shape == (4, 4)
        (result, img) = cg.trace(geometry=model.GeometryEngine.geant4)
        assert isinstance(result, model.TraceOutput)
        assert img.shape == (4, 4)
        result = cg.close()
        assert result


def test_ReverseIndexDict():
    to_volume = []
    to_id = visualize.ReverseIndexDict(to_volume)

    assert to_id["foo"] == 0
    assert to_id["bar"] == 1
    assert to_id["baz"] == 2
    assert to_id["foo"] == 0
    assert to_volume == ["foo", "bar", "baz"]


def test_IdMapper():
    map_ids = visualize.IdMapper()
    (img, vol) = map_ids(np.array([1, 1, 1]), ["foo", "bar"])
    assert_array_equal(img, np.array([0, 0, 0]))
    assert vol == ["bar"]

    (img, vol) = map_ids(np.array([1, 2, 1]), ["foo", "baz", "bar"])
    assert_array_equal(img, np.array([1, 0, 1]))
    assert vol == ["bar", "baz"]

    (img, vol) = map_ids(np.array([0, 1, 2]), ["bar", "foo", "baz"])
    assert_array_equal(img, np.array([0, 2, 1]))
    assert vol == ["bar", "baz", "foo"]

    (img, vol) = map_ids(np.array([-1, 0, -1]), ["foo"])
    assert_array_equal(img, np.array([2, 2, 2]))
    assert_array_equal(img.mask, [True, False, True])
    assert vol == ["bar", "baz", "foo"]


def test_centered_image():
    centered_image = visualize.centered_image
    kwargs = dict()

    # Test case 1: Square image with x projection along -y
    result = centered_image(
        center=[0, 0, 0], xdir=[0, -1, 0], outdir=[0, 0, 1], width=2.0, **kwargs
    )
    assert result.lower_left == [-1.0, 1.0, 0.0]
    assert result.upper_right == [1.0, -1.0, 0.0]
    assert result.rightward == [0.0, -1.0, 0.0]

    # Test case 2: Rectangle image with xdir = [1, 0, 0]
    result = centered_image(
        center=[1, 0, 0],
        xdir=[1, 0, 0],
        outdir=[0, 0, -1],
        width=(2.0, 4.0),
        **kwargs,
    )
    assert result.lower_left == [0.0, 2.0, 0.0]
    assert result.upper_right == [2.0, -2.0, 0.0]