diff --git a/stan/docs/interface/fft.rst b/stan/docs/interface/fft.rst index f17cdda..1ec0c90 100644 --- a/stan/docs/interface/fft.rst +++ b/stan/docs/interface/fft.rst @@ -9,7 +9,7 @@ Utility Functions The following functions, and their two-dimensional analogues, primarily exist as utility functions but may be useful for more complex models. - :stan:func:`gp_rfft` transforms Gaussian process realizations to the Fourier domain and scales the coefficients such that they are white noise under the Gaussian process prior. -- :stan:func:`gp_rfft_log_abs_det_jacobian` evaluates the log absolute determinant of the Jacobian associated with the transformations :stan:func:`gp_rfft`. +- :stan:func:`gp_rfft_log_abs_det_jac` evaluates the log absolute determinant of the Jacobian associated with the transformations :stan:func:`gp_rfft`. Together, these two functions are used by :stan:func:`gp_rfft_lpdf` to evaluate the likelihood. diff --git a/stan/gptools/stan/gptools/fft1.stan b/stan/gptools/stan/gptools/fft1.stan index 66710ab..7333564 100644 --- a/stan/gptools/stan/gptools/fft1.stan +++ b/stan/gptools/stan/gptools/fft1.stan @@ -75,7 +75,7 @@ Evaluate the log absolute determinant of the Jacobian associated with :param n: Size of the real signal. Necessary because the size cannot be inferred from :code:`rfft`. :returns: Log absolute determinant of the Jacobian. */ -real gp_rfft_log_abs_det_jacobian(vector cov_rfft, int n) { +real gp_rfft_log_abs_det_jac(vector cov_rfft, int n) { vector[n %/% 2 + 1] rfft_scale = gp_evaluate_rfft_scale(cov_rfft, n); return - sum(log(rfft_scale[1:n %/% 2 + 1])) -sum(log(rfft_scale[2:(n + 1) %/% 2])) - log(2) * ((n - 1) %/% 2) + n * log(n) / 2; @@ -94,7 +94,7 @@ real gp_rfft_lpdf(vector y, vector loc, vector cov_rfft) { int n = size(y); int nrfft = n %/% 2 + 1; vector[n] z = gp_rfft(y, loc, cov_rfft); - return std_normal_lpdf(z) + gp_rfft_log_abs_det_jacobian(cov_rfft, n); + return std_normal_lpdf(z) + gp_rfft_log_abs_det_jac(cov_rfft, n); } diff --git a/stan/gptools/stan/gptools/fft2.stan b/stan/gptools/stan/gptools/fft2.stan index 9b810fc..d698444 100644 --- a/stan/gptools/stan/gptools/fft2.stan +++ b/stan/gptools/stan/gptools/fft2.stan @@ -140,7 +140,7 @@ Evaluate the log absolute determinant of the Jacobian associated with Fourier transform :code:`cov_rfft2`). :returns: Log absolute determinant of the Jacobian. */ -real gp_rfft2_log_abs_det_jacobian(matrix cov_rfft2, int width) { +real gp_rfft2_log_abs_det_jac(matrix cov_rfft2, int width) { int height = rows(cov_rfft2); int n = width * height; int fftwidth = width %/% 2 + 1; @@ -185,7 +185,7 @@ Evaluate the log probability of a two-dimensional Gaussian process realization i */ real gp_rfft2_lpdf(matrix y, matrix loc, matrix cov_rfft2) { return std_normal_lpdf(to_vector(gp_rfft2(y, loc, cov_rfft2))) - + gp_rfft2_log_abs_det_jacobian(cov_rfft2, cols(y)); + + gp_rfft2_log_abs_det_jac(cov_rfft2, cols(y)); } diff --git a/torch/gptools/torch/fft/fft1.py b/torch/gptools/torch/fft/fft1.py index 4fe22a2..1ccd869 100644 --- a/torch/gptools/torch/fft/fft1.py +++ b/torch/gptools/torch/fft/fft1.py @@ -1,4 +1,4 @@ -from gptools.util.fft import transform_rfft, transform_irfft, evaluate_rfft_log_abs_det_jacobian +from gptools.util.fft import transform_rfft, transform_irfft, evaluate_rfft_log_abs_det_jac from gptools.util.fft.fft1 import _get_rfft_scale import torch as th from torch.distributions import constraints @@ -43,8 +43,8 @@ def _inv_call(self, y: th.Tensor) -> th.Tensor: """ return transform_irfft(y, self.loc, rfft_scale=self.rfft_scale) - def log_abs_det_jacobian(self, x: th.Tensor, y: th.Tensor) -> th.Tensor: - return evaluate_rfft_log_abs_det_jacobian(x.shape[-1], rfft_scale=self.rfft_scale) + def log_abs_det_jac(self, x: th.Tensor, y: th.Tensor) -> th.Tensor: + return evaluate_rfft_log_abs_det_jac(x.shape[-1], rfft_scale=self.rfft_scale) class FourierGaussianProcess1D(th.distributions.TransformedDistribution): diff --git a/torch/gptools/torch/fft/fft2.py b/torch/gptools/torch/fft/fft2.py index 884d4ab..b72bab0 100644 --- a/torch/gptools/torch/fft/fft2.py +++ b/torch/gptools/torch/fft/fft2.py @@ -1,4 +1,4 @@ -from gptools.util.fft import transform_rfft2, transform_irfft2, evaluate_rfft2_log_abs_det_jacobian +from gptools.util.fft import transform_rfft2, transform_irfft2, evaluate_rfft2_log_abs_det_jac from gptools.util.fft.fft2 import _get_rfft2_scale import torch as th from .. import OptionalTensor @@ -44,8 +44,8 @@ def _inv_call(self, y: th.Tensor) -> th.Tensor: """ return transform_irfft2(y, self.loc, rfft2_scale=self.rfft2_scale) - def log_abs_det_jacobian(self, x: th.Tensor, y: th.Tensor) -> th.Tensor: - return evaluate_rfft2_log_abs_det_jacobian(x.shape[-1], rfft2_scale=self.rfft2_scale) + def log_abs_det_jac(self, x: th.Tensor, y: th.Tensor) -> th.Tensor: + return evaluate_rfft2_log_abs_det_jac(x.shape[-1], rfft2_scale=self.rfft2_scale) class FourierGaussianProcess2D(th.distributions.TransformedDistribution): diff --git a/util/gptools/util/fft/__init__.py b/util/gptools/util/fft/__init__.py index 8da5aac..bbd3fbe 100644 --- a/util/gptools/util/fft/__init__.py +++ b/util/gptools/util/fft/__init__.py @@ -1,6 +1,6 @@ -from .fft1 import evaluate_log_prob_rfft, evaluate_rfft_log_abs_det_jacobian, expand_rfft, \ +from .fft1 import evaluate_log_prob_rfft, evaluate_rfft_log_abs_det_jac, expand_rfft, \ transform_irfft, transform_rfft -from .fft2 import evaluate_log_prob_rfft2, evaluate_rfft2_log_abs_det_jacobian, transform_irfft2, \ +from .fft2 import evaluate_log_prob_rfft2, evaluate_rfft2_log_abs_det_jac, transform_irfft2, \ transform_rfft2 from .util import log_prob_stdnorm @@ -8,8 +8,8 @@ __all__ = [ "evaluate_log_prob_rfft", "evaluate_log_prob_rfft2", - "evaluate_rfft_log_abs_det_jacobian", - "evaluate_rfft2_log_abs_det_jacobian", + "evaluate_rfft_log_abs_det_jac", + "evaluate_rfft2_log_abs_det_jac", "expand_rfft", "log_prob_stdnorm", "transform_irfft", diff --git a/util/gptools/util/fft/fft1.py b/util/gptools/util/fft/fft1.py index 9372d93..48cea9b 100644 --- a/util/gptools/util/fft/fft1.py +++ b/util/gptools/util/fft/fft1.py @@ -185,10 +185,10 @@ def evaluate_log_prob_rfft(y: ArrayOrTensor, loc: ArrayOrTensor, *, rfft_scale = _get_rfft_scale(cov_rfft, cov, rfft_scale, y.shape[-1]) rfft = transform_rfft(y, loc, rfft_scale=rfft_scale) return log_prob_stdnorm(rfft).sum(axis=-1) \ - + evaluate_rfft_log_abs_det_jacobian(y.shape[-1], rfft_scale=rfft_scale) + + evaluate_rfft_log_abs_det_jac(y.shape[-1], rfft_scale=rfft_scale) -def evaluate_rfft_log_abs_det_jacobian(size: int, *, cov_rfft: OptionalArrayOrTensor = None, +def evaluate_rfft_log_abs_det_jac(size: int, *, cov_rfft: OptionalArrayOrTensor = None, cov: OptionalArrayOrTensor = None, rfft_scale: OptionalArrayOrTensor = None) -> ArrayOrTensor: """ @@ -202,7 +202,7 @@ def evaluate_rfft_log_abs_det_jacobian(size: int, *, cov_rfft: OptionalArrayOrTe :code:`(..., size // 2 + 1)`. Returns: - log_abs_det_jacobian: Log absolute determinant of the Jacobian. + log_abs_det_jac: Log absolute determinant of the Jacobian. """ imagidx = (size + 1) // 2 rfft_scale = _get_rfft_scale(cov_rfft, cov, rfft_scale, size) diff --git a/util/gptools/util/fft/fft2.py b/util/gptools/util/fft/fft2.py index fba7788..28c60ad 100644 --- a/util/gptools/util/fft/fft2.py +++ b/util/gptools/util/fft/fft2.py @@ -178,7 +178,7 @@ def transform_rfft2(y: ArrayOrTensor, loc: ArrayOrTensor, *, return unpack_rfft2(dispatch[y].fft.rfft2(y - loc) / rfft2_scale, y.shape) -def evaluate_rfft2_log_abs_det_jacobian(width: int, *, cov_rfft2: OptionalArrayOrTensor = None, +def evaluate_rfft2_log_abs_det_jac(width: int, *, cov_rfft2: OptionalArrayOrTensor = None, cov: OptionalArrayOrTensor = None, rfft2_scale: OptionalArrayOrTensor = None) -> ArrayOrTensor: """ @@ -194,7 +194,7 @@ def evaluate_rfft2_log_abs_det_jacobian(width: int, *, cov_rfft2: OptionalArrayO :code:`(..., height, width // 2 + 1)`. Returns: - log_abs_det_jacobian: Log absolute determinant of the Jacobian. + log_abs_det_jac: Log absolute determinant of the Jacobian. """ rfft2_scale = _get_rfft2_scale(cov_rfft2, cov, rfft2_scale, width) height = rfft2_scale.shape[-2] @@ -249,4 +249,4 @@ def evaluate_log_prob_rfft2(y: ArrayOrTensor, loc: ArrayOrTensor, *, rfft2_scale = _get_rfft2_scale(cov_rfft2, cov, rfft2_scale, y.shape[-1]) rfft2 = transform_rfft2(y, loc, rfft2_scale=rfft2_scale) return log_prob_stdnorm(rfft2).sum(axis=(-2, -1)) \ - + evaluate_rfft2_log_abs_det_jacobian(y.shape[-1], rfft2_scale=rfft2_scale) + + evaluate_rfft2_log_abs_det_jac(y.shape[-1], rfft2_scale=rfft2_scale)