Function compute_ifft

utils/strand_util.py:475–502

(input, spatial_size=256)

#inverse of compute_stft
def compute_ifft(input, spatial_size=256):

    #if the input if 4 dimensional and the last dimension 2, then we view it as complex
    if len(input.shape)==4 and input.shape[-1]==2:
        input=torch.view_as_complex(input)

    if len(input.shape)==2:
        #we have a (B,T) input so we can just run stft
         x= torch.fft.irfft(
            input,
        )
    elif len(input.shape)==3:
        #do stft for each dimension
        dim=input.shape[-1]
        x_total=[]
        for i in range(dim):
            input_axis = input[:,:,i]
            x_axis = torch.fft.irfft(
                input_axis,
                n=spatial_size,
            ).unsqueeze(-1)
            x_total.append(x_axis)
        x = torch.cat(x_total,-1)
    else:
        return None
    

    return x


#strand positions should be [nr_strands,nr_verts,3]