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.to(DTYPE)
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        F_int = forward_model_batched(internal_points, board, p_ref=p_ref,norms=norms,k=k, transducer_radius=transducer_radius)
363        bs = torch.cat([bs, F_int], dim=1)
364    
365    if a is not None: #Modified Greens function
366        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)
367        bs += c * f_mod
368    
369    
370    if BM_alpha is not None: #Burton-Miller
371        
372        if BM_mode == 'analytical': 
373            bs_a_grad = torch.stack(forward_model_grad(centres, board, p_ref=p_ref, k=k, transducer_radius=transducer_radius), dim=1)
374            bs_norm_grad = torch.sum(bs_a_grad * mesh_norms.unsqueeze(3), dim=1)
375
376
377            if internal_points is not None: #CHIEF
378                # 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)
379                p_n = internal_points.shape[1]
380                M = board.shape[0]
381                # int_bs_grad = torch.zeros_like(int_bs_grad)[:,0,:]
382                int_bs_grad = torch.zeros((1,p_n,M))
383
384                bs_norm_grad = torch.cat([bs_norm_grad,int_bs_grad], dim=1)
385
386
387            
388            bs = bs - BM_alpha * bs_norm_grad
389        else:
390            bs_grad = (bs-bs_grad)/h
391            bs = bs - BM_alpha * (bs-bs_grad)/h
392
393
394    return bs   
395
396 
397def 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, 
398              a:Tensor=None, c:Tensor=None, internal_points:Tensor=None, smooth_distance:float=0, 
399              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:
400    '''
401    Computes H for the BEM model \n
402    :param scatterer: The mesh used (as a `vedo` `mesh` object)
403    :param board: Transducers to use 
404    :param use_LU: if True computes H with LU decomposition, otherwise solves using standard linear inversion
405    :param use_OLS: if True computes H with OLS, otherwise solves using standard linear inversion
406    :param p_ref: The value to use for p_ref
407    :param norms: Tensor of normals for transduers
408    :param k: wavenumber
409    :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
410    :param internal_points: The internal points to use for CHIEF based BEM
411    :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)
412    :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)
413    :param smooth_distance: amount to add to distances to avoid explosion over small values
414    :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)
415    :param h: finite difference step for Burton-Miller BEM
416    :param BM_alpha: constant alpha to use in Burton-Miller BEM
417    :return H: H
418    '''
419
420    # internal_points = None
421
422    centres = torch.tensor(scatterer.cell_centers().points, dtype=DTYPE, device=device)
423    M = centres.shape[0]
424
425    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)
426
427
428    if internal_points is not None and (internal_points.shape[1] == 3 and internal_points.shape[2] != 3):
429            internal_points = internal_points.permute(0,2,1)
430
431    
432    # if internal_points is not None: print(internal_points.shape)
433    
434
435
436    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)
437
438
439    # print(A.shape, torch.linalg.matrix_rank(A), A.det(), torch.linalg.cond(A))
440    # print(A)
441    # print(A.inverse())
442
443    if use_LU:
444        LU, pivots = torch.linalg.lu_factor(A)
445        H = torch.linalg.lu_solve(LU, pivots, bs)
446    elif use_OLS: 
447       
448        H = torch.linalg.lstsq(A,bs, rcond=1e-6).solution    
449    else:
450         H = torch.linalg.solve(A,bs)
451    
452    # H = H / (1-eta*1j)
453    # exit()
454    
455    H_clipped = H[:,:M,: ]
456
457    # print(torch.linalg.eig(A))
458
459    # exit()
460
461    # print((k*H, A@H, bs), H/bs)
462    # exit()
463
464    if return_components: return H_clipped,A,bs, H
465
466    
467    return H_clipped
468
469
470
471def get_cache_or_compute_H(scatterer:Mesh,board,use_cache_H:bool=True, path:str="Media", 
472                           print_lines:bool=False, cache_name:str|None=None, p_ref = Constants.P_ref, 
473                           norms:Tensor|None=None, method:Literal['OLS','LU', 'INV']='LU', k:float=Constants.k,alphas:float|Tensor = 1, betas:float|Tensor = 0, 
474                           a:Tensor=None, c:Tensor=None, internal_points:Tensor=None, smooth_distance:float=0, 
475                           CHIEF_mode:Literal['square', 'rect'] = 'square', h:float=None, BM_alpha:complex=None, transducer_radius=Constants.radius) -> Tensor:
476    '''
477    Get H using cache system. Expects a folder named BEMCache in `path`\n
478
479    :param scatterer: The mesh used (as a `vedo` `mesh` object)
480    :param board: Transducers to use 
481    :param use_LU: if True computes H with LU decomposition, otherwise solves using standard linear inversion
482    :param use_OLS: if True computes H with OLS, otherwise solves using standard linear inversion
483    :param p_ref: The value to use for p_ref
484    :param norms: Tensor of normals for transduers
485    :param k: wavenumber
486    :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
487    :param internal_points: The internal points to use for CHIEF based BEM
488    :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)
489    :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)
490    :param smooth_distance: amount to add to distances to avoid explosion over small values
491    :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)
492    :param h: finite difference step for Burton-Miller BEM
493    :param BM_alpha: constant alpha to use in Burton-Miller BEM
494    :param use_cache_H: If true uses the cache system, otherwise computes H and does not save it
495    :param path: path to folder containing `BEMCache/ `
496    :param print_lines: if true prints messages detaling progress
497    :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`
498    :return H: H tensor
499    '''
500
501    use_OLS=False
502    use_LU = False
503    
504    if method == "OLS":
505        use_OLS = True
506    elif method == "LU":
507        use_LU = True
508    
509    
510    if use_cache_H:
511        
512        if cache_name is None:
513            ps = locals()
514            ps.__delitem__('scatterer')
515            cache_name = scatterer.filename+"--" + str(ps)
516            cache_name = hashlib.md5(cache_name.encode()).hexdigest()
517        f_name = path+"/BEMCache/"  +  cache_name + ".bin"
518        # print(f_name)
519
520        try:
521            if print_lines: print("Trying to load H at", f_name ,"...")
522            H = pickle.load(open(f_name,"rb")).to(device).to(DTYPE)
523        except FileNotFoundError: 
524            if print_lines: print("Not found, computing H...")
525            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)
526            try:
527                f = open(f_name,"wb")
528            except FileNotFoundError as e:
529                print(e)
530                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")
531            pickle.dump(H,f)
532            f.close()
533    else:
534        if print_lines: print("Computing H...")
535        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)
536
537    return H
538
539def compute_E(scatterer:Mesh, points:Tensor, board:Tensor|None=None, use_cache_H:bool=True, print_lines:bool=False,
540               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, 
541               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:
542    '''
543    Computes E in the BEM model\n
544    :param scatterer: The mesh used (as a `vedo` `mesh` object)
545    :param board: Transducers to use 
546    :param use_LU: if True computes H with LU decomposition, otherwise solves using standard linear inversion
547    :param use_OLS: if True computes H with OLS, otherwise solves using standard linear inversion
548    :param p_ref: The value to use for p_ref
549    :param norms: Tensor of normals for transduers
550    :param k: wavenumber
551    :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
552    :param internal_points: The internal points to use for CHIEF based BEM
553    :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)
554    :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)
555    :param smooth_distance: amount to add to distances to avoid explosion over small values
556    :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)
557    :param h: finite difference step for Burton-Miller BEM
558    :param BM_alpha: constant alpha to use in Burton-Miller BEM
559    :param use_cache_H: If true uses the cache system, otherwise computes H and does not save it
560    :param path: path to folder containing `BEMCache/ `
561    :param print_lines: if true prints messages detaling progress
562    :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`
563
564    :return E: Propagation matrix for BEM E
565
566    ```Python
567    from acoustools.Mesh import load_scatterer
568    from acoustools.BEM import compute_E, propagate_BEM_pressure, compute_H
569    from acoustools.Utilities import create_points, TOP_BOARD
570    from acoustools.Solvers import wgs
571    from acoustools.Visualiser import Visualise
572
573    import torch
574
575    path = "../../BEMMedia"
576    scatterer = load_scatterer(path+"/Sphere-lam2.stl",dy=-0.06,dz=-0.08)
577    
578    p = create_points(N=1,B=1,y=0,x=0,z=0)
579    
580    H = compute_H(scatterer, TOP_BOARD)
581    E = compute_E(scatterer, p, TOP_BOARD,path=path,H=H)
582    x = wgs(p,board=TOP_BOARD,A=E)
583    
584    A = torch.tensor((-0.12,0, 0.12))
585    B = torch.tensor((0.12,0, 0.12))
586    C = torch.tensor((-0.12,0, -0.12))
587
588    Visualise(A,B,C, x, colour_functions=[propagate_BEM_pressure],
589                colour_function_args=[{"scatterer":scatterer,"board":TOP_BOARD,"path":path,'H':H}],
590                vmax=8621, show=True,res=[256,256])
591    ```
592    
593    '''
594    if board is None:
595        board = TOP_BOARD
596
597    if norms is None: #Transducer Norms
598        norms = (torch.zeros_like(board) + torch.tensor([0,0,1], device=device)) * torch.sign(board[:,2].real).unsqueeze(1).to(DTYPE)
599
600    if print_lines: print("H...")
601    
602    if H is None:
603        H = get_cache_or_compute_H(scatterer,board,use_cache_H, path, print_lines,p_ref=p_ref,
604                                   norms=norms, method=H_method, k=k, betas=betas, a=a, c=c, internal_points=internal_points,
605                                     smooth_distance=smooth_distance, CHIEF_mode=CHIEF_mode, h=h, BM_alpha=BM_alpha, transducer_radius=transducer_radius, alphas=alphas).to(DTYPE)
606        
607    if print_lines: print("G...")
608    G = compute_G(points, scatterer, k=k, betas=betas,alphas=alphas, smooth_distance=smooth_distance).to(DTYPE)
609    
610    if print_lines: print("F...")
611    F = forward_model_batched(points,board,p_ref=p_ref,norms=norms, k=k, transducer_radius=transducer_radius).to(DTYPE)  
612    # if a is not None:
613    #     F += c * forward_model_batched(a,board, p_ref=p_ref,norms=norms)
614    
615    if print_lines: print("E...")
616
617    E = F+G@H
618
619
620    torch.cuda.empty_cache()
621    if return_components:
622        return E.to(DTYPE), F.to(DTYPE), G.to(DTYPE), H.to(DTYPE)
623    return E.to(DTYPE)
624
625
626
627
628
629def __get_G_partial(points:Tensor, scatterer:Mesh, board:Tensor|None=None, return_components:bool=False, k=Constants.k) -> tuple[Tensor, Tensor, Tensor]:
630    '''
631    @private
632    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
633    Computes gradient of the G matrix in BEM \n
634    :param points: Points to propagate to
635    :param scatterer: The mesh used (as a `vedo` `mesh` object)
636    :param board: Ignored
637    :param return_components: if true will return the subparts used to compute
638    :return: Gradient of the G matrix in BEM
639    '''
640    #Bk3. Pg. 26
641    # if board is None:
642    #     board = TRANSDUCERS
643
644    areas = get_areas(scatterer)
645    centres = get_centres_as_points(scatterer)
646    normals = get_normals_as_points(scatterer)
647
648
649    N = points.shape[2]
650    M = centres.shape[2]
651
652
653    # points = points.unsqueeze(3).expand(-1,-1,-1,M)
654    # centres = centres.unsqueeze(2).expand(-1,-1,N,-1)
655    points  = points.unsqueeze(3)  # [B, 3, N, 1]
656    centres = centres.unsqueeze(2)  # [B, 3, 1, M]
657
658    diff = (points - centres)
659    diff_square = diff**2
660    distances = torch.sqrt(torch.sum(diff_square, 1))
661    distances_expanded = distances.unsqueeze(1)#.expand((1,3,N,M))
662    distances_expanded_square = distances_expanded**2
663    distances_expanded_cube = distances_expanded**3
664
665    # G  =  e^(ikd) / 4pi d
666    G = areas * torch.exp(1j * k * distances_expanded) / (4*3.1415*distances_expanded)
667
668    #Ga =  [i*da * e^{ikd} * (kd+i) / 4pi d^2]
669
670    #d = distance
671    #da = -(at - a)^2 / d
672    da = diff / distances_expanded
673    kd = k * distances_expanded
674    phase = torch.exp(1j*kd)
675    Ga =  areas * ( (1j*da*phase * (kd + 1j))/ (4*3.1415*distances_expanded_square))
676
677    #P = (ik - 1/d)
678    P = (1j*k - 1/distances_expanded)
679    #Pa = da / d^2 = (diff / d^2) /d
680    Pa = diff / distances_expanded_cube
681
682    #C = (diff \cdot normals) / distances
683
684    nx = normals[:,0]
685    ny = normals[:,1]
686    nz = normals[:,2]
687
688    dx = diff[:,0,:]
689    dy = diff[:,1,:]
690    dz = diff[:,2,:]
691
692    n_dot_d = nx*dx + ny*dy + nz*dz
693
694    C = (n_dot_d) / distances
695
696
697    distance_square = distances**2
698
699
700    Cx = 1/distance_square * (nx * distances - (n_dot_d * dx) / distances)
701    Cy = 1/distance_square * (ny * distances - (n_dot_d * dy) / distances)
702    Cz = 1/distance_square * (nz * distances - (n_dot_d * dz) / distances)
703
704    Cx.unsqueeze_(1)
705    Cy.unsqueeze_(1)
706    Cz.unsqueeze_(1)
707
708    Ca = torch.cat([Cx, Cy, Cz],axis=1)
709
710    grad_G = Ga*P*C + G*P*Ca + G*Pa*C
711
712    grad_G =  -grad_G.to(DTYPE)
713
714    if return_components:
715        return grad_G[:,0,:], grad_G[:,1,:], grad_G[:,2,:], G,P,C,Ga,Pa, Ca
716    
717    return grad_G[:,0,:], grad_G[:,1,:], grad_G[:,2,:]