src.acoustools.BEM.Forward_models

  1import torch
  2from torch import Tensor
  3
  4from vedo import Mesh
  5from typing import Literal
  6import hashlib
  7import pickle
  8
  9
 10import acoustools.Constants as Constants
 11
 12from acoustools.Utilities import device, DTYPE, forward_model_batched, TOP_BOARD, create_points
 13from acoustools.Mesh import get_normals_as_points, get_centres_as_points, get_areas, get_centre_of_mass_as_points
 14
 15from acoustools.Utilities import forward_model_grad
 16
 17
 18
 19
 20
 21def compute_green_derivative(y:Tensor,x:Tensor,norms:Tensor,B:int,N:int,M:int, return_components:bool=False, 
 22                             a:Tensor=None, c:Tensor=None, k=Constants.k, smooth_distance:float=0) -> Tensor:
 23    '''
 24    Computes the derivative of greens function \n
 25    :param y: y in greens function - location of the source of sound
 26    :param x: x in greens function - location of the point to be propagated to
 27    :param norms: norms to y 
 28    :param B: Batch dimension
 29    :param N: size of x
 30    :param M: size of y
 31    :param return_components: if true will return the subparts used to compute the derivative \n
 32    :param k: Wavenumber to use
 33    :param smooth_distance: amount to add to distances to avoid explosion over small values
 34    :param a: position to use for a modified-green's function based BEM formulation (see (18) in `Eliminating the fictitious frequency problem in BEM solutions of the external Helmholtz equation` for more details)
 35    :param c: constant to use for a modified-green's function based BEM formulation (see (18) in `Eliminating the fictitious frequency problem in BEM solutions of the external Helmholtz equation` for more details)
 36    :return: returns the partial derivative of greeens fucntion wrt y
 37    '''
 38    norms= norms.real
 39
 40    vecs = y.real-x.real
 41
 42 
 43    distance = torch.sqrt(torch.sum((vecs)**2,dim=3))
 44    
 45    if a is None: #Were computing with a OR we have no a to begin with
 46        if len(vecs.shape) > 4: #Vecs isnt expanded - we must never have had an a
 47            norms = norms.unsqueeze(4).expand(B,N,-1,-1,1)
 48        else: #vecs included a 
 49            norms = norms.expand(B,N,-1,-1)
 50    else:
 51        norms = norms.expand(B,N,-1,-1)
 52
 53    
 54    # norm_norms = torch.norm(norms,2,dim=3) # === 1x
 55    # vec_norms = torch.norm(vecs,2,dim=3) # === distance?
 56    # print(vec_norms == distance)
 57    angles = (torch.sum(norms*vecs,3) / (distance))
 58
 59    # del norms, vecs
 60    torch.cuda.empty_cache()
 61
 62    # distance = distance + smooth_distance
 63    distance = distance.clamp_min(smooth_distance)
 64
 65
 66    A = 1 * greens(y,x,distance=distance,k=k)
 67    ik_d = (1j*k - 1/(distance))
 68    
 69    del distance
 70    # torch.cuda.empty_cache()
 71
 72    partial_greens = A*ik_d*angles
 73    
 74    # if not return_components:
 75    #     del A,B,angles
 76    torch.cuda.empty_cache()
 77
 78    
 79
 80    if a is not None:
 81        n_a = a.shape[2]
 82        # a = a.permute(0,2,1)
 83        a = a.unsqueeze(1).unsqueeze(2)
 84        a = a.expand(B,N,M,3,n_a).clone()
 85        y = y.unsqueeze(4).expand(B,N,M,3,n_a)
 86        g_mod =  torch.sum(c*compute_green_derivative(y, a, norms, B, N, M,k=k),dim=3) #Allow for multiple a's
 87        partial_greens += g_mod
 88    
 89    
 90    partial_greens[partial_greens.isnan()] = 0
 91    if return_components:
 92        return partial_greens, A,ik_d,angles
 93    
 94
 95    return partial_greens 
 96
 97def greens(y:Tensor,x:Tensor, k:float=Constants.k, distance=None):
 98    '''
 99    Computes greens function for a source at y and a point at x\n
100    :param y: source location
101    :param x: point location
102    :param k: wavenumber
103    :param distance: precomputed distances from y->x
104    :returns greens function:
105    '''
106    if distance is None:
107        vecs = y.real-x.real
108        distance = torch.sqrt(torch.sum((vecs)**2,dim=3)) 
109    green = torch.exp(1j*k*distance) / (4*Constants.pi*distance)
110
111    return green
112
113def compute_G(points: Tensor, scatterer: Mesh, k:float=Constants.k, alphas:float|Tensor=1, betas:float|Tensor = 0, a:Tensor=None, c:Tensor=None, smooth_distance:float=0) -> Tensor:
114    '''
115    Computes G in the BEM model\n
116    :param points: The points to propagate to
117    :param scatterer: The mesh used (as a `vedo` `mesh` object)
118    :param k: wavenumber
119    :param alphas: Absorbance of each element, can be Tensor for element-wise attribution or a number for all elements. If Tensor, should have shape [B,M] where M is the mesh size, B is the batch size and will normally be 1
120    :param betas: Ratio of impedances of medium and scattering material for each element, can be Tensor for element-wise attribution or a number for all elements
121    :param smooth_distance: amount to add to distances to avoid explosion over small values
122    :param a: position to use for a modified-green's function based BEM formulation (see (18) in `Eliminating the fictitious frequency problem in BEM solutions of the external Helmholtz equation` for more details)
123    :param c: constant to use for a modified-green's function based BEM formulation (see (18) in `Eliminating the fictitious frequency problem in BEM solutions of the external Helmholtz equation` for more details)
124    :return G: `torch.Tensor` of G
125    '''
126    torch.cuda.empty_cache()
127    areas = torch.Tensor(scatterer.celldata["Area"]).to(device).real
128    B = points.shape[0]
129    N = points.shape[2]
130    M = areas.shape[0]
131    areas = areas.expand((B,N,-1))
132
133    #Compute the partial derivative of Green's Function
134
135    #Firstly compute the distances from mesh points -> control points
136    centres = torch.tensor(scatterer.cell_centers().points).to(device).real #Uses centre points as position of mesh
137    centres = centres.expand((B,N,-1,-1))
138    
139    # print(points.shape)
140    # p = torch.reshape(points,(B,N,3))
141    p = torch.permute(points,(0,2,1)).real
142    p = torch.unsqueeze(p,2).expand((-1,-1,M,-1))
143
144    #Compute cosine of angle between mesh normal and point
145    # scatterer.compute_normals()
146    # norms = torch.tensor(scatterer.cell_normals).to(device)
147    norms = get_normals_as_points(scatterer,permute_to_points=False).real.expand((B,N,-1,-1))
148
149    # centres_p = get_centres_as_points(scatterer)
150    partial_greens = compute_green_derivative(centres,p,norms, B,N,M, a=a, c=c, k=k,smooth_distance=smooth_distance )
151    
152    if ((type(betas) in [int, float]) and betas != 0) or (type(betas) is Tensor and (betas != 0).any()):  #Either β non 0 and type(β) is number or β is Tensor and any elemenets non 0
153        green = greens(centres, p, k=k) * 1j * k * betas
154        partial_greens += green
155    
156    G = areas * partial_greens
157
158
159    if ((type(alphas) in [int, float]) and alphas != 1) or (type(alphas) is Tensor and (alphas != 1).any()):
160        #Does this need to be in A too?
161        if type(alphas) is Tensor:
162            alphas = alphas.unsqueeze(1)
163            alphas = alphas.expand(B, N, M)
164        vecs = p - centres
165        angle = torch.sum(vecs * norms, dim=3) #Technically not the cosine of the angle - would need to /distance but as we only care about the sign then it doesnt matter
166        angle = angle.real
167        if type(alphas) is Tensor:
168            G[angle>0] = G[angle>0] * alphas[angle>0]
169        else:
170            G[angle>0] = G[angle>0] * alphas
171    
172    return G
173
174
175
176def augment_A_CHIEF(A:Tensor, internal_points:Tensor, CHIEF_mode:Literal['square', 'rect'] = 'square', centres:Tensor=None, norms:Tensor=None, areas:Tensor=None, scatterer:Mesh=None, k:float=Constants.k):
177    '''
178    Augments an A matrix with CHIEF BEM equations\n
179    
180    :param A: A matrix
181    :param internal_points:  The internal points to use for CHIEF based BEM
182    :param CHIEF_mode: Mode for CHIEF -> augment A to be either rectangular (only appended to one axis) or square (appended to both with a 0 block in the corner)
183    :param centres: mesh centres
184    :param norms: mesh normals
185    :param areas: mesh aeras
186    :param scatterer: mesh
187    :param k: wavenumber
188    '''
189     # if internal_points is not None: print(internal_points.shape)
190
191    if internal_points is not None and (internal_points.shape[1] == 3 and internal_points.shape[2] != 3):
192        internal_points = internal_points.permute(0,2,1)
193
194    if areas is None: areas = torch.tensor(scatterer.celldata["Area"], dtype=DTYPE, device=device)
195    if centres is None: centres = torch.tensor(scatterer.cell_centers().points, dtype=DTYPE, device=device)
196    if norms is None: norms = get_normals_as_points(scatterer, permute_to_points=False)
197
198
199
200    P = internal_points.shape[1]
201    M = centres.shape[0]
202
203    m_int = centres.unsqueeze(0).unsqueeze(1)
204    int_p = internal_points.unsqueeze(2)
205    # G_block = greens(m_int,int_p, k=k) 
206    
207    int_norms = norms.unsqueeze(1)
208    G_block = -compute_green_derivative(m_int,int_p, int_norms, 1, P, M, k=k)
209    G_block = G_block * areas[None,None,:] 
210    
211    G_block_t = G_block.mT
212    zero_block = torch.zeros((1, P, P), device=device, dtype=DTYPE)
213
214    
215    # return torch.cat((A, G_block), dim=1)
216    if CHIEF_mode.lower() == 'square':
217        A_aux = torch.cat((A, G_block_t), dim=2)
218        
219        GtZ = torch.cat((G_block, zero_block), dim=2)
220        A = torch.cat((A_aux, GtZ), dim=1)
221    elif CHIEF_mode.lower() == 'rect':
222        A = torch.cat([A, G_block], dim= 1)
223    else:
224        raise RuntimeError(f"Invalid CHIEF Mode {CHIEF_mode}, should be on of [square, rect]")
225
226    return A
227
228
229
230
231def compute_A(scatterer: Mesh, k:float=Constants.k, alphas:float|Tensor = 1, betas:float|Tensor = 0, a:Tensor=None, c:Tensor=None, internal_points:Tensor=None, smooth_distance:float=0, CHIEF_mode:Literal['square', 'rect'] = 'square', h:float=None, BM_alpha:complex=None, BM_mode:str='fd') -> Tensor:
232    '''
233    Computes A for the computation of H in the BEM model\n
234    :param scatterer: The mesh used (as a `vedo` `mesh` object)
235    :param k: wavenumber
236    :param betas: Ratio of impedances of medium and scattering material for each element, can be Tensor for element-wise attribution or a number for all elements
237    :param smooth_distance: amount to add to distances to avoid explosion over small values
238    :param a: position to use for a modified-green's function based BEM formulation (see (18) in `Eliminating the fictitious frequency problem in BEM solutions of the external Helmholtz equation` for more details)
239    :param c: constant to use for a modified-green's function based BEM formulation (see (18) in `Eliminating the fictitious frequency problem in BEM solutions of the external Helmholtz equation` for more details)
240    :param internal_points: The internal points to use for CHIEF based BEM
241    :param CHIEF_mode: Mode for CHIEF -> augment A to be either rectangular (only appended to one axis) or square (appended to both with a 0 block in the corner)
242    :param h: finite difference step for Burton-Miller BEM
243    :param BM_alpha: constant alpha to use in Burton-Miller BEM
244    :param BM_mode: Mode to use for Burton-Miller BEM - should be fd for finite differences
245    
246
247    :return A: A tensor
248    '''
249
250    
251    
252    areas = torch.tensor(scatterer.celldata["Area"], dtype=DTYPE, device=device)
253    centres = torch.tensor(scatterer.cell_centers().points, dtype=DTYPE, device=device)
254    norms = get_normals_as_points(scatterer, permute_to_points=False)
255
256    M = centres.shape[0]
257    # if internal_points is not None:
258        # M = M + internal_points.shape[2]
259
260    if h is not None: #We want to do fin-diff BM
261        if BM_alpha is None: #We are in the grad step
262            centres = centres - h * norms.squeeze(0)
263        else:
264            if BM_mode != 'analytical':
265                Aminus = compute_A(scatterer=scatterer, k=k, betas=betas, a=a, c=c, internal_points=internal_points, smooth_distance=smooth_distance, CHIEF_mode=CHIEF_mode, h=h, BM_alpha=None)
266                Aplus = compute_A(scatterer=scatterer, k=k, betas=betas, a=a, c=c, internal_points=internal_points, smooth_distance=smooth_distance, CHIEF_mode=CHIEF_mode, h=-h, BM_alpha=None)
267
268            else: #This will need to move as h is not required for analytical
269                A_grad = torch.stack(__get_G_partial(centres.unsqueeze(0).permute(0,2,1), scatterer, None, k=k), dim=1)
270                # A_grad = A_grad.permute(0,1,3,2)
271
272                n = norms.permute(0,2,1).unsqueeze(3)
273                A_norm = -1 * torch.sum(A_grad * n , dim=1)
274            
275    m = centres.expand((M, M, 3))
276    m_prime = centres.unsqueeze(1).expand((M, M, 3))
277
278
279    partial_greens = compute_green_derivative(m.unsqueeze_(0), m_prime.unsqueeze_(0), norms, 1, M, M, a=a, c=c, k=k,smooth_distance=smooth_distance)
280    
281
282    if ((type(betas) in [int, float]) and betas != 0) or (isinstance(betas, Tensor) and (betas != 0).any()):
283        green = greens(m, m_prime, k=k) * 1j * k * betas
284        partial_greens += green
285
286    A = -partial_greens * areas
287    A[:, torch.eye(M, dtype=torch.bool, device=device)] = 0.5
288    
289
290    if internal_points is not None: #CHIEF
291
292        A = augment_A_CHIEF(A, internal_points, CHIEF_mode, centres, norms, areas, scatterer,k=k)
293     
294
295    if BM_alpha is not None: #Burton-Miller F.D
296
297        if BM_mode == 'analytical':
298            A_grad = A_norm
299        
300        else:
301            A_grad = (Aplus - Aminus) / (2*h)
302
303
304        A_grad[:, torch.eye(A_grad.shape[2], dtype=torch.bool, device=device)] = 0.5     
305        A = A - BM_alpha * A_grad
306
307        A[:, torch.eye(A_grad.shape[2], dtype=torch.bool, device=device)] = 0.5     
308
309
310    if ((type(alphas) in [int, float]) and alphas != 1) or (type(alphas) is Tensor and (alphas != 1).any()):
311        if type(alphas) is Tensor:
312            alphas = alphas.unsqueeze(1)
313            alphas = alphas.expand(1, M, M)
314        vecs = m_prime - m
315        angle = torch.sum(vecs * norms, dim=3) #Technically not the cosine of the angle - would need to /distance but as we only care about the sign then it doesnt matter
316        angle = angle.real
317        if type(alphas) is Tensor:
318            A[angle>0] = A[angle>0] * alphas[angle>0]
319        else:
320            A[angle>0] = A[angle>0] * alphas
321    
322    return A.to(DTYPE)
323
324 
325def compute_bs(scatterer: Mesh, board:Tensor, p_ref=Constants.P_ref, norms:Tensor|None=None, 
326               a:Tensor=None, c:Tensor=None, k=Constants.k, internal_points:Tensor=None, h:float=None, 
327               BM_alpha:complex=None, BM_mode:str='analytical', transducer_radius = Constants.radius) -> Tensor:
328    '''
329    Computes B for the computation of H in the BEM model\n
330    :param scatterer: The mesh used (as a `vedo` `mesh` object)
331    :param board: Transducers to use 
332    :param p_ref: The value to use for p_ref
333    :param norms: Tensor of normals for transduers
334    :param a: position to use for a modified-green's function based BEM formulation (see (18) in `Eliminating the fictitious frequency problem in BEM solutions of the external Helmholtz equation` for more details)
335    :param c: constant to use for a modified-green's function based BEM formulation (see (18) in `Eliminating the fictitious frequency problem in BEM solutions of the external Helmholtz equation` for more details)
336    :param internal_points: The internal points to use for CHIEF based BEM
337    :param CHIEF_mode: Mode for CHIEF -> augment A to be either rectangular (only appended to one axis) or square (appended to both with a 0 block in the corner)
338    :param h: finite difference step for Burton-Miller BEM
339    :param BM_alpha: constant alpha to use in Burton-Miller BEM
340    :param BM_mode: Mode to use for Burton-Miller BEM - should be analytical
341    :param k: wavenumber
342    :return B: B tensor
343    '''
344
345    if norms is None:
346        norms = (torch.zeros_like(board) + torch.tensor([0,0,1], device=device)) * torch.sign(board[:,2].real).unsqueeze(1).to(DTYPE)
347
348
349    centres = torch.tensor(scatterer.cell_centers().points).to(DTYPE).to(device).T.unsqueeze_(0)
350    if h is not None: #Burton-Miller F.D
351        mesh_norms = get_normals_as_points(scatterer, permute_to_points=True)
352        if BM_alpha is None: #We are in the grad step
353            centres = centres - h * mesh_norms.squeeze(0)
354        elif BM_mode != 'analytical':
355            bs_grad = compute_bs(scatterer=scatterer, board=board, p_ref=p_ref, norms=norms, a=a, c=c, k=k, internal_points=internal_points,h=h, BM_alpha=None, transducer_radius=transducer_radius)
356
357    bs = forward_model_batched(centres,board, p_ref=p_ref,norms=norms,k=k, transducer_radius=transducer_radius) 
358
359
360    if internal_points is not None: #CHIEF
361        F_int = forward_model_batched(internal_points.permute(0,2,1), board, p_ref=p_ref,norms=norms,k=k, transducer_radius=transducer_radius)
362        bs = torch.cat([bs, F_int], dim=1)
363    
364    if a is not None: #Modified Greens function
365        f_mod = torch.sum(forward_model_batched(a,board, p_ref=p_ref,norms=norms, k=k, transducer_radius=transducer_radius), dim=1, keepdim=True)
366        bs += c * f_mod
367    
368    
369    if BM_alpha is not None: #Burton-Miller
370        
371        if BM_mode == 'analytical': 
372            bs_a_grad = torch.stack(forward_model_grad(centres, board, p_ref=p_ref, k=k, transducer_radius=transducer_radius), dim=1)
373            bs_norm_grad = torch.sum(bs_a_grad * mesh_norms.unsqueeze(3), dim=1)
374
375
376            if internal_points is not None: #CHIEF
377                # int_bs_grad = torch.stack(forward_model_grad(internal_points.permute(0,2,1), board, p_ref=p_ref, k=k, transducer_radius=transducer_radius), dim=1)
378                p_n = internal_points.shape[1]
379                M = board.shape[0]
380                # int_bs_grad = torch.zeros_like(int_bs_grad)[:,0,:]
381                int_bs_grad = torch.zeros((1,p_n,M))
382
383                bs_norm_grad = torch.cat([bs_norm_grad,int_bs_grad], dim=1)
384
385
386            
387            bs = bs - BM_alpha * bs_norm_grad
388        else:
389            bs_grad = (bs-bs_grad)/h
390            bs = bs - BM_alpha * (bs-bs_grad)/h
391
392
393    return bs   
394
395 
396def compute_H(scatterer: Mesh, board:Tensor ,use_LU:bool=True, use_OLS:bool = False, p_ref = Constants.P_ref, norms:Tensor|None=None, k:float=Constants.k, alphas:float|Tensor = 1, betas:float|Tensor = 0, 
397              a:Tensor=None, c:Tensor=None, internal_points:Tensor=None, smooth_distance:float=0, 
398              return_components:bool=False, CHIEF_mode:Literal['square', 'rect'] = 'square', h:float=None, BM_alpha:complex=None, transducer_radius = Constants.radius, A:Tensor|None=None, bs:Tensor|None=None) -> Tensor:
399    '''
400    Computes H for the BEM model \n
401    :param scatterer: The mesh used (as a `vedo` `mesh` object)
402    :param board: Transducers to use 
403    :param use_LU: if True computes H with LU decomposition, otherwise solves using standard linear inversion
404    :param use_OLS: if True computes H with OLS, otherwise solves using standard linear inversion
405    :param p_ref: The value to use for p_ref
406    :param norms: Tensor of normals for transduers
407    :param k: wavenumber
408    :param betas: Ratio of impedances of medium and scattering material for each element, can be Tensor for element-wise attribution or a number for all elements
409    :param internal_points: The internal points to use for CHIEF based BEM
410    :param a: position to use for a modified-green's function based BEM formulation (see (18) in `Eliminating the fictitious frequency problem in BEM solutions of the external Helmholtz equation` for more details)
411    :param c: constant to use for a modified-green's function based BEM formulation (see (18) in `Eliminating the fictitious frequency problem in BEM solutions of the external Helmholtz equation` for more details)
412    :param smooth_distance: amount to add to distances to avoid explosion over small values
413    :param CHIEF_mode: Mode for CHIEF -> augment A to be either rectangular (only appended to one axis) or square (appended to both with a 0 block in the corner)
414    :param h: finite difference step for Burton-Miller BEM
415    :param BM_alpha: constant alpha to use in Burton-Miller BEM
416    :return H: H
417    '''
418
419    # internal_points = None
420
421    centres = torch.tensor(scatterer.cell_centers().points, dtype=DTYPE, device=device)
422    M = centres.shape[0]
423
424    if internal_points is not None and (internal_points.shape[1] == 3 and internal_points.shape[2] != 3):
425            internal_points = internal_points.permute(0,2,1)
426
427    
428    # if internal_points is not None: print(internal_points.shape)
429    
430
431
432    if A is None: A = compute_A(scatterer, alphas=alphas, betas=betas, a=a, c=c, k=k,internal_points=internal_points, smooth_distance=smooth_distance, CHIEF_mode=CHIEF_mode, h=h, BM_alpha=BM_alpha)
433    if bs is None: bs = compute_bs(scatterer,board,p_ref=p_ref,norms=norms,a=a,c=c, k=k,internal_points=internal_points, h=h, BM_alpha=BM_alpha, transducer_radius=transducer_radius)
434
435
436    if use_LU:
437        LU, pivots = torch.linalg.lu_factor(A)
438        H = torch.linalg.lu_solve(LU, pivots, bs)
439    elif use_OLS: 
440       
441        H = torch.linalg.lstsq(A,bs, rcond=1e-6).solution    
442    else:
443         H = torch.linalg.solve(A,bs)
444    
445    # H = H / (1-eta*1j)
446
447    # exit()
448    H = H[:,:M,: ]
449    
450    # print(torch.linalg.eig(A))
451
452    # exit()
453
454    # print((k*H, A@H, bs), H/bs)
455    # exit()
456
457    if return_components: return H,A,bs
458    return H
459
460
461
462def get_cache_or_compute_H(scatterer:Mesh,board,use_cache_H:bool=True, path:str="Media", 
463                           print_lines:bool=False, cache_name:str|None=None, p_ref = Constants.P_ref, 
464                           norms:Tensor|None=None, method:Literal['OLS','LU', 'INV']='LU', k:float=Constants.k,alphas:float|Tensor = 1, betas:float|Tensor = 0, 
465                           a:Tensor=None, c:Tensor=None, internal_points:Tensor=None, smooth_distance:float=0, 
466                           CHIEF_mode:Literal['square', 'rect'] = 'square', h:float=None, BM_alpha:complex=None, transducer_radius=Constants.radius) -> Tensor:
467    '''
468    Get H using cache system. Expects a folder named BEMCache in `path`\n
469
470    :param scatterer: The mesh used (as a `vedo` `mesh` object)
471    :param board: Transducers to use 
472    :param use_LU: if True computes H with LU decomposition, otherwise solves using standard linear inversion
473    :param use_OLS: if True computes H with OLS, otherwise solves using standard linear inversion
474    :param p_ref: The value to use for p_ref
475    :param norms: Tensor of normals for transduers
476    :param k: wavenumber
477    :param betas: Ratio of impedances of medium and scattering material for each element, can be Tensor for element-wise attribution or a number for all elements
478    :param internal_points: The internal points to use for CHIEF based BEM
479    :param a: position to use for a modified-green's function based BEM formulation (see (18) in `Eliminating the fictitious frequency problem in BEM solutions of the external Helmholtz equation` for more details)
480    :param c: constant to use for a modified-green's function based BEM formulation (see (18) in `Eliminating the fictitious frequency problem in BEM solutions of the external Helmholtz equation` for more details)
481    :param smooth_distance: amount to add to distances to avoid explosion over small values
482    :param CHIEF_mode: Mode for CHIEF -> augment A to be either rectangular (only appended to one axis) or square (appended to both with a 0 block in the corner)
483    :param h: finite difference step for Burton-Miller BEM
484    :param BM_alpha: constant alpha to use in Burton-Miller BEM
485    :param use_cache_H: If true uses the cache system, otherwise computes H and does not save it
486    :param path: path to folder containing `BEMCache/ `
487    :param print_lines: if true prints messages detaling progress
488    :param method: Method to use to compute H: One of OLS (Least Squares), LU. (LU decomposition). If INV (or anything else) will use `torch.linalg.solve`
489    :return H: H tensor
490    '''
491
492    use_OLS=False
493    use_LU = False
494    
495    if method == "OLS":
496        use_OLS = True
497    elif method == "LU":
498        use_LU = True
499    
500    
501    if use_cache_H:
502        
503        if cache_name is None:
504            ps = locals()
505            ps.__delitem__('scatterer')
506            cache_name = scatterer.filename+"--" + str(ps)
507            cache_name = hashlib.md5(cache_name.encode()).hexdigest()
508        f_name = path+"/BEMCache/"  +  cache_name + ".bin"
509        # print(f_name)
510
511        try:
512            if print_lines: print("Trying to load H at", f_name ,"...")
513            H = pickle.load(open(f_name,"rb")).to(device).to(DTYPE)
514        except FileNotFoundError: 
515            if print_lines: print("Not found, computing H...")
516            H = compute_H(scatterer,board,use_LU=use_LU,use_OLS=use_OLS,norms=norms, k=k, alphas=alphas, betas=betas, a=a, c=c, internal_points=internal_points, smooth_distance=smooth_distance, CHIEF_mode=CHIEF_mode,h=h, BM_alpha=BM_alpha, transducer_radius=transducer_radius)
517            try:
518                f = open(f_name,"wb")
519            except FileNotFoundError as e:
520                print(e)
521                raise FileNotFoundError("AcousTools BEM expects a directory named BEMCache inside of `path' in order to use the cache and this was not found. Check this directory exists")
522            pickle.dump(H,f)
523            f.close()
524    else:
525        if print_lines: print("Computing H...")
526        H = compute_H(scatterer,board, p_ref=p_ref,norms=norms,use_LU=use_LU,use_OLS=use_OLS, k=k, alphas=alphas, betas=betas, a=a, c=c, internal_points=internal_points, smooth_distance=smooth_distance, CHIEF_mode=CHIEF_mode, h=h, BM_alpha=BM_alpha, transducer_radius=transducer_radius)
527
528    return H
529
530def compute_E(scatterer:Mesh, points:Tensor, board:Tensor|None=None, use_cache_H:bool=True, print_lines:bool=False,
531               H:Tensor|None=None,path:str="Media", return_components:bool=False, p_ref = Constants.P_ref, norms:Tensor|None=None, H_method:Literal['OLS','LU', 'INV']=None, 
532               k:float=Constants.k, betas:float|Tensor = 0, alphas:float|Tensor=1, a:Tensor=None, c:Tensor=None, internal_points:Tensor=None, smooth_distance:float=0,  CHIEF_mode:Literal['square', 'rect'] = 'square', h:float=None, BM_alpha:complex=None, transducer_radius=Constants.radius) -> Tensor:
533    '''
534    Computes E in the BEM model\n
535    :param scatterer: The mesh used (as a `vedo` `mesh` object)
536    :param board: Transducers to use 
537    :param use_LU: if True computes H with LU decomposition, otherwise solves using standard linear inversion
538    :param use_OLS: if True computes H with OLS, otherwise solves using standard linear inversion
539    :param p_ref: The value to use for p_ref
540    :param norms: Tensor of normals for transduers
541    :param k: wavenumber
542    :param betas: Ratio of impedances of medium and scattering material for each element, can be Tensor for element-wise attribution or a number for all elements
543    :param internal_points: The internal points to use for CHIEF based BEM
544    :param a: position to use for a modified-green's function based BEM formulation (see (18) in `Eliminating the fictitious frequency problem in BEM solutions of the external Helmholtz equation` for more details)
545    :param c: constant to use for a modified-green's function based BEM formulation (see (18) in `Eliminating the fictitious frequency problem in BEM solutions of the external Helmholtz equation` for more details)
546    :param smooth_distance: amount to add to distances to avoid explosion over small values
547    :param CHIEF_mode: Mode for CHIEF -> augment A to be either rectangular (only appended to one axis) or square (appended to both with a 0 block in the corner)
548    :param h: finite difference step for Burton-Miller BEM
549    :param BM_alpha: constant alpha to use in Burton-Miller BEM
550    :param use_cache_H: If true uses the cache system, otherwise computes H and does not save it
551    :param path: path to folder containing `BEMCache/ `
552    :param print_lines: if true prints messages detaling progress
553    :param method: Method to use to compute H: One of OLS (Least Squares), LU. (LU decomposition). If INV (or anything else) will use `torch.linalg.solve`
554
555    :return E: Propagation matrix for BEM E
556
557    ```Python
558    from acoustools.Mesh import load_scatterer
559    from acoustools.BEM import compute_E, propagate_BEM_pressure, compute_H
560    from acoustools.Utilities import create_points, TOP_BOARD
561    from acoustools.Solvers import wgs
562    from acoustools.Visualiser import Visualise
563
564    import torch
565
566    path = "../../BEMMedia"
567    scatterer = load_scatterer(path+"/Sphere-lam2.stl",dy=-0.06,dz=-0.08)
568    
569    p = create_points(N=1,B=1,y=0,x=0,z=0)
570    
571    H = compute_H(scatterer, TOP_BOARD)
572    E = compute_E(scatterer, p, TOP_BOARD,path=path,H=H)
573    x = wgs(p,board=TOP_BOARD,A=E)
574    
575    A = torch.tensor((-0.12,0, 0.12))
576    B = torch.tensor((0.12,0, 0.12))
577    C = torch.tensor((-0.12,0, -0.12))
578
579    Visualise(A,B,C, x, colour_functions=[propagate_BEM_pressure],
580                colour_function_args=[{"scatterer":scatterer,"board":TOP_BOARD,"path":path,'H':H}],
581                vmax=8621, show=True,res=[256,256])
582    ```
583    
584    '''
585    if board is None:
586        board = TOP_BOARD
587
588    if norms is None: #Transducer Norms
589        norms = (torch.zeros_like(board) + torch.tensor([0,0,1], device=device)) * torch.sign(board[:,2].real).unsqueeze(1).to(DTYPE)
590
591    if print_lines: print("H...")
592    
593    if H is None:
594        H = get_cache_or_compute_H(scatterer,board,use_cache_H, path, print_lines,p_ref=p_ref,
595                                   norms=norms, method=H_method, k=k, betas=betas, a=a, c=c, internal_points=internal_points,
596                                     smooth_distance=smooth_distance, CHIEF_mode=CHIEF_mode, h=h, BM_alpha=BM_alpha, transducer_radius=transducer_radius, alphas=alphas).to(DTYPE)
597        
598    if print_lines: print("G...")
599    G = compute_G(points, scatterer, k=k, betas=betas,alphas=alphas, smooth_distance=smooth_distance).to(DTYPE)
600    
601    if print_lines: print("F...")
602    F = forward_model_batched(points,board,p_ref=p_ref,norms=norms, k=k, transducer_radius=transducer_radius).to(DTYPE)  
603    # if a is not None:
604    #     F += c * forward_model_batched(a,board, p_ref=p_ref,norms=norms)
605    
606    if print_lines: print("E...")
607
608    E = F+G@H
609
610
611    torch.cuda.empty_cache()
612    if return_components:
613        return E.to(DTYPE), F.to(DTYPE), G.to(DTYPE), H.to(DTYPE)
614    return E.to(DTYPE)
615
616
617
618
619
620def __get_G_partial(points:Tensor, scatterer:Mesh, board:Tensor|None=None, return_components:bool=False, k=Constants.k) -> tuple[Tensor, Tensor, Tensor]:
621    '''
622    @private
623    here so it can be used in Burton-Miller BEM above - not ideal for it to be here so this should be refactored at some point
624    Computes gradient of the G matrix in BEM \n
625    :param points: Points to propagate to
626    :param scatterer: The mesh used (as a `vedo` `mesh` object)
627    :param board: Ignored
628    :param return_components: if true will return the subparts used to compute
629    :return: Gradient of the G matrix in BEM
630    '''
631    #Bk3. Pg. 26
632    # if board is None:
633    #     board = TRANSDUCERS
634
635    areas = get_areas(scatterer)
636    centres = get_centres_as_points(scatterer)
637    normals = get_normals_as_points(scatterer)
638
639
640    N = points.shape[2]
641    M = centres.shape[2]
642
643
644    # points = points.unsqueeze(3).expand(-1,-1,-1,M)
645    # centres = centres.unsqueeze(2).expand(-1,-1,N,-1)
646    points  = points.unsqueeze(3)  # [B, 3, N, 1]
647    centres = centres.unsqueeze(2)  # [B, 3, 1, M]
648
649    diff = (points - centres)
650    diff_square = diff**2
651    distances = torch.sqrt(torch.sum(diff_square, 1))
652    distances_expanded = distances.unsqueeze(1)#.expand((1,3,N,M))
653    distances_expanded_square = distances_expanded**2
654    distances_expanded_cube = distances_expanded**3
655
656    # G  =  e^(ikd) / 4pi d
657    G = areas * torch.exp(1j * k * distances_expanded) / (4*3.1415*distances_expanded)
658
659    #Ga =  [i*da * e^{ikd} * (kd+i) / 4pi d^2]
660
661    #d = distance
662    #da = -(at - a)^2 / d
663    da = diff / distances_expanded
664    kd = k * distances_expanded
665    phase = torch.exp(1j*kd)
666    Ga =  areas * ( (1j*da*phase * (kd + 1j))/ (4*3.1415*distances_expanded_square))
667
668    #P = (ik - 1/d)
669    P = (1j*k - 1/distances_expanded)
670    #Pa = da / d^2 = (diff / d^2) /d
671    Pa = diff / distances_expanded_cube
672
673    #C = (diff \cdot normals) / distances
674
675    nx = normals[:,0]
676    ny = normals[:,1]
677    nz = normals[:,2]
678
679    dx = diff[:,0,:]
680    dy = diff[:,1,:]
681    dz = diff[:,2,:]
682
683    n_dot_d = nx*dx + ny*dy + nz*dz
684
685    C = (n_dot_d) / distances
686
687
688    distance_square = distances**2
689
690
691    Cx = 1/distance_square * (nx * distances - (n_dot_d * dx) / distances)
692    Cy = 1/distance_square * (ny * distances - (n_dot_d * dy) / distances)
693    Cz = 1/distance_square * (nz * distances - (n_dot_d * dz) / distances)
694
695    Cx.unsqueeze_(1)
696    Cy.unsqueeze_(1)
697    Cz.unsqueeze_(1)
698
699    Ca = torch.cat([Cx, Cy, Cz],axis=1)
700
701    grad_G = Ga*P*C + G*P*Ca + G*Pa*C
702
703    grad_G =  -grad_G.to(DTYPE)
704
705    if return_components:
706        return grad_G[:,0,:], grad_G[:,1,:], grad_G[:,2,:], G,P,C,Ga,Pa, Ca
707    
708    return grad_G[:,0,:], grad_G[:,1,:], grad_G[:,2,:]