eic-opticks/ana/pvplt.py

0001 #!/usr/bin/env python
0002 """
0003 pvplt.py
0004 ===========
0005
0006 pvplt_simple
0007 mpplt_focus
0008     given xlim, ymin args and FOCUS env-array return new xlim,ylim
0009
0010 pvplt_viewpoint
0011     according to EYE, LOOK, UP, ZOOM, PARA envvars
0012 pvplt_photon
0013     pos, mom, pol plotting
0014 pvplt_plotter
0015 pvplt_arrows
0016 pvplt_lines
0017 pvplt_add_contiguous_line_segments
0018
0019 mpplt_add_contiguous_line_segments
0020
0021
0022
0023
0024 """
0025
0026 import os, logging
0027 log = logging.getLogger(__name__)
0028 import numpy as np
0029
0030 D_ = os.environ.get("D", "3")
0031 MODE_ = os.environ.get("MODE", D_ )
0032 MODE = int(MODE_)
0033 mp = None
0034 pv = None
0035
0036 if MODE in [2,3,-2,-3]:
0037     try:
0038         import matplotlib as mp
0039         import matplotlib.pyplot as mpp
0040         # this works in a varietry of versions
0041         # trying to control window position using  fig.canvas.manager.window
0042         # might need mp.use('TkAgg')  but that crashes
0043         #import matplotlib.pyplot as mp
0044     except ImportError:
0045         mp = None
0046         mpp = None
0047     pass
0048 pass
0049
0050 if MODE in [3,-3]:
0051     try:
0052         import pyvista as pv
0053     except ImportError:
0054         pv = None
0055     pass
0056 pass
0057
0058
0059 if not mp is None:
0060     from matplotlib import collections  as mp_collections
0061     from matplotlib import lines as mp_lines
0062     from matplotlib.patches import Circle, Rectangle, Ellipse
0063 else:
0064     mp_collections = None
0065     Circle = None
0066     Rectangle = None
0067     Ellipse = None
0068 pass
0069
0070 from opticks.ana.eget import efloatlist_, efloatarray_, elookce_, elook_epsilon_, eint_
0071 from opticks.ana.axes import Axes, X,Y,Z
0072
0073 themes = ["default", "dark", "paraview", "document" ]
0074
0075 if not pv is None:
0076     pv.set_plot_theme(themes[1])
0077 pass
0078
0079 GUI = not "NOGUI" in os.environ
0080 SIZE = np.array([1280, 720] )
0081 ASPECT = float(SIZE[0])/float(SIZE[1])  # 1280/720 = 1.7777777777777777
0082
0083
0084 eary_ = lambda ekey, edef:np.array( list(map(float, os.environ.get(ekey,edef).split(","))) )
0085 efloat_ = lambda ekey, edef: float( os.environ.get(ekey,edef) )
0086
0087
0088
0089 XDIST = efloat_("XDIST", "200")
0090 FOCUS = eary_("FOCUS", "0,0,0")
0091 SCALE = efloat_("SCALE", "1")
0092
0093 def pvplt_simple(pl, xyz, label):
0094     """
0095     :param pl: pyvista plotter
0096     :param xyz: (n,3) shaped array of positions
0097     :param label: to place on plot
0098     """
0099     pl.add_text( "pvplt_simple %s " % label, position="upper_left")
0100     pl.add_points( xyz, color="white" )
0101
0102 def mpplt_focus_aspect(aspect=ASPECT, scale=SCALE):
0103     log.info("mpplt_focus_aspect aspect:%s FOCUS:%s scale:%s" % (str(aspect),str(FOCUS), scale))
0104     if np.all(FOCUS == 0): return None, None
0105
0106     center = FOCUS[:2]
0107     extent = FOCUS[2] if len(FOCUS) > 2 else 100
0108     diagonal  = np.array([extent*aspect, extent])
0109     botleft = center - diagonal
0110     topright = center + diagonal
0111     log.info("mpplt_focus_aspect botleft:%s" % str(botleft))
0112     log.info("mpplt_focus_aspect topright:%s" % str(topright))
0113     xlim = np.array([botleft[0], topright[0]])*scale
0114     ylim = np.array([botleft[1], topright[1]])*scale
0115     return xlim, ylim
0116
0117
0118
0119 def mpplt_focus(xlim, ylim):
0120     """
0121     Used from::
0122
0123         sysrap/sframe.py:mp_subplots (as used by tests/CSGSimtraceTest.py)
0124         g4cx/tests/cf_G4CXSimtraceTest.py:main
0125
0126     :param xlim: array of shape (2,)
0127     :param ylim: array of shape (2,)
0128     :return xlim,ylim: restricted by FOCUS envvar
0129     """
0130     aspect = (xlim[1]-xlim[0])/(ylim[1]-ylim[0])
0131     log.info("mpplt_focus xlim:%s ylim:%s FOCUS:%s " % (str(xlim),str(ylim), str(FOCUS)))
0132
0133
0134
0135 def pvplt_viewpoint(pl, reset=False, verbose=False):
0136     """
0137     https://github.com/pyvista/pyvista-support/issues/40
0138
0139     """
0140     eye = eary_("EYE",  "1,1,1.")
0141     look = eary_("LOOK", "0,0,0")
0142     up = eary_("UP", "0,0,1")
0143     zoom = efloat_("ZOOM", "1")
0144
0145     PARA = "PARA" in os.environ
0146     if verbose:
0147         print("pvplt_viewpoint reset:%d PARA:%d " % (reset, PARA))
0148         print(" eye  : %s " % str(eye) )
0149         print(" look : %s " % str(look) )
0150         print(" up   : %s " % str(up) )
0151         print(" zoom : %s " % str(zoom) )
0152     pass
0153     if PARA:
0154         pl.camera.ParallelProjectionOn()
0155     pass
0156     pl.set_focus(    look )
0157     pl.set_viewup(   up )
0158     pl.set_position( eye, reset=reset )
0159     pl.camera.Zoom(zoom)
0160
0161
0162 def pvplt_photon( pl, p, polcol="blue", polscale=1, wscale=False, wcut=True  ):
0163     """
0164     :param p: array of shape (1,4,4)
0165     :param polcol:
0166     :param polscale:
0167     :param wscale: when True, using wavelength(Coeff) as scale for polarization vector
0168     :param wcut: when True, apply selection that wavelength(Coeff) must be greater than 0.
0169     """
0170     assert p.shape == (1,4,4)
0171     pos = p[:,0,:3]
0172     mom = p[:,1,:3]
0173     pol = p[:,2,:3]
0174     wav = p[:,2,3]
0175
0176     if wcut and wav < 1e-6: return
0177
0178     if wscale:
0179         polscale *= wav
0180     pass
0181
0182     pl.add_points( pos, color="magenta", point_size=16.0 )
0183
0184     lmom = np.zeros( (2, 3), dtype=np.float32 )
0185     lmom[0] = pos[0]
0186     lmom[1] = pos[0] + mom[0]
0187
0188     lpol = np.zeros( (2, 3), dtype=np.float32 )
0189     lpol[0] = pos[0]
0190     lpol[1] = pos[0] + pol[0]*polscale
0191
0192     pl.add_lines( lmom, color="red" )
0193     pl.add_lines( lpol, color=polcol )
0194
0195
0196 def pvplt_plotter(label="pvplt_plotter", verbose=False):
0197     if verbose:
0198         print("STARTING PVPLT_PLOTTER ... THERE COULD BE A WINDOW WAITING FOR YOU TO CLOSE")
0199     pass
0200     pl = pv.Plotter(window_size=SIZE*2 )
0201     pvplt_viewpoint(pl, reset=False, verbose=verbose)
0202
0203     if not "NOGRID" in os.environ:
0204         pl.show_grid()
0205     else:
0206         print("pvplt_plotter NOGRID")
0207     pass
0208     TEST = os.environ.get("TEST","")
0209     pl.add_text( "%s %s " % (label,TEST), position="upper_left")
0210     return pl
0211
0212
0213
0214 def mpplt_annotate_fig( fig, label  ):
0215     suptitle = os.environ.get("SUPTITLE",label)
0216
0217     TOF = float(os.environ.get("TOF","0.99"))   # adjust the position of the title, to legibly display 4 lines
0218
0219     if len(suptitle) > 0:
0220         log.debug("suptitle:%s " % (suptitle) )
0221         fig.suptitle(suptitle, family="monospace", va="top", ha='left', x=0.1, y=TOF, fontweight='bold' )
0222     else:
0223         log.debug("no SUPTITLE/label" )
0224     pass
0225
0226
0227 def mpplt_annotate_ax( ax ):
0228     subtitle = os.environ.get("SUBTITLE", "")
0229     thirdline = os.environ.get("THIRDLINE", "")
0230     lhsanno  = os.environ.get("LHSANNO", "")
0231     rhsanno  = os.environ.get("RHSANNO", "")
0232
0233     if len(subtitle) > 0:
0234         log.debug("subtitle:%s " % (subtitle) )
0235         ax.text( 1.05, -0.12, subtitle, va='bottom', ha='right', family="monospace", fontsize=12, transform=ax.transAxes)
0236     else:
0237         log.debug("no SUBTITLE")
0238     pass
0239
0240     if len(thirdline) > 0:
0241         log.debug("thirdline:%s " % (thirdline) )
0242         ax.text(-0.05,  1.02, thirdline, va='bottom', ha='left', family="monospace", fontsize=12, transform=ax.transAxes)
0243     else:
0244         log.debug("no THIRDLINE")
0245     pass
0246
0247     if len(lhsanno) > 0:
0248         log.debug("lhsanno:%s " % (lhsanno) )
0249         ax.text(-0.05,  0.01, lhsanno, va='bottom', ha='left', family="monospace", fontsize=12, transform=ax.transAxes)
0250     else:
0251         log.debug("no lhsanno")
0252     pass
0253
0254     if len(rhsanno) > 0:
0255         rhsanno_pos = efloatarray_("RHSANNO_POS","0.6,0.01")
0256         log.debug("rhsanno:%s " % (rhsanno) )
0257         ax.text(rhsanno_pos[0], rhsanno_pos[1], rhsanno, va='bottom', ha='left', family="monospace", fontsize=12, transform=ax.transAxes)
0258     else:
0259         log.debug("no rhsanno")
0260     pass
0261
0262
0263
0264 def mpplt_plotter(nrows=1, ncols=1, label="", equal=True):
0265     """
0266     ISSUE: Observe that when plotting onto an external screen the layout of the
0267     plot differs : mysteriously this even happens when a plot is created on the
0268     laptop screen and then dragged onto the external monitor. But everything
0269     works fine when the screen capture is done on the laptop screen.
0270     This might be related to retina resolution of the laptop screen.
0271     """
0272     fig, axs = mpp.subplots(nrows=nrows, ncols=ncols, figsize=SIZE/100.) # 100 dpi
0273     if type(axs).__name__ == 'AxesSubplot':axs=np.array([axs], dtype=np.object )
0274
0275     if equal:
0276         for ax in axs:
0277             ax.set_aspect('equal')
0278         pass
0279     pass
0280     mpplt_annotate_fig(fig, label)
0281     for ax in axs:
0282         mpplt_annotate_ax(ax)
0283     pass
0284     return fig, axs
0285
0286
0287 def plotter(label=""):
0288     if MODE == 2:
0289         pl = mpplt_plotter(label=label)
0290         assert type(pl) is tuple and len(pl) == 2  # fig, ax
0291     elif MODE == 3:
0292         pl = pvplt_plotter(label=label)
0293     else:
0294         pl = None
0295     pass
0296     return pl
0297
0298
0299
0300
0301
0302
0303 def pvplt_arrows( pl, pos, vec, color='yellow', factor=0.15 ):
0304     """
0305
0306     glyph.orient
0307         Use the active vectors array to orient the glyphs
0308     glyph.scale
0309         Use the active scalars to scale the glyphs
0310     glyph.factor
0311         Scale factor applied to sclaing array
0312     glyph.geom
0313         The geometry to use for the glyph
0314
0315     """
0316     init_pl = pl == None
0317     if init_pl:
0318         pl = pvplt_plotter(label="pvplt_arrows")
0319     pass
0320     pos_cloud = pv.PolyData(pos)
0321     pos_cloud['vec'] = vec
0322     vec_arrows = pos_cloud.glyph(orient='vec', scale=False, factor=factor )
0323
0324     pl.add_mesh(pos_cloud, render_points_as_spheres=True, show_scalar_bar=False)
0325     pl.add_mesh(vec_arrows, color=color, show_scalar_bar=False)
0326
0327
0328 def pvplt_lines( pl, pos, vec, linecolor='white', pointcolor='white', factor=1.0 ):
0329     """
0330     :param pl: plotter
0331     :param pos: (n,3) array
0332     :param vec: (n,3) array
0333
0334     2025/01 : changed args to use separate line and point color
0335
0336     """
0337     init_pl = pl == None
0338     if init_pl:
0339         pl = pvplt_plotter(label="pvplt_line")
0340     pass
0341     pos_cloud = pv.PolyData(pos)
0342     pos_cloud['vec'] = vec
0343     geom = pv.Line(pointa=(0.0, 0., 0.), pointb=(1.0, 0., 0.),)
0344     vec_lines = pos_cloud.glyph(orient='vec', scale=False, factor=factor, geom=geom)
0345     pl.add_mesh(pos_cloud, color=pointcolor, render_points_as_spheres=True, show_scalar_bar=False)
0346     pl.add_mesh(vec_lines, color=linecolor, show_scalar_bar=False)
0347
0348
0349 def pvplt_frame( pl, sfr, local=False ):
0350     extent = sfr.ce[3]
0351     m2w = sfr.m2w if local == False else np.eye(4)
0352     pvplt_frame_(pl, extent, m2w )
0353
0354 def pvplt_frame_(pl, e, m2w=np.eye(4) ):
0355     """
0356
0357                  (-,+,+)
0358                    +-----------+ (+,+,+)
0359         (-,-,+)    |  (+,-,+)  |
0360             +----------2       |
0361             |      |   |       |
0362             |      |   |       |
0363             |      |   |       |
0364             |      +---|-------+ (+,+,-)      plus_y
0365             |          |
0366             0----------1       minus_y
0367        (-,-,-)      (+,-,-)
0368
0369            -X         +X
0370
0371
0372             Z
0373             |   Y
0374             |  /
0375             | /
0376             |/
0377             O----- X
0378
0379     """
0380     minus_y = np.array( [
0381            [-e,-e,-e],
0382            [+e,-e,-e],
0383            [+e,-e,+e],
0384            [-e,-e,+e],
0385            [-e,-e,-e] ])
0386
0387     pvplt_add_contiguous_line_segments(pl, minus_y, m2w=m2w )
0388
0389     plus_y = minus_y.copy()
0390     plus_y[:,1] = e
0391
0392     pvplt_add_contiguous_line_segments(pl, plus_y, m2w=m2w )
0393
0394
0395
0396 def transform_points( pos3, transform ):
0397     assert len(pos3.shape) == 2 and pos3.shape[1] == 3 and pos3.shape[0] > 0
0398     opos = np.ones_like( pos3, shape=(len(pos3),4 ) )
0399     opos[:,:3] = pos3
0400     tpos = np.dot( opos, transform )
0401     return tpos[:,:3]
0402
0403
0404 def pvplt_add_line_a2b(pl, a, b, color="white", width=1):
0405     lines = np.zeros( (2,3), dtype=np.float32 )
0406     lines[0] = a
0407     lines[1] = b
0408     pl.add_lines( lines, color=color, width=width  )
0409
0410
0411
0412
0413 def pvplt_add_contiguous_line_segments( pl, xpos, point_size=1, point_color="white", line_color="white", m2w=None, render_points_as_spheres=False,  ):
0414     """
0415     :param pl: pyvista plotter
0416     :param xpos: (n,3) array of positions
0417     :param m2w: (4,4) transform array or None
0418
0419     Adds red points at *xpos* and joins them with blue line segments using add_lines.
0420     This has been used only for small numbers of positions such as order less than 10
0421     photon step positions.
0422     """
0423     upos = xpos if m2w is None else transform_points( xpos, m2w )
0424
0425     pl.add_points( upos, color=point_color, render_points_as_spheres=render_points_as_spheres, point_size=point_size )
0426     xseg = contiguous_line_segments(upos)
0427     pl.add_lines( xseg, color=line_color )
0428
0429
0430 def mpplt_add_contiguous_line_segments(ax, xpos, axes, linewidths=2, colors="red", linestyle="dotted", label="label:mpplt_add_contiguous_line_segments", s=5 ):
0431     """
0432     :param ax: matplotlib 2D axis
0433     :param xpos: (n,3) array of positions
0434     :param axes: (2,)  2-tuple identifying axes to pick from the where X=0, Y=1, Z=2
0435     """
0436
0437     if len(xpos) == 0:
0438         log.info("len(xpos) zero : skip ")
0439         return
0440     pass
0441
0442     assert len(axes) == 2
0443     xseg = contiguous_line_segments(xpos[:,:3]).reshape(-1,2,3)  # (n,2,3)
0444     xseg2D = xseg[:,:,axes]   #  (n,2,2)    same as xseg[...,axes]  see https://numpy.org/doc/stable/user/basics.indexing.html
0445
0446     lc = mp_collections.LineCollection(xseg2D, linewidths=linewidths, colors=colors, linestyle=linestyle )
0447     ax.add_collection(lc)
0448
0449     xpos2d = xpos[:,axes]
0450     ax.scatter( xpos2d[:,0], xpos2d[:,1], s=s, label=label )
0451
0452
0453 def mpplt_add_line(ax, a, b, axes ):
0454     """
0455     :param ax:  matplotlib 2D axis
0456     :param a: (3,) xyz position array
0457     :param b: (3,) xyz position array
0458     :param axes: (2,) 2-tuple identifying axes X=0, Y=1, Z=2
0459     """
0460     H,V = axes
0461     l = mp_lines.Line2D([a[H],b[H]], [a[V],b[V]])
0462     ax.add_line(l)
0463
0464
0465 def mpplt_hist(mp, v, bins=100):
0466     fig, ax = mp.subplots()
0467     ax.hist(v, bins=bins )
0468     fig.show()
0469
0470
0471 def get_ellipse_param(elpar):
0472     """
0473     :param elpar_: list of up to 7 float
0474     """
0475     elw = elpar[0]*2 if len(elpar) > 0 else 100
0476     elh = elpar[1]*2 if len(elpar) > 1 else 100
0477     elx = elpar[2]   if len(elpar) > 2 else 0
0478     ely = elpar[3]   if len(elpar) > 3 else 0
0479     ela = elpar[4]   if len(elpar) > 4 else 0.1
0480     ez0 = elpar[5]   if len(elpar) > 5 else 0.
0481     ez1 = elpar[6]   if len(elpar) > 6 else 0.
0482     return elw,elh,elx,ely,ela,ez0,ez1
0483
0484 def mpplt_add_ellipse_(ax, par, opt):
0485     """
0486     :param ax: matplotlib 2D axis
0487     :param elpar: list of up to 5 float
0488
0489       tl            .
0490               .            .
0491
0492       ml  +                      .
0493
0494               .            .
0495       bl  +         .
0496
0497     """
0498     elw,elh,elx,ely,ela,ez0,ez1 = get_ellipse_param(par)
0499     el = Ellipse(xy=[elx,ely], width=elw, height=elh, alpha=ela )
0500     ax.add_artist(el)
0501
0502     #dz =  elh/2+ez0
0503
0504     topleft = (-elw/2, elh/2 + ez0 )
0505     midleft = (-elw/2, ez0 )
0506     botleft = (-elw/2,-elh/2 + ez0 )
0507     ec = "red" if opt.find("red") > -1 else "none"
0508
0509     log.info( " elw/2 %s elh/2 %s ez0 %s botleft %s ec %s opt %s" % (elw/2, elh/2, ez0, str(botleft), ec, opt) )
0510
0511     if opt.find("top") > -1:
0512         clip_bx = Rectangle(midleft,elw,elh/2, facecolor="none", edgecolor=ec )
0513     elif opt.find("bot") > -1:
0514         clip_bx = Rectangle(botleft,elw,elh/2, facecolor="none", edgecolor=ec )
0515     else:
0516         clip_bx = None
0517     pass
0518     if not clip_bx is None:
0519         ax.add_artist(clip_bx)
0520         el.set_clip_path(clip_bx)
0521         ## clipping is selecting part of the ellipse to include (not to cut away)
0522     pass
0523
0524
0525 def mpplt_add_ellipse(ax, ekey ):
0526     assert ekey.startswith("ELLIPSE")
0527     par = efloatlist_(ekey, "0,0,0,0")      #  elw,elh,elx,ely,ela,ez0,ez1
0528     opt = os.environ.get("%s_OPT" % ekey, "")
0529
0530     if len(par) > 0 and par[0] > 0.:
0531         mpplt_add_ellipse_(ax, par, opt)
0532     pass
0533
0534
0535
0536 def get_rectangle_param(repar):
0537     """
0538     :param repar: list of up to 5 float
0539     """
0540     hx = repar[0] if len(repar) > 0 else 100
0541     hy = repar[1] if len(repar) > 1 else 100
0542     cx = repar[2] if len(repar) > 2 else 0
0543     cy = repar[3] if len(repar) > 3 else 0
0544     al = repar[4] if len(repar) > 4 else 0.2
0545     dy = repar[5] if len(repar) > 5 else 0
0546     return hx,hy,cx,cy,al,dy
0547
0548
0549 def mpplt_add_rectangle_(ax, par, opt):
0550     """
0551                                    (cx+hx, cy+hy+dy)
0552            +----------+----------+
0553            |          |          |
0554            |          |(cx,cy+dy)|
0555            +----------+----------+
0556            |          |          |
0557            |          |          |
0558            +----------+----------+
0559        (cx-hx,cy-hy+dy)
0560
0561     """
0562     hx,hy,cx,cy,al,dy = get_rectangle_param(par)
0563     botleft = (cx-hx, cy-hy+dy)
0564     width = hx*2
0565     height = hy*2
0566     bx = Rectangle( botleft, width, height, alpha=al, facecolor="none", edgecolor="red" )
0567     ax.add_artist(bx)
0568
0569
0570 def mpplt_add_rectangle(ax, ekey):
0571     assert ekey.startswith("RECTANGLE")
0572     par = efloatlist_(ekey, "0,0,0,0")  #  hx,hy,cx,cy,al,dy
0573     opt = os.environ.get("%s_OPT" % ekey, "")
0574     mpplt_add_rectangle_(ax, par, opt)
0575
0576
0577 def mpplt_add_shapes(ax):
0578     """
0579     """
0580     mpplt_add_ellipse(ax,"ELLIPSE0")
0581     mpplt_add_ellipse(ax,"ELLIPSE1")
0582     mpplt_add_ellipse(ax,"ELLIPSE2")
0583
0584     mpplt_add_rectangle(ax,"RECTANGLE0")
0585     mpplt_add_rectangle(ax,"RECTANGLE1")
0586     mpplt_add_rectangle(ax,"RECTANGLE2")
0587
0588
0589
0590 def get_from_simtrace_isect( isect, mode="o2i"):
0591     """
0592     :param isect: (4,4) simtrace array item
0593     :param mode: str "o2i" "nrm" "nrm10"
0594     :return step: (2,3) point to point step
0595     """
0596     assert isect.shape == (4,4)
0597
0598     dist = isect[0,3] # simtrace layout assumed, see CSG/tests/SimtraceRerunTest.cc
0599     nrm = isect[0,:3]
0600     pos = isect[1,:3]
0601     ori = isect[2,:3]
0602     mom = isect[3,:3]
0603
0604     step = np.zeros( (2,3), dtype=np.float32 )
0605     if mode == "o2i":
0606         step[0] = ori
0607         step[1] = ori+dist*mom
0608     elif mode == "o2i_XDIST":
0609         step[0] = ori
0610         step[1] = ori+XDIST*mom
0611     elif mode == "nrm":
0612         step[0] = pos
0613         step[1] = pos+nrm
0614     elif mode == "nrm10":
0615         step[0] = pos
0616         step[1] = pos+10*nrm
0617     else:
0618         assert 0
0619     pass
0620     return step
0621
0622
0623 def mpplt_simtrace_selection_line(ax, sts, axes, linewidths=2):
0624     """
0625     :param ax:
0626     :param sts: simtrace_selection array of shape (n,2,4,4)  where n is small eg < 10
0627     :param axes:
0628
0629     The simtrace_selection created in CSG/tests/SimtraceRerunTest.cc
0630     contains pairs of isect the first from normal GPU simtrace and
0631     the second from CPU rerun.
0632
0633     TODO: at dot at pos
0634     """
0635     pass
0636     log.info("mpplt_simtrace_selection_line sts\n%s\n" % repr(sts))
0637
0638     MPPLT_SIMTRACE_SELECTION_LINE = os.environ.get("MPPLT_SIMTRACE_SELECTION_LINE", "o2i,o2i_XDIST,nrm10" )
0639     cfg = MPPLT_SIMTRACE_SELECTION_LINE.split(",")
0640     log.info("MPPLT_SIMTRACE_SELECTION_LINE %s cfg %s " % (MPPLT_SIMTRACE_SELECTION_LINE, str(cfg)))
0641
0642     colors = ["red","blue"]
0643     if sts.ndim in (3,4) and sts.shape[-2:] == (4,4):
0644         for i in range(len(sts)):
0645             jj = list(range(sts.shape[1])) if sts.ndim == 4 else [-1,]
0646             log.info(" jj %s " % str(jj))
0647             for j in jj:
0648                 isect = sts[i,j] if sts.ndim == 4 else sts[i]
0649                 color = colors[j%len(colors)]
0650                 if "o2i" in cfg:
0651                     o2i = get_from_simtrace_isect(isect, "o2i")
0652                     mpplt_add_contiguous_line_segments(ax, o2i, axes, linewidths=linewidths, colors=color, label="o2i", s=10 )
0653                 pass
0654                 if "o2i_XDIST" in cfg:
0655                     o2i_XDIST = get_from_simtrace_isect(isect, "o2i_XDIST")
0656                     mpplt_add_contiguous_line_segments(ax, o2i_XDIST, axes, linewidths=linewidths, colors=color, label="o2i_XDIST"  )
0657                 pass
0658                 if "nrm10" in cfg:
0659                     nrm10 = get_from_simtrace_isect(isect, "nrm10")
0660                     mpplt_add_contiguous_line_segments(ax, nrm10, axes, linewidths=linewidths, colors=color, label="nrm10" )
0661                 pass
0662             pass
0663         pass
0664
0665 def pvplt_simtrace_selection_line(pl, sts):
0666     """
0667     """
0668     print("pvplt_simtrace_selection_line sts\n", sts)
0669     colors = ["red","blue"]
0670     if sts.ndim in (3,4) and sts.shape[-2:] == (4,4):
0671         for i in range(len(sts)):
0672             jj = list(range(sts.shape[1])) if sts.ndim == 4 else [-1,]
0673             for j in jj:
0674                 isect = sts[i,j] if sts.ndim == 4 else sts[i]
0675                 color = colors[j%len(colors)]
0676                 o2i = get_from_simtrace_isect(isect, "o2i")
0677                 nrm10 = get_from_simtrace_isect(isect, "nrm10")
0678                 pvplt_add_contiguous_line_segments(pl, o2i, line_color=color, point_color=color )
0679                 pvplt_add_contiguous_line_segments(pl, nrm10, line_color=color, point_color=color )
0680             pass
0681         pass
0682     pass
0683
0684 def ce_line_segments( ce, axes ):
0685     """
0686     :param ce: center-extent array of shape (4,)
0687     :param axes: tuple of length 2 picking axes from X=0 Y=1 Z=2
0688     :return box_lseg: box line segments of shape (4, 2, 3)
0689
0690
0691        tl            tr = c + [e,e]
0692          +----------+
0693          |          |
0694          |----c-----|
0695          |          |
0696          +----------+
0697        bl            br = c + [e,-e]
0698
0699     """
0700     assert len(axes) == 2
0701
0702     c = ce[:3]
0703     e = ce[3]
0704
0705     h = e*Axes.UnitVector(axes[0])
0706     v = e*Axes.UnitVector(axes[1])
0707
0708     tr = c + h + v
0709     bl = c - h - v
0710     tl = c - h + v
0711     br = c + h - v
0712
0713     box_lseg = np.empty( (4,2,3), dtype=np.float32 )
0714     box_lseg[0] = (tl, tr)
0715     box_lseg[1] = (tr, br)
0716     box_lseg[2] = (br, bl)
0717     box_lseg[3] = (bl, tl)
0718     return box_lseg
0719
0720 def pvplt_parallel_lines(pl, gslim , aa, axes, look ):
0721     """
0722     HMM: seems no simple way to make dotted or dashed lines with pyvista ?
0723     The Chart system which has linestyles seems entirely separate from 3D plotting.
0724
0725     For example with canonical XZ axes with envvar XX=x0,x1,x2
0726     will get 3 parallel vertical lines spanning between the gslim Z-limits
0727     at the provided x positions
0728
0729          +---+------------+
0730          |   |            |
0731          |   |            |    Z
0732          |   |            |    |
0733          +---+------------+    +-- X
0734         x0  x1            x2
0735
0736
0737     gslim {0: array([209.35 , 210.197], dtype=float32), 1: array([-64.597, -64.597], dtype=float32), 2: array([129.514, 129.992], dtype=float32)}
0738     aa    {0: [209.5, 210.0], 1: [], 2: []}
0739     axes  (0, 2)
0740     look  [209.774, -64.59664, 129.752]
0741     """
0742     print("pvplt_parallel_lines")
0743     print("gslim %s " % str(gslim) )
0744     print("aa    %s " % str(aa) )
0745     print("axes  %s " % str(axes) )
0746     print("look  %s " % str(look) )
0747     assert len(axes) == 2
0748     H,V = axes
0749     for i in [X,Y,Z]:
0750         if len(aa[i]) == 0: continue
0751         for a in aa[i]:
0752             lo = look.copy()
0753             hi = look.copy()
0754             lo[i] = a
0755             hi[i] = a
0756             if i == H:  # i:horizontal axis, so create vertical lines
0757                 lo[V] = gslim[V][0]
0758                 hi[V] = gslim[V][1]
0759             elif i == V:  #  i:vertical axis, so create horizontal lines
0760                 lo[H] = gslim[H][0]
0761                 hi[H] = gslim[H][1]
0762             pass
0763             line = pv.Line(lo, hi, resolution=10)
0764             pl.add_mesh(line, color="w")
0765             log.info("i %d pv horizontal  lo %s hi %s " % (i, str(lo), str(hi)))
0766         pass
0767     pass
0768
0769
0770
0771 def mpplt_parallel_lines(ax, bbox, aa, axes, linestyle=None  ):
0772     """
0773     Draws axis parallel line segments in matplotlib and pyvista.
0774     The segments extend across the genstep grid limits.
0775     Lines to draw are configured using comma delimited value lists
0776     in envvars XX, YY, ZZ
0777
0778     :param ax: matplotlib axis
0779     :param bbox: array of shape (3,2)
0780     :param aa: {X:XX, Y:YY, Z:ZZ } dict of values
0781     :param axes: tuple eg (0,2)
0782
0783            +----------------------+
0784            |                      |
0785            |                      |
0786            +----------------------+
0787            |                      |
0788            +----------------------+
0789            |                      |
0790            |                      |   V=Z
0791            +----------------------+
0792
0793              H=X
0794
0795     """
0796     if ax is None: return
0797     H,V = axes
0798     log.info("mpplt_parallel_lines bbox[H] %s bbox[V] %s aa %s " % (str(bbox[H]), str(bbox[V]), str(aa)))
0799     for i in [X,Y,Z]:
0800         if len(aa[i]) == 0: continue
0801         for a in aa[i]:
0802             if V == i:  # vertical ordinate -> horizontal line
0803                 ax.plot( bbox[H], [a,a], linestyle=linestyle )
0804             elif H == i:  # horizontal ordinate -> vertical line
0805                 ax.plot( [a,a], bbox[V], linestyle=linestyle )
0806             pass
0807         pass
0808     pass
0809
0810 def mpplt_parallel_lines_auto(ax, bbox, axes, linestyle=None  ):
0811     aa = {}
0812     aa[X] = efloatlist_("XX")
0813     aa[Y] = efloatlist_("YY")
0814     aa[Z] = efloatlist_("ZZ")
0815     mpplt_parallel_lines(ax, bbox, aa, axes, linestyle=linestyle )
0816
0817
0818
0819
0820 def mpplt_ce(ax, ce, axes, linewidths=2, colors="blue"):
0821     assert len(axes) == 2
0822     box_lseg = ce_line_segments(ce, axes)
0823     box_lseg_2D = box_lseg[:,:,axes]
0824     lc = mp_collections.LineCollection(box_lseg_2D, linewidths=linewidths, colors=colors, label="ce %10.4f" % ce[3])
0825     ax.add_collection(lc)
0826
0827 def pvplt_ce(pl, ce, axes, color="blue"):
0828     assert len(axes) == 2
0829     box_lseg = ce_line_segments(ce, axes)
0830     box_lseg_pv = box_lseg.reshape(-1,3)
0831     pl.add_lines( box_lseg_pv, color=color )
0832
0833
0834 LOOKCE_COLORS = ["red","green", "blue", "cyan", "magenta", "yellow", "black"]
0835 def lookce_color(i):
0836     return LOOKCE_COLORS[i%len(LOOKCE_COLORS)]
0837
0838 def mpplt_ce_multiple(ax, lookce, axes):
0839     for i, ce in enumerate(lookce):
0840         mpplt_ce(ax, ce, axes=axes,colors=lookce_color(i) )
0841     pass
0842
0843 def pvplt_ce_multiple(pl, lookce, axes):
0844     for i,ce in enumerate(lookce):
0845         pvplt_ce(pl, ce, axes=axes, color=lookce_color(i))
0846     pass
0847
0848
0849
0850 def contiguous_line_segments( pos ):
0851     """
0852     :param pos: (n,3) array of positions
0853     :return seg: ( 2*(n-1),3) array of line segments suitable for pl.add_lines
0854
0855     Note that while the pos is assumed to represent a contiguous sequence of points
0856     such as photon step record positions the output line segments *seg*
0857     are all independent so they could be concatenated to draw line sequences
0858     for multiple photons with a single pl.add_lines
0859
0860     For example, a set of three positions (3,3)::
0861
0862         [0, 0, 0], [1, 0, 0], [1, 1, 0]
0863
0864     Would give seg (2,2,3)::
0865
0866          [[0,0,0],[1,0,0]],
0867          [[1,0,0],[1,1,0]]
0868
0869     Which is reshaped to (4,3)::
0870
0871          [[0, 0, 0], [1, 0, 0], [1, 0, 0], [1, 1, 0]])
0872
0873     That is the form needed to represent line segments between the points
0874     with pl.add_lines
0875
0876     """
0877     assert len(pos.shape) == 2 and pos.shape[1] == 3 and pos.shape[0] > 1
0878     num_seg = len(pos) - 1
0879     seg = np.zeros( (num_seg, 2, 3), dtype=pos.dtype )
0880     seg[:,0] = pos[:-1]   # first points of line segments skips the last position
0881     seg[:,1] = pos[1:]    # second points of line segments skipd the first position
0882     return seg.reshape(-1,3)
0883
0884
0885 def test_pvplt_contiguous_line_segments():
0886     pos = np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0]])
0887     seg = pvplt_contiguous_line_segments(pos)
0888     x_seg = np.array([[0, 0, 0], [1, 0, 0], [1, 0, 0], [1, 1, 0]])
0889
0890     print(pos)
0891     print(seg)
0892     assert np.all( x_seg == seg )
0893
0894
0895
0896 def pvplt_check_transverse( mom, pol, assert_transverse=True ):
0897     mom_pol_transverse = np.abs(np.sum( mom*pol , axis=1 )).max()
0898
0899     if mom_pol_transverse > 1e-5:
0900         print("WARNING mom and pol ARE NOT TRANSVERSE mom_pol_transverse %s assert_transverse %d " % ( mom_pol_transverse , assert_transverse ))
0901         if assert_transverse:
0902             assert mom_pol_transverse < 1e-5
0903         pass
0904     else:
0905         print("pvplt_check_transverse  mom_pol_transverse %s " % (mom_pol_transverse  ))
0906     pass
0907
0908
0909
0910
0911 def pvplt_polarized( pl, pos, mom, pol, factor=0.15, assert_transverse=True ):
0912     """
0913     :param pl:
0914     :param pos: shape (n,3)
0915     :param mom: shape (n,3)
0916     :param pol: shape (n,3)
0917     :param factor: scale factor, needs experimentation for visibility
0918
0919
0920     https://docs.pyvista.org/examples/00-load/create-point-cloud.html
0921     https://docs.pyvista.org/examples/01-filter/glyphs.html
0922
0923     Note bizarre issue of mom arrows only in one direction appearing ?
0924
0925     Managed to get them to appear using add_arrows and fiddling with mag
0926     in CSG/tests/CSGFoundry_MakeCenterExtentGensteps_Test.py
0927     """
0928     pvplt_check_transverse(mom, pol, assert_transverse=assert_transverse)
0929
0930     init_pl = pl == None
0931     if init_pl:
0932         pl = pvplt_plotter(label="pvplt_polarized")
0933     pass
0934
0935     pos_cloud = pv.PolyData(pos)
0936     pos_cloud['mom'] = mom
0937     pos_cloud['pol'] = pol
0938     mom_arrows = pos_cloud.glyph(orient='mom', scale=False, factor=factor )
0939     pol_arrows = pos_cloud.glyph(orient='pol', scale=False, factor=factor )
0940
0941     pl.add_mesh(pos_cloud, render_points_as_spheres=True, show_scalar_bar=False)
0942     pl.add_mesh(pol_arrows, color='lightblue', show_scalar_bar=False)
0943     pl.add_mesh(mom_arrows, color='red', show_scalar_bar=False)
0944
0945     if init_pl:
0946         cp = pl.show() if GUI else None
0947     else:
0948         cp = None
0949     pass
0950     return cp
0951
0952 def pvplt_add_points_adaptive( pl, pos, **kwa ):
0953     p3 = None
0954     if pos.ndim == 2:
0955         p3 = pos[:,:3]
0956     elif pos.ndim == 3:
0957         p3 = pos[:,:,:3].reshape(-1,3)
0958     else:
0959         print("pvplt_add_points_adaptive UNEXPECTED pos.ndim : %d " % pos.ndim )
0960     pass
0961     return pvplt_add_points(pl, p3, **kwa )
0962
0963
0964 def pvplt_add_points( pl, pos, **kwa ):
0965     if pos is None or len(pos) == 0:
0966         if VERBOSE: print("pvplt_add_points len(pos) ZERO   %s " % repr(kwa))
0967         return None
0968     pass
0969     #print("pvplt_add_points pos.shape %s kwa %s " % ( str(pos.shape), str(kwa) ))
0970     pos_cloud = pv.PolyData(pos)
0971     pl.add_mesh(pos_cloud, **kwa )
0972     return pos_cloud
0973
0974
0975 def pvplt_make_delta_lines( pos, delta ):
0976     """
0977     :param pos: (n,3) coordinates
0978     :param delta: (n,3) offsets
0979     :return vlin, lines:
0980
0981     vlin
0982         (n,3) vertices of start and end points of the lines
0983
0984     lines
0985         specification with number of points to join (2)
0986         and indices into vlin, eg::
0987
0988              np.array( [[2,100,101], [2,102,103]] ).ravel()
0989
0990     """
0991     assert(len(pos) == len(delta))
0992
0993     num_line = len(pos)
0994     vlin = np.zeros([num_line,2,3])
0995     vlin[:,0] = pos
0996     vlin[:,1] = pos + delta
0997     vlin = vlin.reshape(-1,3)
0998
0999     _lines = np.zeros( (num_line,3), dtype=np.int64 )
1000     _lines[:,0] = 2                           # 2 indices for each line segment
1001     _lines[:,1] = np.arange(0, 2*num_line,2)  # line start index within vlin
1002     _lines[:,2] = 1 + _lines[:,1]               # line end index within vlin
1003     lines = _lines.ravel()
1004
1005     return vlin, lines
1006
1007
1008 def pvplt_add_delta_lines( pl, pos, delta, **kwa ):
1009     """
1010     :param pl: plotter
1011     :param pos: (n,3) coordinates
1012     :param delta: (n,3) offsets
1013     """
1014     vlin, lines = pvplt_make_delta_lines(pos, delta )
1015     pvplt_add_lines(pl, vlin, lines, **kwa )
1016
1017
1018 def pvplt_add_lines( pl, pos, lines, **kwa ):
1019     """
1020     :param pos: float array of shape (n,3)
1021     :param lines: int array of form np.array([[2,0,1],[2,1,2]]).ravel()
1022
1023     Used for adding intersect normals
1024     """
1025     if len(pos) == 0:
1026         if VERBOSE: print("pvplt_add_lines len(pos) ZERO   %s " % repr(kwa))
1027         return None
1028     pass
1029     if VERBOSE: print("pvplt_add_lines pos.shape %s lines.shape %s kwa %s " % ( str(pos.shape), str(lines.shape), str(kwa) ))
1030     pos_lines = pv.PolyData(pos, lines=lines)
1031     pl.add_mesh(pos_lines, **kwa )
1032     return pos_lines
1033
1034 if __name__ == '__main__':
1035     test_pvplt_contiguous_line_segments()
1036 pass
1037