pg-sui 1.0.2.1__py3-none-any.whl → 1.6.8__py3-none-any.whl

pgsui/impute/supervised/base.py

@@ -0,0 +1,339 @@
+import json
+from pathlib import Path
+from typing import Any, Dict, List
+
+import matplotlib.pyplot as plt
+import numpy as np
+from sklearn.exceptions import NotFittedError
+from sklearn.experimental import enable_iterative_imputer  # noqa
+from sklearn.metrics import (
+    accuracy_score,
+    classification_report,
+    f1_score,
+    precision_score,
+    recall_score,
+)
+from snpio.utils.logging import LoggerManager
+
+from pgsui.utils.classification_viz import ClassificationReportVisualizer
+
+
+class BaseImputer:
+    """A base class for supervised, iterative imputer models.
+
+    This class provides a common framework and shared functionality for imputers that use scikit-learn's `IterativeImputer`. It is not intended for direct instantiation. Child classes should inherit from this class and provide a specific estimator model (e.g., RandomForest, GradientBoosting).
+
+    Notes:
+        - A hyperparameter tuning workflow using Optuna.
+        - Standardized data splitting, model training, and evaluation methods.
+        - Utilities for creating output directories and handling model state.
+    """
+
+    def __init__(self, verbose: bool = False, debug: bool = False) -> None:
+        """Initializes the BaseImputer class.
+
+        This class sets up logging and verbosity/debug settings. It also contains methods that all supervised imputers will share.
+
+        Note:
+            Inheriting child classes must define `self.prefix` before calling `super().__init__()`, as it is required for logger initialization.
+
+        Args:
+            verbose (bool): If True, enables detailed logging output. Defaults to False.
+            debug (bool): If True, enables debug mode. Defaults to False.
+        """
+        self.verbose = verbose
+        self.debug = debug
+
+        logman = LoggerManager(
+            __name__, prefix=self.prefix, verbose=self.verbose, debug=self.debug
+        )
+        self.logger = logman.get_logger()
+
+    def _create_model_directories(self, prefix: str, outdirs: List[str]) -> None:
+        """Creates the output directory structure for the imputer.
+
+        This method sets up a standardized folder hierarchy for saving models, plots, metrics, and optimization results, organized by the model's name.
+
+        Args:
+            prefix (str): The prefix for the main output directory.
+            outdirs (List[str]): A list of subdirectories to create (e.g., 'models', 'plots').
+        """
+        base_dir = Path(f"{prefix}_output") / "Supervised"
+        for d in outdirs:
+            subdir = base_dir / d / self.model_name
+            setattr(self, f"{d}_dir", subdir)
+            subdir.mkdir(parents=True, exist_ok=True)
+
+    def _make_class_reports(
+        self,
+        y_true: np.ndarray,
+        y_pred: np.ndarray,
+        metrics: Dict[str, float],
+        y_pred_proba: np.ndarray | None = None,
+        labels: List[str] = ["REF", "HET", "ALT"],
+    ) -> None:
+        """Generate and save detailed classification reports and visualizations.
+
+        3-class (zygosity) or 10-class (IUPAC) depending on `labels` length.
+
+        Args:
+            y_true (np.ndarray): True labels (1D array).
+            y_pred (np.ndarray): Predicted labels (1D array).
+            metrics (Dict[str, float]): Computed metrics.
+            y_pred_proba (np.ndarray | None): Predicted probabilities (2D array). Defaults to None.
+            labels (List[str], optional): Class label names
+                (default: ["REF", "HET", "ALT"] for 3-class).
+        """
+        report_name = "zygosity" if len(labels) == 3 else "iupac"
+        middle = "IUPAC" if report_name == "iupac" else "Zygosity"
+
+        msg = f"{middle} Report (on {y_true.size} total genotypes)"
+        self.logger.info(msg)
+
+        if y_pred_proba is not None:
+            self.plotter_.plot_metrics(
+                y_true,
+                y_pred_proba,
+                metrics,
+                label_names=labels,
+                prefix=report_name,
+            )
+
+        self.plotter_.plot_confusion_matrix(
+            y_true, y_pred, label_names=labels, prefix=report_name
+        )
+
+        self.logger.info(
+            "\n"
+            + classification_report(
+                y_true,
+                y_pred,
+                labels=list(range(len(labels))),
+                target_names=labels,
+                zero_division=0,
+            )
+        )
+
+        report = classification_report(
+            y_true,
+            y_pred,
+            labels=list(range(len(labels))),
+            target_names=labels,
+            zero_division=0,
+            output_dict=True,
+        )
+
+        with open(self.metrics_dir / f"{report_name}_report.json", "w") as f:
+            json.dump(report, f, indent=4)
+
+        viz = ClassificationReportVisualizer(reset_kwargs=self.plotter_.param_dict)
+
+        plots = viz.plot_all(
+            report,
+            title_prefix=f"{self.model_name} {middle} Report",
+            show=getattr(self, "show_plots", False),
+            heatmap_classes_only=True,
+        )
+
+        for name, fig in plots.items():
+            fout = self.plots_dir / f"{report_name}_report_{name}.{self.plot_format}"
+            if hasattr(fig, "savefig"):
+                fig.savefig(fout, dpi=300, facecolor="#111122")
+                plt.close(fig)
+            else:
+                fig.write_html(file=fout.with_suffix(".html"))
+
+        viz._reset_mpl_style()
+
+    def _evaluate_012_and_plot(self, y_true: np.ndarray, y_pred: np.ndarray) -> None:
+        """0/1/2 zygosity report & confusion matrix.
+
+        This method generates a classification report and confusion matrix for genotypes encoded as 0, 1, or 2. If the data is haploid, it treats genotypes 1 and 2 as equivalent (presence of the alternate allele).
+
+        Args:
+            y_true (np.ndarray): True genotypes (0/1/2) for masked
+            y_pred (np.ndarray): Predicted genotypes (0/1/2) for masked
+
+        Raises:
+            NotFittedError: If fit() and transform() have not been called.
+        """
+        labels = [0, 1, 2]
+        # Haploid parity: fold ALT (2) into ALT/Present (1)
+        if self.is_haploid_:
+            y_true[y_true == 2] = 1
+            y_pred[y_pred == 2] = 1
+            labels = [0, 1]
+
+        metrics = {
+            "n_masked_test": int(y_true.size),
+            "accuracy": accuracy_score(y_true, y_pred),
+            "f1": f1_score(
+                y_true, y_pred, average="macro", labels=labels, zero_division=0
+            ),
+            "precision": precision_score(
+                y_true, y_pred, average="macro", labels=labels, zero_division=0
+            ),
+            "recall": recall_score(
+                y_true, y_pred, average="macro", labels=labels, zero_division=0
+            ),
+        }
+
+        metrics.update({f"zygosity_{k}": v for k, v in metrics.items()})
+
+        report_names = ["REF", "HET"] if self.is_haploid_ else ["REF", "HET", "ALT"]
+
+        self.logger.info(
+            f"\n{classification_report(y_true, y_pred, labels=labels, target_names=report_names, zero_division=0)}"
+        )
+
+        report = classification_report(
+            y_true,
+            y_pred,
+            labels=labels,
+            target_names=report_names,
+            zero_division=0,
+            output_dict=True,
+        )
+
+        viz = ClassificationReportVisualizer(reset_kwargs=self.plotter_.param_dict)
+
+        plots = viz.plot_all(
+            report,
+            title_prefix=f"{self.model_name} Zygosity Report",
+            show=getattr(self, "show_plots", False),
+            heatmap_classes_only=True,
+        )
+
+        for name, fig in plots.items():
+            fout = self.plots_dir / f"zygosity_report_{name}.{self.plot_format}"
+            if hasattr(fig, "savefig"):
+                fig.savefig(fout, dpi=300, facecolor="#111122")
+                plt.close(fig)
+            else:
+                fig.write_html(file=fout.with_suffix(".html"))
+
+        viz._reset_mpl_style()
+
+        # Save JSON
+        self._save_report(report, suffix="zygosity")
+
+        # Confusion matrix
+        self.plotter_.plot_confusion_matrix(
+            y_true, y_pred, label_names=report_names, prefix="zygosity"
+        )
+
+    def _evaluate_iupac10_and_plot(
+        self, y_true: np.ndarray, y_pred: np.ndarray
+    ) -> None:
+        """10-class IUPAC report & confusion matrix.
+
+        This method generates a classification report and confusion matrix for genotypes encoded using the 10 IUPAC codes (0-9). The IUPAC codes represent various nucleotide combinations, including ambiguous bases.
+
+        Args:
+            y_true (np.ndarray): True genotypes (0-9) for masked
+            y_pred (np.ndarray): Predicted genotypes (0-9) for masked
+
+        Raises:
+            NotFittedError: If fit() and transform() have not been called.
+        """
+        labels_idx = list(range(10))
+        labels_names = ["A", "C", "G", "T", "W", "R", "M", "K", "Y", "S"]
+
+        metrics = {
+            "accuracy": accuracy_score(y_true, y_pred),
+            "f1": f1_score(
+                y_true, y_pred, average="macro", labels=labels_idx, zero_division=0
+            ),
+            "precision": precision_score(
+                y_true, y_pred, average="macro", labels=labels_idx, zero_division=0
+            ),
+            "recall": recall_score(
+                y_true, y_pred, average="macro", labels=labels_idx, zero_division=0
+            ),
+        }
+        metrics.update({f"iupac_{k}": v for k, v in metrics.items()})
+
+        self.logger.info(
+            f"\n{classification_report(y_true, y_pred, labels=labels_idx, target_names=labels_names, zero_division=0)}"
+        )
+
+        report = classification_report(
+            y_true,
+            y_pred,
+            labels=labels_idx,
+            target_names=labels_names,
+            zero_division=0,
+            output_dict=True,
+        )
+
+        viz = ClassificationReportVisualizer(reset_kwargs=self.plotter_.param_dict)
+
+        plots = viz.plot_all(
+            report,
+            title_prefix=f"{self.model_name} IUPAC Report",
+            show=getattr(self, "show_plots", False),
+            heatmap_classes_only=True,
+        )
+
+        # Reset the style from Optuna's plotting.
+        plt.rcParams.update(self.plotter_.param_dict)
+
+        for name, fig in plots.items():
+            fout = self.plots_dir / f"iupac_report_{name}.{self.plot_format}"
+            if hasattr(fig, "savefig"):
+                fig.savefig(fout, dpi=300, facecolor="#111122")
+                plt.close(fig)
+            else:
+                fig.write_html(file=fout.with_suffix(".html"))
+
+        # Reset the style
+        viz._reset_mpl_style()
+
+        # Save JSON
+        self._save_report(report, suffix="iupac")
+
+        # Confusion matrix
+        self.plotter_.plot_confusion_matrix(
+            y_true, y_pred, label_names=labels_names, prefix="iupac"
+        )
+
+    def _save_report(self, report_dict: Dict[str, float], suffix: str) -> None:
+        """Save classification report dictionary as a JSON file.
+
+        This method saves the provided classification report dictionary to a JSON file in the metrics directory, appending the specified suffix to the filename.
+
+        Args:
+            report_dict (Dict[str, float]): The classification report dictionary to save.
+            suffix (str): Suffix to append to the filename (e.g., 'zygosity' or 'iupac').
+
+        Raises:
+            NotFittedError: If fit() and transform() have not been called.
+        """
+        if not self.is_fit_:
+            msg = "No report to save. Ensure fit() has been called."
+            raise NotFittedError(msg)
+
+        out_fp = self.metrics_dir / f"classification_report_{suffix}.json"
+
+        with open(out_fp, "w") as f:
+            json.dump(report_dict, f, indent=4)
+
+        self.logger.info(f"{self.model_name} {suffix} report saved to {out_fp}.")
+
+    def _save_best_params(self, best_params: Dict[str, Any]) -> None:
+        """Save the best hyperparameters to a JSON file.
+
+        This method saves the best hyperparameters found during hyperparameter tuning to a JSON file in the optimization directory. The filename includes the model name for easy identification.
+
+        Args:
+            best_params (Dict[str, Any]): A dictionary of the best hyperparameters to save.
+        """
+        if not hasattr(self, "parameters_dir"):
+            msg = "Attribute 'parameters_dir' not found. Ensure _create_model_directories() has been called."
+            self.logger.error(msg)
+            raise AttributeError(msg)
+
+        fout = self.parameters_dir / "best_parameters.json"
+
+        with open(fout, "w") as f:
+            json.dump(best_params, f, indent=4)

