API reference¶

def general_ligand_count(self) -> int:
    """How many screening ligands are rows here — what the importer asks before offering
    to shard, since the two libraries cannot coexist."""
    self._require_active_project()
    return general_ligand_count(self.workspace.project_db)

def set_forced_dependencies(self, job_ids: Optional[list[str]]) -> None:
    """Workflow channel: ids the NEXT submit_job call(s) must also depend on. Lets the
    orchestrator inject same-category serialization that the category auto-deps don't cover,
    without every submit API needing a depends_on parameter. Cleared by passing None."""
    self._forced_dependencies = list(job_ids) if job_ids else None

def submit_job(
    self,
    job,
    *,
    params: Any = None,
    executor_name: Optional[str] = None,
    cpu_required: Optional[int] = None,
    max_job_cpu: Optional[int] = None,
    max_inflight_tasks: Optional[int] = None,
    total_chunks: Optional[int] = None,
    depends_on: Optional[list[str]] = None,
    priority: int = 0,
    result_handler_kwargs: Optional[dict] = None,
    auto_depends: bool = True,
) -> str:
    """Submit a job, auto-waiting for active prerequisite jobs by default.

    With auto_depends=True (default) the job is held until every active job of its
    prerequisite categories finishes, so a deferred pipeline can be queued at once.
    Pass auto_depends=False to submit immediately regardless of running jobs.
    Returns the key of the job in the project's journal, which is its id everywhere here.
    """
    self._require_active_project()
    depends_on = list(depends_on or [])
    depends_on += [job_id for job_id in getattr(self, "_forced_dependencies", None) or () if job_id not in depends_on]
    if auto_depends:
        depends_on = self.resolve_job_dependencies(job, depends_on) or []
    workspace = self.workspace
    job_config = self._job_config()
    engine_job, on_error = job.engine_job(handler_kwargs=result_handler_kwargs)
    job_params = job.validate_params(params)
    total = total_chunks if total_chunks is not None else job.count_chunks(job_params, config=job_config)
    job_id = uuid4().hex
    engine_ids = self._engine_ids()
    engine_ids[job_id] = workspace.submit(
        engine_job,
        job_params,
        key=job_id,
        executor=self._run_on(job, executor_name),
        # A job that ended in an earlier session has nothing left to wait for.
        after=[engine_ids[dep] for dep in depends_on if dep in engine_ids],
        total=total,
        cpu=cpu_required,
        window=max_inflight_tasks,
        max_cpu=max_job_cpu,
        priority=priority,
        on_error=on_error,
        **({"thread": True} if job.on_thread is not None and job.on_thread(job_params) else {}),
    )
    # Single funnel for "a job just started" UI feedback (the main window plugs a toast in
    # here). May fire from a worker thread, so the hook must be thread-safe (a Qt signal).
    hook = getattr(self, "on_job_submitted", None)
    if hook is not None:
        try:
            hook(str(getattr(job, "name", "") or "Job"), job_id)
        except Exception:  # noqa: BLE001 - feedback must never break a submit
            pass
    return job_id

def job_destinations(self, job_id: str) -> list[tuple[str, str]]:
    """`run_destinations` for a job on record: where it may be continued now. Nowhere for one
    that cannot be submitted again."""
    try:
        return self.run_destinations(self._recorded_job(job_id)[0])
    except Exception:  # noqa: BLE001 - not on record, or its declaration is gone
        return []

def resume_job(self, job_id: str, executor_name: Optional[str] = None) -> str:
    """Continue a job that was canceled or interrupted, on `executor_name` or where it was.

    A job that is still running is canceled first and waited for, which is how it is moved
    to another executor: not from a thread that cannot wait. It waits for the jobs its kind
    waits for, as when it was submitted.
    """
    self._require_active_project()
    job_id = str(job_id)
    job, was_on = self._recorded_job(job_id)
    executor = self._run_on(job, executor_name or was_on)
    self._job_config()
    engine_ids = self._engine_ids()
    waits = [engine_ids[dep] for dep in self.resolve_job_dependencies(job) or ()
             if dep in engine_ids and dep != job_id]
    engine_ids[job_id] = self.workspace.move(job_id, executor, after=waits)
    return job_id

