# mutcleaner/cleaners/archstabms_1e10_sup5_cleaner.py
from __future__ import annotations
import logging
import pandas as pd
from pathlib import Path
from typing import TYPE_CHECKING
from dataclasses import dataclass, field
from .basic_cleaners import (
read_dataset,
merge_columns,
convert_data_types,
validate_mutations,
filter_and_clean_data,
extract_and_rename_columns,
apply_mutations_to_sequences,
remap_mutation_positions_by_name,
convert_to_mutation_dataset_format,
)
from .archstabms_1e10_custom_cleaners import (
add_wild_type_sequences_by_library,
convert_pairwise_couplings_to_ddg,
)
from .base_config import BaseCleanerConfig
from ..core.pipeline import Pipeline, create_pipeline
from ..core.dataset import MutationDataset
if TYPE_CHECKING:
from typing import Callable, Optional, Tuple, Dict, Union, Any, List
__all__ = [
"ArchStabMS1E10CleanerSup5Config",
"create_archstabms_1e10_sup5_cleaner",
"clean_archstabms_1e10_sup5_dataset",
]
def __dir__() -> List[str]:
return __all__
# Create module logger
logger = logging.getLogger(__name__)
GRB2_SH3_WT_SEQUENCE = (
"TYVQALFDFDPQEDGELGFRRGDFIHVMDNSDPNWWKGACHGQTGMFPRNYVTPVN"
)
SRC_WT_SEQUENCE = (
"MGSNKSKPKDASQRRRSLEPAENVHGAGGGAFPASQTPSKPASADGHRGPSAAFAPAAAEPK"
"LFGGFNSSDTVTSPQRAGPLAGGVTTFVALYDYESRTETDLSFKKGERLQIVNNTEGDWWLA"
"HSLSTGQTGYIPSNYVAPSDSIQAEEWYFGKITRRESERLLLNAENPRGTFLVRESETTKGAY"
"CLSVSDFDNAKGLNVKHYKIRKLDSGGFYITSRTQFNSLQQLVAYYSKHADGLCHRLTTVCPT"
"SKPQTQGLAKDAWEIPRESLRLEVKLGQGCFGEVWMGTWNGTTRVAIKTLKPGTMSPEAFLQ"
"EAQVMKKLRHEKLVQLYAVVSEEPIYIVTEYMSKGSLLDFLKGETGKYLRLPQLVDMAAQIAS"
"GMAYVERMNYVHRDLRAANILVGENLVCKVADFGLARLIEDNEYTARQGAKFPIKWTAPEAAL"
"YGRFTIKSDVWSFGILLTELTTKGRVPYPGMVNREVLDQVERGYRMPCPPECPESLHDLMCQC"
"WRKEPEERPTFEYLQAFLEDYFTSTEPQYQPGENL"
)
[docs]
@dataclass
class ArchStabMS1E10CleanerSup5Config(BaseCleanerConfig):
"""
Configuration class for ArchStabMS1E10 dataset cleaner.
Inherits from BaseCleanerConfig and adds ArchStabMS1E10-specific configuration options.
Simply run `mutcleaner.download_archstabms1e10_source_file()` to download the dataset.
Alternatively, the raw archstabms1e10 file can be obtained from:
- Hugging Face: https://huggingface.co/datasets/xulab-research/MutCleaner/blob/main/ArchStabMS1E10_Epistasis_Dataset/ArchStabMS1E10_Epistasis_Dataset.csv
Attributes
----------
column_mapping : Dict[str, str]
Mapping from source to target column names
filters : Dict[str, Callable]
Filter conditions for ndata cleaning
type_conversions : Dict[str, str]
Data type conversion specifications
label_columns : List[str]
List of score columns to process
primary_label_column : str
Primary score column for the dataset
"""
# column mapping configuration
column_mapping: Dict[str, str] = field(
default_factory=lambda: {
"name": "name",
"id_ref": "mut_info",
"mean_kcal/mol": "label",
"wt_sequence": "wt_sequence",
}
)
# Data filtering configuration
filters: Dict[str, Callable] = field(
default_factory=lambda: {"conf": lambda x: x == True}
)
library_sequences: Dict[int, str] = field(
default_factory=lambda: {
1: GRB2_SH3_WT_SEQUENCE,
2: GRB2_SH3_WT_SEQUENCE,
3: GRB2_SH3_WT_SEQUENCE,
4: SRC_WT_SEQUENCE,
}
)
# Type conversion configuration
type_conversions: Dict[str, str] = field(
default_factory=lambda: {"label": "float"}
)
# Score columns configuration
label_columns: List[str] = field(default_factory=lambda: ["label"])
primary_label_column: str = "label"
# Override default pipeline name
pipeline_name: str = "archstabms1e10_sup5_cleaner pipeline"
[docs]
def validate(self) -> None:
"""Validate ArchStabMS1E10CleanerConfig
Raises
------
ValueError
If configuration is invalid
"""
# Call parent validation
super().validate()
# Validate score columns
if not self.label_columns:
raise ValueError("label_columns cannot be empty")
if self.primary_label_column not in self.label_columns:
raise ValueError(
f"primary_label_column '{self.primary_label_column}' must be in label_columns {self.label_columns}"
)
# Validate column mapping
required_mappings = {"name", "mean_kcal/mol", "id_ref"}
missing = required_mappings - set(self.column_mapping.keys())
if missing:
raise ValueError(f"Missing required column mappings: {missing}")
[docs]
def create_archstabms_1e10_sup5_cleaner(
dataset_or_path: Optional[Union[pd.DataFrame, str, Path]] = None,
config: Optional[
Union[ArchStabMS1E10CleanerSup5Config, Dict[str, Any], str, Path]
] = None,
) -> Pipeline:
"""Create ArchStabMS1E10 dataset cleaning piipeline
Parameters
----------
dataset_or_path : Optional[Union[pd.DataFrame, str, Path]], default=None
Raw dataset DataFrame or file path to archstams 1e10 dataset
config : Optional[Union[ArchStabMS1E10CleanerSup5Config, Dict[str, Any], str, Path]]. default=None
Configuration for the cleaing pipeline. Can be:
- CDNAProtelysisCleanerConfig object
- Dictionary with configuration parameters (merged with defaults)
- Path to JSON configuration file (str or Path)
- None (uses default configuration)
Returns
-------
Pipeline
Pipeline: The cleaning pipeline used
Raises
------
TypeError
If config has invalid type
ValueError
If configuration validation fails
"""
# Handle configuration parameter
if config is None:
final_config = ArchStabMS1E10CleanerSup5Config()
elif isinstance(config, ArchStabMS1E10CleanerSup5Config):
final_config = config
elif isinstance(config, dict):
default_config = ArchStabMS1E10CleanerSup5Config()
final_config = default_config.merge(config)
elif isinstance(config, (str, Path)):
# Load from file
final_config = ArchStabMS1E10CleanerSup5Config.from_json(config)
else:
raise TypeError(
f"config must be ArchStabMS1E10CleanerSup5Config, dict, str, Path or None, got {type(config)}"
)
# Log configuration summary
logger.info(
f"archstabms 1e10 dataset will be cleaned with pipeline: {final_config.pipeline_name}"
)
logger.debug(f"Configuration:\n{final_config.get_summary()}")
try:
# Create pipeline
pipeline = create_pipeline(dataset_or_path, final_config.pipeline_name)
# Add cleaning steps
pipeline = (
pipeline.delayed_then(
filter_and_clean_data,
filters=final_config.filters,
)
.delayed_then(
add_wild_type_sequences_by_library,
library_sequences=final_config.library_sequences,
)
.delayed_then(
merge_columns,
columns_to_merge=["library", "trait_name"],
new_column_name="name",
drop_original=True,
)
.delayed_then(
extract_and_rename_columns,
column_mapping=final_config.column_mapping,
)
.delayed_then(
validate_mutations,
mutation_sep="_",
)
.delayed_then(
remap_mutation_positions_by_name,
position_offsets={"4_Folding": 3},
name_column=final_config.column_mapping.get("name", "name"),
mutation_column=final_config.column_mapping.get("id_ref", "id_ref")
)
.delayed_then(
convert_data_types,
type_conversions=final_config.type_conversions,
)
.delayed_then(
convert_pairwise_couplings_to_ddg,
group_columns=[
final_config.column_mapping.get("name", "name")
],
mutation_column=final_config.column_mapping.get("id_ref", "id_ref"),
label_column=final_config.column_mapping.get("mean_kcal/mol", "mean_kcal/mol"),
)
.delayed_then(
apply_mutations_to_sequences,
mutation_column=final_config.column_mapping.get("id_ref", "id_ref"),
sequence_column="wt_sequence",
)
.delayed_then(
convert_to_mutation_dataset_format,
name_column=final_config.column_mapping.get("name", "name"),
mutation_column=final_config.column_mapping.get("id_ref", "id_ref"),
sequence_column=final_config.column_mapping.get("wt_sequence", "wt_sequence"),
mutated_sequence_column=final_config.column_mapping.get("mut_seq", "mut_seq"),
label_column=final_config.column_mapping.get("mean_kcal/mol", "mean_kcal/mol"),
is_zero_based=True,
)
)
# Create pipeline based on dataset_or_path type
if isinstance(dataset_or_path, (str, Path)):
pipeline.add_delayed_step(read_dataset, 0)
elif not isinstance(dataset_or_path, pd.DataFrame):
raise TypeError(
f"dataset_or_path must be pd.DataFrame or str/Path, got {type(dataset_or_path)}"
)
return pipeline
except Exception as e:
logger.error(f"Error in creating archstabms cleaning pipeline: {str(e)}")
raise RuntimeError(f"Error in creating archstabms cleaning pipeline: {str(e)}")
[docs]
def clean_archstabms_1e10_sup5_dataset(
pipeline: Pipeline,
) -> Tuple[Pipeline, MutationDataset]:
"""Clean ArchStabMS1E10 dataset using configurable pipeline
Parameters
----------
pipeline : Pipeline
ArchStabMS1E10 dataset cleaning pipeline
Returns
-------
Tuple[Pipeline, MutationDataset]
- Pipeline: The cleaned pipeline
- MutationDataset: The cleaned ArchStabMS1E10 dataset
Examples
--------
Use default configuration:
>>> pipeline = create_archstabms_1e10_sup5_cleaner(df) # df is raw ArchStabMS1E10 dataset file
Use partial configuration:
>>> pipeline = create_archstabms_1e10_sup5_cleaner(df, config={
... "column_mapping": {
... "mean_kcal/mol": "label",
... "id_ref": "mutation_name",
... },
... })
Load configuration from file:
>>> pipeline = create_archstabms_1e10_sup5_cleaner(df, config="config.json")
>>> pipeline, dataset = clean_archstabms_1e10_sup5_dataset(pipeline)
"""
try:
# Run pipeline
pipeline.execute()
# Extract results
archstabms_1e10_dataset_df, archstabms_1e10_ref_seq = pipeline.data
archstabms_1e10_dataset = MutationDataset.from_dataframe(
archstabms_1e10_dataset_df, archstabms_1e10_ref_seq
)
logger.info(
f"Successfully cleaned archstabms1e10 dataset: {len(archstabms_1e10_dataset_df)} mutations from {len(archstabms_1e10_ref_seq)} proteins"
)
return pipeline, archstabms_1e10_dataset
except Exception as e:
logger.error(
f"Error in running archstabms1e10 dataset cleaning pipeline: {str(e)}"
)
raise RuntimeError(
f"Error in running archstabms1e10 dataset cleaning pipeline: {str(e)}"
)