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):
 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    transducers = torch.unsqueeze(transducers,2)
 25    transducers = transducers.expand((B,-1,-1,N))
 26    points = torch.unsqueeze(points,1)
 27    points = points.expand((-1,M,-1,-1))
 28
 29    diff =  points - transducers
 30    distances = torch.sqrt(torch.sum(diff**2, 2))
 31    planar_distance= torch.sqrt(torch.sum((diff**2)[:,:,0:2,:],dim=2))
 32    planar_distance = planar_distance +  1e-8
 33
 34    #Partial derivates of bessel function section wrt xyz
 35    sin_theta = torch.divide(planar_distance,distances) 
 36    partialFpartialU = -1* (k**2 * transducer_radius**2)/4 * sin_theta + (k**4 * transducer_radius**4)/48 * sin_theta**3
 37    partialUpartiala = torch.ones_like(diff)
 38    
 39    diff_z = torch.unsqueeze(diff[:,:,2,:],2)
 40    diff_z = diff_z.expand((-1,-1,2,-1))
 41    
 42    denom = torch.unsqueeze((planar_distance*distances**3),2)
 43    denom = denom.expand((-1,-1,2,-1))
 44    # denom[denom == 0] = 1
 45    
 46    partialUpartiala[:,:,0:2,:] = -1 * (diff[:,:,0:2,:] * diff_z**2) / denom
 47    partialUpartiala[:,:,2,:] = (diff[:,:,2,:] * planar_distance) / distances**3
 48
 49    partialFpartialU = torch.unsqueeze(partialFpartialU,2)
 50    partialFpartialU = partialFpartialU.expand((-1,-1,3,-1))
 51    partialFpartialX  = partialFpartialU * partialUpartiala
 52
 53    #Grad of Pref / d(xt,t)
 54    dist_expand = torch.unsqueeze(distances,2)
 55    dist_expand = dist_expand.expand((-1,-1,3,-1))
 56
 57    partialGpartialX = (p_ref.unsqueeze(3) * diff) / dist_expand**3
 58
 59    #Grad of e^ikd(xt,t)
 60    partialHpartialX = 1j * k * (diff / dist_expand) * torch.exp(1j * k * dist_expand)
 61
 62    #Combine
 63    bessel_arg=k*transducer_radius*torch.divide(planar_distance,distances)
 64    F=1-torch.pow(bessel_arg,2)/8+torch.pow(bessel_arg,4)/192
 65    F = torch.unsqueeze(F,2)
 66    F = F.expand((-1,-1,3,-1))
 67
 68    G = p_ref.unsqueeze(3) / dist_expand
 69    H = torch.exp(1j * k * dist_expand)
 70
 71
 72    return F,G,H, partialFpartialX, partialGpartialX, partialHpartialX, partialFpartialU, partialUpartiala
 73
 74def forward_model_grad(points:Tensor, transducers:Tensor|None = None, p_ref=Constants.P_ref, k=Constants.k, transducer_radius=Constants.radius) -> tuple[Tensor]:
 75    '''
 76    Computes the analytical gradient of the piston model\n
 77    :param points: Point position to compute propagation to 
 78    :param transducers: The Transducer array, default two 16x16 arrays 
 79    :return: derivative of forward model wrt x,y,z position
 80
 81    ```Python
 82    from acoustools.Utilities import forward_model_grad
 83
 84    Fx, Fy, Fz = forward_model_grad(points,transducers=board)
 85    Px  = torch.abs(Fx@activations) #gradient wrt x position
 86    Py  = torch.abs(Fy@activations) #gradient wrt y position
 87    Pz  = torch.abs(Fz@activations) #gradient wrt z position
 88
 89    ```
 90    '''
 91    if transducers is None:
 92        transducers=TRANSDUCERS
 93
 94    F,G,H, partialFpartialX, partialGpartialX, partialHpartialX,_,_ = compute_gradients(points, transducers, p_ref=p_ref,k=k, transducer_radius=transducer_radius)
 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) ->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    B = points.shape[0]
