src.acoustools.BEM.Forward_models

  1import torch
  2from torch import Tensor
  3
  4from vedo import Mesh
  5
  6import hashlib
  7import pickle
  8
  9import acoustools.Constants as Constants
 10
 11from acoustools.Utilities import device, DTYPE, forward_model_batched, TOP_BOARD
 12from acoustools.Mesh import get_normals_as_points, board_name, get_centres_as_points
 13
 14def compute_green_derivative(y:Tensor,x:Tensor,norms:Tensor,B:int,N:int,M:int, return_components:bool=False) -> Tensor:
 15    '''
 16    Computes the derivative of greens function \n
 17    :param y: y in greens function - location of the source of sound
 18    :param x: x in greens function - location of the point to be propagated to
 19    :param norms: norms to y 
 20    :param B: Batch dimension
 21    :param N: size of x
 22    :param M: size of y
 23    :param return_components: if true will return the subparts used to compute the derivative \n
 24    :return: returns the partial derivative of greeens fucntion wrt y
 25    '''
 26    norms= norms.real
 27    vecs = y.real-x.real
 28
 29 
 30    distance = torch.sqrt(torch.sum((vecs)**2,dim=3))
 31
 32    norms = norms.expand(B,N,-1,-1)
 33
 34    
 35    # norm_norms = torch.norm(norms,2,dim=3) # === 1
 36    # vec_norms = torch.norm(vecs,2,dim=3) # === distance?
 37    # print(vec_norms == distance)
 38    angles = (torch.sum(norms*vecs,3) / (distance))
 39
 40    del norms, vecs
 41    torch.cuda.empty_cache()
 42
 43    A = ((torch.exp(1j*Constants.k*distance))/(4*torch.pi*distance))
 44    B = (1j*Constants.k - 1/(distance))
 45    
 46    del distance
 47    # torch.cuda.empty_cache()
 48
 49    partial_greens = A*B*angles
 50    
 51    if not return_components:
 52        del A,B,angles
 53    torch.cuda.empty_cache()
 54
 55    
 56    partial_greens[partial_greens.isnan()] = 1
 57
 58
 59    if return_components:
 60        return partial_greens, A,B,angles
 61    
 62    return partial_greens
 63
 64def compute_G(points: Tensor, scatterer: Mesh) -> Tensor:
 65    '''
 66    Computes G in the BEM model\n
 67    :param points: The points to propagate to
 68    :param scatterer: The mesh used (as a `vedo` `mesh` object)
 69    :return G: `torch.Tensor` of G
 70    '''
 71    torch.cuda.empty_cache()
 72    areas = torch.Tensor(scatterer.celldata["Area"]).to(device).real
 73    B = points.shape[0]
 74    N = points.shape[2]
 75    M = areas.shape[0]
 76    areas = areas.expand((B,N,-1))
 77
 78    #Compute the partial derivative of Green's Function
 79
 80    #Firstly compute the distances from mesh points -> control points
 81    centres = torch.tensor(scatterer.cell_centers().points).to(device).real #Uses centre points as position of mesh
 82    centres = centres.expand((B,N,-1,-1))
 83    
 84    # print(points.shape)
 85    # p = torch.reshape(points,(B,N,3))
 86    p = torch.permute(points,(0,2,1)).real
 87    p = torch.unsqueeze(p,2).expand((-1,-1,M,-1))
 88
 89    #Compute cosine of angle between mesh normal and point
 90    # scatterer.compute_normals()
 91    # norms = torch.tensor(scatterer.cell_normals).to(device)
 92    norms = get_normals_as_points(scatterer,permute_to_points=False).real
 93  
 94    partial_greens = compute_green_derivative(centres,p,norms, B,N,M)
 95    
 96    G = areas * partial_greens
 97    return G
 98
 99 
