API Reference

ovwt

read_feature_file(file_path)

Reads a feature file and returns a Polars DataFrame.

Parameters:

Name	Type	Description	Default
file_path	PathLike	Path to the feature file. Supported formats: .parquet, .pq, .csv.	required

Returns:

Type	Description
DataFrame	pl.DataFrame: The feature data as a Polars DataFrame.

Raises:

Type	Description
ValueError	If the file extension is not supported.

Source code in src/ovwt/__init__.py

def read_feature_file(file_path: PathLike) -> pl.DataFrame:
    """
    Reads a feature file and returns a Polars DataFrame.

    Args:
        file_path (PathLike):
            Path to the feature file. Supported formats: .parquet, .pq, .csv.

    Returns:
        pl.DataFrame:
            The feature data as a Polars DataFrame.

    Raises:
        ValueError:
            If the file extension is not supported.
    """
    path = pathlib.Path(file_path)
    suffix = path.suffix.lower()
    if suffix in [".parquet", ".pq"]:
        return pl.read_parquet(path)
    elif suffix == ".csv":
        return pl.read_csv(path)
    else:
        raise ValueError(
            f"Unsupported file format: {suffix!r}. Expected .parquet or .csv"
        )

get_feature_cols(df)

Returns the CellProfiler feature columns from a DataFrame.

Infers feature columns as those whose name starts with an uppercase letter and contains an underscore, matching CellProfiler naming conventions.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame to extract feature column names from.	required

Returns:

Type	Description
list[str]	list[str]: Column names identified as CellProfiler features.

Source code in src/ovwt/__init__.py

def get_feature_cols(df: pl.DataFrame) -> list[str]:
    """
    Returns the CellProfiler feature columns from a DataFrame.

    Infers feature columns as those whose name starts with an uppercase letter
    and contains an underscore, matching CellProfiler naming conventions.

    Args:
        df (pl.DataFrame):
            DataFrame to extract feature column names from.

    Returns:
        list[str]:
            Column names identified as CellProfiler features.
    """
    return [
        col for col in df.columns if len(col) > 0 and col[0].isupper() and "_" in col
    ]

filter_min_cells(data_df, label_col, wt_label, min_cells)

Remove variants (non-wildtype labels) with fewer than min_cells cells.

Wildtype rows are always retained regardless of count.

Parameters:

Name	Type	Description	Default
data_df	DataFrame	Input DataFrame containing a label column.	required
label_col	str	Name of the column holding class labels.	required
wt_label	str	The label value that identifies wildtype cells.	required
min_cells	int	Minimum number of cells a variant must have to be retained.	required

Returns:

Type	Description
DataFrame	pl.DataFrame: DataFrame with under-represented variants removed.

Source code in src/ovwt/__init__.py

def filter_min_cells(
    data_df: pl.DataFrame,
    label_col: str,
    wt_label: str,
    min_cells: int,
) -> pl.DataFrame:
    """
    Remove variants (non-wildtype labels) with fewer than ``min_cells`` cells.

    Wildtype rows are always retained regardless of count.

    Args:
        data_df (pl.DataFrame):
            Input DataFrame containing a label column.
        label_col (str):
            Name of the column holding class labels.
        wt_label (str):
            The label value that identifies wildtype cells.
        min_cells (int):
            Minimum number of cells a variant must have to be retained.

    Returns:
        pl.DataFrame:
            DataFrame with under-represented variants removed.
    """
    variant_counts = (
        data_df.filter(pl.col(label_col) != wt_label).group_by(label_col).len()
    )
    keep_labels = (
        variant_counts.filter(pl.col("len") >= min_cells)
        .get_column(label_col)
        .to_list()
    )
    return data_df.filter(
        (pl.col(label_col) == wt_label) | pl.col(label_col).is_in(keep_labels)
    )

downsample_wildtype(data_df, label_col, wt_label, seed)

Downsample wildtype cells to the count of the largest non-wildtype variant.

Parameters:

Name	Type	Description	Default
data_df	DataFrame	Input DataFrame containing a label column.	required
label_col	str	Name of the column holding class labels.	required
wt_label	str	The label value that identifies wildtype cells.	required
seed	int	Random seed for reproducible sampling.	required

