Source code for src.Libs.Simulation.UFO.UFOTrace

#!/usr/bin/env python3
"""
This file is part of the PAFFrontendSim.

The PAFFrontendSim is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License v3 as published by the Free Software Foundation.

The PAFFrontendSim is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
See the GNU General Public License for more details.

You should have received a copy of the GNU General Public License along with the PAFFrontendSim. If not, see https://www.gnu.org/licenses/.
"""
# ToDo
#############################################################################
# - clean-up code (frequently and regularly, please)
# - improve readability and add more comments (frequently and regularly)
#
#############################################################################
# import libs
#############################################################################
import numpy                     as np
from   copy                      import copy, deepcopy
import matplotlib.pyplot         as plt

_SciPY   = False
try:
    from scipy.interpolate           import griddata
    from scipy.spatial               import ConvexHull
    _SciPy = True
except:
    _SciPy = False

import LogClass
from Simulation.constants        import _ETA, _Tiso, _Impedance, _GaussBeamDiam
from Simulation                  import getColors

# import UFO-libs
from Simulation.UFO.constants    import _WaistContour, _Rings, _Rays
from Simulation.UFO              import getAngles
#############################################################################
# constants
#############################################################################
_CLASSNAME_RAY     = "UFO-Ray"
_CLASSNAME_TRACE   = "UFO-Trace"

