src.acoustools.Force

  1from acoustools.Gorkov import gorkov_fin_diff, get_finite_diff_points_all_axis, get_gorkov_constants
  2from acoustools.Utilities import forward_model_batched, forward_model_grad, forward_model_second_derivative_unmixed, forward_model_second_derivative_mixed, TRANSDUCERS, propagate, DTYPE, device
  3import acoustools.Constants as c
  4
  5import torch
  6from torch import Tensor
  7from types import FunctionType
  8
  9
 10def force_fin_diff(activations:Tensor, points:Tensor, axis:str="XYZ", stepsize:float= 0.000135156253,K1:float|None=None, 
 11                   K2:float|None=None,U_function:FunctionType=gorkov_fin_diff,U_fun_args:dict={}, board:Tensor|None=None, V=c.V, p_ref=c.P_ref,
 12                    k=c.k, transducer_radius = c.radius, transducer_norms = None,
 13                    medium_density=c.p_0, medium_speed = c.c_0, particle_density = c.p_p, particle_speed = c.c_p) -> Tensor:
 14    '''
 15    Returns the force on a particle using finite differences to approximate the derivative of the gor'kov potential\n
 16    :param activations: Transducer hologram
 17    :param points: Points to propagate to
 18    :param axis: string containing `X`, `Y` or `Z` defining the axis to take into account eg `XYZ` considers all 3 axes and `YZ` considers only the y and z-axes
 19    :param stepsize: stepsize to use for finite differences 
 20    :param K1: Value for K1 to be used in the gor'kov computation, see `Holographic acoustic elements for manipulation of levitated objects` for more information
 21    :param K2: Value for K1 to be used in the gor'kov computation, see `Holographic acoustic elements for manipulation of levitated objects` for more information
 22    :param U_function: The function used to compute the gor'kov potential
 23    :param U_fun_args: arguments for `U_function` 
 24    :param board: Transducers to use, if `None` uses `acoustools.Utilities.TRANSDUCERS`
 25    :parm V: Particle volume
 26    :return: Force
 27    '''
 28    B = points.shape[0]
 29    D = len(axis)
 30    N = points.shape[2]
 31
 32    if board is None:
 33        board = TRANSDUCERS
 34
 35    fin_diff_points = get_finite_diff_points_all_axis(points, axis, stepsize)
 36
 37    
 38    U_points = U_function(activations, fin_diff_points, axis=axis, stepsize=stepsize/4 ,K1=K1,K2=K2,**U_fun_args, board=board,V=V, 
 39                            p_ref=p_ref, k=k, transducer_radius=transducer_radius, transducer_norms=transducer_norms,
 40                            medium_density=medium_density, medium_speed=medium_speed, particle_density=particle_density,particle_speed=particle_speed)
 41    U_grads = U_points[:,N:]
 42    split = torch.reshape(U_grads,(B,2,-1))
 43    # print(split)
 44    # print((split[:,0,:] - split[:,1,:]))
 45
 46    # print()
 47
 48    F = -1* (split[:,0,:] - split[:,1,:]) / (2*stepsize)
 49    F = F.reshape(B,3,N).permute(0,2,1)
 50    return F
 51
 52def compute_force(activations:Tensor, points:Tensor,board:Tensor|None=None,return_components:bool=False, V=c.V, p_ref=c.P_ref, 
 53                  transducer_radius=c.radius, k=c.k,
 54                 medium_density=c.p_0, medium_speed = c.c_0, particle_density = c.p_p, particle_speed = c.c_p, transducer_norms=None) -> Tensor | tuple[Tensor, Tensor, Tensor]:
 55    '''
 56    Returns the force on a particle using the analytical derivative of the Gor'kov potential and the piston model\n
 57    :param activations: Transducer hologram
 58    :param points: Points to propagate to
 59    :param board: Transducers to use, if `None` uses `acoustools.Utilities.TRANSDUCERS`
 60    :param return_components: If true returns force as one tensor otherwise returns Fx, Fy, Fz
 61    :param V: Particle volume
 62    :return: force  
 63    '''
 64
 65    #Bk.2 Pg.319
 66
 67    if board is None:
 68        board = TRANSDUCERS
 69
 70    if transducer_norms is None:
 71        transducer_norms = (torch.zeros_like(board) + torch.tensor([0,0,1], device=device)) * torch.sign(board[:,2].real).unsqueeze(1).to(DTYPE)
 72    
 73    F = forward_model_batched(points,transducers=board,p_ref=p_ref, transducer_radius=transducer_radius, k=k, norms=transducer_norms)
 74    Fx, Fy, Fz = forward_model_grad(points,transducers=board,p_ref=p_ref, transducer_radius=transducer_radius, k=k, transducer_norms=transducer_norms)
 75    Fxx, Fyy, Fzz = forward_model_second_derivative_unmixed(points,transducers=board,p_ref=p_ref, transducer_radius=transducer_radius, k=k,transducer_norms=transducer_norms)
 76    Fxy, Fxz, Fyz = forward_model_second_derivative_mixed(points,transducers=board,p_ref=p_ref, transducer_radius=transducer_radius, k=k,transducer_norms=transducer_norms)
 77
 78    p   = (F@activations)
 79    Px  = (Fx@activations)
 80    Py  = (Fy@activations)
 81    Pz  = (Fz@activations)
 82    Pxx = (Fxx@activations)
 83    Pyy = (Fyy@activations)
 84    Pzz = (Fzz@activations)
 85    Pxy = (Fxy@activations)
 86    Pxz = (Fxz@activations)
 87    Pyz = (Fyz@activations)
 88
 89
 90    # grad_p = torch.stack([Px,Py,Pz], dim=2).squeeze(3)
 91    # grad_px = torch.stack([Pxx,Pxy,Pxz])
 92    # grad_py = torch.stack([Pxy,Pyy,Pyz])
 93    # grad_pz = torch.stack([Pxz,Pyz,Pzz])
 94
 95    p_term_x= (p*Px.conj() + p.conj()*Px) 
 96    p_term_y= (p*Py.conj() + p.conj()*Py) 
 97    p_term_z= (p*Pz.conj() + p.conj()*Pz) 
 98
 99    # px_term = Px*grad_px.conj() + Px.conj()*grad_px
