evoseq 0.1.0__py3-none-any.whl

evoseq/__init__.py

@@ -0,0 +1,3 @@
+from .config import load_config, run_from_config
+
+__all__ = ["load_config", "run_from_config"]

evoseq/cli.py

@@ -0,0 +1,21 @@
+import argparse
+
+from .config import run_from_config
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(description="Run EvoSeq from a TOML config.")
+    parser.add_argument("config", help="Path to an EvoSeq TOML config file.")
+    args = parser.parse_args(argv)
+
+    outputs = run_from_config(args.config)
+    print("EvoSeq run completed.")
+    for key, value in outputs.items():
+        if key.endswith("_paths"):
+            print(f"{key}:")
+            for name, path in value.items():
+                print(f"  {name}: {path}")
+
+
+if __name__ == "__main__":
+    main()

evoseq/config.py

@@ -0,0 +1,107 @@
+from pathlib import Path
+import tomllib
+
+from .preprocess import preprocess_files, preprocess_folder
+from .scoring import export_perbase_logprobs, score_pairs_file
+
+
+def _none_if_blank(value):
+    return None if value == "" else value
+
+
+def load_config(path):
+    with open(path, "rb") as fh:
+        return tomllib.load(fh)
+
+
+def run_from_config(path):
+    config = load_config(path)
+
+    project = config.get("project", {})
+    input_dir = project.get("input_dir", project.get("base_dir", "."))
+
+    preprocess_config = config.get("preprocess", {})
+    scoring_config = config.get("scoring", {})
+    perbase_config = config.get("perbase", {})
+
+    outputs = {}
+    if preprocess_config.get("enabled", True):
+        reference_fasta_path = _none_if_blank(preprocess_config.get("reference_fasta_path"))
+        mutant_fasta_path = _none_if_blank(preprocess_config.get("mutant_fasta_path"))
+        manifest_path = preprocess_config.get("manifest_path", "auto")
+
+        if reference_fasta_path and mutant_fasta_path:
+            evo_df, saved = preprocess_files(
+                reference_fasta_path=reference_fasta_path,
+                mutant_fasta_path=mutant_fasta_path,
+                manifest_path=manifest_path,
+                output_dir=_none_if_blank(
+                    preprocess_config.get("output_dir", preprocess_config.get("out_dir"))
+                ),
+                strict_manifest=preprocess_config.get("strict_manifest", False),
+                progress=preprocess_config.get("progress", True),
+            )
+        else:
+            evo_df, saved = preprocess_folder(
+                input_dir=input_dir,
+                output_dir=_none_if_blank(
+                    preprocess_config.get("output_dir", preprocess_config.get("out_dir"))
+                ),
+                manifest_path=manifest_path,
+                reference_fasta_path=reference_fasta_path,
+                mutant_fasta_path=mutant_fasta_path,
+                dataset_type=preprocess_config.get("dataset_type", "auto"),
+                window_size=preprocess_config.get("window_size"),
+                strict_manifest=preprocess_config.get("strict_manifest", False),
+                progress=preprocess_config.get("progress", True),
+            )
+        outputs["preprocess_df"] = evo_df
+        outputs["preprocess_paths"] = saved
+
+    if scoring_config.get("enabled", False):
+        pairs_path = (
+            _none_if_blank(scoring_config.get("pairs_path"))
+            or outputs.get("preprocess_paths", {}).get("pairs")
+        )
+        if not pairs_path:
+            raise ValueError(
+                "scoring.enabled is true, but no pairs_path was provided and "
+                "preprocessing did not produce a pair table."
+            )
+
+        result_df, paths = score_pairs_file(
+            pairs_path=pairs_path,
+            output_dir=_none_if_blank(
+                scoring_config.get("output_dir", scoring_config.get("result_dir"))
+            ),
+            manifest_path=scoring_config.get("manifest_path", "auto"),
+            model_name=scoring_config.get("model_name", "evo2_7b"),
+            device=scoring_config.get("device", "cuda:0"),
+            local_path=_none_if_blank(scoring_config.get("local_path")),
+            batch_size=scoring_config.get("batch_size", 8),
+            force_reload=scoring_config.get("force_reload", False),
+            require_recommended_gpu=scoring_config.get(
+                "require_recommended_gpu",
+                True,
+            ),
+            progress=scoring_config.get("progress", True),
+        )
+        outputs["scoring_df"] = result_df
+        outputs["scoring_paths"] = paths
+
+    if perbase_config.get("enabled", False):
+        output_path = export_perbase_logprobs(
+            fasta_path=perbase_config["fasta_path"],
+            output_path=_none_if_blank(perbase_config.get("output_path")),
+            output_dir=_none_if_blank(perbase_config.get("output_dir")),
+            model_name=perbase_config.get("model_name", "evo2_7b"),
+            device=perbase_config.get("device", "cuda:0"),
+            center=perbase_config.get("center", 4096),
+            half_window=perbase_config.get("half_window", 320),
+            local_path=_none_if_blank(perbase_config.get("local_path")),
+            progress=perbase_config.get("progress", True),
+        )
+        outputs["perbase_path"] = output_path
+
+    outputs["config_path"] = Path(path)
+    return outputs

