"""HMM-backed segmentation from per-CpG state labels to genomic regions."""
from typing import Dict, List, Optional, Tuple
import numpy as np
import pandas as pd
from .utils import (
plot_state_labels,
relabel_by_mean_emission,
resolve_overlay_plot_args,
)
from .helper_classes import (
MethylStateAssignmentMethod,
MethylationStates,
SampleInfo,
)
from .methylseg_hmm import MethylSegHMM
from .methyl_state_analyzer import MethylStateAnalyzer
[docs]
class MethylSegmentor:
"""
Class to handle segmentation of methylation data using HMMs.
Recommend CTHMM for sparse data with variable probe spacing, and the
sticky categorical smoother for dense discrete state-label smoothing.
"""
[docs]
def __init__(
self,
analyzer: MethylStateAnalyzer,
hmm_model: MethylSegHMM,
state_assignment_method: MethylStateAssignmentMethod = MethylStateAssignmentMethod.DEFINITION,
out_dir=".",
random_state: int = 42,
):
"""
Initialize the segmentor with state-assignment and HMM backends.
Parameters
----------
analyzer
State analyzer that prepares emissions and biological labels.
hmm_model
Configured HMM backend used to smooth state observations.
state_assignment_method
Strategy used to obtain state labels before segmentation.
out_dir
Directory for segmentation artifacts and plot outputs.
random_state
Random seed used by stochastic segmentation operations.
"""
self.analyzer = analyzer
self.hmm_model = hmm_model
self.state_assignment_method = MethylStateAssignmentMethod(
state_assignment_method
)
self.out_dir = out_dir
self.random_state = random_state
self.segment_results = {}
self.default_sample_info: SampleInfo | None = None
def _encode_states_for_hmm(self, states: np.ndarray) -> np.ndarray:
numeric_states = MethylationStates.convert_to_numeric(states)
unique_states = np.sort(np.unique(numeric_states))
supported_observations = getattr(
self.hmm_model,
"n_emissions",
getattr(self.hmm_model, "n_states", None),
)
if supported_observations is not None and len(unique_states) > int(
supported_observations
):
raise ValueError(
"Observed state labels contain more distinct categories than the "
"configured HMM can represent."
)
state_to_obs = {
int(state_value): obs_idx
for obs_idx, state_value in enumerate(unique_states.tolist())
}
return np.array(
[state_to_obs[int(state_value)] for state_value in numeric_states],
dtype=int,
)
def _get_state_cutoffs(self) -> Optional[Dict[str, object]]:
return getattr(self.analyzer, "state_cutoffs", None)
def _prepare_emissions(
self,
sample_info: SampleInfo,
chrom: str | None = None,
) -> Tuple[pd.DataFrame, pd.DataFrame]:
meth_data, emissions_df = (
self.analyzer.assigner.prepare_sample_for_clustering(
sample_info=sample_info,
chrom=chrom,
)
)
self.meth_data = meth_data.copy()
self.emissions_df = emissions_df.copy()
return self.meth_data, self.emissions_df
def _derive_sequence_lengths(
self,
meth_data: pd.DataFrame,
chrom: str | None = None,
) -> Optional[List[int]]:
if chrom is not None or len(meth_data) == 0:
return None
chrom_values = meth_data["CpG_chrm"].astype(str).to_numpy()
lengths = []
current_chrom = chrom_values[0]
current_length = 1
for chrom_name in chrom_values[1:]:
if chrom_name == current_chrom:
current_length += 1
else:
lengths.append(current_length)
current_chrom = chrom_name
current_length = 1
lengths.append(current_length)
return lengths if len(lengths) > 1 else None
def _set_hmm_sequence_lengths(
self,
lengths: Optional[List[int]],
) -> None:
if not hasattr(self.hmm_model, "lengths"):
return
self.hmm_model.lengths = (
None if lengths is None else [int(length) for length in lengths]
)
def _segment_sample_discrete_states(
self,
sample_info: SampleInfo,
chrom: str | None = None,
) -> Tuple[np.ndarray, np.ndarray]:
self.assign_states(sample_info, chrom)
states = self._encode_states_for_hmm(
self.meth_data["state"].to_numpy(dtype=int)
)
sequence_lengths = self._derive_sequence_lengths(self.meth_data, chrom=chrom)
self._set_hmm_sequence_lengths(sequence_lengths)
self.hmm_model.create_model()
self.hmm_model.fit(states, sample_info, chrom)
hidden_states = self.hmm_model.predict(states)
readable_states = relabel_by_mean_emission(
hidden_states,
self.emissions_df,
self._get_state_cutoffs(),
self.analyzer.assigner.int_low_cutoff,
self.analyzer.assigner.int_high_cutoff,
self.analyzer.assigner.window_specs,
)
return hidden_states, readable_states
[docs]
def assign_states(
self,
sample_info: SampleInfo,
chrom: str | None = None,
) -> np.ndarray:
"""
Assign coarse methylation states before HMM smoothing.
Parameters
----------
sample_info
Prepared methylation sample to summarize and label.
chrom
Optional chromosome restriction for per-chromosome state assignment.
Returns
-------
tuple
Pair of ``(meth_data, emissions_df)`` cached on the segmentor after
populating ``state`` and ``state_readable`` columns on ``meth_data``.
Raises
------
ValueError
If the configured state-assignment method is unknown.
NotImplementedError
If ``AUTO`` assignment is requested.
"""
meth_data, emissions_df = self._prepare_emissions(
sample_info=sample_info, chrom=chrom
)
if (
self.state_assignment_method.value
== MethylStateAssignmentMethod.DEFINITION.value
):
states = self.analyzer.define_states_by_rules(
sample_info=sample_info,
chrom=chrom,
sample_emissions=emissions_df,
)
elif (
self.state_assignment_method.value
== MethylStateAssignmentMethod.KMEANS.value
):
_, _, _, states = self.analyzer.assigner.apply_kmeans_to_emissions(
emissions_df
)
elif (
self.state_assignment_method.value == MethylStateAssignmentMethod.AUTO.value
):
raise NotImplementedError(
"AUTO state assignment method not implemented yet."
)
else:
raise ValueError(
f"Unknown state assignment method: {self.state_assignment_method}"
)
meth_data = meth_data.copy()
meth_data["state"] = MethylationStates.convert_to_numeric(states)
meth_data["state_readable"] = states
self.meth_data = meth_data
self.emissions_df = emissions_df
return meth_data, emissions_df
[docs]
def segment_sample(
self,
sample_info: SampleInfo | None = None,
chrom: str | None = None,
force_resegment: bool = False,
) -> Tuple[pd.DataFrame, object]:
"""
Segment a sample and refresh probe-level results plus raw regions.
Parameters
----------
sample_info
Prepared sample to segment. When omitted, uses
``default_sample_info``.
chrom
Optional chromosome restriction for per-chromosome segmentation.
force_resegment
If ``True``, ignore cached segmentation results and rerun the HMM.
Returns
-------
tuple
``(meth_data, hmm_model)`` where ``meth_data`` is the segmented
probe-level table and ``hmm_model`` is the fitted backend model.
Returns the segmented probe-level methylation table and fitted HMM
object. Raw contiguous regions are stored on ``self.regions_df``.
"""
if sample_info is None:
sample_info = self.default_sample_info
if sample_info is None:
raise ValueError(
"No sample_info provided and no default_sample_info configured."
)
chrom_segmented_on_sample = (
sample_info.sample_id in self.segment_results
and chrom in self.segment_results[sample_info.sample_id]
)
if not chrom_segmented_on_sample or force_resegment:
hidden_states, readable_states = self._segment_sample_discrete_states(
sample_info=sample_info,
chrom=chrom,
)
self.segment_results.setdefault(sample_info.sample_id, {})
self.segment_results[sample_info.sample_id][chrom] = {
"meth_data": self.meth_data.copy(),
"emissions_df": self.emissions_df.copy(),
"hmm_state": hidden_states,
"hmm_state_readable": readable_states,
}
else:
self.meth_data = self.segment_results[sample_info.sample_id][chrom][
"meth_data"
].copy()
self.emissions_df = self.segment_results[sample_info.sample_id][chrom][
"emissions_df"
].copy()
hidden_states = self.segment_results[sample_info.sample_id][chrom][
"hmm_state"
]
readable_states = self.segment_results[sample_info.sample_id][chrom][
"hmm_state_readable"
]
cache_entry = self.segment_results[sample_info.sample_id][chrom]
# Attach HMM states and refresh the raw contiguous regions.
self.meth_data["hmm_state"] = hidden_states
self.meth_data["hmm_state_readable"] = readable_states
self.regions_df = self.create_regions(
state_col="hmm_state_readable",
region_min_probes=1,
)
cache_entry["meth_data"] = self.meth_data.copy()
cache_entry["regions_df"] = self.regions_df.copy()
# print(f"State relabeling completed in {time.time() - start:.2f} seconds.")
return self.meth_data, self.hmm_model.hmm_model
[docs]
def create_regions(self, state_col="hmm_state_readable", region_min_probes=1):
"""
Create regions (start, end) for contiguous segments of the same state.
Parameters
----------
meth_data : DataFrame
Must contain 'CpG_chrm', 'CpG_beg', 'CpG_end', and state_col.
state_col : str
Column name for the state labels.
region_min_probes : int
Minimum number of probes required to form a region.
Returns
-------
regions_df : DataFrame
Columns: 'CpG_chrm', 'start', 'end', state_col
"""
for col in ["CpG_chrm", "CpG_beg", "CpG_end", state_col]:
if col not in self.meth_data.columns:
raise ValueError(f"Column {col} not found in meth_data.")
regions = []
current_chrom = None
current_state = None
region_start = None
region_end = None
current_meth_sum = 0
current_probe_count = 0
for idx, row in self.meth_data.iterrows():
chrom = row["CpG_chrm"]
state = row[state_col]
beg = row["CpG_beg"]
end = row["CpG_end"]
beta_val = float(row["beta"])
if (chrom != current_chrom) or (state != current_state):
# Save previous region
if (
current_chrom is not None
and current_probe_count >= region_min_probes
):
regions.append(
{
"CpG_chrm": current_chrom,
"start": region_start,
"end": region_end,
"avg_beta": current_meth_sum / current_probe_count,
"probe_count": current_probe_count,
"state": current_state,
}
)
# Start new region
current_chrom = chrom
current_state = state
region_start = beg
region_end = end
current_meth_sum = beta_val
current_probe_count = 1
else:
# Extend current region
region_end = end
current_meth_sum += beta_val
current_probe_count += 1
# Save last region
if current_chrom is not None and current_probe_count >= region_min_probes:
regions.append(
{
"CpG_chrm": current_chrom,
"start": region_start,
"end": region_end,
"avg_beta": current_meth_sum / current_probe_count,
"probe_count": current_probe_count,
"state": current_state,
}
)
regions_df = pd.DataFrame(
regions,
columns=["CpG_chrm", "start", "end", "avg_beta", "probe_count", "state"],
)
self.regions_df = regions_df
return regions_df
[docs]
def regions_to_bed(self, bed_path: str, separate_beds_by_state: bool = False):
"""
Save regions DataFrame to BED file.
Parameters
----------
bed_path
Output path for the BED file. A ``.bed`` suffix is added when it is
missing.
separate_beds_by_state
If ``True``, write one BED per biological state instead of one
combined BED file.
Returns
-------
None
Writes BED file(s) derived from ``self.regions_df``.
"""
if not bed_path.lower().endswith(".bed"):
bed_path += ".bed"
regions_df = self.regions_df.copy()
regions_df["start"] = regions_df["start"].astype(int)
regions_df["end"] = regions_df["end"].astype(int)
if not separate_beds_by_state:
bed_df = regions_df[["CpG_chrm", "start", "end", "state"]].copy()
bed_df.to_csv(bed_path, sep="\t", header=False, index=False)
else:
for state in MethylationStates:
state_df = regions_df[regions_df["state"] == state]
bed_df = state_df[["CpG_chrm", "start", "end", "state"]].copy()
state_bed_path = bed_path.replace(".bed", f"_{state.name}.bed")
bed_df.to_csv(state_bed_path, sep="\t", header=False, index=False)
[docs]
def plot_labels(
self,
sample_info: SampleInfo | None = None,
chrom: str | None = None,
sample_info_removed: pd.DataFrame | None = None,
overlay_regions_df: pd.DataFrame | None = None,
overlay_style: str = "state",
region_start: int | None = None,
region_end: int | None = None,
x_col: str = "CpG_beg",
y_col: str = "beta",
label_title: str | None = None,
show_plot: bool = True,
max_points: int = 120_000,
state_colors: dict | None = None,
):
"""
Plot genomic-position vs beta for HMM labels.
Parameters
----------
sample_info
Sample to segment and plot. When omitted, uses
``default_sample_info``.
chrom
Chromosome to segment and display.
sample_info_removed
Optional table of CpGs removed during preprocessing to show as a
background layer.
overlay_regions_df
Optional region table used to recolor points by overlapping
intervals.
overlay_style
Overlay mode, either ``"state"`` or ``"highlight"``.
region_start
Optional genomic start coordinate for x-axis zooming.
region_end
Optional genomic end coordinate for x-axis zooming.
x_col
Probe-level column used for the x-axis.
y_col
Probe-level column used for the y-axis.
label_title
Optional legend title override.
show_plot
If ``True``, display the Plotly figure immediately.
max_points
Maximum number of plotted points before downsampling.
state_colors
Optional biological-state color overrides.
Returns
-------
plotly.graph_objects.Figure
Interactive beta scatter plot for the resolved HMM labels.
Region args only zoom the x-axis viewport; they do not create a
highlight overlay unless one is passed explicitly.
"""
if chrom is None:
raise ValueError("chrom is required when plotting HMM labels.")
meth_data, _ = self.segment_sample(sample_info=sample_info, chrom=chrom)
resolved_sample_info = (
self.default_sample_info if sample_info is None else sample_info
)
return plot_state_labels(
df_plot=meth_data.copy(),
sample_info=resolved_sample_info,
sample_info_removed=sample_info_removed,
chrom=chrom,
out_dir=self.out_dir,
label_col="hmm_state_readable",
overlay_regions_df=overlay_regions_df,
overlay_style=overlay_style,
region_start=region_start,
region_end=region_end,
x_col=x_col,
y_col=y_col,
label_title=label_title if label_title is not None else "HMM state",
show_plot=show_plot,
max_points=max_points,
state_colors=state_colors,
)