#############################################################################
# classes
#############################################################################
[docs] class UFORay : """ Class describing one geometric beam-ray and offer a method to trace it through an optics """ ##################################### # internal data and init method # init method def __init__ ( self, position = np.asarray( [0., 0., 0.] ), direction = np.asarray( [0., 0., 1.] ), power = 0., log = None ): """ Initialize class. Parameters: position (numpy.array): 1d-array of len 3, indicating the coordinates of the ray starting point direction (numpy.array): 1d-array of len 3, indicating the normal vector of beam propagation at starting point power (float): power of the ray at point of emmission log (LogClass): Logging Class used by the calling application """ if log == None: self.log = LogClass.LogClass(_CLASSNAME_RAY) else: self.log = deepcopy(log) # logging class used for all in and output messages self.log.addLabel(self) self.ray = [ [np.asarray( position ), np.asarray( direction )] ] # list of ray-reflection points and their propagation directions self.rayPower = power # initial power of the ray self.endSurface = -1 # ray ended at this surface-number. < 0 : no ending # internal methods
[docs] def printOut ( self, line, message_type, prefix = True, lineFeed = True ): """ Wrapper for the printOut method of the actual log class """ self.log.printOut(line, message_type, prefix, lineFeed)
[docs] def getRay ( self, cut = None ): """ Extracts an array of the ray reflection points, which connected by lines give the ray-beam-path. Limits can be applied. Parameters: cut (list): of 3-element lists. The first element in the sub-list denotes the axis (0 to 2) and the second element the boundary data, and the third the limit type (> 0 upper, < 0 lower, and == 0 no limit) Returns: tuple (tuple): trace, self.endSurface (numpy array of all 3d ray-reflection points on the track) """ def _cutTraceAxis ( trace, value, axis = 2, direction = 1 ) : """ Cut data points above or below the limit given by value. If the limit is a upper limit is dependent on the direction flag. Parameters: trace (numpy.array): 2d numpy-array with 3 columns giving 3d-points value (float): giving the limit axis (int): giving the axis to limit direction (float): positive indicates a upper limit, negative indicates a lower limit, zero indicates to take all points """ newTrace = [] p1 = trace[0] for idx in range( 1, len( trace ) ) : p2 = trace[idx] if ( direction * p1[axis] > direction * value ) and ( direction * p2[axis] > direction * value ) : p1 = p2 continue if ( direction * p1[axis] <= direction * value ) and ( direction * p2[axis] <= direction * value ) : if len( newTrace ) == 0 : newTrace.append( p1 ) newTrace.append( p2 ) p1 = p2 else : m = ( value - p1[axis] ) / ( p2[axis] - p1[axis] ) p = m * ( p2 - p1 ) + p1 newTrace.append( p ) if ( direction * p1[axis] <= direction * value ) : if len( newTrace ) == 0 : newTrace.append( p1 ) else : newTrace.append( p2 ) p1 = p2 return np.asarray( newTrace ) # main procedure if type(cut) != type([]) : # if no limits given, return point-list return np.asarray(self.ray)[:, 0], self.endSurface for c in cut : # apply all limits trace = _cutTraceAxis( np.asarray(self.ray)[:, 0], c[1], c[0], c[2] ) return trace, self.endSurface
[docs] def getEndpoint ( self , zero = True ): """ Returns a list with the ray endpoint as numpy array. The list is empty if the ray ended during tracing. Parameters: zero (boolean): True = zero intensity rays are included, False = they are excluded. """ if ( ( self.endSurface >= 0 ) or ( ( not zero ) and ( self.rayPower == 0. ) ) ) : return [] return [ [ np.asarray( self.ray[-1][0] ), self.rayPower ] ]
[docs] def traceRay ( self, ufoList, posList = None ) : """ Traces the ray to the surface given by ufo. Parameters: ufoList (list): of UFODataClasses with the surface of the element to approach posList (list): of numpy-arrays with all positions of the point cloude of ufo (to speed up tracing, if None, data will be extracted from ufo -- slow) """ # extract position-arrays of all optical elements if not given if type(posList) == type(None) : # check type of posList and establish one if not given posList = [] # establish list for ufo in ufoList : # loop all optical elements posList.append( ufo.getData( onlyPos = True ) ) # extract position array for i, pos in enumerate( posList ) : # loop all elements # get directions from ray-point to all points of the surface dists = pos - self.ray[-1][0] # calculate connecting vector from ray end-point to all sub-elements norm = np.sqrt(np.einsum('...i,...i', dists, dists)).reshape(-1, 1) # calculate distance (absolute value) of all connections dirs = dists / norm # calculate normalized direction vectors # get direction closest to direction of propagation cosArray = np.dot( dirs , self.ray[-1][1] ) # dot product between ray propagation vector and direction vectors gives cosine between idx = np.argmax( cosArray ) # maximum cosine ==> most parallel vector err = ( 1. - cosArray[idx]**2. ) * norm[idx] **2. # calculate remaining error (ray usually does not hit sub-element fully) if err > np.dot( pos[0] - pos[1], pos[0] - pos[1] ) * 0.51 : # check if error is larger than sub-element spacing / sqrt(2) self.endSurface = i # if so, end ray break max_n = np.asarray( [ ufoList[i].Data[idx].n.x, # extract normal vector of closest sub-element ufoList[i].Data[idx].n.y, ufoList[i].Data[idx].n.z ] ) # here the uncertainty comes in: # 1. we only use the sub-element closesed to the surface <--> ray interception point and do not calculate the correct point # 2. we use the normal vector of the closeses sub-element as normal vector for the reflection and not the interpolated normal vector at point of interception # be awaire, this is only a diagnostic tool and does not provide physical optics data!! lens = False rot = None rotI = None # filter ufo type to get correct direction for different optical elements if ufoList[i].SurfaceType == "Aperture" : # for an aperture, leave ray untouched continue # Lenses (.... needs testing and confirmation .... ) if ufoList[i].SurfaceType == "Lens" : # calculate second surface as nominal geometric reflector with an f of the given focal length if isinstance(ufoList[i].er, (float, int, complex) ) and ( ufoList[i].er > 1. ) : # re-calculate normal vector to match a reflector with the given focal length # get rotation to and from nominal coordinate system rot, rotI, angles = getAngles( zOrientation = [ ufoList[i].Header.n.x, ufoList[i].Header.n.y, ufoList[i].Header.n.z ], xOrientation = [ ufoList[i].Header.o.x, ufoList[i].Header.o.y, ufoList[i].Header.o.z ] ) # calculate vector from lens center to lens reflection point dist = np.asarray( [ pos[idx][0] - ufoList[i].Header.v.x, pos[idx][1] - ufoList[i].Header.v.y, pos[idx][2] - ufoList[i].Header.v.z ] ) # and rotate it to nominal coordinate system (z of the header facing in positive z-direction) dist = np.dot( rot, dist ) d = np.sqrt( dist[0]**2 + dist[1]**2 ) # calculate distance to center a = np.pi/4. - np.arctan( ufoList[i].FocalLength / d ) / 2. # calculate normal vector angle with z-axis b = np.arctan2( dist[1], dist[0] ) # calculate normal vector angle with x-axis # calculate new normal vector max_n = np.asarray( [ -np.sin(a)*np.cos(b) , -np.sin(a)*np.sin(b), -np.cos(a) ] ) # fh = open("test_gt.dat", "a") # fh.write("%.3f %.3f %.3f " %( pos[idx][0] - ufoList[i].Header.v.x, pos[idx][1] - ufoList[i].Header.v.y, pos[idx][2] - ufoList[i].Header.v.z) ) # fh.write("%.3f %.3f %.3f " %( max_n[0], max_n[1], max_n[2] ) ) # and rotate it into the lens coordinate system max_n = np.dot( rotI, max_n ) # fh.write("%.3f %.3f %.3f " %( max_n[0], max_n[1], max_n[2] ) ) # fh.write("%.3f %.3f %.3f " %( dirs[idx][0], dirs[idx][1], dirs[idx][2] ) ) lens = True else : # see first surface as aperture only #flat reflector continue #max_n = np.asarray( [ ufoList[i].Header.n.x, ufoList[i].Header.n.y, ufoList[i].Header.n.z ] ) # calculate new ray direction newDir = dirs[idx] - 2. * np.dot( dirs[idx], max_n ) * max_n # using the normalized connecting vector of the sub-elements as ray propagation vector #newDir = self.ray[-1][1] - 2 * np.dot( self.ray[-1][1], max_n ) * max_n # using the actual ray direction as propagation vector newDir = newDir / np.sqrt(np.dot(newDir, newDir)) # normalize new propagation vector if lens : newDir = np.dot( rot, newDir ) newDir[2] *= -1. newDir = np.dot( rotI, newDir ) # fh.write("%.3f %.3f %.3f\n" %( newDir[0] , newDir[1], newDir[2]) ) # fh.close() # calculate new ray starting point newPos = pos[idx] # using the actual sub-element position as new ray-position. self.ray.append( [ newPos, newDir ] ) # add new ray-point to the ray
[docs] class GeometricTraceData : """ Class for geometric raytracing """ ##################################### # init method def __init__ ( self, objects, log = None ) : """ Initialize class. Parameters: objects (list): of UFODataClasses describing the optical elements in the order as they shall be traced. The first element is taken as transmitting surface log (LogClass): Logging Class used by the calling application """ if log == None: self.log = LogClass.LogClass(_CLASSNAME_TRACE) else: self.log = deepcopy(log) # logging class used for all in and output messages self.log.addLabel(self) self.objects = [] # list of UFO-objects for o in objects : # filter input list of objects for geometric raytracing # if o.SurfaceType == "Aperture" : # continue self.objects.append( o ) self.posList = [] # List of objects to be traces (only the position data) self.rayList = [] # List of available rays self.Power0 = 0. # power value of all rays started (arbitary units) for obj in self.objects[1:] : self.posList.append( obj.getData( onlyPos = True ) ) # internal methods
[docs] def printOut ( self, line, message_type, prefix = True, lineFeed = True ): """ Wrapper for the printOut method of the actual log class. """ self.log.printOut(line, message_type, prefix, lineFeed)
def _cutposList ( self, cut = None ) : """ Cuts data points of the objects outside the limits given by cut. Parameters: cut (list): List of 3-element lists. The first element in the sub-list denotes the axis (0 to 2) and the second element the boundary data, and the third the limit type (> 0 upper, < 0 lower, and == 0 no limit) """ if type(cut) == type(None) : return self.posList cutList = [] for pos in self.posList : newPos = [] for p in pos : take = True for c in cut : if c[2] * p[c[0]] > c[2] * c[1] : take = False if take : newPos.append(p) cutList.append( np.asarray( newPos ) ) return cutList # class methods
[docs] def Trace ( self, startWaist, wavelength , rings = _Rings, rays = _Rays ) : """ Trace a sequence of mirrors (UFODataClasses) with the UFO trace program - geometric ray-tracer - and return the ray-paths as np.array. Parameters: startWaist (float): waist size of the starting waist wavelength (float): wavelength of the optics setup (used to determine the opening angle of the initial beam) Returns: list (list): TODO something is wrong here """ self.printOut("Geometric tracing of %s objects:" % len(self.objects), self.log.FLOW_TYPE) self.printOut(" start-waist : %.2f mm" % startWaist, self.log.DATA_TYPE, False) self.printOut(" wavelength : %.2f mm" % wavelength, self.log.DATA_TYPE, False) self.printOut(" rings of rays : %d" % rings, self.log.DATA_TYPE, False) self.printOut(" rays per ring : %d" % rays, self.log.DATA_TYPE, False) if len(self.objects) < 2 : return [] # self.rayList = [] # clean list of rays self.Power0 = 0. # start-power # calculate ray-list for the start # get start-Object center position centerPos = np.asarray( [ self.objects[0].Header.v.x, self.objects[0].Header.v.y, self.objects[0].Header.v.z] ) centerDir = np.asarray( [ self.objects[0].Header.n.x, self.objects[0].Header.n.y, self.objects[0].Header.n.z] ) # and rotation matrix to direction (z-rotation around beam-axis is discarted)) aX, aY, aZ = self.objects[0].getRot() aX = aX / 180. * np.pi aY = aY / 180. * np.pi aZ = aZ / 180. * np.pi rx = np.asarray([ [1, 0, 0], [0, np.cos(-aX), -np.sin(-aX)], [0, np.sin(-aX), np.cos(-aX)] ]) ry = np.asarray([ [ np.cos(-aY), 0, np.sin(-aY)], [ 0, 1, 0], [-np.sin(-aY), 0, np.cos(-aY)] ]) rz = np.asarray([ [np.cos(-aZ), -np.sin(-aZ), 0.], [np.sin(-aZ), np.cos(-aZ), 0.], [ 0., 0., 1.] ]) rotMatrix = np.dot( rx, np.dot(ry, rz) ) maxBeta = _WaistContour / 2. * wavelength / ( np.pi * startWaist ) # openening angle at beam contour size stepBeta = maxBeta / rings # loop al theta-angles for idxTheta, theta in enumerate( np.arange( 0., maxBeta + 0.5 * stepBeta, stepBeta) ) : self.printOut(" tracing theta-angle of %.1f deg" % ( theta * 180 / np.pi ), self.log.FLOW_TYPE, False, False) radius = idxTheta * _WaistContour * startWaist / ( 2. * rings ) # field radius belonging to the angle theta, _WaistContour is a diameter ==> division by 2 if theta > 0. : thetaRays = radius * rays / startWaist rayPower = np.exp( - 2. * (radius**2.) / (startWaist**2.) ) else : thetaRays = 1 rayPower = 0. phiStep = 2. * np.pi / thetaRays for phi in np.arange( 0., 2. * np.pi - 0.5*phiStep, phiStep ) : # loop 360deg around phi for sign in ( -1., 1) : # loop two beam-orientations posVector = centerPos + np.dot( rotMatrix, radius * np.asarray( [ np.cos( phi ), np.sin( phi ), 0. ] ) ) dirVector = np.dot( rotMatrix, np.asarray( [ -1. * sign * np.sin( phi ) * np.sin( theta ), sign * np.cos( phi ) * np.sin( theta ), np.cos( theta ) ] ) ) ray = UFORay( posVector, dirVector, rayPower ) ray.traceRay( self.objects[1:], self.posList ) self.rayList.append( ray ) self.Power0 += rayPower self.printOut("", self.log.FLOW_TYPE, False) self.printOut(" total number of rays traced: %d" % len( self.rayList ), self.log.DEBUG_TYPE, False)
[docs] def getTraceLine ( self, cut = None ) : """ Extract one continuose line which shows all rays traced. Parameters: cut (list): of 3-element lists. The first element in the sub-list denotes the axis (0 to 2) and the second element the boundary data, and the third the limit type (> 0 upper, < 0 lower, and == 0 no limit) Returns: numpy.array (numpy.array): of points which can be plotted as one line, showing all rays traced. """ self.printOut("Getting trace line", self.log.FLOW_TYPE) traceData = [] for ray in self.rayList : rayData, es = ray.getRay( cut ) traceData += list( rayData ) + list( rayData[::-1] ) return np.asarray( traceData )
[docs] def getPattern ( self, zero = False ) : """ Returns pattern of the ray-end points at the last surface. Parameters: zero (boolean): True = zero intensizy rays are included, False = they are excluded (except center ray as first element). """ self.printOut("Getting end-pattern line", self.log.FLOW_TYPE) # get end-Object center position centerPos = np.asarray( [ self.objects[-1].Header.v.x, self.objects[-1].Header.v.y, self.objects[-1].Header.v.z] ) aX, aY, aZ = self.objects[-1].getRot() aX = aX / 180. * np.pi aY = aY / 180. * np.pi aZ = aZ / 180. * np.pi rx = np.asarray([ [1, 0, 0], [0, np.cos(aX), -np.sin(aX)], [0, np.sin(aX), np.cos(aX)] ]) ry = np.asarray([ [ np.cos(aY), 0, np.sin(aY)], [ 0, 1, 0], [-np.sin(aY), 0, np.cos(aY)] ]) rz = np.asarray([ [np.cos(aZ), -np.sin(aZ), 0.], [np.sin(aZ), np.cos(aZ), 0.], [ 0., 0., 1.] ]) rotMatrix = np.dot( rz, np.dot(ry, rx) ) pattern = self.rayList[0].getEndpoint() intensities = [] for ray in self.rayList[1:] : pattern += ray.getEndpoint( zero ) retArray = [] for p in pattern : p[0] = np.dot( rotMatrix, ( p[0] - centerPos ) ) retArray.append( np.asarray( [ p[0][0], p[0][1], p[0][2], p[1] ] ) ) return np.asarray( retArray )
[docs] def getPower ( self ) : """ Returns the fraction of the starting-power in the final trace pattern. """ self.printOut("Calculate power in resulting pattern", self.log.FLOW_TYPE) power = 0. for ray in self.rayList : if ray.endSurface < 0 : power += ray.rayPower return power / self.Power0
[docs] def plotPattern ( self, figname = "", show = False) : """ Plot the pattern of the rays on the end-surface. Parameters: figname (string): Filename of the figure (if given figure is only shown on screen if show == True show (boolean): boolean flag to show plot on screen even when being written to disk (True) """ self.printOut( "Plotting end-pattern", self.log.FLOW_TYPE ) points = self.getPattern() if len(points) < 3 : return surface = deepcopy( self.objects[-1] ) surface.resetPosition() sPoints = surface.getData( onlyPos = True ) fig, ax = plt.subplots(figsize=(16, 10)) ax.set_aspect('equal', 'box') ax.set_ylabel( "y-position [mm]" ) ax.set_xlabel( "x-position [mm]" ) ax.plot(sPoints[:, 0], sPoints[:, 1], "y.", label = "Surface points" ) ax.plot(points [1:, 0], points [1:, 1], "k.", label = "ray-pattern" ) ax.plot(points [0, 0], points [0, 1], "bo", label = "center-ray" ) ax.legend( loc = 'upper left' ) if figname != "" : fig.savefig(figname) if show : plt.show() else : plt.show() plt.close()
[docs] def plotIntensity ( self, figname = "", show = False) : """ Plot the intensity-pattern of the rays on the end-surface. Parameters: figname (string): Filename of the figure (if given figure is only shown on screen if show == True show (boolean): boolean flag to show plot on screen even when being written to disk (True) """ self.printOut("Plotting intensity-distribution", self.log.FLOW_TYPE) points = self.getPattern( zero = False ) if len(points) < 3 : return [], [], [] x = points[1:, 0] y = points[1:, 1] z = points[1:, 3] z = z / np.amax( z ) if not _SciPy : self.printOut("SciPy-Lib is not available, can't create the plot!", self.log.WARNING_TYPE) return x, y, z # grid the data. xi, yi = np.mgrid[ min(x):max(x):100*1j, min(y):max(y):100*1j] zi = griddata( ( x, y ), z, ( xi, yi ), method='linear' ) surface = deepcopy( self.objects[-1] ) surface.resetPosition() sPoints = surface.getData( onlyPos = True ) hull = ConvexHull( sPoints[:, 0:2] ) fig, ax = plt.subplots( figsize = ( 16, 10) ) ax.set_aspect( 'equal', 'box' ) ax.set_ylabel( "y-position [mm]" ) ax.set_xlabel( "x-position [mm]" ) cf = ax.pcolormesh( xi, yi, zi, cmap = plt.get_cmap( 'rainbow' ) ) ax.plot( points [1:, 0], points [1:, 1], "k.", label = "ray-pattern" ) ax.plot( points [0, 0], points [0, 1], "bo", label = "center-ray" ) addLeg = True for simplex in hull.simplices: if addLeg : ax.plot(sPoints[simplex, 0], sPoints[simplex, 1], 'r-', label = "mirror-contour") addLeg = False else : ax.plot(sPoints[simplex, 0], sPoints[simplex, 1], 'r-') ax.legend( loc = 'upper left' ) fig.colorbar(cf) if figname != "" : fig.savefig(figname) if show : plt.show() else : plt.show() plt.close() return xi, yi, zi
[docs] def plotGeometricTrace ( self, figname = "", show = False, cut = None) : """ Plot the geometrical layout including the mirror surfaces as three 2D projections. Parameters: figname (string): Filename of the figure (if given figure is only shown on screen if show == True show (boolean): flag to show plot on screen even when being written to disk (True) cut (list): of 3-element lists. The first element in the sub-list denotes the axis (0 to 2) and the second element the boundary data, and the third the limit type (> 0 upper, < 0 lower, and == 0 no limit) """ self.printOut("Plot geometric 3D trace in projection including mirrors :", self.log.FLOW_TYPE) trace = self.getTraceLine( cut ) # get trace-line posList = self._cutposList ( cut ) # get object data inside the limits given by cut colorList = getColors ( len( self.objects ) ) # get mirror colors if len(trace) < 2 : return # set up figures fig, (ax0, ax1, ax2) = plt.subplots(nrows=1, ncols=3, figsize=(12, 6)) # plot projection on y-z axis ax0.set_ylabel("z-position [mm]") ax0.set_xlabel("y-position [mm]") ax0.set_aspect('equal', 'box') for idx, data in enumerate(posList) : if len( data ) > 0 : ax0.plot(data[:, 1], data[:, 2], color = colorList[idx], markersize = 1) ax0.plot(trace [:, 1], trace [:, 2], "r-", label = "Beam-path" ) # plot projection on x-z axis ax1.set_ylabel("z-position [mm]") ax1.set_xlabel("x-position [mm]") ax1.set_aspect('equal', 'box') for idx, data in enumerate(posList) : if len( data ) > 0 : ax1.plot(data[:, 0], data[:, 2], color = colorList[idx], markersize = 1) ax1.plot(trace [:, 0], trace [:, 2], "r-") # plot projection on x-y axis ax2.set_ylabel("y-position [mm]") ax2.set_xlabel("x-position [mm]") ax2.set_aspect('equal', 'box') for idx, data in enumerate(posList) : if len( data ) > 0 : ax2.plot(data[:, 1], data[:, 0], color = colorList[idx], markersize = 1) ax2.plot(trace [:, 1], trace [:, 0], "r-") if figname != "" : fig.savefig(figname) if show : plt.show() else : plt.show() plt.close()
[docs] def plotGeometricTrace3D ( self, figname = "", show = False, cut = None ) : """ Plot the geometrical layout including the mirror surfaces as three 2D projections. Parameters: figname (string): Filename of the figure (if given figure is only shown on screen if show == True show (boolean): flag to show plot on screen even when being written to disk (True) cut (list): of 3-element lists. The first element in the sub-list denotes the axis (0 to 2) and the second element the boundary data, and the third the limit type (> 0 upper, < 0 lower, and == 0 no limit) """ self.printOut("Plot geometric 3D trace with mirrors :", self.log.FLOW_TYPE) trace = self.getTraceLine( cut ) # get trace-line posList = self._cutposList ( cut ) # get object data inside the limits given by cut colorList = getColors ( len( self.objects ) ) # get mirror colors # set up figure fig = plt.figure( figsize = ( 16, 10 ) ) ax = fig.add_subplot( projection = '3d' ) ax.view_init( elev = 0, azim = -90 ) # ax.set_aspect('equal') -- not supported in the used version of matplotlib # plot all surfaces for idx, data in enumerate( posList ) : if len( data ) > 0 : ax.plot(data[:, 0], data[:, 1], data[:, 2], color = colorList[idx]) # plot rays for ray in self.rayList : rayData, es = ray.getRay( cut ) if len( rayData ) > 0 : if es < 0 : rayColor = "red" else : rayColor = colorList[es] ax.plot(rayData[:, 0], rayData[:, 1], rayData[:, 2], color = rayColor) # display and / or save figure if figname != "" : fig.savefig(figname) if show : plt.show() else : plt.show() plt.close()