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