100    # py_term = Py*grad_py.conj() + Py.conj()*grad_py
101    # pz_term = Pz*grad_pz.conj() + Pz.conj()*grad_pz
102
103    # K1 = V / (4*c.p_0*c.c_0**2)
104    # K2 = 3*V / (4*(2*c.f**2 * c.p_0))
105    K1, K2 = get_gorkov_constants(V=V, c_0=medium_speed, c_p=particle_speed, p_0=medium_density, p_p=particle_density)
106
107
108    # grad_U = K1 * p_term - K2 * (px_term + py_term + pz_term)
109
110    grad_U_x = K1 * p_term_x - K2 * ((Pxx*Px.conj() + Px*Pxx.conj()) + (Pxy*Py.conj() + Py.conj()*Pxy) + (Pxz*Pz.conj() + Pxz.conj()*Pz))
111    grad_U_y = K1 * p_term_y - K2 * ((Pxy*Px.conj() + Px*Pxy.conj()) + (Pyy*Py.conj() + Py.conj()*Pyy) + (Pyz*Pz.conj() + Pyz.conj()*Pz))
112    grad_U_z = K1 * p_term_z - K2 * ((Pxz*Px.conj() + Px*Pxz.conj()) + (Pyz*Py.conj() + Py.conj()*Pyz) + (Pzz*Pz.conj() + Pzz.conj()*Pz))
113
114    grad_U = torch.stack([grad_U_x, grad_U_y, grad_U_z])
115
116    force = -(grad_U).real.squeeze(3).permute(1,2,0)
117    
118    if return_components:
119        return force[:,:,0], force[:,:,1], force[:,:,2] 
120    else:
121        return force 
122
123    
124def get_force_axis(activations:Tensor, points:Tensor,board:Tensor|None=None, axis:int=2, transducer_radius=c.radius, k=c.k, transducer_norms=None,
125                 medium_density=c.p_0, medium_speed = c.c_0, particle_density = c.p_p, particle_speed = c.c_p) -> Tensor:
126    '''
127    Returns the force in one axis on a particle using the analytical derivative of the Gor'kov potential and the piston model \n
128    Equivalent to `compute_force(activations, points,return_components=True)[axis]` \n 
129
130    :param activations: Transducer hologram
131    :param points: Points to propagate to
132    :param board: Transducers to use if `None` uses `acoustools.Utilities.TRANSDUCERS`
133    :param axis: Axis to take the force in
134    :return: force  
135    '''
136    if board is None:
137        board = TRANSDUCERS
138    forces = compute_force(activations, points,return_components=True, board=board, transducer_radius=transducer_radius, k=k, transducer_norms=transducer_norms,
139                           medium_density=medium_density, medium_speed=medium_speed,particle_density=particle_density, particle_speed=particle_speed)
140    force = forces[axis]
141
142    return force
143
144
145def force_mesh(activations:Tensor, points:Tensor, norms:Tensor, areas:Tensor, board:Tensor, grad_function:FunctionType=forward_model_grad, 
146               grad_function_args:dict={}, F_fun:FunctionType|None=forward_model_batched, F_function_args:dict={},
147               F:Tensor|None=None, Ax:Tensor|None=None, Ay:Tensor|None=None,Az:Tensor|None=None,
148               use_momentum:bool=False, return_components:bool=False, p_ref=c.P_ref, transducer_radius=c.radius, 
149               transducer_norms=None,k=c.k, 
150               medium_density = c.p_0, wave_speed = c.c_0, angular_frequency = c.angular_frequency ) -> Tensor:
151    '''
152    Returns the force on a mesh using a discritised version of Eq. 1 in `Acoustical boundary hologram for macroscopic rigid-body levitation`\n
153    :param activations: Transducer hologram
154    :param points: Points to propagate to
155    :param norms: The normals to the mesh faces
156    :param areas: The areas of the mesh points
157    :param board: Transducers to use 
158    :param grad_function: The function to use to compute the gradient of pressure
159    :param grad_function_args: The argument to pass to `grad_function`
160    :param F_fun: Function to compute F
161    :param F_function_args:Fucntion to compute Grad F
162    :param F: A precomputed forward propagation matrix, if `None` will be computed
163    :param Ax: The gradient of `F` wrt x, if `None` will be computed
164    :param Ay: The gradient of `F` wrt y, if `None` will be computed
165    :param Az: The gradient of `F` wrt z, if `None` will be computed
166    :param use_mpmentum: If true will add the term for momentum advection, for sound hard boundaries should be false
167    :param return_components: If True will return force, momentum_flux (force is still the total force)
168    :return: the force on each mesh element
169    '''
170
171    if F is None:
172        F = F_fun(points=points, transducers=board, p_ref=p_ref, transducer_radius=transducer_radius,norms=transducer_norms, k=k ,**F_function_args)
173    p = propagate(activations,points,board,A=F, p_ref=p_ref, transducer_radius=transducer_radius, k=k)
174    pressure_square = torch.abs(p)**2
175    pressure_time_average = 1/2 * pressure_square
176
177    # return pressure_time_average.expand((1,3,-1)), None 
178    if Ax is None or Ay is None or Az is None:
179        Ax, Ay, Az = grad_function(points=points, transducers=board, p_ref=p_ref, k=k, transducer_radius=transducer_radius, transducer_norms=transducer_norms, **grad_function_args)
180    
181    px = (Ax@activations).squeeze(2).unsqueeze(0)
182    py = (Ay@activations).squeeze(2).unsqueeze(0)
183    pz = (Az@activations).squeeze(2).unsqueeze(0)
184
185    grad  = torch.cat((px,py,pz),dim=1) 
186    velocity = grad /( 1j * medium_density * angular_frequency)
187
188    
189    k0 = 1/( medium_density * wave_speed**2)
190    velocity_time_average = 1/2 * torch.sum(velocity * velocity.conj().resolve_conj(), dim=1, keepdim=True).real 
191
192    # + velocity_time_average / velocity_time_average.max()
193    # pressure_square / pressure_square.max()
194
195    force = ( 0.5 * k0 * pressure_time_average - (medium_density / 2) * velocity_time_average) * norms
196
197    if use_momentum:        
198        momentum = medium_density/2 * (torch.sum(velocity * norms, dim=1, keepdim=True) * velocity.conj().resolve_conj()).real + 0j
199        
200        force += momentum 
201    else:
202        momentum = 0
203    
204    force *= -areas # *0.7
205    # force = torch.real(force) #Im(F) == 0 but needs to be complex till now for dtype compatability
206    # print(torch.sgn(torch.sgn(force) * torch.log(torch.abs(force))) == torch.sgn(force))
207
208    if return_components: 
209        return force, momentum
210    
211    return force
212
213def torque_mesh(activations:Tensor, points:Tensor, norms:Tensor, areas:Tensor, centre_of_mass:Tensor, board:Tensor,force:Tensor|None=None, 
214                grad_function:FunctionType=forward_model_grad,grad_function_args:dict={},F:Tensor|None=None, transducer_norms=None,
215                Ax:Tensor|None=None, Ay:Tensor|None=None,Az:Tensor|None=None, transducer_radius=c.radius, k=c.k, 
216               medium_density = c.p_0, wave_speed = c.c_0, angular_frequency = c.angular_frequency, p_ref=c.P_ref) -> Tensor:
217    '''
218    Returns the torque on a mesh using a discritised version of Eq. 1 in `Acoustical boundary hologram for macroscopic rigid-body levitation`\n
219    :param activations: Transducer hologram
220    :param points: Points to propagate to
221    :param norms: The normals to the mesh faces
222    :param areas: The areas of the mesh points
223    :param centre_of_mass: The position of the centre of mass of the mesh
224    :param board: Transducers to use 
225    :param force: Precomputed force on the mesh faces, if `None` will be computed
226    :param grad_function: The function to use to compute the gradient of pressure
227    :param grad_function_args: The argument to pass to `grad_function`
228    :param F: A precomputed forward propagation matrix, if `None` will be computed
229    :param Ax: The gradient of F wrt x, if `None` will be computed
230    :param Ay: The gradient of F wrt y, if `None` will be computed
231    :param Az: The gradient of F wrt z, if `None` will be computed
232    :return: the force on each mesh element
233    '''
234
235    if force is None:
236        force = force_mesh(activations, points, norms, areas, board,grad_function,grad_function_args,F=F, Ax=Ax, Ay=Ay, Az=Az, transducer_norms=transducer_norms,
237                           p_ref=p_ref, k=k, wave_speed=wave_speed, medium_density=medium_density, transducer_radius=transducer_radius, angular_frequency=angular_frequency)
238    force = force.to(DTYPE)
239    
240    displacement = points - centre_of_mass
241    displacement = displacement.to(DTYPE)
242    torque = torch.linalg.cross(displacement,force,dim=1)
243
244    return torch.real(torque)