Returns:

Type	Description
DataFrame	pl.DataFrame: DataFrame with wildtype rows downsampled. Non-wildtype rows are unchanged. Row order is not guaranteed.

Source code in src/ovwt/__init__.py

def downsample_wildtype(
    data_df: pl.DataFrame,
    label_col: str,
    wt_label: str,
    seed: int,
) -> pl.DataFrame:
    """
    Downsample wildtype cells to the count of the largest non-wildtype variant.

    Args:
        data_df (pl.DataFrame):
            Input DataFrame containing a label column.
        label_col (str):
            Name of the column holding class labels.
        wt_label (str):
            The label value that identifies wildtype cells.
        seed (int):
            Random seed for reproducible sampling.

    Returns:
        pl.DataFrame:
            DataFrame with wildtype rows downsampled. Non-wildtype rows are
            unchanged. Row order is not guaranteed.
    """
    max_variant_count = (
        data_df.filter(pl.col(label_col) != wt_label)
        .group_by(label_col)
        .len()
        .get_column("len")
        .max()
    )
    wt_df = data_df.filter(pl.col(label_col) == wt_label)
    if max_variant_count is not None and len(wt_df) > max_variant_count:
        wt_df = wt_df.sample(n=max_variant_count, seed=seed)
    return pl.concat([data_df.filter(pl.col(label_col) != wt_label), wt_df])

train_test_val_split(data_df, cfg)

Splits the data into an 8:1:1 train/test/validation split.

The split is stratified based on the label column to ensure that the distribution of classes is preserved across the train, test, and validation sets.

Parameters:

Name	Type	Description	Default
data_df	DataFrame	The input data as a Polars DataFrame.	required
cfg	DictConfig	Hydra config. Uses cfg.app.label_col, cfg.app.seed, and cfg.app.feature_cols. If cfg.app.feature_cols is None, feature columns are inferred via get_feature_cols.	required

Returns:

Type	Description
tuple[DataFrame, DataFrame, DataFrame]	tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame]: DataFrames for train, test, and validation sets respectively, each containing the feature columns and the label column.

Source code in src/ovwt/__init__.py

def train_test_val_split(
    data_df: pl.DataFrame,
    cfg: DictConfig,
) -> tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame]:
    """
    Splits the data into an 8:1:1 train/test/validation split.

    The split is stratified based on the label column to ensure that the distribution of
    classes is preserved across the train, test, and validation sets.

    Args:
        data_df (pl.DataFrame):
            The input data as a Polars DataFrame.
        cfg (DictConfig):
            Hydra config. Uses cfg.app.label_col, cfg.app.seed, and
            cfg.app.feature_cols. If cfg.app.feature_cols is None, feature
            columns are inferred via `get_feature_cols`.

    Returns:
        tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame]:
            DataFrames for train, test, and validation sets respectively, each
            containing the feature columns and the label column.
    """
    label_col = cfg.app.label_col
    if cfg.app.feature_cols is not None:
        feature_cols = list(cfg.app.feature_cols)
    else:
        feature_cols = get_feature_cols(data_df)

    select_cols = feature_cols + [label_col]
    data_df = data_df.select(select_cols)
    data_df = data_df.filter(pl.col(label_col).is_not_null())

    if cfg.app.min_cells is not None:
        data_df = filter_min_cells(
            data_df, label_col, cfg.app.wt_label, cfg.app.min_cells
        )

    if cfg.app.downsample_wt:
        data_df = downsample_wildtype(
            data_df, label_col, cfg.app.wt_label, cfg.app.seed
        )

    data_df = data_df.with_row_index("__idx__")
    labels = data_df.get_column(label_col).to_numpy()
    all_idx = data_df.get_column("__idx__").to_numpy()

    train_idx, val_test_idx = sklearn.model_selection.train_test_split(
        all_idx,
        test_size=0.2,
        stratify=labels,
        random_state=cfg.app.seed,
    )

    test_idx, val_idx = sklearn.model_selection.train_test_split(
        val_test_idx,
        test_size=0.5,
        stratify=labels[val_test_idx],
        random_state=cfg.app.seed,
    )

    def select_rows(idx: np.ndarray) -> pl.DataFrame:
        return data_df.filter(pl.col("__idx__").is_in(idx)).select(select_cols)

    return select_rows(train_idx), select_rows(test_idx), select_rows(val_idx)

