smftools 0.2.3__py3-none-any.whl → 0.2.5__py3-none-any.whl

smftools/hmm/__init__.py

@@ -1,20 +1,13 @@
-from .apply_hmm_batched import apply_hmm_batched
-from .calculate_distances import calculate_distances
 from .call_hmm_peaks import call_hmm_peaks
 from .display_hmm import display_hmm
 from .hmm_readwrite import load_hmm, save_hmm
-from .nucleosome_hmm_refinement import refine_nucleosome_calls, infer_nucleosomes_in_large_bound
-from .train_hmm import train_hmm
-
+from .nucleosome_hmm_refinement import infer_nucleosomes_in_large_bound, refine_nucleosome_calls
 
 __all__ = [
-    "apply_hmm_batched",
-    "calculate_distances",
     "call_hmm_peaks",
     "display_hmm",
     "load_hmm",
     "refine_nucleosome_calls",
     "infer_nucleosomes_in_large_bound",
     "save_hmm",
-    "train_hmm"
-]
+]

smftools/hmm/archived/call_hmm_peaks.py

@@ -0,0 +1,121 @@
+def call_hmm_peaks(
+    adata,
+    feature_configs,
+    obs_column='Reference_strand',
+    site_types=['GpC_site', 'CpG_site'],
+    save_plot=False,
+    output_dir=None,
+    date_tag=None,
+    inplace=False
+):
+    """Call peaks from HMM feature layers and annotate AnnData.
+
+    Args:
+        adata: AnnData containing feature layers.
+        feature_configs: Mapping of layer name to peak config.
+        obs_column: Obs column for reference categories.
+        site_types: Site types to summarize around peaks.
+        save_plot: Whether to save peak plots.
+        output_dir: Output directory for plots.
+        date_tag: Optional tag for plot filenames.
+        inplace: Whether to modify AnnData in place.
+
+    Returns:
+        Annotated AnnData with peak masks and summary columns.
+    """
+    import numpy as np
+    import pandas as pd
+    import matplotlib.pyplot as plt
+    from scipy.signal import find_peaks
+
+    if not inplace:
+        adata = adata.copy()
+
+    # Ensure obs_column is categorical
+    if not isinstance(adata.obs[obs_column].dtype, pd.CategoricalDtype):
+        adata.obs[obs_column] = pd.Categorical(adata.obs[obs_column])
+
+    coordinates = adata.var_names.astype(int).values
+    peak_columns = []
+
+    obs_updates = {}
+
+    for feature_layer, config in feature_configs.items():
+        min_distance = config.get('min_distance', 200)
+        peak_width = config.get('peak_width', 200)
+        peak_prominence = config.get('peak_prominence', 0.2)
+        peak_threshold = config.get('peak_threshold', 0.8)
+
+        matrix = adata.layers[feature_layer]
+        means = np.mean(matrix, axis=0)
+        peak_indices, _ = find_peaks(means, prominence=peak_prominence, distance=min_distance)
+        peak_centers = coordinates[peak_indices]
+        adata.uns[f'{feature_layer} peak_centers'] = peak_centers.tolist()
+
+        # Plot
+        plt.figure(figsize=(6, 3))
+        plt.plot(coordinates, means)
+        plt.title(f"{feature_layer} with peak calls")
+        plt.xlabel("Genomic position")
+        plt.ylabel("Mean intensity")
+        for i, center in enumerate(peak_centers):
+            start, end = center - peak_width // 2, center + peak_width // 2
+            plt.axvspan(start, end, color='purple', alpha=0.2)
+            plt.axvline(center, color='red', linestyle='--')
+            aligned = [end if i % 2 else start, 'left' if i % 2 else 'right']
+            plt.text(aligned[0], 0, f"Peak {i}\n{center}", color='red', ha=aligned[1])
+        if save_plot and output_dir:
+            filename = f"{output_dir}/{date_tag or 'output'}_{feature_layer}_peaks.png"
+            plt.savefig(filename, bbox_inches='tight')
+            print(f"Saved plot to {filename}")
+        else:
+            plt.show()
+
+        feature_peak_columns = []
+        for center in peak_centers:
+            start, end = center - peak_width // 2, center + peak_width // 2
+            colname = f'{feature_layer}_peak_{center}'
+            peak_columns.append(colname)
+            feature_peak_columns.append(colname)
+
+            peak_mask = (coordinates >= start) & (coordinates <= end)
+            adata.var[colname] = peak_mask
+
+            region = matrix[:, peak_mask]
+            obs_updates[f'mean_{feature_layer}_around_{center}'] = np.mean(region, axis=1)
+            obs_updates[f'sum_{feature_layer}_around_{center}'] = np.sum(region, axis=1)
+            obs_updates[f'{feature_layer}_present_at_{center}'] = np.mean(region, axis=1) > peak_threshold
+
+            for site_type in site_types:
+                adata.obs[f'{site_type}_sum_around_{center}'] = 0
+                adata.obs[f'{site_type}_mean_around_{center}'] = np.nan
+
+            for ref in adata.obs[obs_column].cat.categories:
+                ref_idx = adata.obs[obs_column] == ref
+                mask_key = f"{ref}_{site_type}"
+                for site_type in site_types:
+                    if mask_key not in adata.var:
+                        continue
+                    site_mask = adata.var[mask_key].values
+                    site_coords = coordinates[site_mask]
+                    region_mask = (site_coords >= start) & (site_coords <= end)
+                    if not region_mask.any():
+                        continue
+                    full_mask = site_mask.copy()
+                    full_mask[site_mask] = region_mask
+                    site_region = adata[ref_idx, full_mask].X
+                    if hasattr(site_region, "A"):
+                        site_region = site_region.A
+                    if site_region.shape[1] > 0:
+                        adata.obs.loc[ref_idx, f'{site_type}_sum_around_{center}'] = np.nansum(site_region, axis=1)
+                        adata.obs.loc[ref_idx, f'{site_type}_mean_around_{center}'] = np.nanmean(site_region, axis=1)
+                    else:
+                        pass
+
+        adata.var[f'is_in_any_{feature_layer}_peak'] = adata.var[feature_peak_columns].any(axis=1)
+        print(f"Annotated {len(peak_centers)} peaks for {feature_layer}")
+
+    adata.var['is_in_any_peak'] = adata.var[peak_columns].any(axis=1)
+    adata.obs = pd.concat([adata.obs, pd.DataFrame(obs_updates, index=adata.obs.index)], axis=1)
+
+    return adata if not inplace else None