100def compute_A(scatterer: Mesh) -> Tensor:
101    '''
102    Computes A for the computation of H in the BEM model\n
103    :param scatterer: The mesh used (as a `vedo` `mesh` object)
104    :return A: A tensor
105    '''
106
107    areas = torch.Tensor(scatterer.celldata["Area"]).to(device)
108
109    centres = torch.tensor(scatterer.cell_centers().points).to(device)
110    m = centres
111    M = m.shape[0]
112    m = m.expand((M,M,3))
113
114    m_prime = m.clone()
115    m_prime = m_prime.permute((1,0,2))
116
117    # norms = torch.tensor(scatterer.cell_normals).to(device)
118    norms = get_normals_as_points(scatterer,permute_to_points=False)
119
120    green = compute_green_derivative(m.unsqueeze_(0),m_prime.unsqueeze_(0),norms,1,M,M)
121    # areas = areas.unsqueeze(0).T.expand((-1,M)).unsqueeze(0)
122    A = green * areas * -1
123    eye = torch.eye(M).to(bool)
124    A[:,eye] = 0.5
125
126    return A.to(DTYPE)
127
128 
129def compute_bs(scatterer: Mesh, board:Tensor) -> Tensor:
130    '''
131    Computes B for the computation of H in the BEM model\n
132    :param scatterer: The mesh used (as a `vedo` `mesh` object)
133    :param board: Transducers to use 
134    :return B: B tensor
135    '''
136    centres = torch.tensor(scatterer.cell_centers().points).to(device).T.unsqueeze_(0)
137    bs = forward_model_batched(centres,board)
138    return bs.to(DTYPE)
139
140 
141def compute_H(scatterer: Mesh, board:Tensor ,use_LU:bool=True, use_OLS:bool = False) -> Tensor:
142    '''
143    Computes H for the BEM model \n
144    :param scatterer: The mesh used (as a `vedo` `mesh` object)
145    :param board: Transducers to use 
146    :param use_LU: if True computes H with LU decomposition, otherwise solves using standard linear inversion
147    :return H: H
148    '''
149    
150    A = compute_A(scatterer)
151    bs = compute_bs(scatterer,board)
152    if use_LU:
153        LU, pivots = torch.linalg.lu_factor(A)
154        H = torch.linalg.lu_solve(LU, pivots, bs)
155    elif use_OLS:
156       
157        H = torch.linalg.lstsq(A,bs).solution    
158    else:
159         H = torch.linalg.solve(A,bs)
160
161    return H
162
163
164
165def get_cache_or_compute_H(scatterer:Mesh,board,use_cache_H:bool=True, path:str="Media", 
166                           print_lines:bool=False, cache_name:str|None=None,use_LU:bool=True) -> Tensor:
167    '''
168    Get H using cache system. Expects a folder named BEMCache in `path`\n
169    :param scatterer: The mesh used (as a `vedo` `mesh` object)
170    :param  board: Transducers to use 
171    :param use_cache_H_grad: If true uses the cache system, otherwise computes H and does not save it
172    :param path: path to folder containing `BEMCache/ `
173    :param print_lines: if true prints messages detaling progress
174    :param use_LU: If true use LU decomopsition to solve for H
175    :return H: H tensor
176    '''
177    
178    if use_cache_H:
179        
180        if cache_name is None:
181            cache_name = scatterer.filename+"--"+ board_name(board)
182            cache_name = hashlib.md5(cache_name.encode()).hexdigest()
183        f_name = path+"/BEMCache/"  +  cache_name + ".bin"
184        # print(f_name)
185
186        try:
187            if print_lines: print("Trying to load H at", f_name ,"...")
188            H = pickle.load(open(f_name,"rb")).to(device).to(DTYPE)
189        except FileNotFoundError: 
190            if print_lines: print("Not found, computing H...")
191            H = compute_H(scatterer,board,use_LU=use_LU)
192            f = open(f_name,"wb")
193            pickle.dump(H,f)
194            f.close()
195    else:
196        if print_lines: print("Computing H...")
197        H = compute_H(scatterer,board)
198
199    return H
200
201def compute_E(scatterer:Mesh, points:Tensor, board:Tensor|None=None, use_cache_H:bool=True, print_lines:bool=False,
202               H:Tensor|None=None,path:str="Media", return_components:bool=False) -> Tensor:
203    '''
204    Computes E in the BEM model\n
205    :param scatterer: The mesh used (as a `vedo` `mesh` object)
206    :param board: Transducers to use, if `None` then `acoustools.Utilities.TOP_BOARD` is used
207    :param use_cache_H_grad: If true uses the cache system, otherwise computes H and does not save it
208    :param print_lines: if true prints messages detaling progress
209    :param H: Precomputed H - if None H will be compute
210    :param path: path to folder containing `BEMCache/`
211    :param return_components: if true will return the subparts used to compute, F,G,H
212    :return E: Propagation matrix for BEM E
213
214    ```Python
215    from acoustools.Mesh import load_scatterer
216    from acoustools.BEM import compute_E, propagate_BEM_pressure, compute_H
217    from acoustools.Utilities import create_points, TOP_BOARD
218    from acoustools.Solvers import wgs
219    from acoustools.Visualiser import Visualise
220
221    import torch
222
223    path = "../../BEMMedia"
224    scatterer = load_scatterer(path+"/Sphere-lam2.stl",dy=-0.06,dz=-0.08)
225    
226    p = create_points(N=1,B=1,y=0,x=0,z=0)
227    
228    H = compute_H(scatterer, TOP_BOARD)
229    E = compute_E(scatterer, p, TOP_BOARD,path=path,H=H)
230    x = wgs(p,board=TOP_BOARD,A=E)
231    
232    A = torch.tensor((-0.12,0, 0.12))
233    B = torch.tensor((0.12,0, 0.12))
234    C = torch.tensor((-0.12,0, -0.12))
235
236    Visualise(A,B,C, x, colour_functions=[propagate_BEM_pressure],
237                colour_function_args=[{"scatterer":scatterer,"board":TOP_BOARD,"path":path,'H':H}],
238                vmax=8621, show=True,res=[256,256])
239    ```
240    
241    '''
242    if board is None:
243        board = TOP_BOARD
244
245    if print_lines: print("H...")
246    
247    if H is None:
248        H = get_cache_or_compute_H(scatterer,board,use_cache_H, path, print_lines).to(DTYPE)
249        
250    if print_lines: print("G...")
251    G = compute_G(points, scatterer).to(DTYPE)
252    
253    if print_lines: print("F...")
254    F = forward_model_batched(points,board).to(DTYPE)
255    
256    if print_lines: print("E...")
257
258    E = F+G@H
259
260    torch.cuda.empty_cache()
261    if return_components:
262        return E.to(DTYPE), F.to(DTYPE), G.to(DTYPE), H.to(DTYPE)
263    return E.to(DTYPE)