train_xgboost(train, val, cfg)

Trains an XGBoost classifier on the provided training data.

Parameters:

Name	Type	Description	Default
train	DataFrame	Training data including feature columns and the label column.	required
val	DataFrame	Validation data including feature columns and the label column.	required
cfg	DictConfig	Hydra config. Uses cfg.app.label_col, cfg.app.wt_label, and cfg.xgboost.num_boost_round, cfg.xgboost.early_stopping_rounds, cfg.xgboost.weigh_samples, and cfg.xgboost.params (passed directly to xgb.train, with objective, eval_metric, and seed added).	required

Returns:

Type	Description
Booster	xgb.Booster: The trained XGBoost booster.

Source code in src/ovwt/__init__.py

def train_xgboost(
    train: pl.DataFrame,
    val: pl.DataFrame,
    cfg: DictConfig,
) -> xgb.Booster:
    """
    Trains an XGBoost classifier on the provided training data.

    Args:
        train (pl.DataFrame):
            Training data including feature columns and the label column.
        val (pl.DataFrame):
            Validation data including feature columns and the label column.
        cfg (DictConfig):
            Hydra config. Uses cfg.app.label_col, cfg.app.wt_label, and
            cfg.xgboost.num_boost_round, cfg.xgboost.early_stopping_rounds,
            cfg.xgboost.weigh_samples, and cfg.xgboost.params (passed directly
            to xgb.train, with objective, eval_metric, and seed added).

    Returns:
        xgb.Booster:
            The trained XGBoost booster.
    """
    label_col = cfg.app.label_col
    wt_label = cfg.app.wt_label

    y_train = convert_labels_to_boolean(
        train.get_column(label_col).to_numpy(), wt_label
    )
    sample_weight = (
        sklearn.utils.compute_sample_weight("balanced", y_train)
        if cfg.xgboost.weigh_samples
        else None
    )

    dtrain = get_dmatrix(train, label_col, wt_label, weight=sample_weight)
    deval = get_dmatrix(val, label_col, wt_label)

    params = dict(cfg.xgboost.params)
    params["objective"] = "binary:logistic"
    params["eval_metric"] = "auc"
    params["seed"] = cfg.app.seed

    return xgb.train(
        params,
        dtrain,
        num_boost_round=cfg.xgboost.num_boost_round,
        evals=[(dtrain, "train"), (deval, "eval")],
        early_stopping_rounds=cfg.xgboost.early_stopping_rounds,
        verbose_eval=True,
    )

test_xgboost(model, train, val, test, cfg)

Computes the train, validation, and test AUC and accuracy.

Parameters:

Name	Type	Description	Default
model	Booster	The trained XGBoost booster to evaluate.	required
train	DataFrame	Training DataFrame including feature and label columns.	required
val	DataFrame	Validation DataFrame including feature and label columns.	required
test	DataFrame	Test DataFrame including feature and label columns.	required
cfg	DictConfig	Hydra config. Uses cfg.app.label_col and cfg.app.wt_label.	required

Returns:

Name	Type	Description
dict	dict	A dictionary with the keys: - "variant": The first non wt_label value in the label column. - "train_auroc": The AUC on the training set. - "train_accuracy": The accuracy on the training set. - "val_auroc": The AUC on the validation set. - "val_accuracy": The accuracy on the validation set. - "test_auroc": The AUC on the test set. - "test_accuracy": The accuracy on the test set.

Source code in src/ovwt/__init__.py