smftools/hmm/call_hmm_peaks.py

@@ -1,106 +1,314 @@
+# FILE: smftools/hmm/call_hmm_peaks.py
+
+from pathlib import Path
+from typing import Any, Dict, Optional, Sequence, Union
+
+from smftools.logging_utils import get_logger
+
+logger = get_logger(__name__)
+
+
 def call_hmm_peaks(
     adata,
-    feature_configs,
-    obs_column='Reference_strand',
-    site_types=['GpC_site', 'CpG_site'],
-    save_plot=False,
-    output_dir=None,
-    date_tag=None,
-    inplace=False
+    feature_configs: Dict[str, Dict[str, Any]],
+    ref_column: str = "Reference_strand",
+    site_types: Sequence[str] = ("GpC", "CpG"),
+    save_plot: bool = False,
+    output_dir: Optional[Union[str, "Path"]] = None,
+    date_tag: Optional[str] = None,
+    inplace: bool = True,
+    index_col_suffix: Optional[str] = None,
+    alternate_labels: bool = False,
 ):
+    """
+    Peak calling over HMM (or other) layers, per reference group and per layer.
+    Writes:
+      - adata.uns["{layer}_{ref}_peak_centers"] = list of centers
+      - adata.var["{layer}_{ref}_peak_{center}"] boolean window masks
+      - adata.obs per-read summaries for each peak window:
+            mean_{layer}_{ref}_around_{center}
+            sum_{layer}_{ref}_around_{center}
+            {layer}_{ref}_present_at_{center} (bool)
+        and per site-type:
+            sum_{layer}_{site}_{ref}_around_{center}
+            mean_{layer}_{site}_{ref}_around_{center}
+      - adata.var["is_in_any_{layer}_peak_{ref}"]
+      - adata.var["is_in_any_peak"] (global)
+    """
+    import matplotlib.pyplot as plt
     import numpy as np
     import pandas as pd
-    import matplotlib.pyplot as plt
     from scipy.signal import find_peaks
+    from scipy.sparse import issparse
 
     if not inplace:
         adata = adata.copy()
 