def resolve_job_dependencies(self, job, explicit: Optional[list[str]] = None) -> Optional[list[str]]:
    """Union of explicit depends_on and the ids of active jobs this one should wait for."""
    category = _job_category(getattr(job, "name", "") or "")
    prereqs = _JOB_PREREQS.get(category or "", ())
    deps: list[str] = list(explicit or [])
    if prereqs or category == "chemistry":
        for status in self.list_jobs(statuses=NON_TERMINAL_JOB_STATUSES):
            active_category = _job_category(status.task_type)
            serial_chemistry = category == "chemistry" and active_category == "chemistry"
            if (active_category in prereqs or serial_chemistry) and status.job_id not in deps:
                deps.append(status.job_id)
    return deps or None

def wait_for_job(self, job_id: str, *, poll_s: float = 0.2) -> JobStatus:
    """Block until the job has ended. `poll_s` is not used: the engine tells."""
    self._require_active_project()
    engine_id = self._engine_ids().get(str(job_id))
    if engine_id is None:  # of an earlier session: it ended with it
        status = self.get_job_status(job_id)
        if status is None:
            raise KeyError(f"Unknown job: {job_id}")
        return status
    return self._job_status_from_snapshot(self.workspace.engine.wait(engine_id))

def watch_job(self, job_id: str, *, poll_s: float = 0.2) -> JobStatus:
    """Display tqdm progress in a terminal or notebook and return the job status."""
    return self.watch_jobs([job_id], poll_s=poll_s)[str(job_id)]

def watch_jobs(self, job_ids: Iterable[str], *, poll_s: float = 0.2) -> dict[str, JobStatus]:
    """Display concurrent tqdm progress for several jobs of this session."""
    self._require_active_project()
    from ms_flow.progress import watch_jobs

    job_ids = [str(item) for item in job_ids]
    engine_ids = self._engine_ids()
    snapshots = watch_jobs(self.workspace.engine, [engine_ids[job_id] for job_id in job_ids], poll_s=poll_s)
    return {job_id: self._job_status_from_snapshot(snapshots[engine_ids[job_id]]) for job_id in job_ids}

def run_destinations(self, job=None) -> list[tuple[str, str]]:
    """Where a job may be sent now, as (executor_name, label): the executors that are
    active, the default one first. With `job`, only the ones that job can run on.
    """
    self._require_active_project()
    rows = [row for row in self.workspace.executors() if row["state"] == "on"]
    # ponytail: `hpc` in the old list of a job is what says its worker needs nothing of
    # this machine but its payload and the files marked in it; one flag when the lists go.
    if job is not None and "hpc" not in job.supported_executors:
        rows = [row for row in rows if row["hosts"] == ["local"]]
    return [(row["name"], row["name"]) for row in sorted(rows, key=lambda row: not row["default"])]

def shard_ligands(
    self,
    files: Iterable[PathLike],
    *,
    shard_size: int | None = None,
    executor_name: str = "compute",
    depends_on: list[str] | None = None,
    molecule_kind: str = "small_molecule",
    prefilter: Mapping[str, Any] | None = None,
) -> list[str]:
    """The `htpvs` import: split the library into shards, materialize nothing.

    Same feed as `load_ligands` — the byte-span splitter — with the record cap read as
    shard size. What changes is the worker: it writes the span back out as a file and
    registers a `screening_shards` row instead of parsing molecules into the project db.
    """
    normalized_files = normalize_files(files, label="ligand files")
    if not normalized_files:
        return []
    check_library_target(self.runtime.workspace.project_db, target=LIBRARY_SHARDS)
    records = (
        default_shard_size(runtime=self.runtime)
        if shard_size is None
        else int(shard_size)
    )
    records = max(1, min(records, MAX_SHARD_RECORDS))
    params: dict[str, Any] = {
        "file_paths": [str(file_path) for file_path in normalized_files],
        "batch_size": records,
        "molecule_kind": str(molecule_kind or "small_molecule"),
    }
    if prefilter:
        params["prefilter"] = dict(prefilter)
    return [
        self._submit_job(
            shard_ligands_job,
            params=params,
            executor_name=executor_name,
            depends_on=depends_on,
            # ramp=False, like the feed: an over-declared total never completes.
            total_chunks=_total_import_chunks(normalized_files, batch_size=records, ramp=False),
            max_inflight_tasks=import_max_inflight(self.runtime),
            # The queue that cuts the shards lives here, in the parent: it needs the
            # project's shard directory and its database, neither of which a job spec
            # knows at declaration time.
            result_handler_kwargs={
                "project_db": self.runtime.workspace.project_db,
                "shard_dir": self.runtime.get_project_resource_path(RESOURCE_SHARDS),
                "shard_size": records,
                "generation_id": self._import_generation_id(),
            },
        )
    ]