def test_xgboost(
    model: xgb.Booster,
    train: pl.DataFrame,
    val: pl.DataFrame,
    test: pl.DataFrame,
    cfg: DictConfig,
) -> dict:
    """
    Computes the train, validation, and test AUC and accuracy.

    Args:
        model (xgb.Booster):
            The trained XGBoost booster to evaluate.
        train (pl.DataFrame):
            Training DataFrame including feature and label columns.
        val (pl.DataFrame):
            Validation DataFrame including feature and label columns.
        test (pl.DataFrame):
            Test DataFrame including feature and label columns.
        cfg (DictConfig):
            Hydra config. Uses cfg.app.label_col and cfg.app.wt_label.

    Returns:
        dict:
            A dictionary with the keys:
                - "variant": The first non wt_label value in the label column.
                - "train_auroc": The AUC on the training set.
                - "train_accuracy": The accuracy on the training set.
                - "val_auroc": The AUC on the validation set.
                - "val_accuracy": The accuracy on the validation set.
                - "test_auroc": The AUC on the test set.
                - "test_accuracy": The accuracy on the test set.
    """
    label_col = cfg.app.label_col
    wt_label = cfg.app.wt_label

    variant = next(
        v for v in train.get_column(label_col).unique().to_list() if v != wt_label
    )

    evaluate_wrapper = functools.partial(
        evaluate, model=model, label_col=label_col, wt_label=wt_label
    )

    train_auroc, train_accuracy = evaluate_wrapper(train)
    val_auroc, val_accuracy = evaluate_wrapper(val)
    test_auroc, test_accuracy = evaluate_wrapper(test)

    return {
        "variant": variant,
        "train_auroc": train_auroc,
        "train_accuracy": train_accuracy,
        "val_auroc": val_auroc,
        "val_accuracy": val_accuracy,
        "test_auroc": test_auroc,
        "test_accuracy": test_accuracy,
    }

evaluate(df, model, label_col, wt_label)

Evaluates an XGBoost model on a dataset, returning AUROC and accuracy.

Parameters:

Name	Type	Description	Default
df	DataFrame	Dataset including feature columns and the label column.	required
model	Booster	The trained XGBoost booster to evaluate.	required
label_col	str	The name of the label column.	required
wt_label	str	The label value corresponding to the wild-type (positive) class.	required

Returns:

Type	Description
tuple[float, float]	tuple[float, float]: A tuple of (AUROC, accuracy), where accuracy is computed at a decision threshold of 0.5.

Source code in src/ovwt/__init__.py

def evaluate(
    df: pl.DataFrame, model: xgb.Booster, label_col: str, wt_label: str
) -> tuple[float, float]:
    """
    Evaluates an XGBoost model on a dataset, returning AUROC and accuracy.

    Args:
        df (pl.DataFrame):
            Dataset including feature columns and the label column.
        model (xgb.Booster):
            The trained XGBoost booster to evaluate.
        label_col (str):
            The name of the label column.
        wt_label (str):
            The label value corresponding to the wild-type (positive) class.

    Returns:
        tuple[float, float]:
            A tuple of (AUROC, accuracy), where accuracy is computed at a
            decision threshold of 0.5.
    """
    dmatrix = get_dmatrix(df, label_col, wt_label)
    y_true = dmatrix.get_label()
    y_prob = model.predict(dmatrix)
    auroc = sklearn.metrics.roc_auc_score(y_true, y_prob)
    accuracy = sklearn.metrics.accuracy_score(y_true, y_prob >= 0.5)

    return auroc, accuracy

get_dmatrix(df, label_col, wt_label, weight=None)

Converts a Polars DataFrame into an XGBoost DMatrix.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame containing feature columns and the label column.	required
label_col	str	The name of the label column.	required
wt_label	str	The label value corresponding to the wild-type (positive) class.	required
weight	Optional[ndarray]	Sample weights. Default is None.	None

Returns:

Type	Description
DMatrix	xgb.DMatrix: The XGBoost DMatrix with boolean labels.

Source code in src/ovwt/__init__.py

def get_dmatrix(
    df: pl.DataFrame,
    label_col: str,
    wt_label: str,
    weight: Optional[np.ndarray] = None,
) -> xgb.DMatrix:
    """
    Converts a Polars DataFrame into an XGBoost DMatrix.

    Args:
        df (pl.DataFrame):
            DataFrame containing feature columns and the label column.
        label_col (str):
            The name of the label column.
        wt_label (str):
            The label value corresponding to the wild-type (positive) class.
        weight (Optional[np.ndarray], optional):
            Sample weights. Default is None.

    Returns:
        xgb.DMatrix:
            The XGBoost DMatrix with boolean labels.
    """
    feature_cols = [col for col in df.columns if col != label_col]
    x = df.select(feature_cols).cast(pl.Float64).to_numpy().copy()
    x[~np.isfinite(x)] = np.nan
    y = convert_labels_to_boolean(df.get_column(label_col).to_numpy(), wt_label)
    return xgb.DMatrix(x, label=y, weight=weight)

