Unsupervised Imputers
=====================

Overview
--------

PG-SUI ships four neural imputers—NLPCA, UBP, Autoencoder, and VAE—that share the :class:`pgsui.impute.unsupervised.base.BaseNNImputer` scaffolding. They operate on 0/1/2 encodings generated by :class:`snpio.analysis.genotype_encoder.GenotypeEncoder`, learn to reconstruct genotypes with missing calls masked as ``-1``, and expose a ``transform()`` method that returns IUPAC strings.

Workflow Highlights
-------------------

- Instantiate with a ``GenotypeData`` object and a typed ``*Config`` (dataclass instance, nested ``dict``, or YAML path). Dot-key overrides (``{"model.latent_dim": 12}``) are applied last.
- ``fit()`` standardises the data split, computes class weights from observed zygosity, optionally tunes hyperparameters with Optuna, then trains the model while caching metrics/plots under ``{prefix}_output/Unsupervised/``.
- ``transform()`` reuses the trained model to impute the full dataset, decodes to IUPAC, and writes genotype-distribution plots for original vs. imputed calls.

Configuration
-------------

Each imputer exposes the same top-level sections: ``io`` (run identity + logging), ``model`` (architecture), ``train`` (optimizer schedule, class weighting, early stopping), ``tune`` (Optuna surface), ``evaluate`` (latent refinement controls), and ``plot`` (figure styling). Presets (``fast``, ``balanced``, ``thorough``) trade accuracy for runtime, and YAML configs may declare a preset before overriding individual fields. Heavy class imbalance can be tempered by ``train.weights_beta`` / ``train.weights_max_ratio``; layer widths are governed by ``model.layer_schedule`` and ``model.layer_scaling_factor``.

.. tip::

   The neural imputers honour ``cfg.io.prefix`` by creating ``models/``, ``plots/``, ``metrics/``, ``optimize/``, and ``parameters/`` subdirectories, mirroring the supervised toolkit.

Config Dataclasses
------------------

Non-linear PCA (Config)
^^^^^^^^^^^^^^^^^^^^^^^

Captures latent-dimension, decoder depth, and latent-optimisation knobs (``model.latent_init``, ``train.lr_input_factor``) used by :class:`pgsui.impute.unsupervised.imputers.nlpca.ImputeNLPCA`.

.. autoclass:: pgsui.data_processing.containers.NLPCAConfig
   :members:
   :noindex:

Standard Autoencoder (Config)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

A standard autoencoder config capturing architecture and training settings for :class:`pgsui.impute.unsupervised.imputers.autoencoder.ImputeAutoencoder`.

.. autoclass:: pgsui.data_processing.containers.AutoencoderConfig
   :members:
   :noindex:

Unsupervised Backpropagation (Config)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

Extends the NLPCA options with UBP-specific presets tuned for class imbalance and latent re-optimisation.

.. autoclass:: pgsui.data_processing.containers.UBPConfig
   :members:
   :noindex:

Variational Autoencoder (Config)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

Adds a ``vae`` section (``kl_beta``, ``kl_warmup``, ``kl_ramp``) to the autoencoder defaults, controlling KL annealing.

.. autoclass:: pgsui.data_processing.containers.VAEConfig
   :members:
   :noindex:

Model Summaries
---------------

Non-linear PCA (ImputeNLPCA)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^

- Encodes genotypes to 0/1/2 (missing → -1) and detects haploid panels, collapsing ALT/HET where appropriate.
- Supports latent initialisation via ``model.latent_init`` (``"random"`` or ``"pca"``) and keeps a learnable latent vector per sample.
- Trains with focal cross-entropy (``model.gamma``) while jointly optimising decoder weights and latent vectors; ``train.lr_input_factor`` scales the latent optimiser.
- ``evaluate.eval_latent_steps`` / ``eval_latent_lr`` govern latent refinement during validation and inference so reconstruction metrics use re-optimised latents.
- Optional Optuna tuning uses the ``tune`` envelope (dataset subsampling, patience, metric selection) before the final training run.

Unsupervised Backpropagation (ImputeUBP)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

- Implements the three-phase schedule from Gashler et al. (decoder warm-up, decoder fine-tune, joint training) with shared focal loss and class-weighting.
- Initial latent vectors come from ``model.latent_init`` and are re-used across phases; ``train.lr_input_factor`` controls the latent optimiser's step size.
- Class weighting is capped by ``train.weights_max_ratio`` to prevent extreme rebalancing on sparse genotypes.
- ``transform()`` re-optimises latent vectors for the full cohort before predicting, then saves IUPAC outputs plus genotype distribution plots.
- Tunable through Optuna with the same search surface as NLPCA (``tune.*`` fields).

Standard Autoencoder (ImputeAutoencoder)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

- Feed-forward encoder/decoder with symmetric hidden layers (``model.layer_schedule``) and dropout; no latent refinement is performed during evaluation.
- Uses focal cross-entropy weighted by observed zygosity (``train.weights_beta``) and early stopping controlled by ``train.early_stop_gen`` / ``train.min_epochs``.
- Optional Optuna tuning searches over latent dimension, dropout, and learning rate while respecting ``tune.max_samples`` / ``tune.max_loci`` caps.
- ``transform()`` applies the trained model to reconstructed logits, fills only previously missing calls, decodes to IUPAC, and plots original vs. imputed counts.

Variational Autoencoder (ImputeVAE)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

- Extends the autoencoder with a KL divergence term whose strength is annealed via ``vae.kl_warmup`` / ``vae.kl_ramp`` to ``vae.kl_beta``.
- Shares the focal cross-entropy schedule with the standard autoencoder (``model.gamma`` ramp) and records the annealed ``beta``/``gamma`` per epoch.
- Hyperparameter tuning reuses the autoencoder search surface while keeping the latent evaluator disabled (VAEs rely on decoder logits during validation).
- ``transform()`` predicts class probabilities across genotypes, fills masked entries with MAP labels, and emits IUPAC arrays with paired distribution plots.

Usage Examples
--------------

.. code-block:: python

   from snpio import VCFReader
   from pgsui import ImputeVAE
   from pgsui.data_processing.containers import VAEConfig

   gdata = VCFReader("cohort.vcf.gz", popmapfile="pops.popmap")

   cfg = VAEConfig.from_preset("balanced")
   cfg.io.prefix = "vae_demo"
   cfg.vae.kl_beta = 1.5
   cfg.tune.enabled = False  # accept preset defaults

   imputer = ImputeVAE(genotype_data=gdata, config=cfg)
   imputer.fit()
   genotypes_iupac = imputer.transform()

CLI usage mirrors the Python API:

.. code-block:: bash

   pg-sui \
      --vcf cohort.vcf.gz \
      --models ImputeNLPCA ImputeUBP \
      --preset thorough \
      --set io.prefix=nlpca_vs_ubp \
      --config configs/nlpca.yaml \
      --sim-strategy random_weighted \
      --sim-prop 0.20

``--sim-strategy`` and ``--sim-prop`` apply to every selected neural model, ensuring the same simulated-missingness mask is used throughout evaluation. Pass ``--simulate-missing`` if you want to rely solely on organically missing calls for a diagnostic run. See :ref:`simulated_missingness` for how each strategy behaves.