evoseq/paths.py

@@ -0,0 +1,43 @@
+from pathlib import Path
+
+
+def common_parent(paths):
+    resolved = [Path(path).expanduser().resolve().parent for path in paths if path]
+    if not resolved:
+        return Path.cwd()
+    if len(resolved) == 1:
+        return resolved[0]
+
+    import os
+
+    return Path(os.path.commonpath([str(path) for path in resolved]))
+
+
+def default_output_dir(kind, *input_paths, base_dir=None):
+    names = {
+        "preprocess": "evoseq_preprocess_output",
+        "scoring": "evoseq_scoring_output",
+        "perbase": "evoseq_perbase_output",
+    }
+    dirname = names.get(kind, f"evoseq_{kind}_output")
+
+    if base_dir:
+        return Path(base_dir) / dirname
+
+    return common_parent(input_paths) / dirname
+
+
+def ensure_output_dir(path, fallback="/content/evoseq_output"):
+    path = Path(path)
+    try:
+        path.mkdir(parents=True, exist_ok=True)
+        test_path = path / ".write_test"
+        test_path.write_text("ok")
+        test_path.unlink(missing_ok=True)
+        return path
+    except OSError as exc:
+        fallback_path = Path(fallback)
+        print(f"Warning: cannot use output directory {path} ({exc}).")
+        print(f"Using fallback output directory: {fallback_path}")
+        fallback_path.mkdir(parents=True, exist_ok=True)
+        return fallback_path

evoseq/preprocess/__init__.py

@@ -0,0 +1,13 @@
+from .pipeline import (
+    prepare_evo2_input,
+    preprocess_files,
+    preprocess_folder,
+    preprocess_from_base_dir,
+)
+
+__all__ = [
+    "prepare_evo2_input",
+    "preprocess_files",
+    "preprocess_folder",
+    "preprocess_from_base_dir",
+]

evoseq/preprocess/discovery.py