convert_labels_to_boolean(labels, wt_label)

Converts an array of labels to boolean values.

Parameters:

Name	Type	Description	Default
labels	ndarray	An array of labels to be converted.	required
wt_label	str	The label corresponding to the positive class (True).	required

Source code in src/ovwt/__init__.py

def convert_labels_to_boolean(labels: np.ndarray, wt_label: str) -> np.ndarray:
    """
    Converts an array of labels to boolean values.

    Args:
        labels (np.ndarray):
            An array of labels to be converted.
        wt_label (str):
            The label corresponding to the positive class (True).
    """
    return labels == wt_label

configure_logging(log_file=None, level='INFO')

Configures root logger with a stdout handler and optional file handler.

No-ops if the root logger already has handlers (e.g. called twice).

Parameters:

Name	Type	Description	Default
log_file	Optional[PathLike]	If provided, log messages are also written to this file.	None
level	str	Logging level name (e.g. "INFO", "DEBUG"). Case-insensitive.	'INFO'

Source code in src/ovwt/__init__.py

def configure_logging(
    log_file: Optional[PathLike] = None,
    level: str = "INFO",
) -> None:
    """
    Configures root logger with a stdout handler and optional file handler.

    No-ops if the root logger already has handlers (e.g. called twice).

    Args:
        log_file (Optional[PathLike]):
            If provided, log messages are also written to this file.
        level (str):
            Logging level name (e.g. ``"INFO"``, ``"DEBUG"``). Case-insensitive.
    """
    if logging.root.handlers:
        return

    handlers: list[logging.Handler] = [logging.StreamHandler(sys.stdout)]

    if log_file is not None:
        handlers.append(logging.FileHandler(log_file))

    logging.basicConfig(
        level=getattr(logging, level.upper()),
        format="%(asctime)s | %(levelname)s | %(message)s",
        handlers=handlers,
    )

log_config(cfg)

Logs the Hydra config at INFO level.

Parameters:

Name	Type	Description	Default
cfg	DictConfig	The Hydra config to log.	required

Source code in src/ovwt/__init__.py

def log_config(cfg: DictConfig) -> None:
    """
    Logs the Hydra config at INFO level.

    Args:
        cfg (DictConfig):
            The Hydra config to log.
    """
    logging.info("Config:\n%s", OmegaConf.to_yaml(cfg))

profile_variant(v, train_all, test_all, val_all, cfg)

Train and evaluate an XGBoost classifier for a single variant vs. wild-type.

Filters each split to rows belonging to v or the wild-type label, trains a model on the training split, and evaluates it on all three splits.

Parameters:

Name	Type	Description	Default
v	str	The variant label to profile.	required
train_all	DataFrame	Full training split (all variants and wild-type).	required
test_all	DataFrame	Full test split.	required
val_all	DataFrame	Full validation split.	required
cfg	DictConfig	Hydra config. Uses cfg.app.label_col, cfg.app.wt_label, and cfg.xgboost settings.	required

Returns:

Type	Description
tuple[dict, Booster]	tuple[dict, xgb.Booster]: A result dict (as returned by test_xgboost) and the trained xgb.Booster.

Source code in src/ovwt/__init__.py