pgsui/impute/supervised/imputers/hist_gradient_boosting.py

@@ -0,0 +1,293 @@
+# Standard library
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Dict, Generator, List
+
+# Third-party
+import numpy as np
+from sklearn.ensemble import HistGradientBoostingClassifier
+from sklearn.exceptions import NotFittedError
+from sklearn.experimental import enable_iterative_imputer  # noqa
+from sklearn.impute import IterativeImputer
+from sklearn.model_selection import train_test_split
+
+# Project
+from snpio.analysis.genotype_encoder import GenotypeEncoder
+from snpio.utils.logging import LoggerManager
+
+from pgsui.data_processing.config import apply_dot_overrides, load_yaml_to_dataclass
+from pgsui.data_processing.containers import (
+    HGBConfig,
+    _HGBParams,
+    _ImputerParams,
+    _SimParams,
+)
+from pgsui.data_processing.transformers import SimGenotypeDataTransformer
+from pgsui.impute.supervised.base import BaseImputer
+from pgsui.utils.plotting import Plotting
+from pgsui.utils.scorers import Scorer
+
+if TYPE_CHECKING:
+    from snpio.read_input.genotype_data import GenotypeData
+
+
+def ensure_hgb_config(config: HGBConfig | Dict | str | None) -> HGBConfig:
+    """Resolve HGB configuration from dataclass, mapping, or YAML path."""
+
+    if config is None:
+        return HGBConfig()
+    if isinstance(config, HGBConfig):
+        return config
+    if isinstance(config, str):
+        return load_yaml_to_dataclass(
+            config, HGBConfig, preset_builder=HGBConfig.from_preset
+        )
+    if isinstance(config, dict):
+        payload = dict(config)
+        preset = payload.pop("preset", None)
+        base = HGBConfig.from_preset(preset) if preset else HGBConfig()
+
+        def _flatten(prefix: str, data: Dict[str, Any], out: Dict[str, Any]) -> None:
+            for key, value in data.items():
+                dotted = f"{prefix}.{key}" if prefix else key
+                if isinstance(value, dict):
+                    _flatten(dotted, value, out)
+                else:
+                    out[dotted] = value
+
+        flat: Dict[str, Any] = {}
+        _flatten("", payload, flat)
+        return apply_dot_overrides(base, flat)
+
+    raise TypeError("config must be an HGBConfig, dict, YAML path, or None.")
+
+
+class ImputeHistGradientBoosting(BaseImputer):
+    """Supervised HGB imputer driven by :class:`HGBConfig`."""
+
+    def __init__(
+        self,
+        genotype_data: "GenotypeData",
+        *,
+        config: HGBConfig | Dict | str | None = None,
+        overrides: Dict | None = None,
+    ) -> None:
+        self.model_name = "ImputeHistGradientBoosting"
+        self.Model = HistGradientBoostingClassifier
+
+        cfg = ensure_hgb_config(config)
+        if overrides:
+            cfg = cfg.apply_overrides(overrides)
+        self.cfg = cfg
+
+        self.genotype_data = genotype_data
+        self.pgenc = GenotypeEncoder(genotype_data)
+
+        self.prefix = cfg.io.prefix
+        self.seed = cfg.io.seed
+        self.n_jobs = cfg.io.n_jobs
+        self.verbose = cfg.io.verbose
+        self.debug = cfg.io.debug
+
+        super().__init__(verbose=self.verbose, debug=self.debug)
+
+        logman = LoggerManager(
+            __name__, prefix=self.prefix, verbose=self.verbose, debug=self.debug
+        )
+        self.logger = logman.get_logger()
+
+        self._create_model_directories(
+            self.prefix, ["models", "plots", "metrics", "optimize", "parameters"]
+        )
+
+        self.plot_format = cfg.plot.fmt
+        if self.plot_format.startswith("."):
+            self.plot_format = self.plot_format.lstrip(".")
+        self.plot_fontsize = cfg.plot.fontsize
+        self.title_fontsize = cfg.plot.fontsize
+        self.plot_dpi = cfg.plot.dpi
+        self.despine = cfg.plot.despine
+        self.show_plots = cfg.plot.show
+
+        self.validation_split = cfg.train.validation_split
+
+        class_weight = getattr(cfg.model, "class_weight", "balanced")
+        self.params = _HGBParams(
+            max_iter=cfg.model.n_estimators,
+            learning_rate=cfg.model.learning_rate,
+            max_depth=cfg.model.max_depth,
+            min_samples_leaf=cfg.model.min_samples_leaf,
+            max_features=cfg.model.max_features,
+            n_iter_no_change=cfg.model.n_iter_no_change,
+            tol=cfg.model.tol,
+            class_weight=class_weight,
+            random_state=self.seed,
+            verbose=self.debug,
+        )
+
+        self.imputer_params = _ImputerParams(
+            n_nearest_features=cfg.imputer.n_nearest_features,
+            max_iter=cfg.imputer.max_iter,
+            random_state=self.seed,
+            verbose=self.verbose,
+        )
+
+        self.sim_params = _SimParams(
+            prop_missing=cfg.sim.prop_missing,
+            strategy=cfg.sim.strategy,
+            missing_val=cfg.sim.missing_val,
+            het_boost=cfg.sim.het_boost,
+            seed=self.seed,
+        )
+
+        self.max_iter = cfg.imputer.max_iter
+        self.n_nearest_features = cfg.imputer.n_nearest_features
+
+        # Will be set in fit()
+        self.is_haploid_: bool | None = None
+        self.num_classes_: int | None = None
+        self.num_features_: int | None = None
+        self.models_: List[HistGradientBoostingClassifier | None] | None = None
+        self.is_fit_: bool = False
+
+    def fit(self) -> "BaseImputer":
+        """Fit the imputer using self.genotype_data with no arguments.
+
+        This method prepares the imputer by splitting the data into training and testing sets, and masking all originally observed genotype entries in the test set to facilitate unbiased evaluation. It does not perform any actual imputation since the RefAllele imputer is deterministic.
+
+        Steps:
+            1) Encode to 0/1/2 with -9/-1 as missing.
+            2) Split samples into train/test.
+            3) Train IterativeImputer on train (convert missing -> NaN).
+            4) Evaluate on test **non-missing positions** (reconstruction metrics) and call your original plotting stack via _make_class_reports().
+
+        Returns:
+            BaseImputer: self.
+        """
+        # Prepare utilities & metadata
+        self.scorers_ = Scorer(
+            prefix=self.prefix, average="macro", verbose=self.verbose, debug=self.debug
+        )
+
+        self.plotter_ = Plotting(
+            self.model_name,
+            prefix=self.prefix,
+            plot_format=self.plot_format,
+            plot_dpi=self.plot_dpi,
+            plot_fontsize=self.plot_fontsize,
+            title_fontsize=self.title_fontsize,
+            despine=self.despine,
+            show_plots=self.show_plots,
+            verbose=self.verbose,
+            debug=self.debug,
+        )
+
+        X_int = self.pgenc.genotypes_012
+        self.X012_ = X_int.astype(float)
+        self.X012_[self.X012_ < 0] = np.nan  # Ensure missing are NaN
+        self.is_haploid_ = np.count_nonzero(self.X012_ == 1) == 0
+        self.num_classes_ = 2 if self.is_haploid_ else 3
+        self.n_samples_, self.n_features_ = X_int.shape
+
+        # Split
+        X_train, X_test = train_test_split(
+            self.X012_,
+            test_size=self.validation_split,
+            random_state=self.seed,
+            shuffle=True,
+        )
+
+        # Simulate missing values on test set.
+        sim_transformer = SimGenotypeDataTransformer(**self.sim_params.to_dict())
+
+        X_test = np.nan_to_num(X_test, nan=-1)  # ensure missing are -1
+        sim_transformer.fit(X_test)
+        X_test_sim, missing_masks = sim_transformer.transform(X_test)
+        sim_mask = missing_masks["simulated"]
+        X_test_sim[X_test_sim < 0] = np.nan  # ensure missing are NaN
+
+        self.model_params_ = self.params.to_dict()
+        self.model_params_["random_state"] = self.seed
+
+        # Train IterativeImputer
+        est = self.Model(**self.model_params_)
+
+        self.imputer_ = IterativeImputer(estimator=est, **self.imputer_params.to_dict())
+
+        self.imputer_.fit(X_train)
+        self.is_fit_ = True
+
+        X_test_imputed = self.imputer_.transform(X_test_sim)
+
+        # Predict on simulated test set
+        y_true_flat = X_test[sim_mask].copy()
+        y_pred_flat = X_test_imputed[sim_mask].copy()
+
+        # Round and clip predictions to valid {0,1,2} or {0,1} if haploid.
+        if self.is_haploid_:
+            y_pred_flat = np.clip(np.rint(y_pred_flat), 0, 1).astype(int, copy=False)
+            y_true_flat = np.clip(np.rint(y_true_flat), 0, 1).astype(int, copy=False)
+        else:
+            y_pred_flat = np.clip(np.rint(y_pred_flat), 0, 2).astype(int, copy=False)
+            y_true_flat = np.clip(np.rint(y_true_flat), 0, 2).astype(int, copy=False)
+
+        # Evaluate (012 / zygosity)
+        self._evaluate_012_and_plot(y_true_flat.copy(), y_pred_flat.copy())
+
+        # Evaluate (IUPAC)
+        encodings_dict = {
+            "A": 0,
+            "C": 1,
+            "G": 2,
+            "T": 3,
+            "W": 4,
+            "R": 5,
+            "M": 6,
+            "K": 7,
+            "Y": 8,
+            "S": 9,
+            "N": -1,
+        }
+
+        y_true_iupac_tmp = self.pgenc.decode_012(y_true_flat)
+        y_pred_iupac_tmp = self.pgenc.decode_012(y_pred_flat)
+        y_true_iupac = self.pgenc.convert_int_iupac(
+            y_true_iupac_tmp, encodings_dict=encodings_dict
+        )
+        y_pred_iupac = self.pgenc.convert_int_iupac(
+            y_pred_iupac_tmp, encodings_dict=encodings_dict
+        )
+        self._evaluate_iupac10_and_plot(y_true_iupac, y_pred_iupac)
+
+        self.best_params_ = self.model_params_
+        self.best_params_.update(self.imputer_params.to_dict())
+        self.best_params_.update(self.sim_params.to_dict())
+        self._save_best_params(self.best_params_)
+
+        return self
+
+    def transform(self) -> np.ndarray:
+        """Impute all samples and return imputed genotypes.
+
+        This method applies the trained imputer to the entire dataset, filling in missing genotype values. It ensures that any remaining missing values after imputation are set to -9, and decodes the imputed 0/1/2 genotypes back to their original format.
+
+        Returns:
+            np.ndarray: (n_samples, n_loci) integers with no -9/-1/NaN.
+
+        Raises:
+            NotFittedError: If fit() has not been called prior to transform().
+        """
+        if not self.is_fit_:
+            msg = "Imputer has not been fit; call fit() before transform()."
+            self.logger.error(msg)
+            raise NotFittedError(msg)
+
+        X = self.X012_.copy()
+        X_imp = self.imputer_.transform(X)
+
+        if np.any(X_imp < 0) or np.isnan(X_imp).any():
+            self.logger.warning("Some imputed values are still missing; setting to -9.")
+            X_imp[X_imp < 0] = -9
+            X_imp[np.isnan(X_imp)] = -9
+
+        return self.pgenc.decode_012(X_imp)