-    # Ensure obs_column is categorical
-    if not isinstance(adata.obs[obs_column].dtype, pd.CategoricalDtype):
-        adata.obs[obs_column] = pd.Categorical(adata.obs[obs_column])
-
-    coordinates = adata.var_names.astype(int).values
-    peak_columns = []
-
-    obs_updates = {}
-
-    for feature_layer, config in feature_configs.items():
-        min_distance = config.get('min_distance', 200)
-        peak_width = config.get('peak_width', 200)
-        peak_prominence = config.get('peak_prominence', 0.2)
-        peak_threshold = config.get('peak_threshold', 0.8)
-
-        matrix = adata.layers[feature_layer]
-        means = np.mean(matrix, axis=0)
-        peak_indices, _ = find_peaks(means, prominence=peak_prominence, distance=min_distance)
-        peak_centers = coordinates[peak_indices]
-        adata.uns[f'{feature_layer} peak_centers'] = peak_centers.tolist()
-
-        # Plot
-        plt.figure(figsize=(6, 3))
-        plt.plot(coordinates, means)
-        plt.title(f"{feature_layer} with peak calls")
-        plt.xlabel("Genomic position")
-        plt.ylabel("Mean intensity")
-        for i, center in enumerate(peak_centers):
-            start, end = center - peak_width // 2, center + peak_width // 2
-            plt.axvspan(start, end, color='purple', alpha=0.2)
-            plt.axvline(center, color='red', linestyle='--')
-            aligned = [end if i % 2 else start, 'left' if i % 2 else 'right']
-            plt.text(aligned[0], 0, f"Peak {i}\n{center}", color='red', ha=aligned[1])
-        if save_plot and output_dir:
-            filename = f"{output_dir}/{date_tag or 'output'}_{feature_layer}_peaks.png"
-            plt.savefig(filename, bbox_inches='tight')
-            print(f"Saved plot to {filename}")
+    if ref_column not in adata.obs:
+        raise KeyError(f"obs column '{ref_column}' not found")
+
+    # Ensure categorical for predictable ref iteration
+    if not pd.api.types.is_categorical_dtype(adata.obs[ref_column]):
+        adata.obs[ref_column] = adata.obs[ref_column].astype("category")
+
+    # Optional: drop duplicate obs columns once to avoid Pandas/AnnData view quirks
+    if getattr(adata.obs.columns, "duplicated", None) is not None:
+        if adata.obs.columns.duplicated().any():
+            adata.obs = adata.obs.loc[:, ~adata.obs.columns.duplicated(keep="first")].copy()
+
+    # Fallback coordinates from var_names
+    try:
+        base_coordinates = adata.var_names.astype(int).values
+    except Exception:
+        base_coordinates = np.arange(adata.n_vars, dtype=int)
+
+    # Output dir
+    if output_dir is not None:
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Build search pool = union of declared HMM layers and actual layers; exclude helper suffixes
+    declared = list(adata.uns.get("hmm_appended_layers", []) or [])
+    search_pool = [
+        layer
+        for layer in declared
+        if not any(s in layer for s in ("_lengths", "_states", "_posterior"))
+    ]
+
+    all_peak_var_cols = []
+
+    # Iterate per reference
+    for ref in adata.obs[ref_column].cat.categories:
+        ref_mask = (adata.obs[ref_column] == ref).values
+        if not ref_mask.any():
+            continue
+
+        # Per-ref coordinate system
+        if index_col_suffix is not None:
+            coord_col = f"{ref}_{index_col_suffix}"
+            if coord_col not in adata.var:
+                raise KeyError(
+                    f"index_col_suffix='{index_col_suffix}' requested, missing var column '{coord_col}' for ref '{ref}'."
+                )
+            coord_vals = adata.var[coord_col].values
+            try:
+                coordinates = coord_vals.astype(int)
+            except Exception:
+                coordinates = np.asarray(coord_vals, dtype=float)
         else:
-            plt.show()
-
-        feature_peak_columns = []
-        for center in peak_centers:
-            start, end = center - peak_width // 2, center + peak_width // 2
-            colname = f'{feature_layer}_peak_{center}'
-            peak_columns.append(colname)
-            feature_peak_columns.append(colname)
-
-            peak_mask = (coordinates >= start) & (coordinates <= end)
-            adata.var[colname] = peak_mask
-
-            region = matrix[:, peak_mask]
-            obs_updates[f'mean_{feature_layer}_around_{center}'] = np.mean(region, axis=1)
-            obs_updates[f'sum_{feature_layer}_around_{center}'] = np.sum(region, axis=1)
-            obs_updates[f'{feature_layer}_present_at_{center}'] = np.mean(region, axis=1) > peak_threshold
-
-            for site_type in site_types:
-                adata.obs[f'{site_type}_sum_around_{center}'] = 0
-                adata.obs[f'{site_type}_mean_around_{center}'] = np.nan
-
-            for ref in adata.obs[obs_column].cat.categories:
-                ref_idx = adata.obs[obs_column] == ref
-                mask_key = f"{ref}_{site_type}"
-                for site_type in site_types:
-                    if mask_key not in adata.var:
-                        continue
-                    site_mask = adata.var[mask_key].values
-                    site_coords = coordinates[site_mask]
-                    region_mask = (site_coords >= start) & (site_coords <= end)
-                    if not region_mask.any():
-                        continue
-                    full_mask = site_mask.copy()
-                    full_mask[site_mask] = region_mask
-                    site_region = adata[ref_idx, full_mask].X
-                    if hasattr(site_region, "A"):
-                        site_region = site_region.A
-                    if site_region.shape[1] > 0:
-                        adata.obs.loc[ref_idx, f'{site_type}_sum_around_{center}'] = np.nansum(site_region, axis=1)
-                        adata.obs.loc[ref_idx, f'{site_type}_mean_around_{center}'] = np.nanmean(site_region, axis=1)
+            coordinates = base_coordinates
+
+        if coordinates.shape[0] != adata.n_vars:
+            raise ValueError(f"Coordinate length {coordinates.shape[0]} != n_vars {adata.n_vars}")
+
+        # Feature keys to consider
+        for feature_key, config in feature_configs.items():
+            # Resolve candidate layers: exact → suffix → direct present
+            candidates = [ln for ln in search_pool if ln == feature_key]
+            if not candidates:
+                candidates = [ln for ln in search_pool if str(ln).endswith(feature_key)]
+            if not candidates and feature_key in adata.layers:
+                candidates = [feature_key]
+
+            if not candidates:
+                logger.warning(
+                    "[call_hmm_peaks] No layers found matching '%s' in ref '%s'. Skipping.",
+                    feature_key,
+                    ref,
+                )
+                continue
+
+            # Hyperparams (sanitized)
+            min_distance = max(1, int(config.get("min_distance", 200)))
+            peak_width = max(1, int(config.get("peak_width", 200)))
+            peak_prom = float(config.get("peak_prominence", 0.2))
+            peak_threshold = float(config.get("peak_threshold", 0.8))
+            rolling_window = max(1, int(config.get("rolling_window", 1)))
+
+            for layer_name in candidates:
+                if layer_name not in adata.layers:
+                    logger.warning(
+                        "[call_hmm_peaks] Layer '%s' not in adata.layers; skipping.",
+                        layer_name,
+                    )
+                    continue
+
+                # Dense layer data
+                L = adata.layers[layer_name]
+                L = L.toarray() if issparse(L) else np.asarray(L)
+                if L.shape != (adata.n_obs, adata.n_vars):
+                    logger.warning(
+                        "[call_hmm_peaks] Layer '%s' has shape %s, expected (%s, %s); skipping.",
+                        layer_name,
+                        L.shape,
+                        adata.n_obs,
+                        adata.n_vars,
+                    )
+                    continue
+
+                # Ref subset
+                matrix = L[ref_mask, :]
+                if matrix.size == 0 or matrix.shape[0] == 0:
+                    continue
+
+                means = np.nanmean(matrix, axis=0)
+                means = np.nan_to_num(means, nan=0.0)
+
+                if rolling_window > 1:
+                    kernel = np.ones(rolling_window, dtype=float) / float(rolling_window)
+                    peak_metric = np.convolve(means, kernel, mode="same")
+                else:
+                    peak_metric = means
+
+                # Peak detection
+                peak_indices, _ = find_peaks(
+                    peak_metric, prominence=peak_prom, distance=min_distance
+                )
+                if peak_indices.size == 0:
+                    logger.info(
+                        "[call_hmm_peaks] No peaks for layer '%s' in ref '%s'.",
+                        layer_name,
+                        ref,
+                    )
+                    continue
+
+                peak_centers = coordinates[peak_indices]
+                adata.uns[f"{layer_name}_{ref}_peak_centers"] = peak_centers.tolist()
+
+                # Plot once per layer/ref
+                fig, ax = plt.subplots(figsize=(6, 3))
+                ax.plot(coordinates, peak_metric, linewidth=1)
+                ax.set_title(f"{layer_name} peaks in {ref}")
+                ax.set_xlabel("Coordinate")
+                ax.set_ylabel(f"Rolling Mean (win={rolling_window})")
+                for i, center in enumerate(peak_centers):
+                    start = center - peak_width // 2
+                    end = center + peak_width // 2
+                    height = peak_metric[peak_indices[i]]
+                    ax.axvspan(start, end, alpha=0.2)
+                    ax.axvline(center, linestyle="--", linewidth=0.8)
+                    x_text, ha = (
+                        (start, "right") if (not alternate_labels or i % 2 == 0) else (end, "left")
+                    )
+                    ax.text(
+                        x_text, height * 0.8, f"Peak {i}\n{center}", ha=ha, va="bottom", fontsize=8
+                    )
+
+                if save_plot and output_dir is not None:
+                    tag = date_tag or "output"
+                    safe_ref = str(ref).replace("/", "_")
+                    safe_layer = str(layer_name).replace("/", "_")
+                    fname = output_dir / f"{tag}_{safe_layer}_{safe_ref}_peaks.png"
+                    fig.savefig(fname, bbox_inches="tight", dpi=200)
+                    logger.info("[call_hmm_peaks] Saved plot to %s", fname)
+                    plt.close(fig)
+                else:
+                    fig.tight_layout()
+                    plt.show()
+
+                # Collect new obs columns; assign once per layer/ref
+                new_obs_cols: Dict[str, np.ndarray] = {}
+                feature_peak_cols = []
+
+                for center in np.asarray(peak_centers).tolist():
+                    start = center - peak_width // 2
+                    end = center + peak_width // 2
+
+                    # var window mask
+                    colname = f"{layer_name}_{ref}_peak_{center}"
+                    feature_peak_cols.append(colname)
+                    all_peak_var_cols.append(colname)
+                    peak_mask = (coordinates >= start) & (coordinates <= end)
+                    adata.var[colname] = peak_mask
+
+                    # feature-layer summaries for reads in this ref
+                    region = matrix[:, peak_mask]  # (n_ref, n_window)
+
+                    mean_col = f"mean_{layer_name}_{ref}_around_{center}"
+                    sum_col = f"sum_{layer_name}_{ref}_around_{center}"
+                    present_col = f"{layer_name}_{ref}_present_at_{center}"
+
+                    for nm, default, dt in (
+                        (mean_col, np.nan, float),
+                        (sum_col, 0.0, float),
+                        (present_col, False, bool),
+                    ):
+                        if nm not in new_obs_cols:
+                            new_obs_cols[nm] = np.full(adata.n_obs, default, dtype=dt)
+
+                    if region.shape[1] > 0:
+                        means_per_read = np.nanmean(region, axis=1)
+                        sums_per_read = np.nansum(region, axis=1)
                     else:
-                        pass
+                        means_per_read = np.full(matrix.shape[0], np.nan, dtype=float)
+                        sums_per_read = np.zeros(matrix.shape[0], dtype=float)
+
+                    new_obs_cols[mean_col][ref_mask] = means_per_read
+                    new_obs_cols[sum_col][ref_mask] = sums_per_read
+                    new_obs_cols[present_col][ref_mask] = (
+                        np.nan_to_num(means_per_read, nan=0.0) > peak_threshold
+                    )
+
+                    # site-type summaries from adata.X, not an AnnData view
+                    Xmat = adata.X
+                    for site_type in site_types:
+                        mask_key = f"{ref}_{site_type}_site"
+                        if mask_key not in adata.var:
+                            continue
+
+                        site_mask = adata.var[mask_key].values.astype(bool)
+                        if not site_mask.any():
+                            continue
+
+                        site_coords = coordinates[site_mask]
+                        site_region_mask = (site_coords >= start) & (site_coords <= end)
+                        sum_site_col = f"sum_{layer_name}_{site_type}_{ref}_around_{center}"
+                        mean_site_col = f"mean_{layer_name}_{site_type}_{ref}_around_{center}"
+
+                        if sum_site_col not in new_obs_cols:
+                            new_obs_cols[sum_site_col] = np.zeros(adata.n_obs, dtype=float)
+                        if mean_site_col not in new_obs_cols:
+                            new_obs_cols[mean_site_col] = np.full(adata.n_obs, np.nan, dtype=float)
+
+                        if not site_region_mask.any():
+                            continue
+
+                        full_mask = np.zeros_like(site_mask, dtype=bool)
+                        full_mask[site_mask] = site_region_mask
+
+                        if issparse(Xmat):
+                            site_region = Xmat[ref_mask][:, full_mask]
+                            site_region = site_region.toarray()
+                        else:
+                            Xnp = np.asarray(Xmat)
+                            site_region = Xnp[np.asarray(ref_mask), :][:, np.asarray(full_mask)]
+
+                        if site_region.shape[1] > 0:
+                            new_obs_cols[sum_site_col][ref_mask] = np.nansum(site_region, axis=1)
+                            new_obs_cols[mean_site_col][ref_mask] = np.nanmean(site_region, axis=1)
+
+                # one-shot assignment to avoid fragmentation
+                if new_obs_cols:
+                    adata.obs = adata.obs.assign(
+                        **{k: pd.Series(v, index=adata.obs.index) for k, v in new_obs_cols.items()}
+                    )
+
+                # per (layer, ref) any-peak
+                any_col = f"is_in_any_{layer_name}_peak_{ref}"
+                if feature_peak_cols:
+                    adata.var[any_col] = adata.var[feature_peak_cols].any(axis=1)
+                else:
+                    adata.var[any_col] = False
 
