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

smftools/config/experiment_config.py

@@ -12,6 +12,7 @@ from smftools.constants import (
     BAM_SUFFIX,
     BARCODE_BOTH_ENDS,
     CONVERSIONS,
+    LOAD_DIR,
     MOD_LIST,
     MOD_MAP,
     REF_COL,
@@ -664,6 +665,8 @@ class ExperimentConfig:
     # General I/O
     input_data_path: Optional[str] = None
     output_directory: Optional[str] = None
+    emit_log_file: Optional[bool] = True
+    log_level: Optional[str] = "INFO"
     fasta: Optional[str] = None
     bam_suffix: str = BAM_SUFFIX
     recursive_input_search: bool = True
@@ -736,6 +739,10 @@ class ExperimentConfig:
     aligner_args: Optional[List[str]] = None
     make_bigwigs: bool = False
     make_beds: bool = False
+    annotate_secondary_supplementary: bool = True
+    samtools_backend: str = "auto"
+    bedtools_backend: str = "auto"
+    bigwig_backend: str = "auto"
 
     # Anndata structure
     reference_column: Optional[str] = REF_COL
@@ -744,6 +751,9 @@ class ExperimentConfig:
     # General Plotting
     sample_name_col_for_plotting: Optional[str] = "Barcode"
     rows_per_qc_histogram_grid: int = 12
+    clustermap_demux_types_to_plot: List[str] = field(
+        default_factory=lambda: ["single", "double", "already"]
+    )
 
     # Preprocessing - Read length and quality filter params
     read_coord_filter: Optional[Sequence[float]] = field(default_factory=lambda: [None, None])
@@ -813,6 +823,9 @@ class ExperimentConfig:
     duplicate_detection_site_types: List[str] = field(
         default_factory=lambda: ["GpC", "CpG", "ambiguous_GpC_CpG"]
     )
+    duplicate_detection_demux_types_to_use: List[str] = field(
+        default_factory=lambda: ["single", "double", "already"]
+    )
     duplicate_detection_distance_threshold: float = 0.07
     hamming_vs_metric_keys: List[str] = field(default_factory=lambda: ["Fraction_C_site_modified"])
     duplicate_detection_keep_best_metric: str = "read_quality"
@@ -824,6 +837,9 @@ class ExperimentConfig:
 
     # Preprocessing - Position QC
     position_max_nan_threshold: float = 0.1
+    mismatch_frequency_range: Sequence[float] = field(default_factory=lambda: [0.05, 0.95])
+    mismatch_frequency_layer: str = "mismatch_integer_encoding"
+    mismatch_frequency_read_span_layer: str = "read_span_mask"
 
     # Spatial Analysis - Clustermap params
     layer_for_clustermap_plotting: Optional[str] = "nan0_0minus1"
@@ -832,6 +848,14 @@ class ExperimentConfig:
     clustermap_cmap_cpg: Optional[str] = "coolwarm"
     clustermap_cmap_a: Optional[str] = "coolwarm"
     spatial_clustermap_sortby: Optional[str] = "gpc"
+    rolling_nn_layer: Optional[str] = "nan0_0minus1"
+    rolling_nn_plot_layer: Optional[str] = "nan0_0minus1"
+    rolling_nn_window: int = 15
+    rolling_nn_step: int = 2
+    rolling_nn_min_overlap: int = 10
+    rolling_nn_return_fraction: bool = True
+    rolling_nn_obsm_key: str = "rolling_nn_dist"
+    rolling_nn_site_types: Optional[List[str]] = None
 
     # Spatial Analysis - UMAP/Leiden params
     layer_for_umap_plotting: Optional[str] = "nan_half"
@@ -880,11 +904,15 @@ class ExperimentConfig:
     accessible_patches: Optional[bool] = True
     cpg: Optional[bool] = False
     hmm_feature_sets: Dict[str, Any] = field(default_factory=dict)
+    hmm_feature_colormaps: Dict[str, Any] = field(default_factory=dict)
     hmm_merge_layer_features: Optional[List[Tuple]] = field(default_factory=lambda: [(None, 60)])
     clustermap_cmap_hmm: Optional[str] = "coolwarm"
     hmm_clustermap_feature_layers: List[str] = field(
         default_factory=lambda: ["all_accessible_features"]
     )
+    hmm_clustermap_length_layers: List[str] = field(
+        default_factory=lambda: ["all_accessible_features"]
+    )
     hmm_clustermap_sortby: Optional[str] = "hmm"
     hmm_peak_feature_configs: Dict[str, Any] = field(default_factory=dict)
 
@@ -903,6 +931,8 @@ class ExperimentConfig:
     invert_adata: bool = False
     bypass_append_binary_layer_by_base_context: bool = False
     force_redo_append_binary_layer_by_base_context: bool = False
+    bypass_append_mismatch_frequency_sites: bool = False
+    force_redo_append_mismatch_frequency_sites: bool = False
     bypass_calculate_read_modification_stats: bool = False
     force_redo_calculate_read_modification_stats: bool = False
     bypass_filter_reads_on_modification_thresholds: bool = False
@@ -1107,7 +1137,7 @@ class ExperimentConfig:
 
         # Demultiplexing output path
         split_dir = merged.get("split_dir", SPLIT_DIR)
-        split_path = output_dir / split_dir
+        split_path = output_dir / LOAD_DIR / split_dir
 
         # final normalization
         if "strands" in merged:
@@ -1194,6 +1224,9 @@ class ExperimentConfig:
         # Final normalization of hmm_feature_sets and canonical local variables
         merged["hmm_feature_sets"] = normalize_hmm_feature_sets(merged.get("hmm_feature_sets", {}))
         hmm_feature_sets = merged.get("hmm_feature_sets", {})
+        hmm_feature_colormaps = merged.get("hmm_feature_colormaps", {})
+        if not isinstance(hmm_feature_colormaps, dict):
+            hmm_feature_colormaps = {}
         hmm_annotation_threshold = merged.get("hmm_annotation_threshold", 0.5)
         hmm_batch_size = int(merged.get("hmm_batch_size", 1024))
         hmm_use_viterbi = bool(merged.get("hmm_use_viterbi", False))
@@ -1208,6 +1241,9 @@ class ExperimentConfig:
         hmm_clustermap_feature_layers = _parse_list(
             merged.get("hmm_clustermap_feature_layers", "all_accessible_features")
         )
+        hmm_clustermap_length_layers = _parse_list(
+            merged.get("hmm_clustermap_length_layers", hmm_clustermap_feature_layers)
+        )
 
         hmm_fit_strategy = str(merged.get("hmm_fit_strategy", "per_group")).strip()
         hmm_shared_scope = _parse_list(merged.get("hmm_shared_scope", ["reference", "methbase"]))
@@ -1228,6 +1264,7 @@ class ExperimentConfig:
 
         # instantiate dataclass
         instance = cls(
+            annotate_secondary_supplementary=merged.get("annotate_secondary_supplementary", True),
             smf_modality=merged.get("smf_modality"),
             input_data_path=input_data_path,
             recursive_input_search=merged.get("recursive_input_search"),
@@ -1254,6 +1291,8 @@ class ExperimentConfig:
             trim=merged.get("trim", TRIM),
             input_already_demuxed=merged.get("input_already_demuxed", False),
             threads=merged.get("threads"),
+            emit_log_file=merged.get("emit_log_file", True),
+            log_level=merged.get("log_level", "INFO"),
             sample_sheet_path=merged.get("sample_sheet_path"),
             sample_sheet_mapping_column=merged.get("sample_sheet_mapping_column"),
             delete_intermediate_bams=merged.get("delete_intermediate_bams", False),
@@ -1264,6 +1303,9 @@ class ExperimentConfig:
             device=merged.get("device", "auto"),
             make_bigwigs=merged.get("make_bigwigs", False),
             make_beds=merged.get("make_beds", False),
+            samtools_backend=merged.get("samtools_backend", "auto"),
+            bedtools_backend=merged.get("bedtools_backend", "auto"),
+            bigwig_backend=merged.get("bigwig_backend", "auto"),
             delete_intermediate_hdfs=merged.get("delete_intermediate_hdfs", True),
             mod_target_bases=merged.get("mod_target_bases", ["GpC", "CpG"]),
             enzyme_target_bases=merged.get("enzyme_target_bases", ["GpC"]),
@@ -1307,6 +1349,9 @@ class ExperimentConfig:
             ),
             reindexing_offsets=merged.get("reindexing_offsets", {None: None}),
             reindexed_var_suffix=merged.get("reindexed_var_suffix", "reindexed"),
+            clustermap_demux_types_to_plot=merged.get(
+                "clustermap_demux_types_to_plot", ["single", "double", "already"]
+            ),
             layer_for_clustermap_plotting=merged.get(
                 "layer_for_clustermap_plotting", "nan0_0minus1"
             ),
@@ -1315,6 +1360,14 @@ class ExperimentConfig:
             clustermap_cmap_cpg=merged.get("clustermap_cmap_cpg", "coolwarm"),
             clustermap_cmap_a=merged.get("clustermap_cmap_a", "coolwarm"),
             spatial_clustermap_sortby=merged.get("spatial_clustermap_sortby", "gpc"),
+            rolling_nn_layer=merged.get("rolling_nn_layer", "nan0_0minus1"),
+            rolling_nn_plot_layer=merged.get("rolling_nn_plot_layer", "nan0_0minus1"),
+            rolling_nn_window=merged.get("rolling_nn_window", 15),
+            rolling_nn_step=merged.get("rolling_nn_step", 2),
+            rolling_nn_min_overlap=merged.get("rolling_nn_min_overlap", 10),
+            rolling_nn_return_fraction=merged.get("rolling_nn_return_fraction", True),
+            rolling_nn_obsm_key=merged.get("rolling_nn_obsm_key", "rolling_nn_dist"),
+            rolling_nn_site_types=merged.get("rolling_nn_site_types", None),
             layer_for_umap_plotting=merged.get("layer_for_umap_plotting", "nan_half"),
             umap_layers_to_plot=merged.get(
                 "umap_layers_to_plot", ["mapped_length", "Raw_modification_signal"]
@@ -1341,6 +1394,7 @@ class ExperimentConfig:
             hmm_emission_adapt_tol=hmm_emission_adapt_tol,
             hmm_dtype=merged.get("hmm_dtype", "float64"),
             hmm_feature_sets=hmm_feature_sets,
+            hmm_feature_colormaps=hmm_feature_colormaps,
             hmm_annotation_threshold=hmm_annotation_threshold,
             hmm_batch_size=hmm_batch_size,
             hmm_use_viterbi=hmm_use_viterbi,
@@ -1349,6 +1403,7 @@ class ExperimentConfig:
             hmm_merge_layer_features=hmm_merge_layer_features,
             clustermap_cmap_hmm=merged.get("clustermap_cmap_hmm", "coolwarm"),
             hmm_clustermap_feature_layers=hmm_clustermap_feature_layers,
+            hmm_clustermap_length_layers=hmm_clustermap_length_layers,
             hmm_clustermap_sortby=merged.get("hmm_clustermap_sortby", "hmm"),
             hmm_peak_feature_configs=hmm_peak_feature_configs,
             footprints=merged.get("footprints", None),
@@ -1384,6 +1439,9 @@ class ExperimentConfig:
             duplicate_detection_site_types=merged.get(
                 "duplicate_detection_site_types", ["GpC", "CpG", "ambiguous_GpC_CpG"]
             ),
+            duplicate_detection_demux_types_to_use=merged.get(
+                "duplicate_detection_demux_types_to_use", ["single", "double", "already"]
+            ),
             duplicate_detection_distance_threshold=merged.get(
                 "duplicate_detection_distance_threshold", 0.07
             ),

smftools/constants.py

@@ -21,7 +21,30 @@ BAM_SUFFIX: Final[str] = ".bam"
 BARCODE_BOTH_ENDS: Final[bool] = False
 REF_COL: Final[str] = "Reference_strand"
 SAMPLE_COL: Final[str] = "Experiment_name_and_barcode"
+SAMPLE: Final[str] = "Sample"
 SPLIT_DIR: Final[str] = "demultiplexed_BAMs"
+H5_DIR: Final[str] = "h5ads"
+DEMUX_TYPE: Final[str] = "demux_type"
+BARCODE: Final[str] = "Barcode"
+REFERENCE: Final[str] = "Reference"
+REFERENCE_STRAND: Final[str] = "Reference_strand"
+REFERENCE_DATASET_STRAND: Final[str] = "Reference_dataset_strand"
+STRAND: Final[str] = "Strand"
+DATASET: Final[str] = "Dataset"
+READ_MISMATCH_TREND: Final[str] = "Read_mismatch_trend"
+READ_MAPPING_DIRECTION: Final[str] = "Read_mapping_direction"
+SEQUENCE_INTEGER_ENCODING: Final[str] = "sequence_integer_encoding"
+SEQUENCE_INTEGER_DECODING: Final[str] = "sequence_integer_decoding"
+MISMATCH_INTEGER_ENCODING: Final[str] = "mismatch_integer_encoding"
+BASE_QUALITY_SCORES: Final[str] = "base_quality_scores"
+READ_SPAN_MASK: Final[str] = "read_span_mask"
+
+LOAD_DIR: Final[str] = "load_adata_outputs"
+PREPROCESS_DIR: Final[str] = "preprocess_adata_outputs"
+SPATIAL_DIR: Final[str] = "spatial_adata_outputs"
+HMM_DIR: Final[str] = "hmm_adata_outputs"
+LATENT_DIR: Final[str] = "latent_adata_outputs"
+LOGGING_DIR: Final[str] = "logs"
 TRIM: Final[bool] = False
 
 _private_conversions = ["unconverted"]
@@ -35,3 +58,45 @@ MOD_MAP: Final[Mapping[str, str]] = _deep_freeze(_private_mod_map)
 
 _private_strands = ("bottom", "top")
 STRANDS: Final[tuple[str, ...]] = _deep_freeze(_private_strands)
+
+MODKIT_EXTRACT_TSV_COLUMN_CHROM: Final[str] = "chrom"
+MODKIT_EXTRACT_TSV_COLUMN_REF_POSITION: Final[str] = "ref_position"
+MODKIT_EXTRACT_TSV_COLUMN_MODIFIED_PRIMARY_BASE: Final[str] = "modified_primary_base"
+MODKIT_EXTRACT_TSV_COLUMN_REF_STRAND: Final[str] = "ref_strand"
+MODKIT_EXTRACT_TSV_COLUMN_READ_ID: Final[str] = "read_id"
+MODKIT_EXTRACT_TSV_COLUMN_CALL_CODE: Final[str] = "call_code"
+MODKIT_EXTRACT_TSV_COLUMN_CALL_PROB: Final[str] = "call_prob"
+
+MODKIT_EXTRACT_MODIFIED_BASE_A: Final[str] = "A"
+MODKIT_EXTRACT_MODIFIED_BASE_C: Final[str] = "C"
+MODKIT_EXTRACT_REF_STRAND_PLUS: Final[str] = "+"
+MODKIT_EXTRACT_REF_STRAND_MINUS: Final[str] = "-"
+
+_private_modkit_extract_call_code_modified = ("a", "h", "m")
+MODKIT_EXTRACT_CALL_CODE_MODIFIED: Final[tuple[str, ...]] = _deep_freeze(
+    _private_modkit_extract_call_code_modified
+)
+_private_modkit_extract_call_code_canonical = ("-",)
+MODKIT_EXTRACT_CALL_CODE_CANONICAL: Final[tuple[str, ...]] = _deep_freeze(
+    _private_modkit_extract_call_code_canonical
+)
+
+MODKIT_EXTRACT_SEQUENCE_BASES: Final[tuple[str, ...]] = _deep_freeze(("A", "C", "G", "T", "N"))
+MODKIT_EXTRACT_SEQUENCE_PADDING_BASE: Final[str] = "PAD"
+_private_modkit_extract_base_to_int: Dict[str, int] = {
+    "A": 0,
+    "C": 1,
+    "G": 2,
+    "T": 3,
+    "N": 4,
+    "PAD": 5,
+}
+MODKIT_EXTRACT_SEQUENCE_BASE_TO_INT: Final[Mapping[str, int]] = _deep_freeze(
+    _private_modkit_extract_base_to_int
+)
+_private_modkit_extract_int_to_base: Dict[int, str] = {
+    value: key for key, value in _private_modkit_extract_base_to_int.items()
+}
+MODKIT_EXTRACT_SEQUENCE_INT_TO_BASE: Final[Mapping[int, str]] = _deep_freeze(
+    _private_modkit_extract_int_to_base
+)

smftools/datasets/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from .datasets import Kissiov_and_McKenna_2025, dCas9_kinetics
 
 __all__ = ["dCas9_kinetics", "Kissiov_and_McKenna_2025"]

smftools/hmm/HMM.py

@@ -3,14 +3,20 @@ from __future__ import annotations
 import ast
 import json
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
-import torch
-import torch.nn as nn
 from scipy.sparse import issparse
 
 from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
+
+if TYPE_CHECKING:
+    import torch as torch_types
+    import torch.nn as nn_types
+
+torch = require("torch", extra="torch", purpose="HMM modeling")
+nn = torch.nn
 
 logger = get_logger(__name__)
 # =============================================================================
@@ -138,6 +144,83 @@ def _safe_int_coords(var_names) -> Tuple[np.ndarray, bool]:
         return np.arange(len(var_names), dtype=int), False
 
 
+def mask_layers_outside_read_span(
+    adata,
+    layers: Sequence[str],
+    *,
+    start_key: str = "reference_start",
+    end_key: str = "reference_end",
+    use_original_var_names: bool = True,
+) -> List[str]:
+    """Mask layer values outside read reference spans with NaN.
+
+    This uses integer coordinate comparisons against either ``adata.var["Original_var_names"]``
+    (when present) or ``adata.var_names``. Values strictly less than ``start_key`` or greater
+    than ``end_key`` are set to NaN for each read.
+
+    Args:
+        adata: AnnData object to modify in-place.
+        layers: Layer names to mask.
+        start_key: obs column holding reference start positions.
+        end_key: obs column holding reference end positions.
+        use_original_var_names: Use ``adata.var["Original_var_names"]`` when available.
+
+    Returns:
+        List of layer names that were masked.
+    """
+    if not layers:
+        return []
+
+    if start_key not in adata.obs or end_key not in adata.obs:
+        raise KeyError(f"Missing {start_key!r} or {end_key!r} in adata.obs.")
+
+    coord_source = adata.var_names
+    if use_original_var_names and "Original_var_names" in adata.var:
+        orig = np.asarray(adata.var["Original_var_names"])
+        if orig.size == adata.n_vars:
+            try:
+                orig_numeric = np.asarray(orig, dtype=float)
+            except (TypeError, ValueError):
+                orig_numeric = None
+            if orig_numeric is not None and np.isfinite(orig_numeric).any():
+                coord_source = orig
+
+    coords, _ = _safe_int_coords(coord_source)
+    if coords.shape[0] != adata.n_vars:
+        raise ValueError("Coordinate source length does not match adata.n_vars.")
+
+    try:
+        starts = np.asarray(adata.obs[start_key], dtype=float)
+        ends = np.asarray(adata.obs[end_key], dtype=float)
+    except (TypeError, ValueError) as exc:
+        raise ValueError("Start/end positions must be numeric.") from exc
+
+    masked = []
+    for layer in layers:
+        if layer not in adata.layers:
+            raise KeyError(f"Layer {layer!r} not found in adata.layers.")
+
+        arr = np.asarray(adata.layers[layer])
+        if not np.issubdtype(arr.dtype, np.floating):
+            arr = arr.astype(float, copy=True)
+
+        for i in range(adata.n_obs):
+            start = starts[i]
+            end = ends[i]
+            if not np.isfinite(start) or not np.isfinite(end):
+                continue
+            start_i = int(start)
+            end_i = int(end)
+            row_mask = (coords < start_i) | (coords > end_i)
+            if row_mask.any():
+                arr[i, row_mask] = np.nan
+
+        adata.layers[layer] = arr
+        masked.append(layer)
+
+    return masked
+
+
 def _logsumexp(x: torch.Tensor, dim: int) -> torch.Tensor:
     """Compute log-sum-exp in a numerically stable way.
 
@@ -1058,6 +1141,8 @@ class BaseHMM(nn.Module):
         uns_key: str = "hmm_appended_layers",
         uns_flag: str = "hmm_annotated",
         force_redo: bool = False,
+        mask_to_read_span: bool = True,
+        mask_use_original_var_names: bool = True,
         device: Optional[Union[str, torch.device]] = None,
         **kwargs,
     ):
@@ -1079,6 +1164,8 @@ class BaseHMM(nn.Module):
             uns_key: .uns key to track appended layers.
             uns_flag: .uns flag to mark annotations.
             force_redo: Whether to overwrite existing layers.
+            mask_to_read_span: Whether to mask appended layers outside read spans.
+            mask_use_original_var_names: Use ``adata.var["Original_var_names"]`` when available.
             device: Device specifier.
             **kwargs: Additional parameters for specialized workflows.
 
@@ -1239,6 +1326,13 @@ class BaseHMM(nn.Module):
                     np.asarray(adata.layers[nm])
                 )
 
+        if mask_to_read_span and appended:
+            mask_layers_outside_read_span(
+                adata,
+                appended,
+                use_original_var_names=mask_use_original_var_names,
+            )
+
         adata.uns[uns_key] = appended
         adata.uns[uns_flag] = True
         return None

smftools/hmm/__init__.py

@@ -1,13 +1,24 @@
-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 infer_nucleosomes_in_large_bound, refine_nucleosome_calls
-
-__all__ = [
-    "call_hmm_peaks",
-    "display_hmm",
-    "load_hmm",
-    "refine_nucleosome_calls",
-    "infer_nucleosomes_in_large_bound",
-    "save_hmm",
-]
+from __future__ import annotations
+
+from importlib import import_module
+
+_LAZY_ATTRS = {
+    "call_hmm_peaks": "smftools.hmm.call_hmm_peaks",
+    "display_hmm": "smftools.hmm.display_hmm",
+    "load_hmm": "smftools.hmm.hmm_readwrite",
+    "save_hmm": "smftools.hmm.hmm_readwrite",
+    "infer_nucleosomes_in_large_bound": "smftools.hmm.nucleosome_hmm_refinement",
+    "refine_nucleosome_calls": "smftools.hmm.nucleosome_hmm_refinement",
+}
+
+
+def __getattr__(name: str):
+    if name in _LAZY_ATTRS:
+        module = import_module(_LAZY_ATTRS[name])
+        attr = getattr(module, name)
+        globals()[name] = attr
+        return attr
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+
+__all__ = list(_LAZY_ATTRS.keys())

smftools/hmm/archived/apply_hmm_batched.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import numpy as np
 import pandas as pd
 import torch

smftools/hmm/archived/calculate_distances.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 # calculate_distances
 
 def calculate_distances(intervals, threshold=0.9):

smftools/hmm/archived/call_hmm_peaks.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 def call_hmm_peaks(
     adata,
     feature_configs,

smftools/hmm/archived/train_hmm.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 def train_hmm(
     data,
     emission_probs=[[0.8, 0.2], [0.2, 0.8]],

smftools/hmm/call_hmm_peaks.py

@@ -1,9 +1,11 @@
-# FILE: smftools/hmm/call_hmm_peaks.py
+from __future__ import annotations
 
+# 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
+from smftools.optional_imports import require
 
 logger = get_logger(__name__)
 
@@ -35,12 +37,13 @@ def call_hmm_peaks(
       - 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
     from scipy.signal import find_peaks
     from scipy.sparse import issparse
 
+    plt = require("matplotlib.pyplot", extra="plotting", purpose="HMM peak plots")
+
     if not inplace:
         adata = adata.copy()
 

smftools/hmm/display_hmm.py

@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
 
 logger = get_logger(__name__)
 
@@ -11,7 +14,7 @@ def display_hmm(hmm, state_labels=["Non-Methylated", "Methylated"], obs_labels=[
         state_labels: Optional labels for states.
         obs_labels: Optional labels for observations.
     """
-    import torch
+    torch = require("torch", extra="torch", purpose="HMM display")
 
     logger.info("**HMM Model Overview**")
     logger.info("%s", hmm)

smftools/hmm/hmm_readwrite.py

@@ -1,3 +1,8 @@
+from __future__ import annotations
+
+from smftools.optional_imports import require
+
+
 def load_hmm(model_path, device="cpu"):
     """
     Reads in a pretrained HMM.
@@ -5,7 +10,7 @@ def load_hmm(model_path, device="cpu"):
     Parameters:
         model_path (str): Path to a pretrained HMM
     """
-    import torch
+    torch = require("torch", extra="torch", purpose="HMM read/write")
 
     # Load model using PyTorch
     hmm = torch.load(model_path)
@@ -20,6 +25,6 @@ def save_hmm(model, model_path):
         model: HMM model instance.
         model_path: Output path for the model.
     """
-    import torch
+    torch = require("torch", extra="torch", purpose="HMM read/write")
 
     torch.save(model, model_path)

smftools/hmm/nucleosome_hmm_refinement.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from smftools.logging_utils import get_logger
 
 logger = get_logger(__name__)

smftools/informatics/__init__.py

@@ -1,41 +1,66 @@
-from .bam_functions import (
-    align_and_sort_BAM,
-    bam_qc,
-    concatenate_fastqs_to_bam,
-    count_aligned_reads,
-    demux_and_index_BAM,
-    extract_base_identities,
-    extract_read_features_from_bam,
-    extract_readnames_from_bam,
-    separate_bam_by_bc,
-    split_and_index_BAM,
-)
-from .basecalling import canoncall, modcall
-from .bed_functions import (
-    _bed_to_bigwig,
-    _plot_bed_histograms,
-    aligned_BAM_to_bed,
-    extract_read_lengths_from_bed,
-)
-from .converted_BAM_to_adata import converted_BAM_to_adata
-from .fasta_functions import (
-    find_conversion_sites,
-    generate_converted_FASTA,
-    get_chromosome_lengths,
-    get_native_references,
-    index_fasta,
-    subsample_fasta_from_bed,
-)
-from .h5ad_functions import add_demux_type_annotation, add_read_length_and_mapping_qc
-from .modkit_extract_to_adata import modkit_extract_to_adata
-from .modkit_functions import extract_mods, make_modbed, modQC
-from .ohe import ohe_batching, ohe_layers_decode, one_hot_decode, one_hot_encode
-from .pod5_functions import basecall_pod5s, fast5_to_pod5, subsample_pod5
-from .run_multiqc import run_multiqc
+from __future__ import annotations
+
+from importlib import import_module
+
+_LAZY_ATTRS = {
+    "_bed_to_bigwig": "smftools.informatics.bed_functions",
+    "_plot_bed_histograms": "smftools.informatics.bed_functions",
+    "add_demux_type_annotation": "smftools.informatics.h5ad_functions",
+    "add_read_tag_annotations": "smftools.informatics.h5ad_functions",
+    "add_read_length_and_mapping_qc": "smftools.informatics.h5ad_functions",
+    "align_and_sort_BAM": "smftools.informatics.bam_functions",
+    "bam_qc": "smftools.informatics.bam_functions",
+    "basecall_pod5s": "smftools.informatics.pod5_functions",
+    "canoncall": "smftools.informatics.basecalling",
+    "concatenate_fastqs_to_bam": "smftools.informatics.bam_functions",
+    "converted_BAM_to_adata": "smftools.informatics.converted_BAM_to_adata",
+    "count_aligned_reads": "smftools.informatics.bam_functions",
+    "demux_and_index_BAM": "smftools.informatics.bam_functions",
+    "extract_base_identities": "smftools.informatics.bam_functions",
+    "extract_mods": "smftools.informatics.modkit_functions",
+    "extract_read_features_from_bam": "smftools.informatics.bam_functions",
+    "extract_read_tags_from_bam": "smftools.informatics.bam_functions",
+    "extract_read_lengths_from_bed": "smftools.informatics.bed_functions",
+    "extract_readnames_from_bam": "smftools.informatics.bam_functions",
+    "fast5_to_pod5": "smftools.informatics.pod5_functions",
+    "find_conversion_sites": "smftools.informatics.fasta_functions",
+    "generate_converted_FASTA": "smftools.informatics.fasta_functions",
+    "get_chromosome_lengths": "smftools.informatics.fasta_functions",
+    "get_native_references": "smftools.informatics.fasta_functions",
+    "index_fasta": "smftools.informatics.fasta_functions",
+    "make_modbed": "smftools.informatics.modkit_functions",
+    "modQC": "smftools.informatics.modkit_functions",
+    "modcall": "smftools.informatics.basecalling",
+    "modkit_extract_to_adata": "smftools.informatics.modkit_extract_to_adata",
+    "decode_int_sequence": "smftools.informatics.sequence_encoding",
+    "encode_sequence_to_int": "smftools.informatics.sequence_encoding",
+    "ohe_batching": "smftools.informatics.ohe",
+    "ohe_layers_decode": "smftools.informatics.ohe",
+    "one_hot_decode": "smftools.informatics.ohe",
+    "one_hot_encode": "smftools.informatics.ohe",
+    "run_multiqc": "smftools.informatics.run_multiqc",
+    "separate_bam_by_bc": "smftools.informatics.bam_functions",
+    "split_and_index_BAM": "smftools.informatics.bam_functions",
+    "subsample_fasta_from_bed": "smftools.informatics.fasta_functions",
+    "subsample_pod5": "smftools.informatics.pod5_functions",
+    "aligned_BAM_to_bed": "smftools.informatics.bed_functions",
+}
+
+
+def __getattr__(name: str):
+    if name in _LAZY_ATTRS:
+        module = import_module(_LAZY_ATTRS[name])
+        attr = getattr(module, name)
+        globals()[name] = attr
+        return attr
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
 
 __all__ = [
     "basecall_pod5s",
     "converted_BAM_to_adata",
+    "decode_int_sequence",
+    "encode_sequence_to_int",
     "subsample_fasta_from_bed",
     "subsample_pod5",
     "fast5_to_pod5",