120    N = points.shape[2]
121    M = transducers.shape[0]
122    
123    transducers = torch.unsqueeze(transducers,2)
124    transducers = transducers.expand((B,-1,-1,N))
125    points = torch.unsqueeze(points,1)
126    points = points.expand((-1,M,-1,-1))
127
128    diff = transducers - points
129    
130    diff_square = diff**2
131    distances = torch.sqrt(torch.sum(diff_square, 2))
132    distances_square = distances ** 2
133    distances_cube = distances ** 3
134    distances_five = distances ** 5
135    
136
137    distances_expanded = distances.unsqueeze(2).expand((B,M,3,N))
138    distances_expanded_square = distances_expanded**2
139    distances_expanded_cube = distances_expanded ** 3
140    
141    planar_distance= torch.sqrt(torch.sum(diff_square[:,:,0:2,:],dim=2))
142    planar_distance = planar_distance + 1e-8
143    planar_distance_square = planar_distance**2
144
145    sin_theta = planar_distance / distances
146    sin_theta_expand = sin_theta.unsqueeze(2).expand((B,M,3,N))
147    sin_theta_expand_square = sin_theta_expand**2
148
149    p_ref_expand = p_ref.unsqueeze(2).expand((B,-1,3,N))
150
151    dx = diff[:,:,0,:]
152    dy = diff[:,:,1,:]
153    dz = diff[:,:,2,:]
154
155    # F = G * H 
156    # G  = Pref * e^(ikd) / d
157    # H = 1 - (kr sin(theta))^2 / 8 + (kr sin(theta))^4 / 192
158
159    G = p_ref * torch.exp(1j * k * distances) / distances
160
161    kr = k * transducer_radius
162    kr_sine = kr*sin_theta
163    H = 1 - ((kr_sine)**2) / 8 + ((kr_sine)**4)/192 
164
165
166    #(a = {x,y,z})
167    #Faa = 2*Ga*Ha + Gaa * H + G * Haa
168
169    #Ga = Pref * [i*da * e^{ikd} * (kd+i) / d^2]
170
171    #d = distance
172    #da = -(at - a)^2 / d
173
174    da = -1 * diff / distances_expanded
175    kd = k * distances_expanded
176    phase = torch.exp(1j*kd)
177    Ga = p_ref_expand.expand((B,-1,3,N)) * ( (1j*da*phase * (kd + 1j))/ (distances_expanded_square))
178
179    #Gaa = Pref * [ -1/d^3 * e^{ikd} * (da^2 * (k^2*d^2 + 2ik*d - 2) + d*daa * (1-ikd))]
180    #daa = distance_bs / d^3
181    # distance_bs = sum(b_t - b)^2 . b = {x,y,z} \ a
182    distance_xy = diff[:,:,0,:] **2 + diff[:,:,1,:] **2
183    distance_xz = diff[:,:,0,:] **2 + diff[:,:,2,:] **2
184    distance_yz = diff[:,:,1,:] **2 + diff[:,:,2,:] **2
185
186    distance_bs = torch.stack([distance_yz,distance_xz,distance_xy], dim =2)
187    daa = distance_bs / distances_expanded_cube
188
189    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)))
190
191    #Ha = (kr)^2/48 * s * sa * ((kr)^2 * s^2 - 12)
192    #s = planar_distance / distance = sin_theta
193    #sb = -1 * (db * dz^2) / (sqrt(dx^2+dy^2) * distance^3). b = {x,y}
194    #sz = (dz * sqrt(dx^2 + dy^2)) / distance^3
195
196    sx = -1 * (dx * dz**2) / (planar_distance * distances_cube)
197    sy = -1 * (dy * dz**2) / (planar_distance * distances_cube)
198    sz = (dz * planar_distance) / distances_cube
199    sa = torch.stack([sx,sy,sz],axis=2)
200    # sa[sa.isnan()] = 1
201
202    Ha = 1/48 * kr**2 * sin_theta_expand * sa * (kr**2 * sin_theta_expand**2 - 12)
203
204    #Haa = 1/48 * (kr)^2 * (3*sa^2 * ((kr)^2 * s^2 - 4 ) + s * saa * ((kr)^2 * s^2 - 12))
205
206    #sbb = [ dz^2 [ -1 * (db^2 * distance ^2) + (planar_distance^2 * distance^2) - 3*db^2 * planar_distance^2]] / planar_distance ^3 * distance ^5
207    #szz = (-1 * planar_distance * (planar_distance^2 - 2*dz^2)) / distances^5  
208    sxx = (dz**2 * (-1 * (dx**2 * distances_square) + (planar_distance_square * distances_square) - 3*dx**2 * planar_distance_square)) / (planar_distance**3 * distances_five)
209    syy = (dz**2 * (-1 * (dy**2 * distances_square) + (planar_distance_square * distances_square) - 3*dy**2 * planar_distance_square)) / (planar_distance**3 * distances_five)
210    szz = ((-1 * planar_distance) * (planar_distance**2 - 2*dz**2)) / distances_five
211    saa = torch.stack([sxx,syy,szz],axis=2)
212    # saa[saa.isnan()] = 0
213
214    Haa = 1/48 * kr**2 * (3*sa**2 * (kr**2 * sin_theta_expand_square- 4) + sin_theta_expand*saa * (kr**2*sin_theta_expand_square - 12))
215
216
217    H_expand = H.unsqueeze(2).expand((B,M,3,N))
218    G_expand = G.unsqueeze(2).expand((B,M,3,N))
219    Faa = 2*Ga*Ha + Gaa*H_expand + G_expand*Haa
220    Faa = Faa.to(device).to(DTYPE)
221
222    
223
224    return Faa[:,:,0,:].permute((0,2,1)), Faa[:,:,1,:].permute((0,2,1)), Faa[:,:,2,:].permute((0,2,1))
225
226def forward_model_second_derivative_mixed(points: Tensor, transducers:Tensor|None = None, p_ref = Constants.P_ref,  k=Constants.k, transducer_radius=Constants.radius)->Tensor:
227    '''
228    Computes the second degree mixed analytical gradient of the piston model\n
229    :param points: Point position to compute propagation to 
230    :param transducers: The Transducer array, default two 16x16 arrays 
231    Returns second degree mixed derivatives of forward model wrt x,y,z position - Pxy, Pxz, Pyz
232    '''
233
234    #Bk.2 Pg.317
235
236    if type(p_ref) == float or type(p_ref) == int:
237        M = transducers.shape[0]
238        p_ref = torch.ones(1,M,1, device=device, dtype=DTYPE) * p_ref
239
240    if transducers is None:
241        transducers= TRANSDUCERS
242
243    B = points.shape[0]
244    N = points.shape[2]
245    M = transducers.shape[0]
246    
247    transducers = torch.unsqueeze(transducers,2)
248    transducers = transducers.expand((B,-1,-1,N))
249    points = torch.unsqueeze(points,1)
250    points = points.expand((-1,M,-1,-1))
251
252    diff = transducers - points
253    
254    diff_square = diff**2
255    distances = torch.sqrt(torch.sum(diff_square, 2))
256    distances_cube = distances ** 3
257    distances_five = distances ** 5
258    
259    distances_expanded = distances.unsqueeze(2).expand((B,M,3,N))
260    distances_expanded_square = distances_expanded**2
261    
262    planar_distance= torch.sqrt(torch.sum(diff_square[:,:,0:2,:],dim=2))
263    planar_distance = planar_distance +  1e-8
264    planar_distance_cube = planar_distance**3
265
266    sin_theta = planar_distance / distances
267    sin_theta_expand = sin_theta.unsqueeze(2).expand((B,M,3,N))
268
269    dx = diff[:,:,0,:]
270    dy = diff[:,:,1,:]
271    dz = diff[:,:,2,:]
272
273    p_ref_expand = p_ref.unsqueeze(2).expand((B,-1,3,N))
274
275    # F = G * H 
276    # G  = Pref * e^(ikd) / d
277    # H = 1 - (kr sin(theta))^2 / 8 + (kr sin(theta))^4 / 192
278
279    G = p_ref * torch.exp(1j * k * distances) / distances
280
281    kr = k * transducer_radius
282    kr_sine = kr*sin_theta
283    H = 1 - ((kr_sine)**2) / 8 + ((kr_sine)**4)/192 
284
285    #(a = {x,y,z})
286
287    #Ga = Pref * [i*da * e^{ikd} * (kd+i) / d^2]
288
289    #d = distance
290    #da = -(at - a)^2 / d
291
292    da = -1 * diff / distances_expanded
293    dax = da[:,:,0,:]
294    day = da[:,:,1,:]
295    daz = da[:,:,2,:]
296
297
298    kd_exp = k * distances_expanded
299    kd = k * distances
300    phase = torch.exp(1j*kd_exp)
301    Ga = p_ref_expand * ( (1j*da*phase * (kd_exp + 1j))/ (distances_expanded_square))
302
303    #Ha = (kr)^2/48 * s * sa * ((kr)^2 * s^2 - 12)
304    #s = planar_distance / distance = sin_theta
305    #sb = -1 * (db * dz^2) / (sqrt(dx^2+dy^2) * distance^3). b = {x,y}
306    #sz = (dz * sqrt(dx^2 + dy^2)) / distance^3
307
308    sx = -1 * (dx * dz**2) / (planar_distance * distances_cube)
309    sy = -1 * (dy * dz**2) / (planar_distance * distances_cube)
310    sz = (dz * planar_distance) / distances_cube
311
312    # sx[sx.isnan()] = 1
313    # sy[sy.isnan()] = 1
314    # sz[sx.isnan()] = 1
315    sa = torch.stack([sx,sy,sz],axis=2)
316    
317
318    Ha = 1/48 * kr**2 * sin_theta_expand * sa * (kr**2 * sin_theta_expand**2 - 12)
319
320    #Gab = P_ref * e^{ikd} * (db * da * ( (kd)^2 + 2ikd - 2) + d * dab * (1-ikd)) / (-1*d^3)
321    #dab = -da*db / d^3
322
323    dxy = -1*dx*dy / distances_cube
324    dxz = -1*dx*dz / distances_cube
325    dyz = -1*dy*dz / distances_cube
326
327
328    Gxy = (p_ref * torch.exp(1j * kd) * (day * dax * (kd**2 + 2*1j*kd - 2) + distances * dxy * (1 - 1j*kd))) / (-1 * distances_cube)
329    Gxz = (p_ref * torch.exp(1j * kd) * (daz * dax * (kd**2 + 2*1j*kd - 2) + distances * dxz * (1 - 1j*kd))) / (-1 * distances_cube)
330    Gyz = (p_ref * torch.exp(1j * kd) * (day * daz * (kd**2 + 2*1j*kd - 2) + distances * dyz * (1 - 1j*kd))) / (-1 * distances_cube)
331
332    #Hab = (kr)^2/ 48 * (3*Sb*Sa * ((kr)^2 S^2 - 4) + S*Sab*((kr)^2 S^2 - 12))
333
334    #Sxy = -dx * dy * dz^2 ( 4 * (dx^2 + dy^2) + dz^2 ) / (dx^2 + dy^2)^(3/2) * d^5
335    #Saz = da * dz * (2 * dx**2 + 2 * dy**2 - dz**2) / (dx^2 + dy^2)^(1/2) * d^5
336
337    Sxy = -1 * dx * dy * dz**2 * (4 * (dx**2 + dy**2) + dz**2) / (planar_distance_cube* distances_five)
338    Sxz = dx * dz * (2 * dx**2 + 2 * dy**2 - dz**2) / (planar_distance * distances_five)
339    Syz = dy * dz * (2 * dx**2 + 2 * dy**2 - dz**2) / (planar_distance * distances_five)
340    # Sxy[Sxy.isnan()] = 1
341    # Sxz[Sxz.isnan()] = 1
342    # Syz[Syz.isnan()] = 1
343
344    Hxy = kr**2 / 48 * (3 * sx * sy * (kr**2 * sin_theta**2 - 4) + sin_theta * Sxy * (kr**2 * sin_theta**2  -12))
345    Hxz = kr**2 / 48 * (3 * sx * sz * (kr**2 * sin_theta**2 - 4) + sin_theta * Sxz * (kr**2 * sin_theta**2  -12))
346    Hyz = kr**2 / 48 * (3 * sy * sz * (kr**2 * sin_theta**2 - 4) + sin_theta * Syz * (kr**2 * sin_theta**2  -12))
347
348
349    #Fab = Ga*Hb + Gb*Ha + Gab * H + G * Hab
350
351    Gx = Ga[:,:,0,:]
352    Gy = Ga[:,:,1,:]
353    Gz = Ga[:,:,2,:]
354
355    Hx = Ha[:,:,0,:]
356    Hy = Ha[:,:,1,:]
357    Hz = Ha[:,:,2,:]
358
359
360    Fxy = Gx * Hy + Gy * Hx + Gxy * H + G*Hxy
361    Fxz = Gx * Hz + Gz * Hx + Gxz * H + G*Hxz
362    Fyz = Gy * Hz + Gz * Hy + Gyz * H + G*Hyz
363
364    Fxy = Fxy.to(device).to(DTYPE)
365    Fxz = Fxz.to(device).to(DTYPE)
366    Fyz = Fyz.to(device).to(DTYPE)
367
368    return Fxy.permute((0,2,1)), Fxz.permute((0,2,1)), Fyz.permute((0,2,1))