src.acoustools.Utilities.Propagators

  1import torch
  2from torch import Tensor
  3
  4from acoustools.Utilities.Boards import TRANSDUCERS
  5from acoustools.Utilities.Utilities import is_batched_points
  6from acoustools.Utilities.Forward_models import forward_model_batched, forward_model
  7from acoustools.Utilities.Piston_model_gradients import forward_model_grad
  8from acoustools.Utilities.Setup import device, DTYPE
  9import acoustools.Constants as c
 10
 11from types import FunctionType
 12
 13
 14    
 15def propagate(activations: Tensor, points: Tensor,board: Tensor|None=None, A:Tensor|None=None, p_ref=c.P_ref, norms:Tensor=None,k=c.k, transducer_radius=c.radius) -> Tensor:
 16    '''
 17    Propagates a hologram to target points\n
 18    :param activations: Hologram to use
 19    :param points: Points to propagate to
 20    :param board: The Transducer array, default two 16x16 arrays
 21    :param A: The forward model to use, if None it is computed using `forward_model_batched`. Default:`None`
 22    :return: point activations
 23
 24    ```Python
 25    from acoustools.Solvers import iterative_backpropagation
 26    from acoustools.Utilities import create_points, propagate
 27
 28    p = create_points(2,1)
 29    x = iterative_backpropagation(p)
 30    
 31    p = p.squeeze(0)
 32    x = iterative_backpropagation(p)
 33    print(propagate(x,p))
 34    ```
 35    '''
 36    if board is None:
 37        board  = TRANSDUCERS
 38    if norms is None:
 39        norms = (torch.zeros_like(board) + torch.tensor([0,0,1], device=device)) * torch.sign(board[:,2].real).unsqueeze(1).to(DTYPE)
 40    batch = is_batched_points(points)
 41
 42    if A is None:
 43        if len(points.shape)>2:
 44            A = forward_model_batched(points,board, p_ref=p_ref,norms=norms,k=k, transducer_radius=transducer_radius).to(device)
 45        else:
 46            A = forward_model(points,board, p_ref=p_ref,norms=norms,k=k, transducer_radius=transducer_radius).to(device)
 47    prop = A@activations
 48    if batch:
 49        prop = torch.squeeze(prop, 2)
 50    return prop
 51
 52def propagate_abs(activations: Tensor, points: Tensor,board:Tensor|None=None, A:Tensor|None=None, A_function:FunctionType=None, A_function_args:dict={}, p_ref=c.P_ref, norms:Tensor=None,k=c.k, transducer_radius=c.radius) -> Tensor:
 53    '''
 54    Propagates a hologram to target points and returns pressure - Same as `torch.abs(propagate(activations, points,board, A))`\n
 55    :param activations: Hologram to use
 56    :param points: Points to propagate to
 57    :param board: The Transducer array, default two 16x16 arrays
 58    :param A: The forward model to use, if None it is computed using `forward_model_batched`. Default:`None`
 59    :return: point pressure
 60
 61    ```Python
 62    from acoustools.Solvers import iterative_backpropagation
 63    from acoustools.Utilities import create_points, propagate_abs
 64
 65    p = create_points(2,1)
 66    x = iterative_backpropagation(p)
 67    
 68    p = p.squeeze(0)
 69    x = iterative_backpropagation(p)
 70    print(propagate_abs(x,p))
 71    ```
 72    '''
 73    if board is None:
 74        board = TRANSDUCERS
 75    if A_function is not None:
 76        A = A_function(points, board, norms=norms, **A_function_args)
 77
 78    out = propagate(activations, points,board,A=A,p_ref=p_ref,norms=norms,k=k, transducer_radius=transducer_radius)
 79    return torch.abs(out)
 80
 81def propagate_abs_normalised(activations: Tensor, points: Tensor,board:Tensor|None=None, A:Tensor|None=None, A_function:FunctionType=None, A_function_args:dict={}, p_ref=c.P_ref, norms:Tensor=None,k=c.k, transducer_radius=c.radius) -> Tensor:
 82    '''
 83    Propagates a hologram to target points and returns pressure normalised from [0,1] - Same as `torch.abs(propagate(activations, points,board, A))`\n
 84    :param activations: Hologram to use
 85    :param points: Points to propagate to
 86    :param board: The Transducer array, default two 16x16 arrays
 87    :param A: The forward model to use, if None it is computed using `forward_model_batched`. Default:`None`
 88    :return: point pressure
 89
 90    ```Python
 91    from acoustools.Solvers import iterative_backpropagation
 92    from acoustools.Utilities import create_points, propagate_abs
 93
 94    p = create_points(2,1)
 95    x = iterative_backpropagation(p)
 96    
 97    p = p.squeeze(0)
 98    x = iterative_backpropagation(p)
 99    print(propagate_abs(x,p))
100    ```
101    '''
102    if board is None:
103        board = TRANSDUCERS
104    if A_function is not None:
105        A = A_function(points, board, norms=norms, **A_function_args)
106
107    out = propagate(activations, points,board,A=A,p_ref=p_ref,norms=norms,k=k, transducer_radius=transducer_radius)
108    img = torch.abs(out)
109    return img / torch.max(img)
110
111def propagate_phase(activations:Tensor, points:Tensor,board:Tensor|None=None, A:Tensor|None=None,p_ref=c.P_ref, norms:Tensor=None,k=c.k, transducer_radius=c.radius) -> Tensor:
112    '''
113    Propagates a hologram to target points and returns phase - Same as `torch.angle(propagate(activations, points,board, A))`\n
114    :param activations: Hologram to use
115    :param points: Points to propagate to
116    :param board: The Transducer array, default two 16x16 arrays
117    :param A: The forward model to use, if None it is computed using `forward_model_batched`. Default:`None`
118    :return: point phase
119
120    ```Python
121    from acoustools.Solvers import iterative_backpropagation
122    from acoustools.Utilities import create_points, propagate_phase
123
124    p = create_points(2,1)
125    x = iterative_backpropagation(p)
126    
127    p = p.squeeze(0)
128    x = iterative_backpropagation(p)
129    print(propagate_phase(x,p))
130    ```
131    '''
132    if board is None:
133        board = TRANSDUCERS
134    out = propagate(activations, points,board,A=A,p_ref=p_ref,norms=norms,k=k, transducer_radius=transducer_radius)
135    return torch.angle(out)
136
137
138def propagate_velocity_potential(activations: Tensor, points: Tensor,board: Tensor|None=None, A:Tensor|None=None, 
139                                 density = c.p_0, angular_frequency = c.angular_frequency,p_ref=c.P_ref, norms:Tensor=None,k=c.k) -> Tensor:
140    '''
141    Propagates a hologram to velocity potential at points\n
142    :param activations: Hologram to use
143    :param points: Points to propagate to
144    :param board: The Transducer array, default two 16x16 arrays
145    :param A: The forward model to use, if None it is computed using `forward_model_batched`. Default:`None`
146    :return: point velocity potential'
147    '''
148
149    pressure = propagate(activations, points, board, A=A,p_ref=p_ref,norms=norms,k=k)
150    velocity_potential = pressure / (1j * density * angular_frequency)
151
152    return velocity_potential
153
154def propagate_pressure_grad(activations: Tensor, points: Tensor,board: Tensor|None=None, Fx=None, Fy=None, Fz=None, cat=True,p_ref=c.P_ref, norms:Tensor=None,k=c.k, transducer_radius=c.radius):
155    '''
156    Propagates a hologram to pressure gradient at points\n
157    :param activations: Hologram to use
158    :param points: Points to propagate to
159    :param board: The Transducer array, default two 16x16 arrays
160    :param Fx: The forward model to us for Fx, if None it is computed using `forward_model_grad`. Default:`None`
161    :param Fy: The forward model to us for Fy, if None it is computed using `forward_model_grad`. Default:`None`
162    :param Fz: The forward model to us for Fz, if None it is computed using `forward_model_grad`. Default:`None`
163    :return: point velocity potential'
164    '''
165    
166    if Fx is None or Fy is None or Fz is None:
167        _Fx,_Fy,_Fz = forward_model_grad(points, board,norms=norms,k=k, transducer_radius=transducer_radius)
168        if Fx is None: Fx = _Fx
169        if Fy is None: Fy = _Fy
170        if Fz is None: Fz = _Fz
171    
172    Px = Fx@activations
173    Py = Fy@activations
174    Pz = Fz@activations
175
176    if cat: 
177        grad = torch.cat([Px, Py, Pz], dim=2)
178        return grad
179    return Px, Py, Pz
180
181
182def propagate_velocity(activations: Tensor, points: Tensor,board: Tensor|None=None, Fx=None, Fy=None, Fz=None, 
183                                 density = c.p_0, angular_frequency = c.angular_frequency, cat=True,p_ref=c.P_ref, norms:Tensor=None,k=c.k, transducer_radius=c.radius):
184    '''
185    Propagates a hologram to velocity at points\n
186    :param activations: Hologram to use
187    :param points: Points to propagate to
188    :param board: The Transducer array, default two 16x16 arrays
189    :param Fx: The forward model to us for Fx, if None it is computed using `forward_model_grad`. Default:`None`
190    :param Fy: The forward model to us for Fy, if None it is computed using `forward_model_grad`. Default:`None`
191    :param Fz: The forward model to us for Fz, if None it is computed using `forward_model_grad`. Default:`None`
192    :return: point velocity potential'
193    '''
194    
195    pressure_grads = propagate_pressure_grad(activations, points,board, Fx, Fy, Fz, cat=False,p_ref=p_ref,norms=norms,k=k, transducer_radius=transducer_radius)
196    alpha = 1/(1j * density * angular_frequency)
197    velocity = [alpha * i for i in pressure_grads]
198    if cat: velocity = torch.cat(velocity)
199    return velocity
200
201def propagate_velocity_real(activations: Tensor, points: Tensor,board: Tensor|None=None, Fx=None, Fy=None, Fz=None, 
202                                 density = c.p_0, angular_frequency = c.angular_frequency, cat=True,
203                                 p_ref=c.P_ref, norms:Tensor=None,k=c.k, transducer_radius=c.radius):
204    '''
205    Propagates a hologram to velocity's real component at points\n
206    :param activations: Hologram to use
207    :param points: Points to propagate to
208    :param board: The Transducer array, default two 16x16 arrays
209    :param Fx: The forward model to us for Fx, if None it is computed using `forward_model_grad`. Default:`None`
210    :param Fy: The forward model to us for Fy, if None it is computed using `forward_model_grad`. Default:`None`
211    :param Fz: The forward model to us for Fz, if None it is computed using `forward_model_grad`. Default:`None`
212    :return: point velocity potential'
213    '''
214    vel = [i.real for i in propagate_velocity(activations, points,board, Fx, Fy, Fz, density, angular_frequency, cat=False,p_ref=p_ref,norms=norms,k=k, transducer_radius=transducer_radius)]
215    if cat: vel = torch.cat(vel, dim=2)
216    return vel
217
218def propagate_speed(activations: Tensor, points: Tensor,board: Tensor|None=None, Fx=None, Fy=None, Fz=None, 
219                                 density = c.p_0, angular_frequency = c.angular_frequency,p_ref=c.P_ref, norms:Tensor=None,k=c.k, transducer_radius=c.radius):
220    '''
221    Propagates a hologram to speed at points\n
222    :param activations: Hologram to use
223    :param points: Points to propagate to
224    :param board: The Transducer array, default two 16x16 arrays
225    :param Fx: The forward model to us for Fx, if None it is computed using `forward_model_grad`. Default:`None`
226    :param Fy: The forward model to us for Fy, if None it is computed using `forward_model_grad`. Default:`None`
227    :param Fz: The forward model to us for Fz, if None it is computed using `forward_model_grad`. Default:`None`
228    :return: point velocity potential'
229    '''
230    
231    velocity = propagate_velocity(activations, points,board, Fx, Fy, Fz, density, angular_frequency,p_ref=p_ref,norms=norms,k=k, transducer_radius=transducer_radius)
232    speeds = []
233    for vel in velocity:
234        speeds.append(torch.abs(vel))
235    speed = 0
236    for spd in speeds:
237        speed += torch.square(spd)
238    speed = torch.sqrt(speed)
239    return speed
240
241def propagate_laplacian_helmholtz(activations: Tensor, points: Tensor,board: Tensor|None=None, A:Tensor|None=None, k=c.k,p_ref=c.P_ref, norms:Tensor=None, transducer_radius=c.radius) -> Tensor:
242    p = propagate(activations=activations, points=points, board=board,A=A,p_ref=p_ref,norms=norms,k=k, transducer_radius=transducer_radius)
243    return -1 * p * k
244
245
246def propagate_signed_pressure_abs(activations: Tensor, points: Tensor,board:Tensor|None=None, A:Tensor|None=None, A_function:FunctionType=None, A_function_args:dict={}, p_ref=c.P_ref, norms:Tensor=None,k=c.k, transducer_radius=c.radius) -> Tensor:
247    '''
248    Propagates a hologram to target points and returns pressure times sign of phase \n
249    :param activations: Hologram to use
250    :param points: Points to propagate to
251    :param board: The Transducer array, default two 16x16 arrays
252    :param A: The forward model to use, if None it is computed using `forward_model_batched`. Default:`None`
253    :return: point pressure
254
255    ```Python
256    from acoustools.Solvers import iterative_backpropagation
257    from acoustools.Utilities import create_points, propagate_abs
258
259    p = create_points(2,1)
260    x = iterative_backpropagation(p)
261    
262    p = p.squeeze(0)
263    x = iterative_backpropagation(p)
264    print(propagate_abs(x,p))
265    ```
266    '''
267    if board is None:
268        board = TRANSDUCERS
269    if A_function is not None:
270        A = A_function(points, board, **A_function_args)
271
272    out = propagate(activations, points,board,A=A,p_ref=p_ref,norms=norms,k=k, transducer_radius=transducer_radius)
273    return torch.abs(out) * torch.sign(torch.angle(out))