"""
Functions to define commonly used function forms
"""
import numpy as np
from scipy.stats import norm
from typing import Union, Callable, Iterator, Tuple
__author__ = "Sam Hunt"
__created__ = "27/10/2020"
[docs]
def f_tophat(x: np.ndarray, centre: float, width: float) -> np.ndarray:
"""
Evaluate a top hat function along x
:param x: coordinate array
:param centre: x position of centre top hat
:param width: FWHM width of top hat
:return: top hat function
"""
y = np.ones(x.shape)
y[x < centre - width / 2] = 0
y[x > centre + width / 2] = 0
return y
[docs]
def f_gaussian(x: np.ndarray, centre: float, width: float) -> np.ndarray:
"""
Evaluate a gaussian function along x
:param x: coordinate array
:param centre: mean of distribution
:param width: FWHM of distribution
:return: gaussian function
"""
std = width / (2 * np.sqrt(2 * np.log(2)))
return norm.pdf(x, centre, std) * np.sqrt(2 * np.pi) * std
[docs]
def f_triangle(x: np.ndarray, centre: float, width: float) -> np.ndarray:
"""
Evaluate a top hat function along x
:param x: coordinate array
:param centre: x position of start of triangle
:param width: half width of base of triangle (i.e. FWHM)
:return: top hat function
:rtype: numpy.ndarray
"""
y = np.zeros(x.shape)
start = centre - width
stop = centre + width
first_half = np.logical_and(start < x, x <= centre)
y[first_half] = (x[first_half] - start) / (centre - start)
second_half = np.logical_and(centre < x, x < stop)
y[second_half] = (stop - x[second_half]) / (stop - centre)
return y
[docs]
def f_normalised(f: Callable, x: np.ndarray, *args, high_res_sampling: float = 0.001):
"""
Return normalised f along x
:param f: defining function
:param x: coordinate array
:param args: f parameters
:param high_res_sampling: (default 0.001) sampling to evaluate area of function to normalise by
:return: normalised function
"""
y = f(x, *args)
x_highres = np.arange(min(x), max(x), high_res_sampling)
y_highres = f(x_highres, *args)
y /= np.trapz(y_highres, x_highres)
return y
[docs]
def repeat_f(
f: Callable,
centres: np.ndarray,
widths: np.ndarray,
normalise: bool = False,
x_sampling: float = 0.01,
xlim_width: float = 3.0,
) -> Tuple[np.ndarray, np.ndarray]:
"""
Evaluates repeating functions along a coordinate axis
.. note::
Defines all function on common x coordinates, so may use a lot of memory if many bands defined with x sampling.
For a lower memory solution try **matheo.utils.function_def.iter_fs()**..
:param f: function to repeat
:param centres: function centres
:param widths: function widths
:param normalise: (default True) switch to define if area of return SRFs should be normalised to 1
:param x_sampling: sampling along function coordinates
:param xlim_width: (default 3) multiple function widths to define function over
:return: evaluated functions
:return: evaluated function coordinates
"""
fs = RepeatingFuncUtil(
f=f,
centres=centres,
widths=widths,
normalise=normalise,
x_sampling=x_sampling,
xlim_width=xlim_width,
)
return fs.return_fs()
[docs]
def iter_f(
f: Callable,
centres: np.ndarray,
widths: np.ndarray,
normalise: bool = False,
x_sampling: float = 0.01,
xlim_width: float = 3.0,
) -> Iterator:
"""
Returns iterator to evaluate repeating functions along a coordinate axis
.. note::
Offers a lower memory solution to **matheo.utils.function_def.return_fs()**.
:param f: function to repeat
:param centres: distribution centres
:param widths: distribution widths
:param normalise: (default True) switch to define if area of return SRFs should be normalised to 1
:param x_sampling: sampling along function coordinates
:param xlim_width: (default 3) multiple function widths to define function over
:return: repeating function iterator
"""
fs = RepeatingFuncUtil(
f=f,
centres=centres,
widths=widths,
normalise=normalise,
x_sampling=x_sampling,
xlim_width=xlim_width,
)
return iter(fs)
[docs]
class RepeatingFuncUtil:
"""
Helper class to define repeating functions along a coordinate axis
:param f: function to repeat
:param centres: distribution centres
:param widths: distribution widths
:param normalise: (default True) switch to define if area of return SRFs should be normalised to 1
:param x_sampling: sampling along function coordinates
:param xlim_width: (default 3) multiple function widths to define function over
"""
def __init__(
self,
f: Callable,
centres: np.ndarray,
widths: np.ndarray,
normalise: bool = False,
x_sampling: float = 0.01,
xlim_width: float = 3.0,
):
# Set attributes from arguments
self.f = f
self.centres = centres
self.widths = widths
self.normalise = normalise
self.x_sampling = x_sampling
self.xlim_width = xlim_width
[docs]
def return_fs(self) -> Tuple[np.ndarray, np.ndarray]:
"""
Evaluates all repeating functions along a coordinate axis
:return: evaluated functions
:return: evaluated function coordinates
"""
x = np.arange(
min(self.centres - self.xlim_width * self.widths),
max(self.centres + self.xlim_width * self.widths + 1),
self.x_sampling,
)
# Evaluate function
ys = np.zeros((len(self.centres), len(x)))
for i_band, (centre, width) in enumerate(zip(self.centres, self.widths)):
if self.normalise:
ys[i_band, :] = f_normalised(self.f, x, centre, width)
else:
ys[i_band, :] = self.f(x, centre, width)
return ys, x
def __iter__(self):
# Define counter
self.i = 0
return self
def __next__(self) -> Tuple[np.ndarray, np.ndarray]:
"""
Returns ith function
:return: evaluated function
:return: evaluated function coordinates
"""
# Iterate through bands
if self.i < len(self.centres):
centre = self.centres[self.i]
width = self.widths[self.i]
x = np.arange(
centre - self.xlim_width * width,
centre + self.xlim_width * width + 1,
self.x_sampling,
)
if self.normalise:
y = f_normalised(self.f, x, centre, width)
else:
y = self.f(x, centre, width)
# Update counter
self.i += 1
return y, x
else:
raise StopIteration
if __name__ == "__main__":
pass