def run_ligand_pipeline(
    self,
    steps: Sequence[str | tuple[str, dict[str, Any]]],
    *,
    ligands: MoleculeSetRef | MoleculeScope | int | None = None,
    structure_source: str = "current",
    batch_size: int = 128,
    executor_name: str = DEFAULT_LOCAL_CPU_EXECUTOR,
    depends_on: list[str] | None = None,
    wait: bool = False,
    shard_state: str = ShardState.PENDING,
) -> str | JobStatus:
    """Run several chemistry steps in one pass over the active screening library.

    `run_ligand_pipeline(["standardize", "protonate", "generate_3d"])` reads and writes every
    ligand once; the three single-step jobs read and rewrite the whole library three times.

    With no explicit row scope, a sharded library dispatches whole shards automatically.
    Passing ``ligands`` always selects the row adapter for curated/reference molecules.

    ponytail: the single-step methods below validate their own arguments (pH range, method
    names, tool installs). Here the step params go through as given and the worker raises.
    """
    normalized = normalize_steps(steps)
    if not normalized:
        raise ValueError("run_ligand_pipeline needs at least one step.")
    for name, _params in normalized:
        if name not in LIGAND_STEPS:
            raise ValueError(f"Unsupported ligand chemistry operation: {name}")
    normalized = [
        (name, self._protonate_params(**params) if name == "protonate" else params)
        for name, params in normalized
    ]
    if ligands is None:
        self.runtime._require_active_project()
        if has_shards(self.runtime.workspace.project_db):
            return self.run_shard_pipeline(
                normalized,
                state=shard_state,
                executor_name=executor_name,
                depends_on=depends_on,
                wait=wait,
            )
    return self._submit_ligand_operation(
        [[name, params] for name, params in normalized],
        ligands=ligands,
        params={},
        structure_source=structure_source,
        batch_size=batch_size,
        executor_name=executor_name,
        depends_on=depends_on,
        wait=wait,
    )

def run_shard_pipeline(
    self,
    steps: Sequence[str | tuple[str, dict[str, Any]]],
    *,
    state: str = ShardState.PENDING,
    params: dict[str, Any] | None = None,
    executor_name: str = DEFAULT_LOCAL_CPU_EXECUTOR,
    depends_on: list[str] | None = None,
    wait: bool = False,
) -> str | JobStatus:
    """The same steps as `run_ligand_pipeline`, over the shards of an `htpvs` project.

    One shard per chunk, one file in and one file out. The run writes into a new
    generation and only switches to it if it finishes, so a crashed run changes nothing
    and the retry simply starts over from the generation that is still the library.
    """
    normalized = normalize_steps(steps)
    if not normalized:
        raise ValueError("run_shard_pipeline needs at least one step.")
    for name, _params in normalized:
        if name not in LIGAND_STEPS:
            raise ValueError(f"Unsupported ligand chemistry operation: {name}")
    normalized = [
        (name, self._protonate_params(**params) if name == "protonate" else params)
        for name, params in normalized
    ]
    if any(name == "conformers" for name, _params in normalized):
        raise ValueError(
            "Sharded conformer ensembles are not supported: one shard record currently "
            "stores one molecular structure. Generate conformers after materialization."
        )
    self.runtime._require_active_project()
    project_db = self.runtime.workspace.project_db
    # Opened here, before anything runs: the feed needs its output directory and the
    # result handler needs its id, and neither can invent one of its own without the two
    # disagreeing about where the generation lives.
    label = "+".join(name for name, _ in normalized)
    shards_root = Path(self.runtime.get_project_resource_path(RESOURCE_SHARDS))
    generation_id = create_generation(project_db, step=label, shard_dir=shards_root)
    job_id = self.runtime.submit_job(
        shard_chemistry_job,
        params=ShardChemistryJobParams(
            operation=[[name, step_params] for name, step_params in normalized],
            params=dict(params or {}),
            state=state,
            generation_id=generation_id,
            output_dir=str(shards_root / f"gen_{generation_id:04d}_{label}"),
        ).model_dump(mode="python"),
        executor_name=executor_name,
        depends_on=depends_on,
        result_handler_kwargs={"project_db": project_db, "generation_id": generation_id},
    )
    return self.runtime.wait_for_job(job_id) if wait else job_id

