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