PlantAI-dev/tools/create_rnaseq_pipeline_minimal.py

#!/usr/bin/env python3
"""Create a small deterministic RNA-seq fixture for align -> quant -> DE validation."""
from __future__ import annotations

import gzip
import random
import shutil
from pathlib import Path

try:
    import yaml  # type: ignore
except Exception as exc:  # pragma: no cover
    raise SystemExit(f"PyYAML is required: {exc}")


PROJECT_ROOT = Path(__file__).resolve().parents[1]
BASE = PROJECT_ROOT / "项目记录" / "rnaseqUpdate" / "testdata" / "rnaseq_pipeline_minimal"
READ_LEN = 50
FRAG_LEN = 160
GENE_LEN = 240
GENE_COUNT = 120


def revcomp(seq: str) -> str:
    return seq.translate(str.maketrans("ACGTN", "TGCAN"))[::-1]


def write_fastq(path: Path, records: list[tuple[str, str]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with gzip.open(path, "wt", encoding="utf-8") as handle:
        for name, seq in records:
            handle.write(f"@{name}\n{seq}\n+\n{'I' * len(seq)}\n")


def poisson_count(rng: random.Random, lam: float) -> int:
    if lam < 30:
        threshold = 2.718281828459045 ** (-lam)
        product = 1.0
        k = 0
        while product > threshold:
            k += 1
            product *= rng.random()
        return max(0, k - 1)
    return max(0, int(round(rng.gauss(lam, lam ** 0.5))))


def expression_count(rng: random.Random, gene_idx: int, sample: str, replicate: int) -> int:
    base = 14 + ((gene_idx * 17) % 31)
    is_treatment = sample.startswith("TRT")
    if gene_idx <= 20:
        scale = 3.2 if is_treatment else 1.0
    elif gene_idx <= 40:
        scale = 0.45 if is_treatment else 2.4
    elif gene_idx <= 55:
        scale = 1.8 if is_treatment else 1.1
    elif gene_idx <= 75:
        scale = 0.75 if is_treatment else 1.35
    else:
        scale = 1.0
    sample_factor = {
        "CK_1": 0.92,
        "CK_2": 1.08,
        "CK_3": 1.01,
        "TRT_1": 0.95,
        "TRT_2": 1.12,
        "TRT_3": 1.04,
    }[sample]
    replicate_factor = 1.0 + ((gene_idx * (replicate + 5)) % 13 - 6) / 45.0
    mean = max(4.0, base * scale * sample_factor * replicate_factor)
    dispersion = 0.18 + ((gene_idx * 7) % 17) / 35.0
    gamma_mean = rng.gammavariate(1.0 / dispersion, mean * dispersion)
    return max(3, poisson_count(rng, gamma_mean))


def main() -> int:
    if BASE.exists():
        shutil.rmtree(BASE)
    input_dir = BASE / "input"
    final_dir = BASE / "workspace" / "job_minimal_pipeline" / "final"
    input_dir.mkdir(parents=True, exist_ok=True)
    final_dir.mkdir(parents=True, exist_ok=True)

    rng = random.Random(20260628)
    gene_positions: list[tuple[str, int, int, str]] = []
    chrom_parts: list[str] = []
    cursor = 1
    for idx in range(1, GENE_COUNT + 1):
        spacer = "N" * 80
        chrom_parts.append(spacer)
        cursor += len(spacer)
        gene_id = f"Gene{idx:03d}"
        seq = "".join(rng.choice("ACGT") for _ in range(GENE_LEN))
        start = cursor
        end = cursor + len(seq) - 1
        gene_positions.append((gene_id, start, end, seq))
        chrom_parts.append(seq)
        cursor = end + 1
    genome = "".join(chrom_parts)

    (input_dir / "Test.genome.fasta").write_text(f">chr1\n{genome}\n", encoding="utf-8")
    with open(input_dir / "Test.gtf", "w", encoding="utf-8") as handle:
        for gene_id, start, end, _seq in gene_positions:
            handle.write(
                f'chr1\tsynthetic\texon\t{start}\t{end}\t.\t+\t.\tgene_id "{gene_id}"; transcript_id "{gene_id}.t1";\n'
            )

    samples = {
        "CK_1": ("CK", 1),
        "CK_2": ("CK", 2),
        "CK_3": ("CK", 3),
        "TRT_1": ("TRT", 1),
        "TRT_2": ("TRT", 2),
        "TRT_3": ("TRT", 3),
    }
    for sample, (_condition, replicate) in samples.items():
        r1_records: list[tuple[str, str]] = []
        r2_records: list[tuple[str, str]] = []
        for gene_idx, (_gene_id, _start, _end, seq) in enumerate(gene_positions, start=1):
            count = expression_count(rng, gene_idx, sample, replicate)
            max_offset = len(seq) - FRAG_LEN
            for read_idx in range(count):
                offset = (gene_idx * 11 + replicate * 7 + read_idx * 13) % (max_offset + 1)
                fragment = seq[offset : offset + FRAG_LEN]
                # Reverse/dUTP-like paired-end layout: mate 1 aligns to the reverse strand,
                # mate 2 aligns to the forward strand of the plus-strand feature.
                r1 = revcomp(fragment[-READ_LEN:])
                r2 = fragment[:READ_LEN]
                name = f"{sample}_Gene{gene_idx:03d}_{read_idx:04d}"
                r1_records.append((f"{name}/1", r1))
                r2_records.append((f"{name}/2", r2))
        write_fastq(final_dir / f"{sample}_1.clean.fastq.gz", r1_records)
        write_fastq(final_dir / f"{sample}_2.clean.fastq.gz", r2_records)

    payload = {
        "species": {"code": "Test", "speciesName": "Synthetic pipeline test", "keggOrganism": ""},
        "rnaseq": {"outputPrefix": "Test", "library_strandness": "reverse"},
        "samples": {
            sample: {
                "run": sample,
                "r1": f"{sample}_1.clean.fastq.gz",
                "r2": f"{sample}_2.clean.fastq.gz",
                "condition": condition,
                "replicate": replicate,
                "library_strandness": "reverse",
                "group_role": "control" if condition == "CK" else "compare",
                "include": True,
            }
            for sample, (condition, replicate) in samples.items()
        },
        "groups": {sample: condition for sample, (condition, _replicate) in samples.items()},
        "experiment": {"design": "two_group", "control_group": "CK", "treatment_groups": ["TRT"]},
        "compare": [["TRT", "CK"]],
        "de": {
            "strategy": "auto",
            "design": "two_group",
            "control_group": "CK",
            "treatment_groups": ["TRT"],
            "pairwise_modes": ["control"],
            "run_pairwise": True,
            "run_lrt": False,
        },
    }
    (input_dir / "samples.yaml").write_text(
        yaml.safe_dump(payload, allow_unicode=True, sort_keys=False),
        encoding="utf-8",
    )
    print(BASE)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())