def profile_variant(
    v: str,
    train_all: pl.DataFrame,
    test_all: pl.DataFrame,
    val_all: pl.DataFrame,
    cfg: DictConfig,
) -> tuple[dict, xgb.Booster]:
    """
    Train and evaluate an XGBoost classifier for a single variant vs. wild-type.

    Filters each split to rows belonging to ``v`` or the wild-type label, trains
    a model on the training split, and evaluates it on all three splits.

    Args:
        v (str):
            The variant label to profile.
        train_all (pl.DataFrame):
            Full training split (all variants and wild-type).
        test_all (pl.DataFrame):
            Full test split.
        val_all (pl.DataFrame):
            Full validation split.
        cfg (DictConfig):
            Hydra config. Uses cfg.app.label_col, cfg.app.wt_label, and
            cfg.xgboost settings.

    Returns:
        tuple[dict, xgb.Booster]:
            A result dict (as returned by ``test_xgboost``) and the trained
            ``xgb.Booster``.
    """
    keep = pl.col(cfg.app.label_col).is_in([v, cfg.app.wt_label])
    train, test, val = (
        train_all.filter(keep),
        test_all.filter(keep),
        val_all.filter(keep),
    )
    logging.info(
        "Subset sizes — train: %d, val: %d, test: %d",
        len(train),
        len(val),
        len(test),
    )
    model = train_xgboost(train, val, cfg)
    result = test_xgboost(model, train, val, test, cfg)
    logging.info(
        "Results for '%s': train_auroc=%.4f, val_auroc=%.4f, test_auroc=%.4f",
        v,
        result["train_auroc"],
        result["val_auroc"],
        result["test_auroc"],
    )
    return result, model

main(cfg)

Trains and evaluates one XGBoost classifier per variant vs. wild-type.

Performs a single stratified 8:1:1 train/test/val split on the full dataset, then for each unique non-wild-type label trains an XGBoost model on the rows belonging to that variant or the wild-type. Results are written to results.csv and trained models are pickled to models.pkl in out_dir.

Parameters:

Name	Type	Description	Default
cfg	DictConfig	Hydra config with two groups: - cfg.app: feature_file, label_col, wt_label, out_dir - cfg.xgboost: XGBoost hyperparameters and training options	required

Source code in src/ovwt/__init__.py

@hydra.main(config_path="pkg://ovwt.conf", config_name="config", version_base=None)
def main(cfg: DictConfig) -> None:
    """
    Trains and evaluates one XGBoost classifier per variant vs. wild-type.

    Performs a single stratified 8:1:1 train/test/val split on the full
    dataset, then for each unique non-wild-type label trains an XGBoost model
    on the rows belonging to that variant or the wild-type. Results are written
    to ``results.csv`` and trained models are pickled to ``models.pkl`` in
    ``out_dir``.

    Args:
        cfg (DictConfig):
            Hydra config with two groups:
                - cfg.app: feature_file, label_col, wt_label, out_dir
                - cfg.xgboost: XGBoost hyperparameters and training options
    """
    out_dir = pathlib.Path(cfg.app.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    configure_logging(out_dir / "ovwt.log", level=cfg.app.log_level)
    log_config(cfg)

    logging.info("Reading feature file: %s", cfg.app.feature_file)
    feature_df = read_feature_file(cfg.app.feature_file)
    train_all, test_all, val_all = train_test_val_split(feature_df, cfg)
    unique_vars = train_all.get_column(cfg.app.label_col).unique().to_list()
    variants = [v for v in unique_vars if v != cfg.app.wt_label]

    logging.info("Found %d variant(s) to profile", len(variants))
    logging.info(
        "Split sizes — train: %d, val: %d, test: %d",
        len(train_all),
        len(val_all),
        len(test_all),
    )

    if cfg.app.save_splits:
        for name, split_df in (
            ("train", train_all),
            ("test", test_all),
            ("val", val_all),
        ):
            split_path = out_dir / f"{name}.parquet"
            split_df.write_parquet(split_path)
            logging.info("Wrote %s split to %s", name, split_path)

    results = []
    models = {}

    for v in variants:
        logging.info("Training model for variant '%s' vs. '%s'", v, cfg.app.wt_label)
        try:
            result, model = profile_variant(v, train_all, test_all, val_all, cfg)
        except Exception:
            logging.warning(
                "Failed to profile variant '%s', skipping:\n%s",
                v,
                traceback.format_exc(),
            )
            continue
        results.append(result)
        models[v] = model

    results_df = pl.DataFrame(results)
    results_path = out_dir / "results.csv"
    results_df.write_csv(results_path)
    logging.info("Results written to %s", results_path)

    models_path = out_dir / "models.pkl"
    logging.info("Writing models to %s", models_path)
    with open(models_path, "wb") as f:
        pickle.dump(models, f)

    logging.info("Done")