src.acoustools.Utilities.Piston_model_gradients

  1from acoustools.Utilities.Boards import TRANSDUCERS
  2from acoustools.Utilities.Setup import device, DTYPE
  3import acoustools.Constants as Constants
  4
  5import torch
  6from torch import Tensor
  7
  8
  9def compute_gradients(points, transducers = TRANSDUCERS, p_ref = Constants.P_ref, k=Constants.k, transducer_radius = Constants.radius, transducer_norms=None):
 10    '''
 11    @private
 12    Computes the components to be used in the analytical gradient of the piston model, shouldnt be useed use `forward_model_grad` to get the gradient \\
 13    `points` Point position to compute propagation to \\
 14    `transducers` The Transducer array, default two 16x16 arrays \\
 15    Returns (F,G,H, partialFpartialX, partialGpartialX, partialHpartialX, partialFpartialU, partialUpartiala)
 16    '''
 17    B = points.shape[0]
 18    N = points.shape[2]
 19    M = transducers.shape[0]
 20
 21    if type(p_ref) == float or type(p_ref) == int:
 22        p_ref = torch.ones(1,M,1, device=device, dtype=DTYPE) * p_ref
 23    
 24    if type(k) == float or type(k) == int:
 25        k = torch.ones(1,M,1, device=device, dtype=DTYPE) * k
 26    k_expanded = k.unsqueeze(3)
 27
 28    if transducer_norms is None:
 29        transducer_norms = (torch.zeros_like(transducers) + torch.tensor([0,0,1], device=device)) * torch.sign(transducers[:,2].real).unsqueeze(1).to(DTYPE)
 30    transducer_norms_exp = transducer_norms.unsqueeze(0).unsqueeze(3)
 31    
 32    transducers = torch.unsqueeze(transducers,2)
 33    transducers = transducers.expand((B,-1,-1,N))
 34    points = torch.unsqueeze(points,1)
 35    points = points.expand((-1,M,-1,-1))
 36
 37    diff =  points - transducers
 38    diff = diff + 0j
 39    distances = torch.sqrt(torch.sum(diff**2, 2)) 
 40
 41    sin_theta = torch.norm(torch.cross(diff, transducer_norms_exp, dim=2),2, dim=2) / distances
 42    partialFpartialU = -1* (k**2 * transducer_radius**2)/4 * sin_theta + (k**4 * transducer_radius**4)/48 * sin_theta**3
 43    
 44    
 45    partialUpartiala = torch.sum(diff * transducer_norms_exp, dim=2) / distances #d/dx sin(x) = cos(x)
 46    partialUpartiala = partialUpartiala.unsqueeze(2)
 47
 48    partialFpartialU = torch.unsqueeze(partialFpartialU,2)
 49    partialFpartialU = partialFpartialU.expand((-1,-1,3,-1))
 50    partialFpartialX  = partialFpartialU * partialUpartiala
 51
 52    #Grad of Pref / d(xt,t)
 53    dist_expand = torch.unsqueeze(distances,2)
 54    dist_expand = dist_expand.expand((-1,-1,3,-1))
 55
 56    partialGpartialX = (p_ref.unsqueeze(3) * diff) / dist_expand**3
 57
 58    #Grad of e^ikd(xt,t)
 59    partialHpartialX = 1j * k_expanded * (diff / dist_expand) * torch.exp(1j * k_expanded * dist_expand)
 60
 61    #Combine
 62    bessel_arg=k*transducer_radius*sin_theta
 63    F=1-torch.pow(bessel_arg,2)/8+torch.pow(bessel_arg,4)/192
 64    F = torch.unsqueeze(F,2)
 65    F = F.expand((-1,-1,3,-1))
 66
 67    G = p_ref.unsqueeze(3) / dist_expand
 68    H = torch.exp(1j * k_expanded * dist_expand)
 69
 70
 71    return F,G,H, partialFpartialX, partialGpartialX, partialHpartialX, partialFpartialU, partialUpartiala
 72
 73def forward_model_grad(points:Tensor, transducers:Tensor|None = None, p_ref=Constants.P_ref, k=Constants.k, transducer_radius=Constants.radius, transducer_norms=None) -> tuple[Tensor]:
 74    '''
 75    Computes the analytical gradient of the piston model\n
 76    :param points: Point position to compute propagation to 
 77    :param transducers: The Transducer array, default two 16x16 arrays 
 78    :return: derivative of forward model wrt x,y,z position
 79
 80    ```Python
 81    from acoustools.Utilities import forward_model_grad
 82
 83    Fx, Fy, Fz = forward_model_grad(points,transducers=board)
 84    Px  = torch.abs(Fx@activations) #gradient wrt x position
 85    Py  = torch.abs(Fy@activations) #gradient wrt y position
 86    Pz  = torch.abs(Fz@activations) #gradient wrt z position
 87
 88    ```
 89    '''
 90    if transducers is None:
 91        transducers=TRANSDUCERS
 92    
 93
 94    F,G,H, partialFpartialX, partialGpartialX, partialHpartialX,_,_ = compute_gradients(points, transducers, p_ref=p_ref,k=k, transducer_radius=transducer_radius, transducer_norms=transducer_norms)
 95    derivative = G*(H*partialFpartialX + F*partialHpartialX) + F*H*partialGpartialX
 96    derivative = derivative.to(device).to(DTYPE) # minus here to match f.d -> not 100% sure why its needed
 97
 98
 99    return derivative[:,:,0,:].permute((0,2,1)), derivative[:,:,1,:].permute((0,2,1)), derivative[:,:,2,:].permute((0,2,1))
