foscat.unet_2_d_from_healpix_params
===================================

.. py:module:: foscat.unet_2_d_from_healpix_params


Attributes
----------

.. autoapisummary::

   foscat.unet_2_d_from_healpix_params.nside


Classes
-------

.. autoapisummary::

   foscat.unet_2_d_from_healpix_params.PlanarUNet


Functions
---------

.. autoapisummary::

   foscat.unet_2_d_from_healpix_params.fit


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

.. py:class:: PlanarUNet(*, in_nside: int, n_chan_in: int, chanlist: List[int], KERNELSZ: int = 3, task: Literal['regression', 'segmentation'] = 'regression', out_channels: int = 1, final_activation: Optional[Literal['none', 'sigmoid', 'softmax']] = None, device: Optional[torch.device | str] = None, down_type: Optional[Literal['mean', 'max']] = 'max', dtype: Literal['float32', 'float64'] = 'float32', head_reduce: Literal['mean', 'learned'] = 'mean')

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


   U-Net 2D (images HxW) mirroring the parameterization of the HealpixUNet.

   Key compat points with HealpixUNet:
     - Same constructor fields: in_nside, n_chan_in, chanlist, KERNELSZ, task,
       out_channels, final_activation, device, down_type, dtype, head_reduce.
     - Two convs per level (encoder & decoder), GroupNorm + ReLU after each conv.
     - Downsampling by factor 2 at each level; upsampling mirrors back.
     - Head produces `out_channels` with optional BN and final activation.

   Differences vs sphere version:
     - Operates on regular 2D images of size (3*in_nside, 4*in_nside).
     - Standard Conv2d instead of custom spherical stencil.
     - No gauges (G=1 implicit) and no cell_ids.

   Shapes
   ------
   Input  : (B, C_in,  3*in_nside, 4*in_nside)
   Output : (B, C_out, 3*in_nside, 4*in_nside)

   Constraints
   -----------
   `in_nside` must be divisible by 2**depth, where depth == len(chanlist).


   .. py:attribute:: in_nside


   .. py:attribute:: n_chan_in


   .. py:attribute:: chanlist


   .. py:attribute:: KERNELSZ
      :value: 3



   .. py:attribute:: task
      :value: 'regression'



   .. py:attribute:: out_channels
      :value: 1



   .. py:attribute:: down_type
      :value: 'max'



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



   .. py:attribute:: head_reduce
      :value: 'mean'



   .. py:attribute:: device


   .. py:attribute:: skips_channels
      :type:  List[int]
      :value: []



   .. py:attribute:: encoder


   .. py:attribute:: upconvs


   .. py:attribute:: decoder


   .. py:attribute:: head_conv


   .. py:attribute:: head_bn


   .. py:attribute:: head_mixer


   .. py:method:: to_tensor(x)


   .. py:method:: to_numpy(x)


   .. py:method:: forward(x: torch.Tensor) -> torch.Tensor

      x: (B, C_in, H, W) with H=3*in_nside, W=4*in_nside



   .. py:method:: predict(x: torch.Tensor, batch_size: int = 8) -> torch.Tensor


.. py:function:: fit(model: torch.nn.Module, x_train: Union[torch.Tensor, numpy.ndarray], y_train: Union[torch.Tensor, numpy.ndarray], *, n_epoch: int = 10, view_epoch: int = 10, batch_size: int = 16, lr: float = 0.001, weight_decay: float = 0.0, clip_grad_norm: Optional[float] = None, verbose: bool = True, optimizer: Literal['ADAM', 'LBFGS'] = 'ADAM') -> dict

   Training loop mirroring `healpix_unet_torch.fit`, adapted for 2D images.

   - Fixed inputs: tensors/ndarrays of the same size (B, C, H, W) with H=3*nside, W=4*nside
   - Perte: MSE (regression) / BCE(BCEWithLogits si final_activation='none') / CrossEntropy (multiclasses)
   - Optimiseur: ADAM ou LBFGS avec closure
   - Logs: renvoie {"loss": history}


.. py:data:: nside
   :value: 32