def predict_receptors(
        self,
        *,
        receptors: MoleculeSetRef | MoleculeScope | int | None = None,
        model: str = FAST_MODEL,
        force: bool = False,
        depends_on: list[str] | None = None,
        wait: bool = False,
) -> str | JobStatus:
    """Predict protein structures from sequence with the ESMFold API (thread executor).

    Only proteins without 3D by default; `force=True` predicts again and repoints current_path
    to the new model. The token comes from ESM_API_KEY and is never stored.
    """
    self.runtime._require_active_project()
    if not api_token():
        raise RuntimeError(f"{ESM_TOKEN_ENV} is not set: paste your Biohub API token first.")
    receptor_set_ref = None if receptors is None or is_molecule_scope(receptors) else ensure_molecule_set_ref(self.runtime, receptors, name="chemistry_receptor_predict_input")
    receptor_scope = scope_payload(receptors) if is_molecule_scope(receptors) else {}
    receptor_filters = dict(receptor_scope.get("filters") or {})
    if "molecule_type" not in receptor_filters:
        receptor_filters["is_receptor"] = True
    if receptor_scope.get("limit") is not None:
        receptor_filters["_limit"] = receptor_scope.get("limit")
    params = ReceptorPredictionJobParams(
        model=model,
        force=bool(force),
        receptor_set_id=None if receptor_set_ref is None else int(receptor_set_ref.id),
        receptor_filters=receptor_filters if force else {**receptor_filters, "has_3d": False},
    )
    project_db = self.runtime.workspace.project_db
    total_items = db_count(project_db, prediction_scope_spec(params))
    # An empty scope with proteins in it means they all have 3D already. An empty scope with
    # no proteins at all may be an import still running: submit and let the job wait for it.
    if total_items == 0 and not force and db_count(project_db, prediction_scope_spec(params.model_copy(update={"receptor_filters": receptor_filters}))):
        raise ValueError("Every protein in scope already has a 3D structure; tick Force to predict again.")
    job_id = self.runtime.submit_job(
        receptor_prediction_job,
        params=params.model_dump(mode="python"),
        executor_name="thread",
        depends_on=depends_on,
        total_chunks=total_items or None,
    )
    if wait:
        return self.runtime.wait_for_job(job_id)
    return job_id

def details(self, complex_ref: ComplexRef | int) -> Details | None:
    """Load the complete Details-panel projection for a receptor-ligand pair."""
    self.runtime._require_active_project()
    complex_id = int(complex_ref.id if isinstance(complex_ref, ComplexRef) else complex_ref)
    with self.runtime.workspace.project_db.get_session() as session:
        pair = session.get(ComplexRecord, complex_id)
        if pair is None:
            return None
        receptor = session.get(MoleculeRecord, int(pair.receptor_molecule_id or 0))
        ligand = session.get(MoleculeRecord, int(pair.ligand_molecule_id or 0))
        activity = (
            session.get(LigandActivity, int(pair.activity_id or 0))
            if int(pair.activity_id or 0) > 0
            else None
        )
    return self.Details(complex=pair, receptor=receptor, ligand=ligand, activity=activity)

