healpix_analyse.fft_local ========================= .. py:module:: healpix_analyse.fft_local .. autoapi-nested-parse:: Fast local 2-D FFTs for data sampled on HEALPix cells. The HEALPix cell centres are projected onto a square gnomonic grid. The geometry is computed once by :class:`LocalFFT`; subsequent projections and transforms use differentiable PyTorch scatter/gather operations. The FFT itself is exactly invertible on the projected grid. The default HEALPix -> grid -> HEALPix round trip is approximate because it uses fast bilinear gridding in both directions. Classes ------- .. autoapisummary:: healpix_analyse.fft_local.LocalFFT Functions --------- .. autoapisummary:: healpix_analyse.fft_local.fft healpix_analyse.fft_local.ifft healpix_analyse.fft_local.ps Module Contents --------------- .. py:class:: LocalFFT(cell_ids: ArrayLike, level: int, *, ellipsoid: str = 'sphere', max_patch_radius_deg: float = 10.0, pixel_size_rad: Optional[float] = None, grid_size: Optional[int] = None, norm: str = 'ortho', dtype: torch.dtype = torch.float32, device: Optional[Union[str, torch.device]] = None) Bases: :py:obj:`torch.nn.Module` A reusable local flat-sky FFT for a set of NESTED HEALPix cells. :Parameters: * **cell_ids** (:py:class:`array-like`, :py:class:`shape [N]`) -- Unique HEALPix cell identifiers in NESTED ordering. * **level** (:py:class:`int`) -- HEALPix level, with ``nside = 2**level``. * **ellipsoid** (:py:class:`str`, *default* ``"sphere"``) -- Geometry passed to ``healpix_geo`` when locating cell centres. * **max_patch_radius_deg** (:py:class:`float`, *default* ``10``) -- Maximum angular distance of a cell centre from the spherical patch centre. A single gnomonic plane is rejected above this radius. * **pixel_size_rad** (:py:class:`float`, *optional*) -- Grid spacing in gnomonic coordinates. The default is the square root of the HEALPix pixel area, ``sqrt(pi/3) / nside``. * **grid_size** (:py:class:`int`, *optional*) -- Side of the square projected grid. By default, the smallest power of two that covers the patch with a one-pixel margin is used. * **norm** (``{"forward", "backward", "ortho"}``, *default* ``"ortho"``) -- Normalisation passed to :func:`torch.fft.fft2` and ``ifft2``. * **dtype** (:py:class:`torch.dtype`, *default* :py:class:`torch.float32`) -- Real dtype used for geometry and real-valued input maps. * **device** (:py:class:`torch.device` or :py:class:`str`, *optional*) -- Device of the cached projection geometry. CUDA is selected when it is available, otherwise CPU. .. admonition:: Notes The tangent frame is built from three-dimensional unit vectors, not from longitude/latitude differences. It is therefore well defined at both geographic poles and across the 0/360-degree meridian. For torch inputs, gradients propagate through projection, FFT, IFFT and back-projection with respect to the data. The discrete HEALPix geometry is fixed and is not differentiable with respect to ``cell_ids``. .. py:attribute:: level .. py:attribute:: nside .. py:attribute:: ellipsoid :value: '' .. py:attribute:: max_patch_radius_deg .. py:attribute:: norm :value: 'ortho' .. py:attribute:: dtype :value: Ellipsis .. py:attribute:: cdtype :value: Ellipsis .. py:attribute:: n_cells .. py:attribute:: pixel_size_rad .. py:attribute:: grid_size :value: 0 .. py:attribute:: grid_shape .. py:attribute:: centre_lon_deg .. py:attribute:: centre_lat_deg .. py:attribute:: patch_radius_deg .. py:property:: device :type: torch.device Current device of the cached geometry buffers. .. py:property:: coverage_mask :type: torch.Tensor Boolean ``[grid_size, grid_size]`` mask reached by input cells. .. py:property:: grid_density :type: torch.Tensor Accumulated interpolation weight on the projected square grid. .. py:method:: project(data: ArrayLike) -> ArrayLike Bilinearly scatter HEALPix values onto the cached square grid. .. py:method:: unproject(grid: ArrayLike) -> ArrayLike Bilinearly sample a projected grid at the HEALPix cell centres. .. py:method:: fft(data: ArrayLike) -> ArrayLike Project HEALPix data and compute its unshifted two-dimensional FFT. .. py:method:: ifft(spectrum: ArrayLike, *, real_output: bool = True) -> ArrayLike Compute the inverse FFT and sample it at the input HEALPix cells. This is a fast approximate inverse of :meth:`fft`. Set ``real_output=False`` when transforming genuinely complex-valued data. .. py:method:: ps(spectrum: ArrayLike, *, n_bins: Optional[int] = None, include_corners: bool = False) -> tuple[ArrayLike, ArrayLike] Compute a circularly averaged 1-D power spectrum from a 2-D FFT. :Parameters: * **spectrum** (:py:class:`numpy.ndarray` or :py:class:`torch.Tensor`, :py:class:`shape [...`, :py:class:`grid_size`, :py:class:`grid_size]`) -- Unshifted Fourier coefficients, normally returned by :meth:`fft`. * **n_bins** (:py:class:`int`, *optional*) -- Number of equally spaced radial-frequency bins. The default is ``grid_size // 2`` when excluding corners. * **include_corners** (:py:class:`bool`, *default* :py:obj:`False`) -- If False, stop at the axis Nyquist frequency so every annulus is fully contained in the Fourier grid. If True, include the corner modes up to ``sqrt(2)`` times the axis Nyquist frequency; outer annuli are then only partially sampled. :returns: * **frequency** (:py:class:`numpy.ndarray` or :py:class:`torch.Tensor`, :py:class:`shape [n_bins]`) -- Bin centres in cycles per gnomonic tangent unit. With the default grid spacing these are approximately cycles per radian. * **power** (:py:class:`numpy.ndarray` or :py:class:`torch.Tensor`, :py:class:`shape [...`, :py:class:`n_bins]`) -- Mean value of ``abs(spectrum)**2`` in each circular annulus. .. admonition:: Notes No additional FFT normalisation, mean removal, mask correction or apodization is applied. These choices belong to the spectral-analysis workflow and should be performed before calling this method when required. .. py:method:: extra_repr() -> str Return the extra representation of the module. To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable. .. py:function:: fft(cell_ids: ArrayLike, level: int, data: ArrayLike, *, return_transform: bool = False, **transform_kwargs) Build a :class:`LocalFFT` and compute a local 2-D FFT. Use ``return_transform=True`` to return ``(spectrum, transform)``. Reusing that transform for IFFT and additional maps avoids rebuilding geometry. .. py:function:: ifft(spectrum: ArrayLike, transform: Optional[LocalFFT] = None, *, cell_ids: Optional[ArrayLike] = None, level: Optional[int] = None, real_output: bool = True, **transform_kwargs) -> ArrayLike Compute a local inverse FFT using cached or newly built geometry. Prefer passing the :class:`LocalFFT` returned by :func:`fft`. Alternatively provide both ``cell_ids`` and ``level`` to rebuild the deterministic grid. .. py:function:: ps(spectrum: ArrayLike, transform: Optional[LocalFFT] = None, *, pixel_size_rad: Optional[float] = None, n_bins: Optional[int] = None, include_corners: bool = False) -> tuple[ArrayLike, ArrayLike] Compute a circular 1-D power spectrum from a square unshifted 2-D FFT. Pass a :class:`LocalFFT` to use its grid spacing and shape validation. For a standalone Fourier array, provide ``pixel_size_rad``; it defaults to 1, in which case frequencies are expressed in cycles per grid-spacing unit.