@@ -0,0 +1,108 @@
+from pathlib import Path
+
+
+FASTA_SUFFIXES = {".fa", ".fasta", ".fna"}
+
+
+def _score_candidate(path, kind, dataset_type="auto", window_size=None):
+    name = path.name.lower()
+    score = 0
+
+    if kind == "reference":
+        score += 5 if "reference" in name else 0
+        score += 3 if "_ref" in name or "ref_" in name else 0
+    else:
+        score += 5 if "mutant" in name else 0
+        score += 3 if "_mut" in name or "mut_" in name else 0
+
+    if dataset_type and dataset_type != "auto":
+        aliases = {
+            "positive": ["pos", "positive"],
+            "negative": ["neg", "negative"],
+        }.get(dataset_type, [dataset_type])
+        score += 4 if any(alias in name for alias in aliases) else 0
+
+    if window_size:
+        score += 2 if str(window_size) in name else 0
+
+    score -= 4 if "evo2_all" in name else 0
+    score -= 3 if "output" in str(path).lower() else 0
+    return score
+
+
+def infer_dataset_type(base_dir):
+    name = Path(base_dir).name.lower()
+    if "neg" in name or "negative" in name:
+        return "negative"
+    if "pos" in name or "positive" in name:
+        return "positive"
+    return "auto"
+
+
+def discover_manifest(base_dir, manifest_path="auto"):
+    if manifest_path in (None, False):
+        return None
+    if manifest_path != "auto":
+        return Path(manifest_path)
+
+    base_dir = Path(base_dir)
+    candidates = list((base_dir / "data").glob("manifest*.tsv"))
+    candidates += list(base_dir.glob("manifest*.tsv"))
+    return candidates[0] if candidates else None
+
+
+def discover_fasta_pair(
+    base_dir,
+    reference_fasta_path=None,
+    mutant_fasta_path=None,
+    dataset_type="auto",
+    window_size=None,
+):
+    if reference_fasta_path and mutant_fasta_path:
+        return Path(reference_fasta_path), Path(mutant_fasta_path)
+
+    base_dir = Path(base_dir)
+    search_dirs = [base_dir / "data", base_dir]
+    fasta_paths = []
+    for search_dir in search_dirs:
+        if search_dir.exists():
+            fasta_paths.extend(
+                p for p in search_dir.iterdir() if p.suffix.lower() in FASTA_SUFFIXES
+            )
+
+    if dataset_type == "auto":
+        dataset_type = infer_dataset_type(base_dir)
+
+    if not reference_fasta_path:
+        refs = [
+            p for p in fasta_paths
+            if "reference" in p.name.lower() or "_ref" in p.name.lower()
+        ]
+        refs = sorted(
+            refs,
+            key=lambda p: _score_candidate(p, "reference", dataset_type, window_size),
+            reverse=True,
+        )
+        reference_fasta_path = refs[0] if refs else None
+
+    if not mutant_fasta_path:
+        muts = [
+            p for p in fasta_paths
+            if "mutant" in p.name.lower() or "_mut" in p.name.lower()
+        ]
+        muts = sorted(
+            muts,
+            key=lambda p: _score_candidate(p, "mutant", dataset_type, window_size),
+            reverse=True,
+        )
+        mutant_fasta_path = muts[0] if muts else None
+
+    if not reference_fasta_path or not mutant_fasta_path:
+        available = ", ".join(str(p) for p in sorted(fasta_paths)) or "none"
+        raise FileNotFoundError(
+            "Could not discover reference/mutant FASTA files. "
+            "Pass reference_fasta_path and mutant_fasta_path explicitly. "
+            f"Available FASTA files: {available}"
+        )
+
+    return Path(reference_fasta_path), Path(mutant_fasta_path)

evoseq/preprocess/export.py

@@ -0,0 +1,80 @@
+import pandas as pd
+
+from .fasta import wrap_sequence
+
+
+def sanitize_header_value(value):
+    if pd.isna(value):
+        return "NA"
+
+    return str(value).replace(" ", "_").replace("\n", "_").replace("\t", "_")
+
+
+def get_row_value(row, name, default="NA"):
+    return getattr(row, name) if hasattr(row, name) else default
+
+
+def make_fasta_header(row, allele):
+    seq_len = row.ref_len if allele == "ref" else row.mut_len
+
+    fields = [
+        row.id,
+        f"allele={allele}",
+        f"gene={sanitize_header_value(get_row_value(row, 'gene'))}",
+        f"variant={sanitize_header_value(get_row_value(row, 'variant'))}",
+        f"hgvs={sanitize_header_value(get_row_value(row, 'hgvs'))}",
+        f"ann={sanitize_header_value(get_row_value(row, 'annotation'))}",
+        f"type={sanitize_header_value(get_row_value(row, 'variant_type'))}",
+        f"len={seq_len}",
+    ]
+
+    return "|".join(fields)
+
+
+def write_fasta_from_df(table, path, allele):
+    seq_col = "ref_seq" if allele == "ref" else "mut_seq"
+
+    with open(path, "w") as f:
+        for row in table.itertuples(index=False):
+            header = make_fasta_header(row, allele)
+            seq = getattr(row, seq_col)
+
+            f.write(f">{header}\n")
+            f.write(wrap_sequence(seq) + "\n")
+
+
+def export_evo2_input(evo_input_df, out_dir):
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    pairs_path = out_dir / "evo2_pairs.tsv"
+    pair_path = out_dir / "evo2_pair.tsv"
+    ref_path = out_dir / "evo2_reference.fa"
+    mut_path = out_dir / "evo2_mutant.fa"
+    all_path = out_dir / "evo2_all.fa"
+
+    evo_input_df.to_csv(pairs_path, sep="\t", index=False)
+    evo_input_df.to_csv(pair_path, sep="\t", index=False)
+
+    write_fasta_from_df(evo_input_df, ref_path, "ref")
+    write_fasta_from_df(evo_input_df, mut_path, "mut")
+
+    with open(all_path, "w") as fout:
+        for row in evo_input_df.itertuples(index=False):
+            ref_header = make_fasta_header(row, "ref")
+            mut_header = make_fasta_header(row, "mut")
+
+            fout.write(f">{ref_header}\n")
+            fout.write(wrap_sequence(row.ref_seq) + "\n")
+            fout.write(f">{mut_header}\n")
+            fout.write(wrap_sequence(row.mut_seq) + "\n")
+
+    return {
+        "pairs": pairs_path,
+        "pair_tsv": pair_path,
+        "reference": ref_path,
+        "reference_fasta": ref_path,
+        "mutant": mut_path,
+        "mutant_fasta": mut_path,
+        "all": all_path,
+        "all_fasta": all_path,
+    }