def set_activity(
    self,
    *,
    ligand_id: int,
    endpoint: str,
    value: float,
    unit: str = "",
    source: str = "",
    description: str = "",
    replace: bool = True,
) -> ActivityRecord:
    """Create or edit a ligand's activity for an endpoint (activity_type). With
    replace=True (default) it overwrites existing values for that (ligand, endpoint)."""
    self.runtime._require_active_project()
    endpoint_text = str(endpoint or "").strip()
    if not endpoint_text:
        raise ValueError("set_activity requires endpoint.")
    with self.runtime.workspace.project_db.get_session() as session:
        molecule = session.get(MoleculeRecord, int(ligand_id))
        if molecule is None:
            raise ValueError(f"Molecule {ligand_id} does not exist.")
        if replace:
            existing = session.exec(
                select(ActivityRecord)
                .where(ActivityRecord.molecule_id == int(ligand_id))
                .where(ActivityRecord.activity_type == endpoint_text)
            ).all()
            for row in existing:
                session.delete(row)
        record = ActivityRecord(
            molecule_id=int(ligand_id),
            value=float(value),
            unit=str(unit or ""),
            activity_type=endpoint_text,
            description=str(description or ""),
            source=str(source or ""),
        )
        session.add(record)
        molecule.has_activity = True
        session.add(molecule)
        session.commit()
        session.refresh(record)
        return record

def delete_activity(self, *, ligand_id: int, endpoint: str | None = None) -> int:
    """Delete a ligand's activities (all, or only for one endpoint). Returns count removed."""
    self.runtime._require_active_project()
    with self.runtime.workspace.project_db.get_session() as session:
        statement = select(ActivityRecord).where(ActivityRecord.molecule_id == int(ligand_id))
        if endpoint is not None:
            statement = statement.where(ActivityRecord.activity_type == str(endpoint))
        rows = session.exec(statement).all()
        for row in rows:
            session.delete(row)
        remaining = session.exec(
            select(ActivityRecord).where(ActivityRecord.molecule_id == int(ligand_id))
        ).first()
        molecule = session.get(MoleculeRecord, int(ligand_id))
        if molecule is not None:
            molecule.has_activity = remaining is not None
            session.add(molecule)
        session.commit()
        return len(rows)

def activity_rows(self, *, endpoint: str | None = None) -> list[dict[str, Any]]:
    """Activities joined with the ligand name, for the editor table:
    [{molecule_id, name, value, unit}] ordered by ligand id."""
    self.runtime._require_active_project()
    with self.runtime.workspace.project_db.get_session() as session:
        statement = select(ActivityRecord)
        if endpoint is not None:
            statement = statement.where(ActivityRecord.activity_type == str(endpoint))
        rows = list(session.exec(statement.order_by(ActivityRecord.molecule_id)))
        names = {
            int(m.id): str(m.name or "")
            for m in session.exec(
                select(MoleculeRecord).where(MoleculeRecord.id.in_([int(r.molecule_id) for r in rows] or [0]))
            )
        }
    return [
        {"molecule_id": int(r.molecule_id), "name": names.get(int(r.molecule_id), ""),
         "value": float(r.value), "unit": str(r.unit or "")}
        for r in rows
    ]

def list_endpoints(self) -> list[str]:
    """Distinct activity endpoints present in the project."""
    self.runtime._require_active_project()
    with self.runtime.workspace.project_db.get_session() as session:
        return sorted({
            str(r) for r in session.exec(select(ActivityRecord.activity_type)).all() if str(r or "").strip()
        })

