wc_ntuple_tutorial/plotting_3d.py at main · leehagaman/wc_ntuple_tutorial · 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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots

# making detector boundary points

tpc_min_x = -1.
tpc_max_x = 254.3
tpc_min_y = -115.
tpc_max_y = 117.
tpc_min_z = 0.6
tpc_max_z = 1036.4


def generate_box_wireframe(x_min, x_max, y_min, y_max, z_min, z_max):
    """Generate wireframe lines for a 3D box"""

    # Define the 8 vertices of the box
    vertices = np.array([
        [x_min, y_min, z_min],  # 0
        [x_max, y_min, z_min],  # 1
        [x_max, y_max, z_min],  # 2
        [x_min, y_max, z_min],  # 3
        [x_min, y_min, z_max],  # 4
        [x_max, y_min, z_max],  # 5
        [x_max, y_max, z_max],  # 6
        [x_min, y_max, z_max],  # 7
    ])

    # Define the 12 edges of the box
    edges = [
        # Bottom face
        [0, 1], [1, 2], [2, 3], [3, 0],
        # Top face
        [4, 5], [5, 6], [6, 7], [7, 4],
        # Vertical edges
        [0, 4], [1, 5], [2, 6], [3, 7]
    ]

    return vertices, edges

def plot_event(event_index, points_dic):

    # Generate wireframe for detector boundaries
    detector_vertices, detector_edges = generate_box_wireframe(tpc_min_x, tpc_max_x, tpc_min_y, tpc_max_y, tpc_min_z, tpc_max_z)

    x_width = tpc_max_x - tpc_min_x
    expanded_detector_vertices, expanded_detector_edges = generate_box_wireframe(
        tpc_min_x - x_width, tpc_max_x + x_width, tpc_min_y, tpc_max_y, tpc_min_z, tpc_max_z
    )

    # Extract reco neutrino vertex for camera positioning
    reco_vtx = None
    if "wc_reco_nu_vtx" in points_dic:
        vtx = points_dic["wc_reco_nu_vtx"][0]
        # Convert from detector coordinates (x, y, z) to plot coordinates (z, x, y)
        reco_vtx = (vtx[2], vtx[0], vtx[1])

    fig = make_subplots(rows=1, cols=1, specs=[[{'type': 'scene'}]])

    # Add expanded detector boundary wireframe (for camera positioning)
    for edge in expanded_detector_edges:
        start_vertex = expanded_detector_vertices[edge[0]]
        end_vertex = expanded_detector_vertices[edge[1]]

        fig.add_trace(go.Scatter3d(
            x=[start_vertex[2], end_vertex[2]],
            y=[start_vertex[0], end_vertex[0]],
            z=[start_vertex[1], end_vertex[1]],
            mode='lines',
            line=dict(
                color='black',
                width=1
            ),
            opacity=0,
            showlegend=False,
            visible=True
        ))

    # Add detector boundary wireframe
    for edge in detector_edges:
        start_vertex = detector_vertices[edge[0]]
        end_vertex = detector_vertices[edge[1]]

        fig.add_trace(go.Scatter3d(
            x=[start_vertex[2], end_vertex[2]],
            y=[start_vertex[0], end_vertex[0]],
            z=[start_vertex[1], end_vertex[1]],
            mode='lines',
            line=dict(
                color='black',
                width=2,
            ),
            name='TPC Boundary',
            showlegend=False
        ))

    for name, values in points_dic.items():

        points, color, cmap, size, visible_by_default = values

        if points.shape == (3,):
            fig.add_trace(go.Scatter3d(
                x=[points[2]],
                y=[points[0]],
                z=[points[1]],
                mode='markers',
                marker=dict(size=size, color=color, opacity=1, colorscale=cmap),
                name=name,
                visible=visible_by_default
            ))
        else:
            fig.add_trace(go.Scatter3d(
                x=points[:, 2],
                y=points[:, 0],
                z=points[:, 1],
                mode='markers',
                marker=dict(size=size, color=color, opacity=1, colorscale=cmap),
                name=name,
                visible=visible_by_default
            ))

    # Configure camera to focus on reco neutrino vertex
    camera_config = dict(
        eye=dict(x=-1.5, y=-1.5, z=1.5)
    )

    if reco_vtx is not None:
        # Set camera center to the reco vertex and position eye closer
        camera_config = dict(
            center=dict(x=0, y=0, z=0),  # Center on the vertex
            eye=dict(x=-0.8, y=-0.8, z=0.8)  # Closer viewing angle
        )

    fig.update_layout(
        scene=dict(
            xaxis_title='z',
            yaxis_title='x',
            zaxis_title='y',
            aspectratio=dict(
                x=5,
                y=3,
                z=1
            ),
        ),
        width=1500,
        height=900,
        autosize=False,
        scene_camera=camera_config
    )

    fig.show(renderer="browser")