evoseq/preprocess/fasta.py

@@ -0,0 +1,29 @@
+from Bio import SeqIO
+
+
+def read_fasta_as_dict(path):
+    records = {}
+    with open(path) as handle:
+        for rec in SeqIO.parse(handle, "fasta"):
+            records[rec.id] = str(rec.seq).upper()
+    return records
+
+
+def check_fasta_pair(ref_records, mut_records):
+    ref_ids = set(ref_records)
+    mut_ids = set(mut_records)
+
+    ref_only = ref_ids - mut_ids
+    mut_only = mut_ids - ref_ids
+    common = ref_ids & mut_ids
+
+    if ref_only:
+        raise ValueError(f"Reference-only IDs exist: { list(sorted(ref_only))[:5] }")
+    if mut_only:
+        raise ValueError(f"Mutant-only IDs exist: { list(sorted(mut_only))[:5] }")
+
+    return common
+
+
+def wrap_sequence(seq, width=80):
+    return "\n".join(seq[i : i + width] for i in range(0, len(seq), width))

evoseq/preprocess/manifest.py

@@ -0,0 +1,65 @@
+import numpy as np
+import pandas as pd
+
+
+def read_manifest(path):
+    df = pd.read_csv(path, sep="\t")
+    unnamed = [c for c in df.columns if str(c).startswith("Unnamed:")]
+    if unnamed:
+        df = df.drop(columns=unnamed)
+    return df
+
+
+def summarize_manifest(df):
+    def nunique_or_zero(column):
+        return df[column].nunique() if column in df.columns else 0
+
+    return {
+        "rows": len(df),
+        "columns": df.shape[1],
+        "unique_record_id": nunique_or_zero("record_id"),
+        "unique_samples": nunique_or_zero("sample"),
+        "unique_genes": nunique_or_zero("gene"),
+        "unique_hgvs": nunique_or_zero("hgvs"),
+        "unique_spdi": nunique_or_zero("spdi"),
+    }
+
+
+def aggregate_manifest(df, record_id_col="record_id"):
+    if record_id_col not in df.columns:
+        raise ValueError(f"Manifest is missing required column: {record_id_col}")
+
+    grouped = df.groupby(record_id_col, dropna=False)
+    out = grouped.size().rename("n_manifest_rows").reset_index()
+
+    if record_id_col != "record_id":
+        out = out.rename(columns={record_id_col: "record_id"})
+
+    optional_first = [
+        "gene",
+        "hgvs",
+        "annotation",
+        "chrom",
+        "pos1",
+        "ref",
+        "alt",
+        "spdi",
+    ]
+    for column in optional_first:
+        if column in df.columns:
+            values = grouped[column].first().reset_index(drop=True)
+            out[f"{column}_manifest" if column == "gene" else column] = values
+
+    if "sample" in df.columns:
+        out["samples_joined"] = grouped["sample"].agg(
+            lambda x: ";".join(sorted(set(map(str, x.dropna()))))
+        ).reset_index(drop=True)
+
+    if "note" in df.columns:
+        out["note"] = grouped["note"].agg(
+            lambda x: ";".join(sorted(set(map(str, x.dropna()))))
+            if x.notna().any()
+            else np.nan
+        ).reset_index(drop=True)
+
+    return out