def load_activities(
    self,
    file: PathLike,
    *,
    endpoint: str,
    value_key: str = "value",
    match_by: str = "auto",
    key_column: str | None = None,
    unit: str = "",
    unit_key: str | None = None,
    transform: str | None = None,
    kind: str = "continuous",
    source: str = "",
    replace: bool = True,
    delimiter: str | None = None,
    import_missing: bool = False,
    executor_name: str = "compute",
) -> dict[str, Any]:
    """Bulk-load activities from a CSV/TSV.

    match_by: 'auto' (default — prefer InChIKey, then SMILES, then name), 'inchikey', 'smiles'
    or 'name'. Structural keys match the stored representations (InChIKey is robust to SMILES
    formatting). transform (e.g. 'pIC50'/'pKi') converts a concentration value+unit to
    pX = -log10(M) — the modeling endpoint — keeping the raw value in the description.
    replace=True wipes prior activities for the matched ligands+endpoint first (idempotent reload).
    import_missing=True first imports (via the real ligand importer) any CSV row whose molecule
    isn't in the project yet, so a fresh CSV becomes molecules+activities in one call.
    """
    self.runtime._require_active_project()
    file_path = Path(file).expanduser().resolve()
    if not file_path.exists():
        raise FileNotFoundError(file_path)
    if import_missing:
        self._import_missing_molecules(
            file_path, match_by=match_by, key_column=key_column,
            delimiter=delimiter, executor_name=executor_name,
        )
    endpoint_text = str(endpoint or "").strip()
    if not endpoint_text:
        raise ValueError("load_activities requires endpoint.")
    if delimiter is None:
        delimiter = "\t" if file_path.suffix.lower() == ".tsv" else ","
    loaded = skipped_missing = skipped_invalid = 0
    with self.runtime.workspace.project_db.get_session() as session:
        with file_path.open("r", encoding="utf-8", newline="") as handle:
            reader = csv.DictReader(handle, delimiter=delimiter)
            if reader.fieldnames is None:
                raise ValueError(f"Activity file has no header: {file_path}")
            mode, key_col = self._detect_match(reader.fieldnames, match_by, key_column)
            name_to_id: dict[str, int] = {}
            inchikey_to_id: dict[str, int] = {}
            smiles_to_id: dict[str, int] = {}
            if mode == "name":
                name_to_id = {
                    str(r.name): int(r.id)
                    for r in session.exec(select(MoleculeRecord).where(MoleculeRecord.is_ligand.is_(True)))
                    if r.id is not None
                }
            else:
                inchikey_to_id, smiles_to_id = self._structural_ligand_index(session)

            rows_to_add: list[tuple[int, float, str, str]] = []  # (lid, value, unit, description)
            for raw in reader:
                key = str(raw.get(key_col) or "").strip()
                if mode == "name":
                    lid = name_to_id.get(key)
                elif mode == "inchikey":
                    lid = inchikey_to_id.get(key.upper())
                else:  # smiles — canonicalize the CSV smiles, prefer InChIKey identity
                    canon, ikey = _structural_keys_from_smiles(key)
                    lid = (inchikey_to_id.get(ikey) if ikey else None) or (smiles_to_id.get(canon) if canon else None)
                if lid is None:
                    skipped_missing += 1
                    continue
                try:
                    raw_value = float(str(raw.get(value_key) or "").strip())
                except ValueError:
                    skipped_invalid += 1
                    continue
                row_unit = str((raw.get(unit_key) if unit_key else unit) or "").strip()
                if transform:
                    transformed = _to_pchem(raw_value, row_unit)
                    if transformed is None:
                        skipped_invalid += 1
                        continue
                    rows_to_add.append((lid, transformed, transform, f"raw={raw_value} {row_unit}".strip()))
                else:
                    rows_to_add.append((lid, raw_value, row_unit, ""))

        # 'auto' -> categorical when every value is a 0/1 label, else continuous.
        effective_kind = str(kind or "continuous")
        if effective_kind == "auto":
            effective_kind = "categorical" if rows_to_add and all(
                v in (0.0, 1.0) for _, v, *_ in rows_to_add
            ) else "continuous"

        touched = {lid for lid, *_ in rows_to_add}
        if replace and touched:
            for stale in session.exec(
                select(ActivityRecord)
                .where(ActivityRecord.activity_type == endpoint_text)
                .where(ActivityRecord.molecule_id.in_(touched))
            ).all():
                session.delete(stale)
        for lid, value, value_unit, description in rows_to_add:
            session.add(
                ActivityRecord(
                    molecule_id=lid,
                    value=value,
                    unit=value_unit,
                    activity_type=endpoint_text,
                    kind=effective_kind,
                    source=str(source or file_path.name),
                    description=description,
                )
            )
            loaded += 1
        for lid in touched:
            molecule = session.get(MoleculeRecord, lid)
            if molecule is not None:
                molecule.has_activity = True
                session.add(molecule)
        session.commit()
    return {
        "file": str(file_path),
        "endpoint": endpoint_text,
        "kind": effective_kind,
        "match_by": mode,
        "key_column": key_col,
        "transform": transform or "",
        "loaded": loaded,
        "skipped_missing_ligand": skipped_missing,
        "skipped_invalid_value": skipped_invalid,
    }