GsfDebugger/src/drawers.py

0001 import numpy as np
0002 import pandas as pd
0003
0004 detector_color = "grey"
0005
0006
0007 class ViewDrawer:
0008     """
0009     Base class for a drawer that encapsulates drawing the state with a certain view (x-y, z-r, ...)
0010     """
0011
0012     def draw_detector(self, ax):
0013         ax = self.draw_detector_impl(ax)
0014         ax.set_title("{}-{} plot".format(self.coor0, self.coor1))
0015         ax.set_ylabel(self.coor1)
0016         ax.set_xlabel(self.coor0)
0017         return ax
0018
0019     def plot(self, ax, values, title_appendix="", **kwargs):
0020         ax.plot(self.get_coor0(values), self.get_coor1(values), **kwargs)
0021         ax.set_title("{}-{} plot: {}".format(self.coor0, self.coor1, title_appendix))
0022         return ax
0023
0024     def scatter(self, ax, values, title_appendix="", **kwargs):
0025         ax.scatter(self.get_coor0(values), self.get_coor1(values), **kwargs)
0026         return ax
0027
0028     def annotate(self, ax, position_3d, text):
0029         l = np.array([position_3d])
0030         position_2d = (self.get_coor0(l), self.get_coor1(l))
0031         ax.annotate(text, position_2d)
0032         return ax
0033
0034     def set_title(self, ax, title_appendix=""):
0035         ax.set_title("{}-{} plot: {}".format(self.coor0, self.coor1, title_appendix))
0036         return ax
0037
0038
0039 class ZRDrawer(ViewDrawer):
0040     """
0041     Base class for a z-r drawer
0042     """
0043
0044     def __init__(self):
0045         self.coor0 = "z"
0046         self.coor1 = "r"
0047
0048     def get_coor0(self, values):
0049         return values[:, 2]
0050
0051     def get_coor1(self, values):
0052         return np.hypot(values[:, 0], values[:, 1])
0053
0054
0055 class XYDrawer(ViewDrawer):
0056     """
0057     Base class for a x-y drawer
0058     """
0059
0060     def __init__(self):
0061         self.coor0 = "x"
0062         self.coor1 = "y"
0063
0064     def get_coor0(self, values):
0065         return values[:, 0]
0066
0067     def get_coor1(self, values):
0068         return values[:, 1]
0069
0070
0071 class XZDrawer(ViewDrawer):
0072     """
0073     Base class for a x-z drawer
0074     """
0075
0076     def __init__(self):
0077         self.coor0 = "x"
0078         self.coor1 = "z"
0079
0080     def get_coor0(self, values):
0081         return values[:, 0]
0082
0083     def get_coor1(self, values):
0084         return values[:, 2]
0085
0086
0087 # Helper to unify code for drawers not involve radius
0088 class CsvCartesianDrawer:
0089     def __init__(self, detector_csv, centers, bounds=None):
0090         super().__init__()
0091         d = pd.read_csv(detector_csv)
0092         d["cr"] = np.hypot(d["cx"], d["cy"])
0093
0094         l = len(d)
0095
0096         # millimeter precision should be enough
0097         for c in centers:
0098             d[c] = d[c].astype(int)
0099
0100         d = d.drop_duplicates(centers).copy()
0101         print("INFO CsvCartesianDrawer", centers, "dup drop", l, "->", len(d))
0102         self.points = d[centers].to_numpy()
0103
0104         # assumes x axis is telescope axis
0105         self.simple_telescope = bounds is not None
0106         if self.simple_telescope:
0107             self.bounds = d[bounds].to_numpy()
0108
0109     def draw_detector_impl(self, ax):
0110         if not self.simple_telescope:
0111             ax.scatter(self.points[:, 0], self.points[:, 1], color=detector_color, s=1)
0112         else:
0113             for point, bound in zip(self.points, self.bounds):
0114                 ax.plot(
0115                     [point[0], point[0]],
0116                     [point[1] + bound[0], point[1] + bound[1]],
0117                     color=detector_color,
0118                     zorder=0.5,
0119                 )
0120
0121         return ax
0122
0123
0124 class CsvXYDrawer(XYDrawer):
0125     def __init__(self, detector_csv, assume_telescope=False):
0126         super().__init__()
0127         if assume_telescope:
0128             self.drawer = CsvCartesianDrawer(
0129                 detector_csv, ["cx", "cy"], ["bound_param1", "bound_param3"]
0130             )
0131         else:
0132             self.drawer = CsvCartesianDrawer(detector_csv, ["cx", "cy"])
0133
0134     def draw_detector_impl(self, ax):
0135         return self.drawer.draw_detector_impl(ax)
0136
0137
0138 class CsvXZDrawer(XZDrawer):
0139     def __init__(self, detector_csv, assume_telescope=False):
0140         super().__init__()
0141         if assume_telescope:
0142             self.drawer = CsvCartesianDrawer(
0143                 detector_csv, ["cx", "cz"], ["bound_param0", "bound_param2"]
0144             )
0145         else:
0146             self.drawer = CsvCartesianDrawer(detector_csv, ["cx", "cz"])
0147
0148     def draw_detector_impl(self, ax):
0149         return self.drawer.draw_detector_impl(ax)
0150
0151
0152 class CsvZRDrawer(ZRDrawer):
0153     def __init__(self, detector_csv):
0154         super().__init__()
0155         self.drawer = CsvCartesianDrawer(detector_csv, ["cz", "cr"])
0156
0157     def draw_detector_impl(self, ax):
0158         return self.drawer.draw_detector_impl(ax)