eic-opticks/ana/gplt.py

0001 #!/usr/bin/env python
0002 """
0003 gplt.py : 2d debugging/presentation of solids with rotational symmetry
0004 ========================================================================
0005
0006 Make connection between GDML parsing and the simple 2d matplotlib plotting
0007 of xplt.py while avoiding the need for manual translation as done in the x018.py etc..
0008
0009 ::
0010
0011     ipython --pdb -i gplt.py
0012
0013
0014
0015 """
0016
0017 import os, sys, argparse, logging
0018 import numpy as np, math
0019 import matplotlib.pyplot as plt
0020 import matplotlib.lines as mlines
0021
0022 log = logging.getLogger(__name__)
0023 sys.path.insert(0, os.path.expanduser("~"))  # assumes $HOME/opticks
0024
0025 from opticks.analytic.GDML import GDML, odict
0026 from opticks.ana.gargs import GArgs
0027
0028
0029 class GPlot(object):
0030     """
0031     GPlot
0032     ------
0033
0034     2d plotting small pieces of GDML defined geometry
0035
0036     """
0037     def __init__(self, lv, args):
0038         self.root = lv
0039         self.args = args
0040
0041     def plot(self, ax, recurse=True, **kwa):
0042         log.debug("kwa %s " % kwa)
0043         self.plot_r(self.root, ax, recurse=recurse, depth=0, **kwa )
0044
0045     def plot_r(self, lv0, ax, recurse, depth, **kwa):
0046         """
0047         Suspect this is assuming no offsets at pv/lv level
0048         """
0049         pvs = lv0.physvol
0050         indent = "   " * depth
0051         log.debug("[%2d] %s %4d %s " % (depth, indent,len(pvs),lv0.name))
0052
0053         s = lv0.solid
0054         color = self.args.color[lv0.local_index]
0055         kwa.update(color=color)
0056
0057         sh = s.as_shape(**kwa)
0058
0059         for pt in sh.patches():
0060             log.debug("pt %s" % pt)
0061             ax.add_patch(pt)
0062         pass
0063
0064         if recurse and ( depth < self.args.maxdepth or self.args.maxdepth == -1):
0065             for pv in pvs:
0066                 lv = pv.volume
0067                 self.plot_r( lv, ax, recurse, depth+1)
0068             pass
0069         else:
0070             pass
0071         pass
0072
0073
0074     @classmethod
0075     def MakeFig(cls, plt, lv, args, recurse=True):
0076         """
0077         With recurse True all subvolumes are drawn onto the same canvas
0078         """
0079
0080         fig = plt.figure(figsize=args.size)
0081         fig.suptitle(lv.local_prefix)
0082
0083         plt.title("gplt : %s " % lv.local_title)
0084
0085         ax = fig.add_subplot(111)
0086
0087         ax.set_xlim(args.xlim)
0088         ax.set_ylim(args.ylim)
0089
0090         gp = cls( lv, args)
0091         gp.plot(ax, recurse=recurse)
0092
0093         return fig, ax
0094
0095     @classmethod
0096     def MakeFigX(cls, plt, lvs, args, recurse=True):
0097         """
0098         With recurse True all subvolumes are drawn onto the same canvas
0099         """
0100         ny = 1
0101         nx = len(lvs)
0102         fig, axs = plt.subplots(ny, nx, **kwa )
0103         assert axs.shape == (nx)
0104
0105         for i in range(len(lvs)):
0106             lv = lvs[i]
0107             ax = axs[i]
0108             gp = cls( lv, args)
0109             gp.plot(ax, recurse=True)
0110         pass
0111
0112
0113     @classmethod
0114     def MultiFig(cls, plt, lvs, args):
0115         """
0116         Separate canvas for each LV
0117         """
0118         axs = []
0119         for lv in lvs.values():
0120             print(lv)
0121             fig, ax = cls.MakeFig(plt, lv, args, recurse=False)
0122             fig.show()
0123             axs.append(ax)
0124         pass
0125         return axs
0126
0127
0128     @classmethod
0129     def SubPlotsFig(cls, plt, lvsl, args, combiZoom=False, zoomlimits=None):
0130         """
0131         :param plt:
0132         :param lvsl: list containing one or more lvs lists of lv
0133
0134         A list of lists for lvsl is used to allow comparisons between two
0135         versions of the parsed GDML.
0136
0137         All volumes on one page via subplots
0138         """
0139
0140         if combiZoom:
0141             assert not zoomlimits is None
0142         pass
0143         shorten_title_ = getattr(args, 'shorten_title_', lambda t:t)
0144
0145         if len(lvsl) == 1:
0146             cf = False
0147             lvs = lvsl[0]
0148         elif len(lvsl) == 2:
0149             cf = True
0150             lvs0 = lvsl[0]
0151             lvs1 = lvsl[1]
0152             assert len(lvs0) == len(lvs1), (len(lvs0), len(lvs1))
0153             lvs = lvs0
0154         pass
0155
0156         n_lvs = len(lvs)
0157
0158         if n_lvs == 3:
0159             ny, nx = 2, 2
0160         else:
0161             ny, nx = 2, n_lvs//2
0162         pass
0163
0164         log.info("SubFig ny:%d nx:%d n_lvs:%d" % (ny,nx,n_lvs) )
0165
0166         kwa = dict()
0167         if not combiZoom:
0168             kwa["sharex"] = True
0169             kwa["sharey"] = True
0170         pass
0171         kwa["figsize"] = args.size
0172
0173         cx = 2
0174         if combiZoom:
0175             izz = range(2)
0176             fig, axs = plt.subplots(1, cx, **kwa )
0177         else:
0178             izz = range(1)
0179             fig, axs = plt.subplots(ny, nx, **kwa )
0180         pass
0181
0182         suptitle = lvs[0].local_prefix if cf == False else "%s cf %s " % (lvs0[0].local_prefix, lvs1[0].local_prefix)
0183         fig.suptitle(suptitle, fontsize=args.suptitle_fontsize)
0184
0185
0186         for iz in izz:
0187             iv = 0
0188             for iy in range(ny):
0189                 for ix in range(nx):
0190                     if iv < len(lvs):
0191
0192                         if combiZoom:
0193                             ax = axs if cx == 1 else axs[iz]
0194                         elif len(axs.shape) == 1:
0195                             ax = axs[iy]
0196                         elif len(axs.shape) == 2:
0197                             ax = axs[iy,ix]
0198                         pass
0199
0200                         ax.set_xlim(args.xlim)
0201                         ax.set_ylim(args.ylim)
0202                         ax.set_aspect('equal')
0203
0204                         if combiZoom and iz == 1:
0205                             zoomlimits(ax)
0206                         pass
0207
0208                         lv = lvs[iv]
0209                         title = lv.local_title if cf == False else "%s cf %s" % (shorten_title_(lvs0[iv].local_title), shorten_title_(lvs1[iv].local_title))
0210
0211                         if combiZoom and iv == 0:
0212                             log.info(title)
0213                             ax.set_title(title, fontsize=10)
0214                         else:
0215                             ax.set_title(title, fontsize=10)
0216                         pass
0217
0218                         if cf == False:
0219                             gp = cls( lv, args)
0220                             gp.plot(ax, recurse=False)
0221                         else:
0222                             gp0 = cls( lvs0[iv], args)
0223                             gp0.plot(ax, recurse=False)
0224                             gp1 = cls( lvs1[iv], args)
0225                             gp1.plot(ax, recurse=False, linestyle="dotted")
0226                         pass
0227                     pass
0228                     iv += 1
0229                 pass
0230             pass
0231         return fig, axs
0232
0233
0234
0235
0236
0237 def add_line(ax, p0, p1, **kwa):
0238     x0,y0 = p0
0239     x1,y1 = p1
0240     l = mlines.Line2D([x0, x1], [y0, y1], *kwa )
0241     ax.add_line(l)
0242
0243
0244 def pmt_annotate( ax, pmt):
0245     """
0246     """
0247     bulbneck, endtube = pmt.first, pmt.second
0248
0249     bulb = bulbneck.first
0250     neck = bulbneck.second
0251
0252     assert bulb.is_primitive, bulb
0253     assert bulb.__class__.__name__ == 'Ellipsoid', bulb
0254
0255     # position applies offset to second boolean constituent
0256     ztub = bulbneck.position.z  # neck offset
0257     add_line(ax, [-300,ztub], [300,ztub])
0258     add_line( ax, [0,-400], [0,400] )
0259
0260     if not neck.__class__.__name__ == 'Polycone':
0261         tube = neck.first
0262         torus = neck.second
0263         ztor = bulbneck.position.z + neck.position.z  # absolute neck offset
0264
0265         rtor = torus.rtor
0266         add_line(ax, [rtor, -400], [rtor, 400])
0267         add_line(ax, [-rtor, -400], [-rtor, 400])
0268
0269         add_line( ax, [-rtor,ztor], [0,0] )
0270         add_line( ax, [rtor,ztor], [0,0] )
0271         add_line(ax, [-300,ztor], [300,ztor] )
0272     pass
0273
0274
0275
0276
0277 if __name__ == '__main__':
0278     args = GArgs.parse(__doc__)
0279     lvx = args.lvname(1)
0280     lvx = "lAddition"
0281
0282     g = GDML.parse(args.gdmlpath(0))
0283     g.smry()
0284
0285     lv = g.find_one_volume(lvx)
0286     #s = lv.solid
0287     #s.sub_traverse()
0288
0289     log.info( "lv %r" % lv )
0290
0291     lvs = g.get_traversed_volumes( lv, maxdepth=args.maxdepth )
0292
0293     plt.ion()
0294
0295     fig, ax = GPlot.MakeFig(plt, lv, args, recurse=True)  # all volumes together
0296
0297     ax.set_aspect('equal')
0298
0299     if lvx == "lAddition":
0300         ax.set_xlim(-500,500)
0301         ax.set_ylim(-150,50)
0302         ax.plot( [-500,500], [0,0], linestyle="dotted", color="blue" )
0303     pass
0304
0305
0306     fig.show()
0307
0308     combpath = args.figpath("CombinedFig")
0309
0310     if lvx == "lAddition":
0311         combpath = "/tmp/lAddition_uni_acrylic3.png"
0312     pass
0313     log.info("saving to %s " % combpath)
0314     fig.savefig(combpath)
0315
0316     #axs = GPlot.MultiFig(plt, lvs, args)
0317
0318 if 0:
0319     fig, axs = GPlot.SubPlotsFig(plt, [lvs], args)
0320     fig.show()
0321     fig.savefig(args.figpath("SplitFig"))
0322
0323
0324 if 0:
0325     pmt = lvs(-1).solid
0326     pmt_annotate( ax, pmt)
0327
0328     bulbneck, endtube = pmt.first, pmt.second
0329     bulb, neck = bulbneck.first, bulbneck.second
0330     neck_offset = bulbneck.position
0331
0332     neckz = neck_offset.z
0333     add_line( ax, [-300,neckz], [300,neckz] )
0334
0335