-        adata.var[f'is_in_any_{feature_layer}_peak'] = adata.var[feature_peak_columns].any(axis=1)
-        print(f"Annotated {len(peak_centers)} peaks for {feature_layer}")
+                logger.info(
+                    "[call_hmm_peaks] Annotated %s peaks for layer '%s' in ref '%s'.",
+                    len(peak_centers),
+                    layer_name,
+                    ref,
+                )
 
-    adata.var['is_in_any_peak'] = adata.var[peak_columns].any(axis=1)
-    adata.obs = pd.concat([adata.obs, pd.DataFrame(obs_updates, index=adata.obs.index)], axis=1)
+    # global any-peak across all layers/refs
+    if all_peak_var_cols:
+        adata.var["is_in_any_peak"] = adata.var[all_peak_var_cols].any(axis=1)
 
-    return adata if not inplace else None
+    return None if inplace else adata

smftools/hmm/display_hmm.py

@@ -1,18 +1,31 @@
+from smftools.logging_utils import get_logger
+
+logger = get_logger(__name__)
+
+
 def display_hmm(hmm, state_labels=["Non-Methylated", "Methylated"], obs_labels=["0", "1"]):
+    """Log a summary of HMM transition and emission parameters.
+
+    Args:
+        hmm: HMM object with edges and distributions.
+        state_labels: Optional labels for states.
+        obs_labels: Optional labels for observations.
+    """
     import torch
-    print("\n**HMM Model Overview**")
-    print(hmm)
 
-    print("\n**Transition Matrix**")
+    logger.info("**HMM Model Overview**")
+    logger.info("%s", hmm)
+
+    logger.info("**Transition Matrix**")
     transition_matrix = torch.exp(hmm.edges).detach().cpu().numpy()
     for i, row in enumerate(transition_matrix):
         label = state_labels[i] if state_labels else f"State {i}"
         formatted_row = ", ".join(f"{p:.6f}" for p in row)
-        print(f"{label}: [{formatted_row}]")
+        logger.info("%s: [%s]", label, formatted_row)
 
-    print("\n**Emission Probabilities**")
+    logger.info("**Emission Probabilities**")
     for i, dist in enumerate(hmm.distributions):
         label = state_labels[i] if state_labels else f"State {i}"
         probs = dist.probs.detach().cpu().numpy()
         formatted_emissions = {obs_labels[j]: probs[j] for j in range(len(probs))}
-        print(f"{label}: {formatted_emissions}")
+        logger.info("%s: %s", label, formatted_emissions)

smftools/hmm/hmm_readwrite.py

@@ -1,16 +1,25 @@
-def load_hmm(model_path, device='cpu'):
+def load_hmm(model_path, device="cpu"):
     """
     Reads in a pretrained HMM.
-    
+
     Parameters:
         model_path (str): Path to a pretrained HMM
     """
     import torch
+
     # Load model using PyTorch
     hmm = torch.load(model_path)
-    hmm.to(device) 
+    hmm.to(device)
     return hmm
 
+
 def save_hmm(model, model_path):
+    """Save a pretrained HMM to disk.
+
+    Args:
+        model: HMM model instance.
+        model_path: Output path for the model.
+    """
     import torch
-    torch.save(model, model_path)
+
+    torch.save(model, model_path)