healpix_analyse.fft_conv
========================

.. py:module:: healpix_analyse.fft_conv

.. autoapi-nested-parse::

   Fast large-kernel convolution on local HEALPix patches.

   The input cells are projected with :class:`healpix_analyse.fft_local.LocalFFT`,
   convolved on its gnomonic grid with a zero-padded 2-D FFT, then sampled back at
   the original HEALPix cell centres.  All data-path operations are implemented
   with PyTorch and remain differentiable.



Classes
-------

.. autoapisummary::

   healpix_analyse.fft_conv.HealPixFFTConv


Module Contents
---------------

.. py:class:: HealPixFFTConv(level: int, in_channels: int, out_channels: int, kernel_sz: int = 33, *, cell_ids, ellipsoid: str = 'sphere', max_patch_radius_deg: float = 10.0, pixel_size_rad: Optional[float] = None, grid_size: Optional[int] = None, use_norm: bool = False, dtype: torch.dtype = torch.float32, device: Optional[Union[str, torch.device]] = None, cache_kernel_fft: bool = True)

   Bases: :py:obj:`torch.nn.Module`


   Large learned convolution on a local NESTED HEALPix patch.

   :Parameters: * **level** (:py:class:`int`) -- Grid4Earth/HEALPix level.  Internally, ``nside = 2**level``.
                * **in_channels, out_channels** (:py:class:`int`) -- Input and output feature channels.
                * **kernel_sz** (:py:class:`int`, *default* ``33``) -- Positive odd side length of the learned 2-D tangent-plane kernel.
                * **cell_ids** (:py:class:`array-like`) -- Unique NESTED cells forming one local patch.  A single gnomonic plane
                  cannot represent a full sphere.
                * **ellipsoid** (:py:class:`str`, *default* ``"sphere"``) -- Geometry passed to :class:`LocalFFT`.
                * **max_patch_radius_deg** (:py:class:`float`, *default* ``10``) -- Reject patches extending farther from their spherical centre.
                * **pixel_size_rad, grid_size** (**optional**) -- Projected-grid controls forwarded to :class:`LocalFFT`.
                * **use_norm** (:py:class:`bool`, *default* :py:obj:`False`) -- Apply GroupNorm followed by ReLU after back-projection.
                * **dtype, device** (**optional**) -- PyTorch execution options.  CUDA is selected when available.
                * **cache_kernel_fft** (:py:class:`bool`, *default* :py:obj:`True`) -- Cache the transformed kernel during no-gradient inference.  Training
                  always recomputes it so gradients reach the spatial kernel.

   .. admonition:: Notes

      The convolution is linear rather than circular: both the projected data
      and kernel are zero-padded to at least ``grid_size + kernel_sz - 1`` in
      each direction.  This prevents values wrapping across opposite sides of
      the local patch.  Values near the patch boundary still see zero outside
      the supplied domain; provide a halo and crop it when boundary accuracy is
      important.
      
      The learned weight follows the existing ``HealPixConv`` channel order and
      has shape ``[C_in, C_out, kernel_sz, kernel_sz]``.


   .. py:attribute:: level


   .. py:attribute:: nside


   .. py:attribute:: in_channels


   .. py:attribute:: out_channels


   .. py:attribute:: kernel_sz
      :value: 33



   .. py:attribute:: use_norm
      :value: False



   .. py:attribute:: dtype
      :value: Ellipsis



   .. py:attribute:: cache_kernel_fft
      :value: True



   .. py:attribute:: transform


   .. py:attribute:: n_cells


   .. py:attribute:: grid_size
      :value: 0



   .. py:attribute:: fft_shape


   .. py:attribute:: weight


   .. py:attribute:: bias


   .. py:property:: device
      :type: torch.device


      Current device of weights and projection buffers.


   .. py:property:: cell_ids
      :type: torch.Tensor


      NESTED cell identifiers in input/output order.


   .. py:method:: train(mode: bool = True)

      Set the module in training mode.

      This has an effect only on certain modules. See the documentation of
      particular modules for details of their behaviors in training/evaluation
      mode, i.e., whether they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
      etc.

      Args:
          mode (bool): whether to set training mode (``True``) or evaluation
                       mode (``False``). Default: ``True``.

      Returns:
          Module: self



   .. py:method:: set_kernel(weight, bias=None, requires_grad: bool = False)

      Set a spatial kernel using ``[C_in,C_out,K,K]`` or flattened form.



   .. py:method:: forward(x: ArrayLike) -> ArrayLike

      Apply the FFT convolution to ``[N]``, ``[B,N]`` or ``[B,C,N]``.



   .. 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.