src.acoustools.Visualiser

  1import torch
  2from acoustools.Utilities import propagate_abs, device, TRANSDUCERS, create_points
  3import matplotlib.pyplot as plt
  4from mpl_toolkits.mplot3d import art3d
  5import matplotlib.colors as clrs
  6import matplotlib.cm as cm
  7import matplotlib.colors as mcolors
  8from mpl_toolkits.axes_grid1 import make_axes_locatable
  9from matplotlib.widgets import MultiCursor
 10
 11import matplotlib.animation as animation
 12import matplotlib  as mpl
 13
 14from torch import Tensor
 15from types import FunctionType
 16from typing import Literal
 17from vedo import Mesh
 18
 19
 20
 21
 22
 23def Visualise(A:Tensor,B:Tensor,C:Tensor,activation:Tensor,points:list[Tensor]|Tensor=[],
 24              colour_functions:list[FunctionType]|None=[propagate_abs], colour_function_args:list[dict]|None=None, 
 25              res:tuple[int]=(200,200), cmaps:list[str]=[], add_lines_functions:list[FunctionType]|None=None, 
 26              add_line_args:list[dict]|None=None,vmin:int|list[int]|None=None,vmax:int|list[int]|None=None, 
 27              matricies:Tensor|list[Tensor]|None = None, show:bool=True,block:bool=True, clr_labels:list[str]|None=None, depth:int=2, link_ax:str|list|None='all',
 28              cursor:bool=False, arrangement:tuple|None = None, arangement:tuple|None = None, titles:list[str]|None=None, call_abs=False, norm_axes=None ) -> None:
 29    '''
 30    Visualises any number of fields generated from activation to the plane ABC and arranges them in a (1,N) grid \n
 31    :param A: Position of the top left corner of the image
 32    :param B: Position of the top right corner of the image
 33    :param C: Position of the bottom left corner of the image
 34    :param activation: The transducer activation to use
 35    :param points: List of point positions to add crosses for each plot. Positions should be given in their position in 3D
 36    :param colour_functions: List of function to call at each position for each plot. Should return a value to colour the pixel at that position. Default `acoustools.Utilities.propagate_abs`. \n
 37        - If colour_functions is `difference` or `ratio` then the result will be the difference or ratio between existing images. To control which images are used `colour_function_args` should have a parameter called `ids`.
 38    :param colour_function_args: The arguments to pass to `colour_functions`
 39    :param res: Number of pixels as a tuple (X,Y). Default (200,200)
 40    :param cmaps: The cmaps to pass to plot
 41    :param add_lines_functions: List of functions to extract lines and add to the image
 42    :param add_line_args: List of parameters to add to `add_lines_functions`
 43    :param vmin: Minimum value to use across all plots
 44    :param vmax: MAximum value to use across all plots
 45    :param matricies: precomputed matricies to plot
 46    :param show: If True will call `plt.show(block=block)` else does not. Default True
 47    :param block: Will be passed to `plot.show(block=block)`. Default True
 48    :param clr_label: Label for colourbar
 49    :param depth: Number of times to tile image
 50    :param link_ax: Axes to link colourbar of `'all'` to link all axes. To unlink all axis, pass one of ['none', False, None]
 51    :param cursor: If `True` will show cursor across plots
 52    :param arangement: Arangment of subplots 
 53    :param title: Titles for each subplot
 54    :param call_abs: if True will call torch.abs on the image 
 55    :param norm_axes: List of Axes to Normalise
 56
 57    ```Python
 58    from acoustools.Utilities import create_points, add_lev_sig
 59    from acoustools.Solvers import wgs
 60    from acoustools.Visualiser import Visualise
 61
 62    import torch
 63
 64    p = create_points(1,1,x=0,y=0,z=0)
 65    x = wgs(p)
 66    x = add_lev_sig(x)
 67
 68    A = torch.tensor((-0.09,0, 0.09))
 69    B = torch.tensor((0.09,0, 0.09))
 70    C = torch.tensor((-0.09,0, -0.09))
 71    normal = (0,1,0)
 72    origin = (0,0,0)
 73
 74    Visualise(A,B,C, x, points=p)
 75    ```
 76    '''
 77
 78    axs =[]
 79    results = []
 80    lines = []
 81    
 82    if arangement is None and arrangement is not None:
 83        arangement = arrangement
 84
 85
 86    if colour_function_args is None and colour_functions is not None:
 87        colour_function_args = [{}]*len(colour_functions)
 88    
 89    if type(activation) is not list :
 90        activation = [activation] *len(colour_functions)
 91    
 92    if type(A) != list:
 93        A = [A] * len(colour_functions)
 94    if type(B) != list:
 95        B = [B] * len(colour_functions)
 96    if type(C) != list:
 97        C = [C] * len(colour_functions)
 98    
 99    if colour_functions is not None:
100        for i,colour_function in enumerate(colour_functions):
101            if type(colour_function) is not str:
102                if depth > 0:
103                    result = Visualise_single_blocks(A[i],B[i],C[i],activation[i],colour_function, colour_function_args[i], res, depth=depth)
104                else:
105                    result = Visualise_single(A[i],B[i],C[i],activation[i],colour_function, colour_function_args[i], res)
106                results.append(result)
107            
108                if add_lines_functions is not None:
109                    if add_lines_functions[i] is not None:
110                        lines.append(add_lines_functions[i](**add_line_args[i]))
111                    else:
112                        lines.append(None)
113            else:  
114                i1, i2 = colour_function_args[i]['ids'] if (i<len(colour_function_args) and 'ids' in colour_function_args[i]) else (0,1)
115                img1 = results[i1]
116                img2 = results[i2]
117
118                if colour_function.lower() == 'difference' or colour_function.lower() == 'diff' or colour_function.lower() == '-':
119                    image = (img2 - img1)
120                if colour_function.lower() == 'ratio' or colour_function.lower() == 'rat' or colour_function.lower() == '/':
121                    image = (img2 / img1)
122
123                post_function= colour_function_args[i]['post_function'] if (i<len(colour_function_args) and 'post_function' in colour_function_args[i]) else None
124                if post_function is not None:
125                    image = post_function(image)
126
127                results.append(image)
128
129    
130    else:
131        for i,mat in enumerate(matricies):
132            result = mat
133            results.append(result)
134        
135            if add_lines_functions is not None:
136                if add_lines_functions[i] is not None:
137                    lines.append(add_lines_functions[i](**add_line_args[i]))
138                else:
139                    lines.append(None)
140
141    v_min = vmin
142    v_max = vmax
143    
144    if type(vmax) is list:
145        v_max = vmax[0]
146    
147    if type(vmin) is list:
148        v_min = vmin[0]
149    
150    norms = {}
151    
152    if link_ax == 'all' or link_ax == True:
153        norm = mcolors.Normalize(vmin=v_min, vmax=v_max)
154        for i in range(len(results)):
155            norms[i] = norm
156    elif link_ax == 'none' or link_ax is None or link_ax == False:
157        for i in range(len(results)):
158            mn = vmin[i] if (vmin is not None and i < len(vmin)) else None
159            mx = vmax[i] if (vmax is not None and i < len(vmax)) else None
160            norms[i] = mcolors.Normalize(vmin=mn, vmax=mx)
161    
162    else:
163        if type(link_ax[0]) == list or type(link_ax[0]) == tuple:
164            for group in link_ax:
165                norm = mcolors.Normalize(vmin=v_min, vmax=v_max)
166                for i in range(len(results)):
167                    if i in group: 
168                        norms[i] = norm 
169                    elif i not in norms: 
170                        norms[i] = None
171
172        else:
173            norm = mcolors.Normalize(vmin=v_min, vmax=v_max)
174            group = link_ax
175            for i in range(len(results)):
176                if i in group: 
177                    norms[i] = norm
178                elif i not in norms: 
179                    norms[i] = None
180
181
182    for i in range(len(results)):
183        if len(cmaps) > 0:
184            cmap = cmaps[i]
185        else:
186            cmap = 'hot'
187
188        length = len(colour_functions) if colour_functions is not None else len(matricies)
189        if arangement is None:
190            arangement = (1,length)
191        if i > 0:
192            ax = plt.subplot(*arangement,i+1, sharex = ax, sharey=ax)
193        else:
194            ax = plt.subplot(*arangement,i+1)
195        
196        if titles is not None:
197            t = titles[i]
198            if t is not None:
199                ax.set_title(t)
200
201        axs.append(ax)
202        im = results[i]
203        if call_abs: im = torch.abs(im)
204        
205
206        # if v_min is None:
207        #     v_min = torch.min(im)
208        # if v_max is None:
209        #     v_max = torch.max(im)
210        
211
212        # print(vmax,vmin)
213        
214        if clr_labels is None:
215            clr_label = 'Pressure (Pa)'
216        else:
217            clr_label = clr_labels[i]
218        
219        
220        if norm_axes is not None and i in norm_axes:
221            im = im / torch.max(im)
222        img = plt.imshow(im.cpu().detach().numpy(),cmap=cmap,norm=norms[i])
223        plt.yticks([])
224        plt.xticks([])
225
226        if len(points) >0:
227            
228            if type(points) == list:
229                points  = torch.concatenate(points,axis=0)
230            pts_pos = get_point_pos(A[i],B[i],C[i],points.real,res)
231            # print(pts_pos)
232            pts_pos_t = torch.stack(pts_pos).T
233
234            ax.scatter(pts_pos_t[1],pts_pos_t[0],marker="x")
235
236        
237
238        if im.shape[2] == 1:
239            divider = make_axes_locatable(ax)
240            cax = divider.append_axes("right", size="5%", pad=0.05)
241            
242            plt.colorbar(label=clr_label,cax=cax)
243        else:
244            divider = make_axes_locatable(ax)
245            cax = divider.append_axes("right", size="0%", pad=0.05)
246            cax.set_xticks([])
247            cax.set_yticks([])
248
249
250        if add_lines_functions is not None:
251            AB = torch.tensor([B[0] - A[0], B[1] - A[1], B[2] - A[2]])
252            AC = torch.tensor([C[0] - A[0], C[1] - A[1], C[2] - A[2]])
253            # print(AB,AC)
254            norm_x = AB
255            norm_y = AC
256            AB = AB[AB!=0] / res[0]
257            AC = AC[AC!=0] / res[1]
258            # AC = AC / torch.abs(AC)
259            # print(AB,AC)
260            if lines[i] is not None:
261                for con in lines[i]:
262                    xs = [con[0][0]/AB + res[0]/2, con[1][0]/AB + res[0]/2] #Convert real coordinates to pixels - number of steps in each direction
263                    ys = [con[0][1]/AC + res[1]/2, con[1][1]/AC + res[1]/2] #Add res/2 as 0,0,0 in middle of real coordinates not corner of image
264                    # print(xs,ys)
265                    plt.plot(xs,ys,color = "blue")
266        
267       
268    
269    fig = plt.gcf()
270     
271    c = MultiCursor(fig.canvas, axs, color='b',lw=0.5, horizOn=True, vertOn=True)
272
273    multi_event_id = fig.canvas.mpl_connect('motion_notify_event', c.onmove)
274        
275    def press(event):
276        nonlocal multi_event_id
277        if event.key == 'z':
278            if c.visible:
279                fig.canvas.mpl_disconnect(multi_event_id)
280                fig.canvas.draw_idle()
281                c.visible = False
282                c.active = True
283            else: 
284                multi_event_id = fig.canvas.mpl_connect('motion_notify_event', c.onmove)
285                fig.canvas.draw_idle()
286                for line in c.vlines + c.hlines:  
287                    line.set_visible(True)
288                c.visible = True
289        if event.key == 'x' and c.visible:
290            c.active = not c.active
291            
292    fig.canvas.mpl_connect('key_press_event', press)
293
294    if not cursor:
295        fig.canvas.mpl_disconnect(multi_event_id)
296        fig.canvas.draw_idle()
297        c.visible = False
298        c.active = True
299
300
301    
302    if show:
303        plt.show(block=block)
304    else:
305        return fig
306    
307
308def get_point_pos(A:Tensor,B:Tensor,C:Tensor, points:Tensor, res:tuple[int]=(200,200),flip:bool=True) -> list[int]:
309    '''
310    converts point positions in 3D to pixel locations in the plane defined by ABC\n
311    :param A: Position of the top left corner of the image
312    :param B: Position of the top right corner of the image
313    :param C: Position of the bottom left corner of the image
314    :param res: Number of pixels as a tuple (X,Y). Default (200,200)
315    :param flip: Reverses X and Y directions. Default True
316    :return: List of point positions
317    '''
318    AB = torch.tensor([B[0] - A[0], B[1] - A[1], B[2] - A[2]])
319    AC = torch.tensor([C[0] - A[0], C[1] - A[1], C[2] - A[2]])
320
321    ab_dir = AB!=0
322    ac_dir = AC!=0
323
324    step_x = AB / res[0]
325    step_y = AC / res[1]
326
327    if points.shape[2] > 1:
328        points = torch.split(points.squeeze().T,1)
329        points = [pt.squeeze() for pt in points]
330    # print(points)
331
332    pts_norm = []
333
334    for pt in points:
335        Apt =  torch.tensor([pt[0] - A[0], pt[1] - A[1], pt[2] - A[2]])
336        px = Apt / step_x
337        py = Apt / step_y
338        pt_pos = torch.zeros((2))
339        if not flip:
340            pt_pos[0]= torch.round(px[ab_dir])
341            pt_pos[1]=torch.round(py[ac_dir])
342        else:
343            pt_pos[1]= torch.round(px[ab_dir])
344            pt_pos[0]=torch.round(py[ac_dir])
345        
346        pts_norm.append(pt_pos)
347
348   
349
350    return pts_norm
351
352def get_image_positions(A:Tensor,B:Tensor,C:Tensor, res:tuple[int]=(200,200)):
353    '''
354    Gets res[0] x res[1] points in the plane defined by ABC
355    :param A: Position of the top left corner of the image
356    :param B: Position of the top right corner of the image
357    :param C: Position of the bottom left corner of the image
358    :param res: Number of pixels as a tuple (X,Y). Default (200,200)
359    :returnns positions: The positions of pixels
360    '''
361    AB = torch.tensor([B[0] - A[0], B[1] - A[1], B[2] - A[2]]).to(device)
362    AC = torch.tensor([C[0] - A[0], C[1] - A[1], C[2] - A[2]]).to(device)
363
364    step_x = AB / res[0]
365    step_y = AC / res[1]
366
367    positions = torch.zeros((1,3,res[0]*res[1])).to(device)
368
369    for i in range(0,res[0]):
370        for j in range(res[1]):
371            positions[:,:,i*res[1]+j] = A + step_x * i + step_y * j
372    
373    return positions
374
375
376def Visualise_single(A:Tensor,B:Tensor,C:Tensor,activation:Tensor,
377                     colour_function:FunctionType=propagate_abs, colour_function_args:dict={}, 
378                     res:tuple[int]=(200,200), flip:bool=True) -> Tensor:
379    '''
380    Visalises field generated from activation to the plane ABC
381    :param A: Position of the top left corner of the image
382    :param B: Position of the top right corner of the image
383    :param C: Position of the bottom left corner of the image
384    :param activation: The transducer activation to use
385    :param colour_function: Function to call at each position. Should return a numeric value to colour the pixel at that position. Default `acoustools.Utilities.propagate_abs`
386    :param colour_function_args: The arguments to pass to `colour_function`
387    :param res: Number of pixels as a tuple (X,Y). Default (200,200)
388    :param flip: Reverses X and Y directions. Default True
389    :return: Tensor of values of propagated field
390    '''
391    if len(activation.shape) < 3:
392        activation = activation.unsqueeze(0)
393    
394
395    
396    positions = get_image_positions(A,B,C,res)
397   
398    
399    # print(positions.shape)
400    # print(colour_function_args)
401    field_val = colour_function(activations=activation,points=positions,**colour_function_args)
402    if len(field_val.shape) < 3:
403        field_val.unsqueeze_(2)
404    results = []
405    for i in range(field_val.shape[2]):
406        result = torch.reshape(field_val[:,:,i], res)
407        results.append(result)
408    result = torch.stack(results,dim=2)
409
410    if flip:
411        result = torch.rot90(torch.fliplr(result))
412    
413    
414    return result
415
416def force_quiver(points: Tensor, U:Tensor,V:Tensor,norm:tuple[int]=(0,0,0), ylims:int|None=None, xlims:int|None=None,
417                 log:bool=False,show:bool=True,colour:str|None=None, reciprocal:bool = False, block:bool=True, scale:float=1) -> None:
418    '''
419    Plot the force on a mesh as a quiver plot\n
420    :param points: The centre of the mesh faces
421    :param U: Force in first axis
422    :param V: Force in second axis
423    :param norm:
424    :param ylims: limit of y axis
425    :param zlims: limit of x axis
426    :param log: if `True` take the log of the values before plotting
427    :param show: if `True` call `plt.show()`
428    :param colour: colour of arrows
429    :param reciprocal: if `True` plot reciprocal of values
430    :param block: passed into `plt.show`
431    :param scale: value to scale arrows by - note will pass 1/scale to matplotlib
432    '''
433
434    B = points.shape[0]
435    N = points.shape[2]
436    
437    # if len(points) > 0:
438    #     pts_pos = get_point_pos(A,B,C,points,res)
439    
440    mask  = ~(torch.tensor(norm).to(bool))
441    points = points[:,mask,:]
442    # points=torch.reshape(points,(B,2,-1))
443    
444
445    xs = points[:,0,:].cpu().detach().numpy()[0]
446    ys = points[:,1,:].cpu().detach().numpy()[0]
447
448
449    if log:
450        U = torch.sign(U) * torch.abs(torch.log(torch.abs(U)))   
451        V = torch.sign(V) * torch.abs(torch.log(torch.abs(V))) 
452    
453    if reciprocal:
454        U = 1/U
455        V = 1/V
456    
457
458    plt.quiver(xs, ys, U.cpu().detach().numpy(),V.cpu().detach().numpy(),color = colour,linewidths=0.01,scale=1/scale)
459    plt.axis('equal')
460
461
462    if ylims is not None:
463        plt.ylim(ylims[0],ylims[1])
464    
465    if xlims is not None:
466        plt.xlim(xlims[0],xlims[1])
467    
468    if show:
469        plt.show(block=block)
470    
471
472
473def force_quiver_3d(points:Tensor, U:Tensor,V:Tensor,W:Tensor, scale:float=1, show:bool=True, ax:mpl.axes.Axes|None=None) ->None:
474    '''
475    Plot the force on a mesh in 3D
476    :param points: The centre of the mesh faces
477    :param U: Force in first axis
478    :param V: Force in second axis
479    :param W: Force in third axis
480    :param scale: value to scale result by
481    :param show: If True will call `plt.show()`
482    :param ax: Axis to use 
483    '''
484    
485    if ax is None: ax = plt.figure().add_subplot(projection='3d')
486    ax.quiver(points[:,0,:].cpu().detach().numpy(), points[:,1,:].cpu().detach().numpy(), points[:,2,:].cpu().detach().numpy(), U.cpu().detach().numpy()* scale, V.cpu().detach().numpy()* scale, W.cpu().detach().numpy()* scale)
487    
488    if show: plt.show()
489    else: return ax
490
491
492
493
494def Visualise_mesh(mesh:Mesh, colours:Tensor|None=None, points:Tensor|None=None, p_pressure:Tensor|None=None,
495                   vmax:int|None=None,vmin:int|None=None, show:bool=True, subplot:int|plt.Axes|None=None, fig:plt.Figure|None=None, 
496                   buffer_x:int=0, buffer_y:int = 0, buffer_z:int = 0, equalise_axis:bool=False, elev:float=-45, azim:float=45, 
497                   clamp:bool=False) ->None:
498    '''
499    Plot a mesh in 3D and colour the mesh faces
500    :param mesh: Mesh to plot
501    :param colours: Colours for each face
502    :param points: Positions of points to also plot
503    :param p_pressure: Values to colour points with
504    :param vmax: Maximum colour to plot
505    :param vmin: Minimum colour to plot
506    :param show: If `True` call `plot.show()`
507    :param subplot: Optionally use existing subplot
508    :param fig: Optionally use existing fig
509    :param buffer_x: Amount of whitesapce to add in x direction
510    :param buffer_y: Amount of whitesapce to add in y direction
511    :param buffer_z: Amount of whitesapce to add in z direction
512    :param equalise_axis: If `True` call `ax.set_aspect('equal')`
513    :param elev: elevation angle
514    :param azim: azimuth angle
515    :param clamp: if True will clamp values in colours to vmax and vmin
516    '''
517
518    xmin,xmax, ymin,ymax, zmin,zmax = mesh.bounds()
519    
520    if type(colours) is torch.Tensor:
521        colours=colours.flatten()
522    
523    if clamp:
524        colours = torch.clamp(colours,vmin,vmax)
525
526
527    v = mesh.vertices
528    f = torch.tensor(mesh.cells)
529
530    if fig is None:
531        fig = plt.figure()
532    
533    if subplot is None:
534        ax = fig.add_subplot(projection="3d")
535    else:
536        ax = fig.add_subplot(subplot,projection="3d")
537
538    # norm = plt.Normalize(C.min(), C.max())
539    # colors = plt.cm.viridis(norm(C))
540
541    if vmin is None and colours is not None:
542        vmin = torch.min(colours).item()
543        if p_pressure is not None and p_pressure < vmin:
544            vmin = p_pressure
545    
546    if vmax is None and colours is not None:
547        vmax = torch.max(colours).item()
548        if p_pressure is not None and p_pressure > vmax:
549            vmax = p_pressure
550
551    norm = clrs.Normalize(vmin=vmin, vmax=vmax, clip=True)
552    mapper = cm.ScalarMappable(norm, cmap=cm.hot)
553
554    if points is not None:
555        if p_pressure is not None:
556            p_c = mapper.to_rgba(p_pressure.squeeze().cpu().detach())
557        else:
558            p_c = 'blue'
559        points = points.cpu().detach()
560        ax.scatter(points[:,0],points[:,1],points[:,2],color=p_c)
561
562    if colours is not None:
563        colour_mapped = []
564        for c in colours:
565            colour_mapped.append(mapper.to_rgba(c.cpu().detach()))
566    else:
567        colour_mapped=None
568
569    pc = art3d.Poly3DCollection(v[f], edgecolor="black", linewidth=0.01, facecolors=colour_mapped)
570    plt_3d = ax.add_collection(pc)
571
572
573    mappable = cm.ScalarMappable(cmap=cm.hot, norm=norm)
574    mappable.set_array(colour_mapped)  # Link the data to the ScalarMappable
575
576    # Add the color bar to the figure
577    cbar = fig.colorbar(mappable, ax=ax, shrink=0.6)
578    cbar.set_label('Face Value')
579
580
581    scale = mesh.vertices.flatten()
582    ax.auto_scale_xyz(scale, scale, scale)
583    
584    
585    if not equalise_axis:
586        ax.set_xlim([xmin - buffer_x, xmax +  buffer_x])
587        ax.set_ylim([ymin - buffer_y, ymax + buffer_y])
588        ax.set_zlim([zmin - buffer_z, zmax + buffer_z])
589    else:
590        ax.set_aspect('equal')
591
592
593    ax.view_init(elev=elev, azim=azim)
594
595
596
597    if show:
598        plt.show()
599    else:
600        return ax
601    
602
603
604
605def Visualise_line(A:Tensor,B:Tensor,x:Tensor, F:Tensor|None=None,points:Tensor|None=None,steps:int = 1000, 
606                   board:Tensor|None=None, propagate_fun:FunctionType = propagate_abs, propagate_args:dict={}, show:bool=True) -> None:
607    '''
608    Plot the field across a line from A->B\n
609    :param A: Start of line
610    :param B: End of line
611    :param x: Hologram
612    :param F: Optionally, propagation matrix
613    :param points: Optionally, pass the points on line AB instead of computing them
614    :param steps: Number of points along line
615    :param board: Transducers to use
616    :param propagate_fun: Function to use to propagate hologram
617    :propagate_args: arguments for `propagate_fun`
618    :show: If `True` call `plt.show()`
619    '''
620    if board is None:
621        board = TRANSDUCERS
622    
623    if points is None:
624        AB = B-A
625        step = AB / steps
626        points = []
627        for i in range(steps):
628            p = A + i*step
629            points.append(p.unsqueeze(0))
630        
631        points = torch.stack(points, 2).to(device)
632    
633    
634    pressure = propagate_fun(activations=x,points=points, board=board,**propagate_args)
635    if show:
636        plt.plot(pressure.detach().cpu().flatten())
637        plt.show()
638    else:
639        return pressure
640       
641
642def ABC(size:float, plane:Literal['xz', 'yz', 'xy'] = 'xz', origin:Tensor|tuple=None) -> tuple[Tensor]:
643    '''
644    Get ABC values for visualisation
645    * A top right corner
646    * B top right corner
647    * C bottom left corner
648    :param size: The size of the window
649    :param plane: Plane, one of 'xz' 'yz' 'xy'
650    :param origin: The centre of the view window 
651    :return: A,B,C 
652    '''
653    if origin is None:
654        origin = torch.tensor((0,0,0), device=device)
655    if type(origin) == tuple or type(origin) == list:
656        origin = torch.tensor(origin).real
657    
658    origin = origin.squeeze().real
659
660    
661    if plane == 'xz':
662        A = origin+torch.tensor((-1,0, 1), device=device) * size
663        B = origin+torch.tensor((1,0, 1), device=device)* size
664        C = origin+torch.tensor((-1,0, -1), device=device)* size
665    
666    if plane == 'yz':
667        A = origin+torch.tensor((0,-1, 1), device=device) * size
668        B = origin+torch.tensor((0,1, 1), device=device)* size
669        C = origin+torch.tensor((0,-1, -1), device=device)* size
670    
671    if plane == 'xy':
672        A = origin+torch.tensor((-1,1, 0), device=device) * size
673        B = origin+torch.tensor((1, 1,0), device=device)* size
674        C = origin+torch.tensor((-1, -1,0), device=device)* size
675    
676    
677
678    return A.to(device), B.to(device), C.to(device)
679
680def ABC_2_tiles(A:Tensor,B:Tensor,C:Tensor):
681    '''
682    Split ABC defined region into 4 tiles
683    * A top right corner
684    * B top right corner
685    * C bottom left corner
686    
687    '''
688    A1 = A
689    B1 = A + (B-A)/2
690    C1 = A+ (C-A)/2
691
692    A2 = A+ (B-A)/2
693    B2 = B
694    C2 = A+ ((B-A)/2 + (C-A)/2)
695
696    A3 = A+ (C-A)/2
697    B3 = A + ((B-A)/2 + (C-A)/2)
698    C3 = C
699
700    A4 = A + ((B-A)/2 + (C-A)/2)
701    B4 = A + (B-A)+(C-A)/2
702    C4 = A + ((B-A)/2 + (C-A))
703
704    return (A1,B1,C1), (A2,B2,C2), (A3,B3,C3), (A4,B4,C4)
705
706def combine_tiles(t1:Tensor,t2:Tensor,t3:Tensor,t4:Tensor):
707    '''
708    Combines subimages into a larger image, used in `Visualise_single_blocks`
709    :param t1: Top left image
710    :param t2: Top right image
711    :param t3: Bottom left image
712    :param t4: Bottom right image
713    '''
714    top = torch.cat([t1,t2],dim=1)
715    bottom = torch.cat([t3,t4],dim=1)
716
717    return torch.cat([top,bottom],dim=0)
718
719def Visualise_single_blocks(A:Tensor,B:Tensor,C:Tensor,activation:Tensor,
720                     colour_function:FunctionType=propagate_abs, colour_function_args:dict={}, 
721                     res:tuple[int]=(200,200), flip:bool=True, depth=2) -> Tensor:
722    '''
723    Visalises field generated from activation to the plane ABC in a slightly nicer memory efficient way by chunking into tiles
724    :param A: Position of the top left corner of the image
725    :param B: Position of the top right corner of the image
726    :param C: Position of the bottom left corner of the image
727    :param activation: The transducer activation to use
728    :param colour_function: Function to call at each position. Should return a numeric value to colour the pixel at that position. Default `acoustools.Utilities.propagate_abs`
729    :param colour_function_args: The arguments to pass to `colour_function`
730    :param res: Number of pixels as a tuple (X,Y). Default (200,200)
731    :param flip: Reverses X and Y directions. Default True
732    :param depth: Number of times to chunk
733    :return: Tensor of values of propagated field
734    '''
735
736    tiles = ABC_2_tiles(A,B,C)
737
738    new_res = (int(res[0]/2), int(res[1]/2))
739
740    ims = []
741
742
743    for (nA,nB,nC) in tiles:
744        if depth == 1:
745            im = Visualise_single(nA,nB,nC,activation,colour_function=colour_function, colour_function_args=colour_function_args, res=new_res, flip=flip)
746        else:
747            im = Visualise_single_blocks(nA,nB,nC,activation,colour_function=colour_function, colour_function_args=colour_function_args, res=new_res, flip=flip, depth = depth-1)
748        ims.append(im)
749        # torch.cuda.empty_cache()
750
751    im = combine_tiles(*ims)
752    return im
753
754
755def animate_lcode(pth, ax:mpl.axes.Axes|None=None, fig:plt.Figure=None, skip:int=1, show:bool=False, 
756                  fname:str='', extruder:Tensor|None = None, xlims:tuple[float]|None=None, 
757                  ylims:tuple[float]|None=None, zlims:tuple[float]|None=None, dpi:int=100, interval:int = 1, 
758                  legend:bool=True, title:bool=True) -> None:
759    '''
760    Reads a .lcode file and produces a gif of the simulation of the result of that lcode\n
761    :param pth: Path to the .lcode file
762    :param ax: Axis to use, if None will create new
763    :param fig: Figure to use, if None will create new
764    :param skip: Number of instructions to skip per animation frame, default 1 (no skipping)
765    :param show: If true will call plt.show()
766    :param: fname: Name of file to save to
767    :param extruder: If not None the position of the extruder to plot as Tensor
768    :param xlims: Tuple of xlims, if None will use  (-0.12,0.12)
769    :param ylims: Tuple of ylims, if None will use  (-0.12,0.12)
770    :param zlims: Tuple of zlims, if None will use  (-0.12,0.12)
771    :param dpi: dpi to use when saving gif
772    :param inetrval: Time to wait between frames
773    :param legend: If True will add figure legend
774    :param title: If True will add figure title
775    '''
776
777    if fig is None: fig = plt.figure()
778    if ax is None: ax = fig.add_subplot(projection='3d')
779    
780
781    point_commands = ['L0','L1','L2','L3']
782
783    frames = []
784    printed_points = [[],]
785
786    functions = {}
787    in_function = None
788
789    
790
791    with open(pth,'r') as file:
792        lines = file.readlines()
793        LINES = len(lines)
794        for i,line in enumerate(lines):
795            print(f"Line {i}/{LINES}", end='\r')
796            line = line.replace(';','').rstrip()
797            split = line.split(':')
798            cmd = split[0]
799            
800            if cmd == 'function':
801                name = split[1]
802                functions[name] = []
803                in_function = name
804            
805            elif cmd == 'end':
806                name = split[1]
807                in_function = None
808            
809            elif cmd.startswith('F'):
810                frame_points = functions[cmd]
811                for frame in frame_points:
812                    frames.append(frame)
813                    frame_printed_points = printed_points[-1].copy()
814                    printed_points.append(frame_printed_points)
815
816
817
818            elif cmd in point_commands:
819                points = split[1:]
820                ps = []
821                for point in points:
822                    ps.append(point.split(','))
823
824                frame_points = [[float(p) for p in pt] for pt in ps]
825                frames.append(frame_points)
826
827                if in_function is not None:
828                    functions[in_function].append(frame_points)
829            
830            frame_printed_points = printed_points[-1].copy()
831            if cmd == 'C1':
832                frame_printed_points.append(frames[-1].copy())
833            
834            printed_points.append(frame_printed_points)
835
836
837    if extruder is not None:
838        if type(extruder) is bool:
839            extruder = create_points(1,1,0,-0.04, 0.04)
840        ex_x = extruder[:,0].detach().cpu()
841        ex_y = extruder[:,1].detach().cpu()
842        ex_z = extruder[:,2].detach().cpu()
843
844
845    FRAMES = int(len(frames) / skip)
846    # FRAMES = 100
847
848    if xlims is None:
849        xlims = (-0.12,0.12)
850    
851    if ylims is None:
852        ylims = (-0.12,0.12)
853    
854    if zlims is None:
855        zlims = (-0.12,0.12)
856    
857
858    def traverse(index):
859        index = index*skip
860        ax.clear()
861        print(f"Line {index/skip}/{FRAMES}", end='\r')
862        for pt in frames[index]:
863            ax.scatter(*pt, label='Trap')
864        
865        printed_xs = [i for i in [[p[0] for p in pt] for pt in printed_points[index]]]
866        printed_ys = [i for i in [[p[1] for p in pt] for pt in printed_points[index]]]
867        printed_zs = [i for i in [[p[2] for p in pt] for pt in printed_points[index]]]
868
869        ax.scatter(printed_xs,printed_ys, printed_zs, label='Printed', edgecolor='black')
870
871        ax.set_ylim(xlims)
872        ax.set_xlim(ylims)
873        ax.set_zlim(zlims)
874
875        
876        if extruder is not None: ax.scatter(ex_x, ex_y, ex_z, label='Extruder')
877
878        if legend: ax.legend()
879        if title: ax.set_title(f'Location: {index}')
880        
881
882
883    ani = animation.FuncAnimation(fig, traverse, frames=FRAMES, interval=interval)
884    if show: plt.show()
885
886    if fname == '':
887        fname = 'Results.gif'
888    ani.save(fname, writer='imagemagick', dpi = dpi)