ModulePlottingMethods

  1"""
  2Adds plotting and visualisation methods to various classes from ARTcore.
  3Most of the code is simply moved here from the class definitions previously.
  4
  5Created in July 2024
  6
  7@author: André Kalouguine + Stefan Haessler + Anthony Guillaume
  8"""
  9# %%
 10import numpy as np
 11import matplotlib.pyplot as plt
 12from mpl_toolkits.mplot3d import Axes3D
 13import matplotlib.patches as patches
 14import pyvista as pv
 15import pyvistaqt as pvqt
 16import colorcet as cc
 17
 18import ARTcore.ModuleSupport as msup
 19import ARTcore.ModuleMirror as mmir
 20import ARTcore.ModuleGeometry as mgeo
 21import ARTcore.ModuleMask as mmask
 22from ARTcore.ModuleGeometry import Point, Vector, Origin
 23import itertools
 24
 25# %% Adding support drawing for mirror projection
 26
 27def SupportRound_ContourSupport(self, Figure):
 28    """Draw support contour in MirrorProjection plots."""
 29    axe = Figure.add_subplot(111, aspect="equal")
 30    axe.add_patch(patches.Circle((0, 0), self.radius, alpha=0.08))
 31    return axe
 32
 33msup.SupportRound._ContourSupport = SupportRound_ContourSupport
 34
 35def SupportRoundHole_ContourSupport(self, Figure):
 36    """Draws support contour in MirrorProjection plots."""
 37    axe = Figure.add_subplot(111, aspect="equal")
 38    axe.add_patch(patches.Circle((0, 0), self.radius, alpha=0.08))
 39    axe.add_patch(patches.Circle((self.centerholeX, self.centerholeY), self.radiushole, color="white", alpha=1))
 40    return axe
 41
 42msup.SupportRoundHole._ContourSupport = SupportRoundHole_ContourSupport
 43
 44def SupportRectangle_ContourSupport(self, Figure):
 45    """Draws support contour in MirrorProjection plots."""
 46    axe = Figure.add_subplot(111, aspect="equal")
 47    axe.add_patch(patches.Rectangle((-self.dimX * 0.5, -self.dimY * 0.5), self.dimX, self.dimY, alpha=0.08))
 48    return axe
 49
 50msup.SupportRectangle._ContourSupport = SupportRectangle_ContourSupport
 51
 52def SupportRectangleHole_ContourSupport(self, Figure):
 53    """Draws support contour in MirrorProjection plots."""
 54    axe = Figure.add_subplot(111, aspect="equal")
 55    axe.add_patch(patches.Rectangle((-self.dimX * 0.5, -self.dimY * 0.5), self.dimX, self.dimY, alpha=0.08))
 56    axe.add_patch(patches.Circle((self.centerholeX, self.centerholeY), self.radiushole, color="white", alpha=1))
 57    return axe
 58
 59msup.SupportRectangleHole._ContourSupport = SupportRectangleHole_ContourSupport
 60
 61def SupportRectangleRectHole_ContourSupport(self, Figure):
 62    """Draws support contour in MirrorProjection plots."""
 63    axe = Figure.add_subplot(111, aspect="equal")
 64    axe.add_patch(patches.Rectangle((-self.dimX * 0.5, -self.dimY * 0.5), self.dimX, self.dimY, alpha=0.08))
 65    axe.add_patch(
 66        patches.Rectangle(
 67            (-self.holeX * 0.5 + self.centerholeX, -self.holeY * 0.5 + self.centerholeY),
 68            self.holeX,
 69            self.holeY,
 70            color="white",
 71            alpha=1,
 72        )
 73    )
 74    return axe
 75
 76msup.SupportRectangleRectHole._ContourSupport = SupportRectangleRectHole_ContourSupport
 77
 78# %% Support sampling and meshing 
 79
 80def sample_support(Support, Npoints):
 81    """
 82    This function samples a regular grid on the support and filters out the points that are outside of the support.
 83    For that it uses the _estimate_size method of the support.
 84    It then generates a regular grid over that area and filters out the points that are outside of the support.
 85    If less than half of the points are inside the support, it will increase the number of points by a factor of 2 and try again.
 86    """
 87    if hasattr(Support, "size"):
 88        size = Support.size
 89    else:
 90        size_x = Support._estimate_size()
 91        size = np.array([2*size_x, 2*size_x])
 92    enough = False
 93    while not enough:
 94        X = np.linspace(-size[0] * 0.5, size[0] * 0.5, int(np.sqrt(Npoints)))
 95        Y = np.linspace(-size[1] * 0.5, size[1] * 0.5, int(np.sqrt(Npoints)))
 96        X, Y = np.meshgrid(X, Y)
 97        Points = np.array([X.flatten(), Y.flatten()]).T
 98        Points = [p for p in Points if p in Support]
 99        if len(Points) > Npoints * 0.5:
100            enough = True
101        else:
102            Npoints *= 2
103    return mgeo.PointArray(Points)
104
105def sample_z_values(Support, Npoints):
106    Npoints_init = Npoints
107    if hasattr(Support, "size"):
108        size = Support.size
109    else:
110        size_x = Support._estimate_size()
111        size = np.array([2*size_x, 2*size_x])
112    enough = False
113    while not enough:
114        X = np.linspace(-size[0] * 0.5, size[0] * 0.5, int(np.sqrt(Npoints)))
115        Y = np.linspace(-size[1] * 0.5, size[1] * 0.5, int(np.sqrt(Npoints)))
116        X, Y = np.meshgrid(X, Y)
117        Points = np.array([X.flatten(), Y.flatten()]).T
118        Points = [p for p in Points if p in Support]
119        if len(Points) > Npoints_init * 0.5:
120            enough = True
121        else:
122            Npoints *= 2
123    Points = np.array([X.flatten(), Y.flatten()]).T
124    mask = np.array([p in Support for p in Points]).reshape(X.shape)
125    return X,Y,mask
126
127
128# %% Mesh/pointcloud generation for different mirror/mask types
129
130# For the mirrors that are already implemented in ARTcore
131# we can use the method _zfunc to render the mirror surface. 
132# For that we sample the support and then use the _zfunc method to get the z-values of the surface.
133# We then use the pyvista function add_mesh to render the surface.
134
135def _RenderMirror(Mirror, Npoints=10000):
136    """
137    This function renders a mirror in 3D.
138    It samples the support of the mirror and uses the _zfunc method to get the z-values of the surface.
139    It draws a small sphere at each point of the support and connects them to form a surface.
140    """
141    Points = sample_support(Mirror.support, Npoints)
142    Points += Mirror.r0[:2]
143    Z = Mirror._zfunc(Points)
144    Points = mgeo.PointArray([Points[:, 0], Points[:, 1], Z]).T
145    Points = Points.from_basis(*Mirror.basis)
146    mesh = pv.PolyData(Points)
147    return mesh
148
149
150def _RenderMirrorSurface(Mirror, Npoints=10000):
151    """
152    This function renders a mirror in 3D.
153    It samples the support of the mirror and uses the _zfunc method to get the z-values of the surface.
154    It draws a small sphere at each point of the support and connects them to form a surface.
155    """
156    X,Y,mask = sample_z_values(Mirror.support, Npoints)
157    shape = X.shape
158    X += Mirror.r0[0]
159    Y += Mirror.r0[1]
160    Z = Mirror._zfunc(mgeo.PointArray([X.flatten(), Y.flatten()]).T).reshape(X.shape)
161    #Z += Mirror.r0[2]
162    Z[~mask] = np.nan
163    Points = mgeo.PointArray([X.flatten(), Y.flatten(), Z.flatten()]).T
164    Points = Points.from_basis(*Mirror.basis)
165    X,Y,Z = Points.T
166    X = X.reshape(shape)
167    Y = Y.reshape(shape)
168    Z = Z.reshape(shape)
169    mesh = pv.StructuredGrid(X, Y, Z)
170    return mesh
171
172mmir.MirrorPlane._Render = _RenderMirror
173mmir.MirrorParabolic._Render = _RenderMirror
174mmir.MirrorCylindrical._Render = _RenderMirror
175mmir.MirrorEllipsoidal._Render = _RenderMirror
176mmir.MirrorSpherical._Render = _RenderMirror
177mmir.MirrorToroidal._Render = _RenderMirror
178
179mmask.Mask._Render = _RenderMirror
180
181mmir.MirrorPlane._Render = _RenderMirrorSurface
182mmir.MirrorParabolic._Render = _RenderMirrorSurface
183mmir.MirrorCylindrical._Render = _RenderMirrorSurface
184mmir.MirrorEllipsoidal._Render = _RenderMirrorSurface
185mmir.MirrorSpherical._Render = _RenderMirrorSurface
186mmir.MirrorToroidal._Render = _RenderMirrorSurface
187
188mmask.Mask._Render = _RenderMirrorSurface
189
190# %% Optical element rendering
191def _RenderOpticalElement(fig, OE, OEpoints, draw_mesh = False, color="blue", index=0):
192    """
193    This function renders the optical elements in the 3D plot.
194    It receives a pyvista figure handle, and draws the optical element on it.
195    """
196    mesh = OE._Render(OEpoints)
197    fig.add_mesh(mesh, color = color, name = f"{OE.description} {index}")    
198
199def _RenderDetector(fig, Detector, size = 40, name = "Focus", detector_meshes = None):
200    """
201    Unfinished
202    """
203    Points = [np.array([-size/2,size/2,0]),np.array([size/2,size/2,0]),np.array([size/2,-size/2,0]),np.array([-size/2,-size/2,0])]
204    Points = mgeo.PointArray(Points)
205    Points = Points.from_basis(*Detector.basis)
206    Rect = pv.Rectangle(Points[:3])
207    if detector_meshes is not None:
208        detector_meshes += [Rect]
209    fig.add_mesh(Rect, color="green", name=f"Detector {name}")
210
211# %% Support rendering using the sdf method
212
213def _RenderSupport(Support, Npoints=10000):
214    """
215    This function renders a support in 3D.
216    """
217    X,Y,mask = sample_z_values(Support, Npoints)
218    shape = X.shape
219    Z = np.zeros_like(X)
220    Z[~mask] = np.nan
221    mesh = pv.StructuredGrid(X, Y, Z)
222    return mesh
223    
224msup.SupportRound._Render = _RenderSupport
225msup.SupportRoundHole._Render = _RenderSupport
226msup.SupportRectangle._Render = _RenderSupport
227msup.SupportRectangleHole._Render = _RenderSupport
228msup.SupportRectangleRectHole._Render = _RenderSupport
229
230
231# %% Standalone optical element rendering
232
233def _RenderMirror(Mirror, Npoints=1000, draw_support = False, draw_points = False, draw_vectors = True , recenter_support = True):
234    """
235    This function renders a mirror in 3D.
236    """
237    mesh = Mirror._Render(Npoints)
238    p = pv.Plotter()
239    p.add_mesh(mesh)
240    if draw_support:
241        support = Mirror.support._Render()
242        if recenter_support:
243            support.translate(Mirror.r0,inplace=True)
244        p.add_mesh(support, color="gray", opacity=0.5)
245    
246    if draw_vectors:
247        # We draw the important vectors of the optical element
248        # For that, if we have a "vectors" attribute, we use that
249        #  (a dictionary with the vector names as keys and the colors as values)
250        # Otherwise we use the default vectors: "support_normal", "majoraxis"
251        if hasattr(Mirror, "vectors"):
252            vectors = Mirror.vectors
253        else:
254            vectors = {"support_normal_ref": "red", "majoraxis_ref": "blue"}
255        for vector, color in vectors.items():
256            if hasattr(Mirror, vector):
257                p.add_arrows(Mirror.r0, 10*getattr(Mirror, vector), color=color)
258    
259    if draw_points:
260        # We draw the important points of the optical element
261        # For that, if we have a "points" attribute, we use that
262        #  (a dictionary with the point names as keys and the colors as values)
263        # Otherwise we use the default points: "centre_ref"
264        if hasattr(Mirror, "points"):
265            points = Mirror.points
266        else:
267            points = {"centre_ref": "red"}
268        for point, color in points.items():
269            if hasattr(Mirror, point):
270                p.add_mesh(pv.Sphere(radius=1, center=getattr(Mirror, point)), color=color)
271    else:
272        p.add_mesh(pv.Sphere(radius=1, center=Mirror.r0), color="red")
273    p.show()
274    return p
275
276mmir.MirrorPlane.render = _RenderMirror
277mmir.MirrorParabolic.render = _RenderMirror
278mmir.MirrorParabolic.vectors = {
279    "support_normal": "red",
280    "majoraxis": "blue",
281    "towards_focusing_ref": "green"
282}
283mmir.MirrorParabolic.points = {
284    "centre_ref": "red",
285    "focus_ref": "green"
286}
287mmir.MirrorCylindrical.render = _RenderMirror
288mmir.MirrorEllipsoidal.render = _RenderMirror
289mmir.MirrorEllipsoidal.vectors = {
290    "support_normal_ref": "red",
291    "majoraxis_ref": "blue",
292    "towards_image_ref": "green",
293    "towards_object_ref": "green",
294    "centre_normal_ref": "purple"}
295
296mmir.MirrorSpherical.render = _RenderMirror
297mmir.MirrorToroidal.render = _RenderMirror
298
299mmask.Mask.render = _RenderMirror
300
301# %% Ray rendering
302
303def _RenderRays(RayListHistory, EndDistance, maxRays=150):
304    """
305    Generates a list of Pyvista-meshes representing the rays in RayListHistory.
306    This can then be employed to render the rays in a Pyvista-figure.
307
308    Parameters
309    ----------
310        RayListHistory : list(list(Ray))
311            A list of lists of objects of the ModuleOpticalRay.Ray-class.
312
313        EndDistance : float
314            The rays of the last ray bundle are drawn with a length given by EndDistance (in mm).
315
316        maxRays : int, optional
317            The maximum number of rays to render. Rendering all the traced rays is a insufferable resource hog
318            and not required for a nice image. Default is 150.
319    
320    Returns
321    -------
322        meshes : list(pvPolyData)
323            List of Pyvista PolyData objects representing the rays
324    """
325    meshes = []
326    # Ray display
327    for k in range(len(RayListHistory)):
328        x = []
329        y = []
330        z = []
331        if k != len(RayListHistory) - 1:
332            knums = list(
333                map(lambda x: x.number, RayListHistory[k])
334            )  # make a list of all ray numbers that are still in the game
335            if len(RayListHistory[k + 1]) > maxRays:
336                rays_to_render = np.random.choice(RayListHistory[k + 1], maxRays, replace=False)
337            else:
338                rays_to_render = RayListHistory[k + 1]
339
340            for j in rays_to_render:
341                indx = knums.index(j.number)
342                i = RayListHistory[k][indx]
343                Point1 = i.point
344                Point2 = j.point
345                x += [Point1[0], Point2[0]]
346                y += [Point1[1], Point2[1]]
347                z += [Point1[2], Point2[2]]
348
349        else:
350            if len(RayListHistory[k]) > maxRays:
351                rays_to_render = np.random.choice(RayListHistory[k], maxRays, replace=False)
352            else:
353                rays_to_render = RayListHistory[k]
354
355            for j in rays_to_render:
356                Point = j.point
357                Vector = j.vector
358                x += [Point[0], Point[0] + Vector[0] * EndDistance]
359                y += [Point[1], Point[1] + Vector[1] * EndDistance]
360                z += [Point[2], Point[2] + Vector[2] * EndDistance]
361        points = np.column_stack((x, y, z))
362        meshes += [pv.line_segments_from_points(points)]
363    return meshes