100
101
102def forward_model_second_derivative_unmixed(points:Tensor, transducers:Tensor|None = None, p_ref = Constants.P_ref, k=Constants.k, transducer_radius=Constants.radius, transducer_norms=None) ->Tensor:
103    '''
104    Computes the second degree unmixed analytical gradient of the piston model\n
105    :param points: Point position to compute propagation to 
106    :param transducers: The Transducer array, default two 16x16 arrays 
107    :return: second degree unmixed derivatives of forward model wrt x,y,z position Pxx, Pyy, Pzz
108    '''
109
110    #See Bk.2 Pg.314
111
112    if transducers is None:
113        transducers= TRANSDUCERS
114    
115    if type(p_ref) == float or type(p_ref) == int:
116        M = transducers.shape[0]
117        p_ref = torch.ones(1,M,1, device=device, dtype=DTYPE) * p_ref
118    
119    if type(k) == float or type(k) == int:
120        k = torch.ones(1,M,1, device=device, dtype=DTYPE) * k
121    k_expanded = k.unsqueeze(3)
122
123    if transducer_norms is None:
124        transducer_norms = (torch.zeros_like(transducers) + torch.tensor([0,0,1], device=device)) * torch.sign(transducers[:,2].real).unsqueeze(1).to(DTYPE)
125    transducer_norms_exp = transducer_norms.unsqueeze(0).unsqueeze(3)
126
127    B = points.shape[0]
128    N = points.shape[2]
129    M = transducers.shape[0]
130    
131    transducers = torch.unsqueeze(transducers,2)
132    transducers = transducers.expand((B,-1,-1,N))
133    points = torch.unsqueeze(points,1)
134    points = points.expand((-1,M,-1,-1))
135
136    diff = transducers - points + 0j
137    
138    diff_square = diff**2
139    distances = torch.sqrt(torch.sum(diff_square, 2))
140
141    distances_expanded = distances.unsqueeze(2).expand((B,M,3,N))
142    distances_expanded_square = distances_expanded**2
143    distances_expanded_cube = distances_expanded ** 3
144    
145
146    # sin_theta = planar_distance / distances
147    sin_theta = torch.norm(torch.cross(diff, transducer_norms_exp, dim=2),2, dim=2) / distances
148    sin_theta_expand = sin_theta.unsqueeze(2).expand((B,M,3,N))
149    sin_theta_expand_square = sin_theta_expand**2
150
151    p_ref_expand = p_ref.unsqueeze(2).expand((B,-1,3,N))
152
153    dx = diff[:,:,0,:]
154    dy = diff[:,:,1,:]
155    dz = diff[:,:,2,:]
156
157    # F = G * H 
158    # G  = Pref * e^(ikd) / d
159    # H = 1 - (kr sin(theta))^2 / 8 + (kr sin(theta))^4 / 192
160
161    G = p_ref * torch.exp(1j * k * distances) / distances
162
163    kr = k * transducer_radius
164    kr_expand = k_expanded * transducer_radius
165    kr_sine = kr*sin_theta
166    H = 1 - ((kr_sine)**2) / 8 + ((kr_sine)**4)/192 
167
168
169    #(a = {x,y,z})
170    #Faa = 2*Ga*Ha + Gaa * H + G * Haa
171
172    #Ga = Pref * [i*da * e^{ikd} * (kd+i) / d^2]
173
174    #d = distance
175    #da = -(at - a)^2 / d
176
177    da = -1 * diff / distances_expanded #Gradient of distances
178    kd = k_expanded * distances_expanded
179    phase = torch.exp(1j*kd)
180    Ga = p_ref_expand.expand((B,-1,3,N)) * ( (1j*da*phase * (kd + 1j))/ (distances_expanded_square))
181
182    #Gaa = Pref * [ -1/d^3 * e^{ikd} * (da^2 * (k^2*d^2 + 2ik*d - 2) + d*daa * (1-ikd))]
183    #daa = distance_bs / d^3
184    # distance_bs = sum(b_t - b)^2 . b = {x,y,z} \ a
185    distance_xy = diff[:,:,0,:] **2 + diff[:,:,1,:] **2
186    distance_xz = diff[:,:,0,:] **2 + diff[:,:,2,:] **2
187    distance_yz = diff[:,:,1,:] **2 + diff[:,:,2,:] **2
188
189    distance_bs = torch.stack([distance_yz,distance_xz,distance_xy], dim =2)
190    daa = distance_bs / distances_expanded_cube #Second Gradient of distance function
191
192    Gaa =  p_ref_expand * (-1/distances_expanded_cube * torch.exp(1j*kd) * (da**2 * (kd**2 + 2*1j*kd - 2) + distances_expanded *daa * (1-1j * kd)))
193
194    cos_theta = torch.sum(diff * transducer_norms_exp, dim=2) / distances
195    cos_theta_exp = cos_theta.unsqueeze(2)
196
197    sa = da * cos_theta_exp
198
199    Ha = 1/48 * kr_expand**2 * sin_theta_expand * sa * (kr_expand**2 * sin_theta_expand**2 - 12)
200
201    saa = daa * cos_theta_exp - da**2 * sin_theta_expand
202    # saa[saa.isnan()] = 0
203
204    Haa = 1/48 * kr_expand**2 * (3*sa**2 * (kr_expand**2 * sin_theta_expand_square- 4) + sin_theta_expand*saa * (kr_expand**2*sin_theta_expand_square - 12))
205
206    H_expand = H.unsqueeze(2).expand((B,M,3,N))
207    G_expand = G.unsqueeze(2).expand((B,M,3,N))
208    Faa = 2*Ga*Ha + Gaa*H_expand + G_expand*Haa
209    Faa = Faa.to(device).to(DTYPE)
210
211    
212
213    return Faa[:,:,0,:].permute((0,2,1)), Faa[:,:,1,:].permute((0,2,1)), Faa[:,:,2,:].permute((0,2,1))
214
215def forward_model_second_derivative_mixed(points: Tensor, transducers:Tensor|None = None, p_ref = Constants.P_ref,  k=Constants.k, transducer_radius=Constants.radius, transducer_norms=None)->Tensor:
216    '''
217    Computes the second degree mixed analytical gradient of the piston model\n
218    :param points: Point position to compute propagation to 
219    :param transducers: The Transducer array, default two 16x16 arrays 
220    Returns second degree mixed derivatives of forward model wrt x,y,z position - Pxy, Pxz, Pyz
221    '''
222
223    #Bk.2 Pg.317
224
225 
226
227    if transducers is None:
228        transducers= TRANSDUCERS
229
230    if type(p_ref) == float or type(p_ref) == int:
231        M = transducers.shape[0]
232        p_ref = torch.ones(1,M,1, device=device, dtype=DTYPE) * p_ref
233
234    if transducer_norms is None:
235        transducer_norms = (torch.zeros_like(transducers) + torch.tensor([0,0,1], device=device)) * torch.sign(transducers[:,2].real).unsqueeze(1).to(DTYPE)
236    transducer_norms_exp = transducer_norms.unsqueeze(0).unsqueeze(3)
237
238    B = points.shape[0]
239    N = points.shape[2]
240    M = transducers.shape[0]
241
242    if type(k) == float or type(k) == int:
243        k = torch.ones(1,M,1, device=device, dtype=DTYPE) * k
244    k_expanded = k.unsqueeze(3)
245    
246    transducers = torch.unsqueeze(transducers,2)
247    transducers = transducers.expand((B,-1,-1,N))
248    points = torch.unsqueeze(points,1)
249    points = points.expand((-1,M,-1,-1))
250
251    diff = transducers - points + 0j
252    
253    diff_square = diff**2
254    distances = torch.sqrt(torch.sum(diff_square, 2))
255    distances_cube = distances ** 3
256    
257    distances_expanded = distances.unsqueeze(2).expand((B,M,3,N))
258    distances_expanded_square = distances_expanded**2
259    
260    sin_theta = torch.norm(torch.cross(diff, transducer_norms_exp, dim=2),2, dim=2) / distances
261    sin_theta_expand = sin_theta.unsqueeze(2).expand((B,M,3,N))
262
263    dx = diff[:,:,0,:]
264    dy = diff[:,:,1,:]
265    dz = diff[:,:,2,:]
266
267    p_ref_expand = p_ref.unsqueeze(2).expand((B,-1,3,N))
268
269
270    G = p_ref * torch.exp(1j * k * distances) / distances
271
272    kr = k * transducer_radius
273    kr_expanded = k_expanded * transducer_radius
274    kr_sine = kr*sin_theta
275    H = 1 - ((kr_sine)**2) / 8 + ((kr_sine)**4)/192 
276
277    da = -1 * diff / distances_expanded
278    dax = da[:,:,0,:]
279    day = da[:,:,1,:]
280    daz = da[:,:,2,:]
281
282
283    kd_exp = k_expanded * distances_expanded
284    kd = k * distances
285    phase = torch.exp(1j*kd_exp)
286    Ga = p_ref_expand * ( (1j*da*phase * (kd_exp + 1j))/ (distances_expanded_square))
287
288
289    cos_theta = torch.sum(diff * transducer_norms_exp, dim=2) / distances
290    cos_theta_exp = cos_theta.unsqueeze(2)
291    sa = da * cos_theta_exp
292    sx = sa[:,:,0]
293    sy = sa[:,:,1]
294    sz = sa[:,:,2]
295
296
297    Ha = 1/48 * kr_expanded**2 * sin_theta_expand * sa * (kr_expanded**2 * sin_theta_expand**2 - 12)
298
299    dxy = -1*dx*dy / distances_cube
300    dxz = -1*dx*dz / distances_cube
301    dyz = -1*dy*dz / distances_cube
302
303
304    Gxy = (p_ref * torch.exp(1j * kd) * (day * dax * (kd**2 + 2*1j*kd - 2) + distances * dxy * (1 - 1j*kd))) / (-1 * distances_cube)
305    Gxz = (p_ref * torch.exp(1j * kd) * (daz * dax * (kd**2 + 2*1j*kd - 2) + distances * dxz * (1 - 1j*kd))) / (-1 * distances_cube)
306    Gyz = (p_ref * torch.exp(1j * kd) * (day * daz * (kd**2 + 2*1j*kd - 2) + distances * dyz * (1 - 1j*kd))) / (-1 * distances_cube)
307
308
309    Sxy = dxy * cos_theta - dx*dy*sin_theta
310    Sxz = dxz * cos_theta - dx*dz*sin_theta
311    Syz = dyz * cos_theta - dy*dz*sin_theta
312
313
314    Hxy = kr**2 / 48 * (3 * sx * sy * (kr**2 * sin_theta**2 - 4) + sin_theta * Sxy * (kr**2 * sin_theta**2  -12))
315    Hxz = kr**2 / 48 * (3 * sx * sz * (kr**2 * sin_theta**2 - 4) + sin_theta * Sxz * (kr**2 * sin_theta**2  -12))
316    Hyz = kr**2 / 48 * (3 * sy * sz * (kr**2 * sin_theta**2 - 4) + sin_theta * Syz * (kr**2 * sin_theta**2  -12))
317
318
319    #Fab = Ga*Hb + Gb*Ha + Gab * H + G * Hab
320
321    Gx = Ga[:,:,0,:]
322    Gy = Ga[:,:,1,:]
323    Gz = Ga[:,:,2,:]
324
325    Hx = Ha[:,:,0,:]
326    Hy = Ha[:,:,1,:]
327    Hz = Ha[:,:,2,:]
328
329
330    Fxy = Gx * Hy + Gy * Hx + Gxy * H + G*Hxy
331    Fxz = Gx * Hz + Gz * Hx + Gxz * H + G*Hxz
332    Fyz = Gy * Hz + Gz * Hy + Gyz * H + G*Hyz
333
334    Fxy = Fxy.to(device).to(DTYPE)
335    Fxz = Fxz.to(device).to(DTYPE)
336    Fyz = Fyz.to(device).to(DTYPE)
337
338    return Fxy.permute((0,2,1)), Fxz.permute((0,2,1)), Fyz.permute((0,2,1))