"""Convolution functions and kernel producing functions.
Here, we group all convolution functions and kernel functions. A lot of the
convolution functions are brief wrappers around Astropy's convolution.
All three dimensionalities are covered.
"""
# isort: split
# Import required to remove circular dependencies from type checking.
from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from lezargus.library import hint
# isort: split
import astropy.convolution
import astropy.modeling
import numpy as np
from lezargus.library import logging
[docs]
def _check_array_dimensionality(array: hint.NDArray, dimensions: int) -> bool:
"""Check if the array has the expected number of dimensions.
This function checks if the array has the correction number of
dimensions. Of course, the expected dimensions are different so this
function is more a wrapper around the logging message and it
serves as a basic check.
Parameters
----------
array : ndarray
The array that we are testing if it has the same number of
dimensions.
dimensions : int
The number of expected dimensions the array should have.
Returns
-------
valid_dimensionality : bool
If True, the array has the expected dimensionality, as input.
"""
# We just use Numpy shape and the like, type conversion.
array = np.asarray(array)
dimensions = int(dimensions)
# Checking.
valid_dimensionality = len(array.shape) == dimensions
if not valid_dimensionality:
logging.error(
error_type=logging.InputError,
message=(
f"Input array has wrong dimensions, shape {array.shape};"
f" expected dimensionality {dimensions}."
),
)
return valid_dimensionality
[docs]
def _check_kernel_dimensionality(kernel: hint.NDArray, dimensions: int) -> bool:
"""Check if the kernel has the expected number of dimensions.
Same function as :py:meth:`_check_array_dimensionality`, just different
error message.
Parameters
----------
kernel : ndarray
The kernel that we are testing if it has the same number of
dimensions.
dimensions : int
The number of expected dimensions the kernel should have.
Returns
-------
valid_dimensionality : bool
If True, the kernel has the expected dimensionality, as input.
"""
# We just use Numpy shape and the like, type conversion.
kernel = np.asarray(kernel)
dimensions = int(dimensions)
# Checking.
valid_dimensionality = len(kernel.shape) == dimensions
if not valid_dimensionality:
logging.error(
error_type=logging.InputError,
message=(
f"Input kernel has wrong dimensions, shape {kernel.shape};"
f" expected dimensionality {dimensions}."
),
)
return valid_dimensionality
[docs]
def _check_array_kernel_variable_stack(
array: hint.NDArray,
kernel_stack: hint.NDArray,
axis: int | tuple[int],
) -> bool:
"""Check if the kernel stack and array have the exact slice count.
For variable kernels, we need to make sure that there are enough kernels
in the kernel stack for each slice of the array which are being convolved.
The axes which are variable (and not dynamic, i.e. changing with
the convolution axis) are checked to be the same size in the array and
the kernel stack.
Parameters
----------
array : ndarray
The array which would be convolved and which we are checking is
compatible with the kernel stack.
kernel_stack : ndarray
The kernel stack which holds all of the kernels that are being used.
We are checking if it is compatible with the kernel stack.
axis : int, tuple[int]
Either a single axis index, or a tuple of axis indexes which the
kernel is variable. The axis or axes should not be the same as the
convolution axis or axes.
Returns
-------
valid_stack : bool
If True, the kernel stack has the correct amount of kernels for the
array, for the axes which are variable.
"""
# We repackage the axis, if it is a single value, it is the same as just
# checking one so the code below can be the same.
if isinstance(axis, int):
check_axes = (axis,)
elif isinstance(axis, list | tuple | np.ndarray):
# It is already a tuple-like.
check_axes = tuple(int(axisdex) for axisdex in axis)
else:
check_axes = axis
logging.error(
error_type=logging.InputError,
message=(
f"Cannot parse axis input {axis} to a list of axis indexes to"
" check."
),
)
# Just loop over each of the axes, checking if they have the same size.
# We assume a good stack at first.
valid_stack = True
mismatched_axes = []
for axisdex in check_axes:
# We need to make sure that there is enough kernels in the stack for
# each of the array slices.
array_stack_count = array.shape[axisdex]
kernel_stack_count = kernel_stack.shape[axisdex]
if array_stack_count != kernel_stack_count:
# The axes are mismatched, we record the mismatch here to report
# later.
valid_stack = False
mismatched_axes.append(axisdex)
# Is the stack still good?
if not valid_stack:
# Nope, we report the invalid stack. There is no one to one relation
# for the array and kernel.
logging.error(
error_type=logging.InputError,
message=(
f"Array shape {array.shape} incompatible with kernel stack"
f" {kernel_stack.shape}. Mismatch in {mismatched_axes} indexed"
" axes."
),
)
# All done.
return valid_stack
[docs]
def _static_astropy_convolve(
array: hint.NDArray,
kernel: hint.NDArray,
) -> hint.NDArray:
"""Use Astropy to convolve the array provided the kernel.
The Astropy convolve function only can convolve up to 3D, and they
determine it based on the array and kernel dimensionality. We attempt
to do an FFT convolution, but, should it fail, we fall back to
discrete convolution.
Parameters
----------
array : ndarray
The array we are convolving by the kernel.
kernel : ndarray
The kernel we are using to convolve.
Returns
-------
convolved : ndarray
The result of the convolution.
"""
# The array and kernel should have the same dimensionality, and it
# cannot be more than 3. We base it on the array as that is the more
# fundamental part.
dimensionality = len(array.shape)
_check_array_dimensionality(array=array, dimensions=dimensionality)
_check_kernel_dimensionality(kernel=kernel, dimensions=dimensionality)
# Checking if Astropy's convolve function can handle it.
max_astropy_convolve_dimensionality = 3
if dimensionality > max_astropy_convolve_dimensionality:
logging.critical(
critical_type=logging.NotSupportedError,
message=(
"Astropy convolve only supports up to 3D arrays, input"
f" dimensionality is {dimensionality}."
),
)
# FFT convolution uses complex numbers. We want to keep the same
# numerical precision as the input type. We can expand this to 192-bit
# and 256-bit, but, it is likely not needed.
if array.dtype.itemsize * 2 <= np.complex64(None).itemsize:
complex_data_type = np.complex64
elif array.dtype.itemsize * 2 <= np.complex128(None).itemsize:
complex_data_type = np.complex128
else:
complex_data_type = complex
# There are two ways that the convolution can happen, either via FFT
# or via discrete convolution. It is always faster to do it via FFT
# but we fall back to discrete convolution if we run out of memory.
try:
convolved = astropy.convolution.convolve_fft(
array,
kernel=kernel,
boundary="wrap",
complex_dtype=complex_data_type,
nan_treatment="interpolate",
normalize_kernel=True,
preserve_nan=True,
allow_huge=True,
)
except MemoryError:
# There is not enough memory for an FFT version, using discrete
# instead.
# We give some warning first.
logging.warning(
warning_type=logging.MemoryFullWarning,
message=(
"Attempting a FFT convolution of a spectra with shape"
f" {array.shape} with kernel shape {kernel.shape} requires"
" too much memory."
),
)
logging.warning(
warning_type=logging.AlgorithmWarning,
message=(
"Discrete convolution attempted as an alternative to FFT"
" convolution due to memory issues; expect long execution time."
),
)
# Discrete convolution.
convolved = astropy.convolution.convolve(
array,
kernel=kernel,
boundary="extend",
nan_treatment="interpolate",
normalize_kernel=True,
preserve_nan=True,
)
# All done.
return convolved
[docs]
def static_1d_with_1d(
array: hint.NDArray,
kernel: hint.NDArray,
) -> hint.NDArray:
"""Convolve a 1D array using a static 1D kernel.
Parameters
----------
array : ndarray
The 1D array data which we will convolve.
kernel : ndarray
The 1D kernel that we are using to convolve.
Returns
-------
convolved : ndarray
The convolved 1D array data.
"""
# We check that both the array and the kernel has the proper
# dimensionality.
array_dimensionality = 1
kernel_dimensionality = 1
_check_array_dimensionality(
array=array,
dimensions=array_dimensionality,
)
_check_kernel_dimensionality(
kernel=kernel,
dimensions=kernel_dimensionality,
)
# We do the convolution.
convolved = _static_astropy_convolve(array=array, kernel=kernel)
return convolved
[docs]
def static_2d_with_2d(
array: hint.NDArray,
kernel: hint.NDArray,
) -> hint.NDArray:
"""Convolve a 2D array using a static 2D kernel.
Parameters
----------
array : ndarray
The 2D array data which we will convolve.
kernel : ndarray
The 2D kernel that we are using to convolve.
Returns
-------
convolved : ndarray
The convolved 2D array data.
"""
# We check that both the array and the kernel has the proper
# dimensionality.
array_dimensionality = 2
kernel_dimensionality = 2
_check_array_dimensionality(
array=array,
dimensions=array_dimensionality,
)
_check_kernel_dimensionality(
kernel=kernel,
dimensions=kernel_dimensionality,
)
# We do the convolution.
convolved = _static_astropy_convolve(array=array, kernel=kernel)
return convolved
[docs]
def static_3d_with_1d_over_z(
array: hint.NDArray,
kernel: hint.NDArray,
) -> hint.NDArray:
"""Convolve a 3D array using a 1D kernel, over the z dimension.
This convolution convolves 1D slices of the 3D array. The convolution
itself then is a 1D array being convolved with a 1D kernel. We take slices
of the last dimension (z), iterating over the 1st and 2nd dimension.
Parameters
----------
array : ndarray
The 3D array data which we will convolve.
kernel : ndarray
The 1D kernel that we are using to convolve over the z axis of the
array.
Returns
-------
convolved : ndarray
The convolved 3D array data.
"""
# We check that both the array and the kernel has the proper
# dimensionality.
array_dimensionality = 3
kernel_dimensionality = 1
_check_array_dimensionality(
array=array,
dimensions=array_dimensionality,
)
_check_kernel_dimensionality(
kernel=kernel,
dimensions=kernel_dimensionality,
)
# Applying the convolution. This really is just a repeated process of
# 1D convolutions across both other axes. We create the resulting
# array and fill it in with the results of the convolutions.
convolved = np.empty_like(array)
for coldex in np.arange(array.shape[0]):
for rowdex in np.arange(array.shape[1]):
convolved[coldex, rowdex, :] = static_1d_with_1d(
array=array[coldex, rowdex, :],
kernel=kernel,
)
# All done.
return convolved
[docs]
def static_3d_with_2d_over_xy(
array: hint.NDArray,
kernel: hint.NDArray,
) -> hint.NDArray:
"""Convolve a 3D array using a 2D kernel, over the x-y plane.
This convolution convolves 2D slices of the 3D array. The convolution
itself then is a 2D array being convolved with a 3D kernel. A full
3D array and 3D kernel convolution is not done here.
Parameters
----------
array : ndarray
The 3D array data which we will convolve.
kernel : ndarray
The 2D kernel that we are using to convolve over the x-y plane
of the array.
Returns
-------
convolved : ndarray
The convolved 3D array data.
"""
# We check that both the array and the kernel has the proper
# dimensionality.
array_dimensionality = 3
kernel_dimensionality = 2
_check_array_dimensionality(
array=array,
dimensions=array_dimensionality,
)
_check_kernel_dimensionality(
kernel=kernel,
dimensions=kernel_dimensionality,
)
# Applying the convolution. This really is just a repeated process of
# 2D convolutions across the x-y plane. We create the resulting
# array and fill it in with the results of the convolutions.
convolved = np.empty_like(array)
for index in np.arange(array.shape[2]):
convolved[:, :, index] = static_2d_with_2d(
array=array[:, :, index],
kernel=kernel,
)
# All done.
return convolved
[docs]
def variable_3d_with_2d_over_xy(
array: hint.NDArray,
kernel_stack: hint.NDArray,
) -> hint.NDArray:
"""Convolve a 3D array using a variable 2D kernel, over the x-y plane.
Like py:func:`static_3d_with_2d_over_xy`, this convolution convolves
2D slices of the 3D array. However, the kernel here is variable in
in the z dimension.
Parameters
----------
array : ndarray
The 3D array data which we will convolve.
kernel_stack : ndarray
The 2D kernel stack that we are using to convolve over the x-y
plane of the array. Each slice of the stack should correspond to
the kernel for the slice of the array.
Returns
-------
convolved : ndarray
The convolved 3D array data.
"""
# It is best to work with arrays.
kernel_stack = np.asarray(kernel_stack)
# We check that both the array and the kernel has the proper
# dimensionality. The kernel stack is extra by one because of the
# stack axis.
array_dimensionality = 3
kernel_dimensionality = 3
_check_array_dimensionality(array=array, dimensions=array_dimensionality)
_check_kernel_dimensionality(
kernel=kernel_stack,
dimensions=kernel_dimensionality,
)
# We also check if the array and kernel stack is compatible for variable
# convolution along the varying axis.
varying_axis = 2
is_valid_stack = _check_array_kernel_variable_stack(
array=array,
kernel_stack=kernel_stack,
axis=varying_axis,
)
if not is_valid_stack:
logging.warning(
warning_type=logging.AlgorithmWarning,
message=(
"Kernel stack mismatch,"
f" {kernel_stack.shape[varying_axis]} kernels for"
f" {array.shape[varying_axis]} array slices."
),
)
# Applying the convolution. This really is just a repeated process of
# 2D convolutions across the x-y plane, except we also iterate the
# kernel stack.
convolved = np.empty_like(array)
for index in np.arange(array.shape[2]):
convolved[:, :, index] = static_2d_with_2d(
array=array[:, :, index],
kernel=kernel_stack[:, :, index],
)
# All done.
return convolved
[docs]
def kernel_1d_gaussian(
shape: tuple | int,
stddev: float,
) -> hint.NDArray:
"""Return a 1D Gaussian convolution kernel.
We normalize the kernel via the amplitude of the Gaussian
function as a whole for maximal precision: volume = 1. The `stddev` must
be expressed in pixels.
Parameters
----------
shape : tuple | int
The shape of the 1D kernel, in pixels. If a single value (i.e. a size
value instead), we attempt convert it to a shape-like value.
stddev : float
The standard deviation of the Gaussian, in pixels.
Returns
-------
gaussian_kernel : ndarray
The discrete kernel array.
"""
# We need to make sure we can handle odd inputs of the standard
# deviation, just in case.
if stddev < 0:
logging.error(
error_type=logging.InputError,
message=f"Gaussian stddev {stddev}is negative, not physical.",
)
elif np.isclose(stddev, 0):
logging.warning(
warning_type=logging.AlgorithmWarning,
message=(
f"Gaussian stddev is {stddev}, about zero; kernel is basically"
" a delta-function."
),
)
# We need to determine the shape. If it is a single value we attempt to
# interpret it. Granted, we only need a size, but we keep a shape as the
# input to align it better with the 2D kernel functions.
if isinstance(shape, list | tuple) and len(shape) == 1:
# All good.
size = shape[0]
elif isinstance(shape, int | np.number):
size = shape
else:
logging.error(
error_type=logging.InputError,
message=(
f"Kernel shape input {shape} type {type(shape)} is not a 1D"
" array shape."
),
)
size = shape
# Regardless, the center of the array is considered to be the center of
# the Gaussian function.
center = (size - 1) / 2
# The actual input array to the Gaussian function.
input_ = np.arange(size, dtype=int)
# The normalization constant is really just the area of the Gaussian.
norm_constant = 1 / (stddev * np.sqrt(2 * np.pi))
# Deriving the kernel and computing it.
gaussian1d = astropy.modeling.models.Gaussian1D(
amplitude=norm_constant,
mean=center,
stddev=stddev,
)
gaussian_kernel = gaussian1d(input_)
# All done.
return gaussian_kernel
[docs]
def kernel_1d_gaussian_resolution(
shape: tuple | int,
template_wavelength: hint.NDArray | float,
base_resolution: float | None = None,
target_resolution: float | None = None,
base_resolving_power: float | None = None,
target_resolving_power: float | None = None,
reference_wavelength: float | None = None,
) -> hint.NDArray:
"""Gaussian 1D kernel adapted for resolution convolution conversions.
This function is a wrapper around a normal 1D Gaussian kernel. Instead
of specifying the standard deviation, we calculate the approximate
required standard deviation needed to down-sample a base resolution to
some target resolution. We accept both resolution values or resolving
power values for the calculation; but we default to resolution based
determination if possible.
Parameters
----------
shape : tuple | int
The shape of the 1D kernel, in pixels. If a single value (i.e. a size
value instead), we attempt convert it to a shape-like value.
template_wavelength : ndarray or float
An example wavelength array which this kernel will be applied to. This
is required to convert the physical standard deviation value calculated
from the resolution/resolving power to one of length in pixels/points.
If an array, we try and compute the conversion factor. If a float,
that is the conversion factor of wavelength per pixel.
base_resolution : float, default = None
The base resolution that we are converting from. Must be provided
along with `target_resolution` for the resolution mode.
target_resolution : float, default = None
The target resolution we are converting to. Must be provided
along with `base_resolution` for the resolution mode.
base_resolving_power : float, default = None
The base resolving power that we are converting from. Must be provided
along with `target_resolving_power` and `reference_wavelength` for the
resolving power mode.
target_resolving_power : float, default = None
The target resolving power that we are converting from. Must be
provided along with `base_resolving_power` and `reference_wavelength`
for the resolving power mode.
reference_wavelength : float, default = None
The reference wavelength used to convert from resolving power to
resolution. Must be provided along with `base_resolving_power` and
`target_resolving_power` for the resolving power mode.
Returns
-------
resolution_kernel : ndarray
The Gaussian kernel with the appropriate parameters to convert from
the base resolution to the target resolution with a convolution.
"""
# We support two different modes of computing the kernel. Toggle is based
# on what parameters are provided. We switch here.
resolution_mode = (
base_resolution is not None and target_resolution is not None
)
resolving_mode = (
base_resolving_power is not None
and target_resolving_power is not None
and reference_wavelength is not None
)
# Determining which, and based on which, we determine the determine the
# standard deviation for the Gaussian. However, the standard deviation
# value determined here is a physical length, not one in pixels/points.
phys_fwhm = -1
if resolution_mode and resolving_mode:
# If we have both modes, the program cannot decide between both.
# Though we default to resolution based modes, it is still problematic.
logging.error(
error_type=logging.InputError,
message=(
"Both resolution mode and resolving mode information was"
" provided for kernel determination. Mode cannot be determined."
),
)
phys_fwhm = np.sqrt(target_resolution**2 - base_resolution**2)
elif resolution_mode:
# Resolution mode, we determine the standard deviation from the
# provided resolutions.
phys_fwhm = np.sqrt(target_resolution**2 - base_resolution**2)
elif resolving_mode:
# Resolving mode, we determine the standard deviation from the
# provided resolving power and root wavelength.
phys_fwhm = reference_wavelength * (
(base_resolving_power**2 - target_resolving_power**2)
/ (base_resolving_power * target_resolving_power)
)
logging.warning(
warning_type=logging.DeprecationWarning,
message=(
"Resolving power kernel generation should be replaced with"
" resolutions computed from resolving power."
),
)
else:
# No mode could be found usable. The inputs seem to be quite wrong.
# This is equivalent to TypeError missing argument, hence a critical
# failure.
logging.critical(
critical_type=logging.InputError,
message=(
"Kernel calculation mode could not be determined. Resolution"
f" mode values: base, {base_resolution}; target:"
f" {target_resolution}. Resolving mode values: base,"
f" {base_resolving_power}; target, {target_resolving_power};"
f" wavelength, {reference_wavelength}."
),
)
# Before we continue, we need to make sure that the FWHM is reasonable.
# If it is not reasonable, we try and find potential problems to warn
# about.
if phys_fwhm <= 0:
# If the resolutions are the same, it basically leads to a
# delta function.
if resolution_mode and np.isclose(target_resolution, base_resolution):
logging.warning(
warning_type=logging.AlgorithmWarning,
message=(
f"Target resolution {target_resolution} and base resolution"
f" {base_resolution} is the same. Possible delta kernel."
),
)
# Similar, if the resolving powers are the same, it basically leads to a
# delta function.
if resolving_mode and np.isclose(
target_resolving_power,
base_resolving_power,
):
logging.warning(
warning_type=logging.AlgorithmWarning,
message=(
f"Target resolving power {target_resolving_power} and base"
f" resolving power {base_resolving_power} is the same."
" Possible delta kernel."
),
)
# Converting to standard deviation.
fwhm_std_const = 2 * np.sqrt(2 * np.log(2))
phys_stddev = phys_fwhm / fwhm_std_const
# We convert the physical standard deviation into a standard deviation of
# pixels (or points in general). We assume a wavelength spacing
# based on the average spacing of the provided wavelength.
if isinstance(template_wavelength, float | int | np.number):
convert_factor = template_wavelength
else:
convert_factor = np.nanmean(
template_wavelength[1:] - template_wavelength[:-1],
)
# Converting
stddev = phys_stddev / convert_factor
# With the standard deviation known, we can compute the kernel using the
# Gaussian kernel creator.
resolution_kernel = kernel_1d_gaussian(shape=shape, stddev=stddev)
# All done.
return resolution_kernel
[docs]
def kernel_2d_gaussian(
shape: tuple,
x_stddev: float,
y_stddev: float,
rotation: float,
) -> hint.NDArray:
"""Return a 2D Gaussian convolution kernel.
We normalize the kernel via the amplitude of the Gaussian
function as a whole for maximal precision: volume = 1. We require the
input of the shape of the kernel to allow for `x_stddev` and `y_stddev`
to be expressed in pixels to keep it general. By definition, the center
of the Gaussian kernel is in the center of the array.
Parameters
----------
shape : tuple
The shape of the 2D kernel, in pixels.
x_stddev : float
The standard deviation of the Gaussian in the x direction, in pixels.
y_stddev : float
The standard deviation of the Gaussian in the y direction, in pixels.
rotation : float
The rotation angle, increasing counterclockwise, in radians.
Returns
-------
gaussian_kernel : ndarray
The discrete kernel array.
"""
# The center of the array given by the shape is defined as just the center
# of it. However, we need to take into account off-by-one errors.
try:
nrow, ncol = shape
except ValueError:
logging.critical(
critical_type=logging.InputError,
message=(
"The 2D kernel shape cannot be determined from input shape:"
f" {shape}"
),
)
cen_row = (nrow - 1) / 2
cen_col = (ncol - 1) / 2
# The normalization constant is provided as volume itself.
norm_constant = 1 / (2 * np.pi * x_stddev * y_stddev)
# The mesh grid used to evaluate the Gaussian function to derive the kernel.
xx, yy = np.meshgrid(np.arange(ncol, dtype=int), np.arange(nrow, dtype=int))
# Deriving the kernel and computing it.
gaussian2d = astropy.modeling.models.Gaussian2D(
amplitude=norm_constant,
x_mean=cen_col,
y_mean=cen_row,
x_stddev=x_stddev,
y_stddev=y_stddev,
theta=rotation,
)
gaussian_kernel = gaussian2d(xx, yy)
return gaussian_kernel