mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-10 05:22:41 +02:00
36 lines
1.4 KiB
Python
36 lines
1.4 KiB
Python
from __future__ import division, print_function, absolute_import
|
|
import numpy as np
|
|
import torch as t
|
|
|
|
all = ['gaussian']
|
|
|
|
from CDTools.tools import cmath
|
|
|
|
|
|
def gaussian(shape, amplitude, sigma, center = None):
|
|
"""Returns an array with a centered gaussian
|
|
|
|
Takes in the shape, amplitude, and standard deviation of a gaussian
|
|
and returns a torch tensor with values corresponding to a two-dimensional
|
|
gaussian function
|
|
|
|
Note that [0, 0] is taken to be at the upper left corner of the array.
|
|
Default is centered at ((shape[0]-1)/2, (shape[1]-1)/2)) because x and y are zero-indexed.
|
|
|
|
Args:
|
|
shape (array_like) : A 1x2 array-like object specifying the dimensions of the output array in the form (i shape, j shape)
|
|
amplitude (float or int): The amplitude the gaussian to simulate
|
|
sigma (array_like): A 1x2 array-like object specifying the i- and j- standard deviation of the gaussian in the form (i stdev, j stdev)
|
|
center (array_like) : Optional 1x2 array-like object specifying the location of the center of the gaussian (i center, j center)
|
|
|
|
Returns:
|
|
torch.Tensor : The real-valued gaussian array
|
|
"""
|
|
if center is None:
|
|
center = ((shape[0]-1)/2, (shape[1]-1)/2)
|
|
|
|
i, j = np.mgrid[:shape[0], :shape[1]]
|
|
result = amplitude*np.exp(-( (i-center[0])**2 / (2 * sigma[0]**2) )
|
|
-( (j-center[1])**2 / (2 * sigma[1]**2) ))
|
|
return cmath.complex_to_torch(result)
|