src.acoustools.Utilities.Utilities

  1from torch import Tensor
  2import torch
  3from typing import Literal, Iterable
  4
  5
  6def is_batched_points(points:Tensor) -> bool:
  7    '''
  8    :param points: `Tensor` of points
  9    :return: `True` is points has a batched shape
 10    '''
 11    if len(points.shape)> 2 :
 12        return True
 13    else:
 14        return False
 15    
 16
 17
 18def permute_points(points: Tensor,index: int,axis:int=0) -> Tensor:
 19    '''
 20    Permutes axis of a tensor \n
 21    :param points: Tensor to permute
 22    :param index: Indexes describing order to perumte to 
 23    :param axis: Axis to permute. Default `0`
 24    :return: permuted points
 25    '''
 26    if axis == 0:
 27        return points[index,:,:,:]
 28    if axis == 1:
 29        return points[:,index,:,:]
 30    if axis == 2:
 31        return points[:,:,index,:]
 32    if axis == 3:
 33        return points[:,:,:,index]
 34
 35
 36def convert_to_complex(matrix: Tensor) -> Tensor:
 37    '''
 38    Comverts a real tensor of shape `B x M x N` to a complex tensor of shape `B x M/2 x N` 
 39    :param matrix: Matrix to convert
 40    :return: converted complex tensor
 41    '''
 42    # B x 1024 x N (real) -> B x N x 512 x 2 -> B x 512 x N (complex)
 43    matrix = torch.permute(matrix,(0,2,1))
 44    matrix = matrix.view((matrix.shape[0],matrix.shape[1],-1,2))
 45    matrix = torch.view_as_complex(matrix.contiguous())
 46    return torch.permute(matrix,(0,2,1))
 47
 48
 49
 50
 51def return_matrix(x,y,mat=None):
 52    '''
 53    @private
 54    Returns value of parameter `mat` - For compatibility with other functions
 55    '''
 56    return mat
 57
 58
 59def get_convert_indexes(n:int=512, single_mode:Literal['bottom','top']='bottom') -> Tensor:
 60    '''
 61    Gets indexes to swap between transducer order for acoustools and OpenMPD for two boards\n
 62    Use: `row = row[:,FLIP_INDEXES]` and invert with `_,INVIDX = torch.sort(IDX)` 
 63    :param n: number of Transducers
 64    :param single_mode: When using only one board is that board a top or bottom baord. Default bottom
 65    :return: Indexes
 66    '''
 67
 68    indexes = torch.arange(0,n)
 69    # # Flip top board
 70    # if single_mode.lower() == 'top':
 71    #     indexes[:256] = torch.flip(indexes[:256],dims=[0])
 72    # elif single_mode.lower() == 'bottom':
 73    #     indexes[:256] = torch.flatten(torch.flip(torch.reshape(indexes[:256],(16,-1)),dims=[1]))
 74
 75    indexes[:256] = torch.flip(indexes[:256],dims=[0])
 76    
 77    if n > 256:
 78        indexes[256:] = torch.flatten(torch.flip(torch.reshape(indexes[256:],(16,-1)),dims=[1]))
 79    
 80    return indexes
 81
 82def batch_list(iterable:Iterable, batch:int=32):
 83    '''
 84    Split an iterable into batch sized pieces\n
 85    :param iterable: The iterable to batch
 86    :param batch: The size to batch
 87    ```
 88    x = range(100)
 89    for b in batch_list(x):
 90        print(b)
 91    ```
 92    '''
 93    i = 0
 94    while i <= len(iterable):
 95        if i + batch <= len(iterable):
 96            yield iterable[i:i+batch]
 97        else:
 98            yield iterable[i:]
 99        i += batch
100
101def get_rows_in(a_centres, b_centres, expand = True):
102    '''
103    Takes two tensors and returns a mask for `a_centres` where a value of true means that row exists in `b_centres` \\
104    Asssumes in form 1x3xN -> returns mask over dim 1\\
105    `a_centres` Tensor of points to check for inclusion in `b_centres` \\
106    `b_centres` Tensor of points which may or maynot contain some number of points in `a_centres`\\
107    `expand` if True returns mask as `1x3xN` if False returns mask as `1xN`. Default: True\\
108    Returns mask for all rows in `a_centres` which are in `b_centres`
109    '''
110
111    M = a_centres.shape[2] #Number of total elements
112    R = b_centres.shape[2] #Number of elements in b
113
114    a_reshape = torch.unsqueeze(a_centres,3).expand(-1, -1, -1, R)
115    b_reshape = torch.unsqueeze(b_centres,2).expand(-1, -1, M, -1)
116
117    mask = b_reshape == a_reshape
118    mask = mask.all(dim=1).any(dim=2)
119
120    if expand:
121        return mask.unsqueeze(1).expand(-1,3,-1)
122    else:
123        return mask