diff --git a/.gitignore b/.gitignore index d871bad3..146a4e91 100644 --- a/.gitignore +++ b/.gitignore @@ -53,6 +53,9 @@ benchmarks/.env.* profiling/.env profiling/.env.* !profiling/.env.example +scarf/agent/.env +scarf/agent/.env.* +!scarf/agent/.env.example /profiling/config.toml /profiling/config.*.toml !/profiling/config.example.toml diff --git a/pyproject.toml b/pyproject.toml index 6bbc748e..c917eeb8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,9 @@ dependencies = [ ] [project.optional-dependencies] +agent = [ + "pydantic-ai-slim[openai]", +] extra = [ "anndata>=0.12", "kneed>=0.8", diff --git a/scarf/agent/.env.example b/scarf/agent/.env.example new file mode 100644 index 00000000..9ead698e --- /dev/null +++ b/scarf/agent/.env.example @@ -0,0 +1,15 @@ +# Copy to .env and fill in values: +# cp scarf/agent/.env.example scarf/agent/.env +# +# Local Ollama (no API key needed): +# OLLAMA_BASE_URL=http://localhost:11434/v1 +# OLLAMA_MODEL=qwen3.5:4b +# +# Ollama Cloud: +# OLLAMA_BASE_URL=https://ollama.com/v1 +# OLLAMA_API_KEY=your-key-here +# OLLAMA_MODEL=qwen3.5:4b + +OLLAMA_BASE_URL=http://localhost:11434/v1 +OLLAMA_MODEL=qwen3.5:4b +# OLLAMA_API_KEY= diff --git a/scarf/agent/__init__.py b/scarf/agent/__init__.py new file mode 100644 index 00000000..8c5e3a5d --- /dev/null +++ b/scarf/agent/__init__.py @@ -0,0 +1,39 @@ +"""Optional grounded decision helpers for Scarf workflows.""" + +from .characterize_covariates import ( + CovariateCharacterization, + characterize_covariates, +) +from .characterize_features import ( + FeatureCharacterization, + characterize_features, +) +from .decide import DecisionValidationError, decide +from .ingest import IngestResult, detect_format, ingest +from .runtime import check_runtime, load_env +from .types import ( + Decision, + EvidenceItem, + NeedsInput, + StageResult, + StageStatus, +) + +__all__ = [ + "CovariateCharacterization", + "Decision", + "DecisionValidationError", + "EvidenceItem", + "FeatureCharacterization", + "IngestResult", + "NeedsInput", + "StageResult", + "StageStatus", + "characterize_covariates", + "characterize_features", + "check_runtime", + "decide", + "detect_format", + "ingest", + "load_env", +] diff --git a/scarf/agent/_deps.py b/scarf/agent/_deps.py new file mode 100644 index 00000000..b248e2dc --- /dev/null +++ b/scarf/agent/_deps.py @@ -0,0 +1,25 @@ +"""Lazy loaders for optional agent dependencies.""" + +from typing import Any + +AGENT_INSTALL_HINT = ( + "Scarf agent support requires the agent extra. " + "Install with: pip install 'scarf[agent]' " + "or: uv sync --extra agent" +) + + +def require_pydantic() -> tuple[Any, Any]: + try: + from pydantic import BaseModel, Field + except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + return BaseModel, Field + + +def require_pydantic_ai() -> Any: + try: + import pydantic_ai + except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + return pydantic_ai diff --git a/scarf/agent/characterize_covariates.py b/scarf/agent/characterize_covariates.py new file mode 100644 index 00000000..74ade7a8 --- /dev/null +++ b/scarf/agent/characterize_covariates.py @@ -0,0 +1,1137 @@ +"""Characterize cell covariates and study-design confounding.""" + +import re +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, Literal, cast + +import numpy as np +import pandas as pd + +from ..metadata.queries import ( + PartitionDigest, + column_constant_within, + column_partition_digest, + columns_same_partition, + reduce_observation_units, +) +from ..metrics.association import directional_mapping, report_confounding +from ..storage.types import as_zarr_array, as_zarr_group +from ._deps import AGENT_INSTALL_HINT +from .decide import DecisionValidationError, decide +from .types import Decision, EvidenceItem, StageStatus + +try: + from pydantic import BaseModel, Field +except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + +__all__ = [ + "CovariateCharacterization", + "characterize_covariates", +] + +Domain = Literal["biological", "technical", "design", "ignore", "unknown"] +ColumnKind = Literal["categorical", "continuous"] + +_DOMAINS = frozenset({"biological", "technical", "design", "ignore", "unknown"}) +# Only these domains reach the design table, so only they are worth collapsing. +_ANALYSED = frozenset({"biological", "technical", "design"}) +_KINDS = frozenset({"categorical", "continuous"}) +_RESERVED_COLUMNS = frozenset({"I", "ids", "names"}) +_EMBEDDING_TOKENS = ( + "umap", + "pca", + "tsne", + "scvi", + "latent", + "phate", + "forceatlas", + "diffmap", + "diffusionmap", + "diffusion", +) +_SHORT_EMBEDDING_PARTS = frozenset({"fa", "dm", "pc"}) +_INDEXED_NAME = re.compile(r"(?P.+?)[-_]?(?P\d+)") +_ONTOLOGY_SUFFIX = "_ontology_term_id" +_CATEGORICAL_MAX_LEVELS = 50 +_SAMPLE_LEVELS = 8 +_CONTEXT_LIMIT = 1200 +_ASSOCIATION_FLOOR = 0.1 +_DROP_REASONS = { + "dropAssayStat": "Scarf assay statistic column", + "dropProvenance": "analysis-linked column", + "dropEmbedding": "embedding-style column", + "dropConstant": "single-level column", +} + +_DOMAIN_EVIDENCE = [ + EvidenceItem( + id="domain:biological", + label="biological", + summary="Biology of interest such as disease, sex, genotype, treatment", + ), + EvidenceItem( + id="domain:technical", + label="technical", + summary="Technical handling such as batch, chemistry, sequencing run", + ), + EvidenceItem( + id="domain:design", + label="design", + summary="Sampling design such as donor, sample, replicate, subject", + ), + EvidenceItem( + id="domain:ignore", + label="ignore", + summary="Identifiers, QC metrics, clusters, or other non-design labels", + ), + EvidenceItem( + id="domain:unknown", + label="unknown", + summary="Cannot classify from available evidence", + ), +] +_COEFFICIENT_EVIDENCE = [ + EvidenceItem( + id="coefficient:yes", + label="yes", + summary="Treat this biological column as a coefficient of interest", + ), + EvidenceItem( + id="coefficient:no", + label="no", + summary="Do not treat this biological column as a coefficient of interest", + ), +] + + +class CovariateCharacterization(BaseModel): + status: StageStatus + auditLog: list[dict[str, Any]] = Field(default_factory=list) + actions: list[str] = Field(default_factory=list) + notes: list[str] = Field(default_factory=list) + decisions: list[dict[str, Any]] = Field(default_factory=list) + columns: list[dict[str, Any]] = Field(default_factory=list) + coefficients: list[dict[str, Any]] = Field(default_factory=list) + technicalNesting: list[dict[str, Any]] = Field(default_factory=list) + confounding: list[dict[str, Any]] = Field(default_factory=list) + + +@dataclass(frozen=True, slots=True) +class _ColumnProfile: + kind: ColumnKind + summary: str + digest: PartitionDigest + + +@dataclass +class _Run: + """Mutable state shared by the stage steps.""" + + store: Any + cell_key: str + n_rows: int + context: str + model: Any | None + profiles: dict[str, _ColumnProfile] = field(default_factory=dict) + domains: dict[str, Domain] = field(default_factory=dict) + audit: list[dict[str, Any]] = field(default_factory=list) + actions: list[str] = field(default_factory=list) + decisions: list[dict[str, Any]] = field(default_factory=list) + + def note(self, *, kind: str, detail: str, **fields: Any) -> None: + self.audit.append({"kind": kind, "detail": detail, **fields}) + + def summary(self, name: str) -> str: + return self.profiles[name].summary + + def kind(self, name: str) -> ColumnKind: + return self.profiles[name].kind + + def digest(self, name: str) -> PartitionDigest: + return self.profiles[name].digest + + def ask( + self, + *, + task: str, + question: str, + evidence: Sequence[EvidenceItem], + column: str | None = None, + ) -> Decision | None: + """Run one grounded decision, or return None when it cannot be asked.""" + if self.model is None or len(evidence) < 2: + return None + try: + decision = decide(model=self.model, question=question, evidence=evidence) + except DecisionValidationError as exc: + self.note( + kind="decisionInvalid", + detail=str(exc), + task=task, + column=column, + ) + return None + record: dict[str, Any] = { + "task": task, + "selectedId": decision.selectedId, + "rationale": decision.rationale, + "evidenceIds": list(decision.evidenceIds), + } + if column is not None: + record["column"] = column + self.decisions.append(record) + return decision + + +def _is_embedding_column(name: str) -> bool: + match = _INDEXED_NAME.fullmatch(name) + if match is None: + return False + parts = [part for part in re.split(r"[-_]+", match.group("stem").lower()) if part] + compact = "".join(parts) + if any(token in compact for token in _EMBEDDING_TOKENS): + return True + return any(part in _SHORT_EMBEDDING_PARTS for part in parts) + + +def _has_source_artifact(store: Any, column: str) -> bool: + try: + cell_data = as_zarr_group(store.zw["cellData"], name="cellData") + if column not in cell_data: + return False + attrs = as_zarr_array(cell_data[column], name=column).attrs + except (KeyError, TypeError, ValueError): + return False + return isinstance(attrs.get("source_artifact"), dict) + + +def _infer_kind(values: np.ndarray) -> ColumnKind: + if ( + values.dtype == object + or np.issubdtype(values.dtype, np.str_) + or np.issubdtype(values.dtype, np.bool_) + ): + return "categorical" + try: + numeric = np.asarray(values, dtype=float) + except (TypeError, ValueError): + return "categorical" + finite = numeric[np.isfinite(numeric)] + if finite.size == 0 or not bool(np.all(np.mod(finite, 1) == 0)): + return "continuous" + limit = min(_CATEGORICAL_MAX_LEVELS, max(2, len(values) // 20)) + return "categorical" if int(np.unique(finite).size) <= limit else "continuous" + + +def _summarize(values: np.ndarray, kind: ColumnKind) -> str: + series = pd.Series(values) + missing = int(series.isna().sum()) + if kind == "categorical": + levels = series.dropna().astype(str).value_counts() + top = ", ".join( + f"{level}={int(count)}" + for level, count in levels.head(_SAMPLE_LEVELS).items() + ) + return f"categorical levels={levels.shape[0]} missing={missing} top=[{top}]" + numeric = pd.to_numeric(series, errors="coerce").to_numpy(dtype=float, copy=False) + finite = numeric[np.isfinite(numeric)] + if finite.size == 0: + return f"continuous missing={missing} finite=0" + return ( + f"continuous missing={missing} min={float(finite.min()):.4g} " + f"max={float(finite.max()):.4g} mean={float(finite.mean()):.4g}" + ) + + +def _digest_key(digest: PartitionDigest) -> tuple[bytes, int, int]: + return (digest.digest, digest.nLevels, digest.nMissing) + + +def _profile_column( + store: Any, + name: str, + *, + cell_key: str, + kind: ColumnKind | None = None, +) -> _ColumnProfile: + values = store.cells.fetch(name, key=cell_key) + resolved_kind = kind or _infer_kind(values) + summary = _summarize(values, resolved_kind) + digest = column_partition_digest(store.cells, name, cell_key=cell_key) + return _ColumnProfile(kind=resolved_kind, summary=summary, digest=digest) + + +def _triage_columns( + store: Any, + *, + cell_key: str, + exclude: set[str], +) -> tuple[list[str], list[tuple[str, str]]]: + """Split cell columns into model candidates and deterministic drops.""" + assay_prefixes = tuple(f"{name}_" for name in store.assay_names) + candidates: list[str] = [] + dropped: list[tuple[str, str]] = [] + for name in store.cells.columns: + if name in _RESERVED_COLUMNS or name == cell_key or name in exclude: + continue + if name.startswith(assay_prefixes): + dropped.append((name, "dropAssayStat")) + elif _has_source_artifact(store, name): + dropped.append((name, "dropProvenance")) + elif _is_embedding_column(name): + dropped.append((name, "dropEmbedding")) + else: + candidates.append(name) + return candidates, dropped + + +def _collapse_ontology_aliases( + store: Any, + columns: Sequence[str], + profiles: Mapping[str, _ColumnProfile], + *, + cell_key: str, +) -> tuple[list[str], dict[str, list[str]], list[dict[str, Any]]]: + """Collapse ``x`` with ``x_ontology_term_id`` when partitions match. + + Arbitrary identical partitions are kept apart: perfect confounding between + biology and batch is a finding to report, not an alias to drop. + """ + present = set(columns) + aliases: dict[str, list[str]] = {} + dropped: set[str] = set() + notes: list[dict[str, Any]] = [] + for name in columns: + if not name.endswith(_ONTOLOGY_SUFFIX): + continue + base = name[: -len(_ONTOLOGY_SUFFIX)] + if base not in present or {name, base} & dropped: + continue + if _digest_key(profiles[name].digest) != _digest_key(profiles[base].digest): + continue + same, correspondence = columns_same_partition( + store.cells, + base, + name, + cell_key=cell_key, + ) + if not same: + continue + aliases.setdefault(base, []).append(name) + dropped.add(name) + note: dict[str, Any] = { + "kind": "ontologyAlias", + "detail": f"Collapsed ontology alias {name} onto {base}", + "representative": base, + "aliases": [name], + } + if correspondence: + note["levels"] = correspondence + notes.append(note) + return [name for name in columns if name not in dropped], aliases, notes + + +def _bounded_context(study_context: str | None) -> str: + text = (study_context or "").strip() + return text if len(text) <= _CONTEXT_LIMIT else text[: _CONTEXT_LIMIT - 3] + "..." + + +def _validate_directions( + directions: Mapping[str, Any], + available: set[str], +) -> list[str]: + errors: list[str] = [] + + def check_names(key: str, names: Any) -> list[str] | None: + if not isinstance(names, Sequence) or isinstance(names, str | bytes): + errors.append(f"{key} must be a sequence of column names") + return None + unknown = sorted(set(names) - available) + if unknown: + errors.append(f"{key} cites unknown columns: {unknown}") + return list(names) + + for key, allowed in (("columnKinds", _KINDS), ("columnDomains", _DOMAINS)): + mapping = directions.get(key) + if mapping is None: + continue + if not isinstance(mapping, Mapping): + errors.append(f"{key} must be a mapping") + continue + check_names(key, list(mapping)) + invalid = sorted({str(value) for value in mapping.values()} - allowed) + if invalid: + errors.append(f"{key} has unsupported values: {invalid}") + + for key in ("coefficientsOfInterest", "excludeColumns"): + names = directions.get(key) + if names is not None: + check_names(key, names) + + units = directions.get("unitsOfInference") + if units is None: + return errors + if not isinstance(units, Mapping): + errors.append("unitsOfInference must be a mapping") + return errors + for coefficient, unit_map in units.items(): + if coefficient not in available: + errors.append(f"unitsOfInference cites unknown coefficient {coefficient!r}") + continue + if not isinstance(unit_map, Mapping): + errors.append(f"unitsOfInference[{coefficient!r}] must be a mapping") + continue + for unit_key in ("observationUnit", "independentUnit"): + unit_name = unit_map.get(unit_key) + if unit_name is not None and unit_name not in available: + errors.append( + f"unitsOfInference[{coefficient!r}].{unit_key} " + f"cites unknown column {unit_name!r}" + ) + return errors + + +def _choose_representative( + run: _Run, + members: Sequence[str], + correspondence: str, +) -> str | None: + decision = run.ask( + task="equivalentColumns", + question=( + f"Columns {', '.join(members)} assign every cell to the same groups. " + f"Their labels line up as {correspondence}. Decide whether they record " + "one variable under different labels, and if so whose labels to keep." + ), + evidence=[ + EvidenceItem( + id="equivalent:distinct", + label="distinct variables", + summary="Different variables that happen to coincide in this dataset", + ), + *( + EvidenceItem( + id=f"equivalent:{name}", + label=name, + summary=run.summary(name), + ) + for name in members + ), + ], + ) + selected = ( + None if decision is None else decision.selectedId.removeprefix("equivalent:") + ) + if selected in set(members): + run.actions.append(f"equivalentColumns:{selected}") + return selected + run.note( + kind="equivalentKeptApart", + detail=( + f"{', '.join(members)} share one partition but were not collapsed " + + ("(no model decision)" if decision is None else "(judged distinct)") + ), + columns=list(members), + levels=correspondence, + ) + return None + + +def _collapse_equivalent_columns( + run: _Run, + candidates: Sequence[str], + aliases: dict[str, list[str]], +) -> list[str]: + """Collapse categorical columns that cut the cells into identical groups. + + An identical partition is also what perfect confounding looks like, so a + class is only eligible when every member carries the same analysis domain, + and the representative is a model judgement rather than a name rule. + Extends ``aliases`` in place. + """ + classes: dict[tuple[bytes, int, int], list[str]] = {} + for name in candidates: + if run.kind(name) != "categorical" or run.domains[name] not in _ANALYSED: + continue + classes.setdefault(_digest_key(run.digest(name)), []).append(name) + + # Store column order is not stable, and members of a class are + # interchangeable, so order them here to keep prompts and notes reproducible. + dropped: set[str] = set() + for members in sorted(sorted(group) for group in classes.values()): + if len(members) < 2: + continue + representative_name = members[0] + verified: list[str] = [representative_name] + correspondence = "" + for other in members[1:]: + same, corr = columns_same_partition( + run.store.cells, + representative_name, + other, + cell_key=run.cell_key, + ) + if not same: + continue + verified.append(other) + correspondence = corr + if len(verified) < 2: + continue + members = verified + domains = sorted({run.domains[name] for name in members}) + if len(domains) > 1: + run.note( + kind="equivalentAcrossDomains", + detail=( + f"{', '.join(members)} share one partition across domains " + f"{domains}; kept apart as perfect confounding" + ), + columns=list(members), + domains=domains, + levels=correspondence, + ) + continue + representative = _choose_representative(run, members, correspondence) + if representative is None: + continue + others = [name for name in members if name != representative] + aliases.setdefault(representative, []).extend(others) + dropped.update(others) + run.note( + kind="equivalentColumns", + detail=f"Collapsed {', '.join(others)} onto {representative}", + representative=representative, + aliases=others, + levels=correspondence, + ) + return [name for name in candidates if name not in dropped] + + +def _assign_domain(run: _Run, name: str, directed: Mapping[str, Domain]) -> Domain: + if name in directed: + domain = directed[name] + run.actions.append(f"domain:{name}->{domain} (directions)") + return domain + decision = run.ask( + task="columnDomain", + column=name, + question=( + f"Assign a domain for cell metadata column {name}. " + f"Column summary: {run.summary(name)}. " + f"Study context: {run.context or 'none provided'}." + ), + evidence=_DOMAIN_EVIDENCE, + ) + if decision is None: + run.note( + kind="domainUnknown", + detail=f"Left {name} as unknown domain", + column=name, + ) + return "unknown" + selected = decision.selectedId.removeprefix("domain:") + if selected not in _DOMAINS: + run.note( + kind="domainUnknown", + detail=f"Unsupported domain {selected!r} returned for {name}", + column=name, + ) + return "unknown" + run.actions.append(f"domain:{name}->{selected}") + return cast(Domain, selected) + + +def _select_coefficients( + run: _Run, + *, + biological: Sequence[str], + directed: set[str], +) -> list[str]: + selected: list[str] = [] + for name in biological: + if name in directed: + selected.append(name) + run.actions.append(f"coefficient:{name} (directions)") + continue + decision = run.ask( + task="coefficientOfInterest", + column=name, + question=( + f"Should biological column {name} be a coefficient of interest? " + f"Column summary: {run.summary(name)}. " + f"Study context: {run.context or 'none provided'}." + ), + evidence=_COEFFICIENT_EVIDENCE, + ) + if decision is None: + run.note( + kind="coefficientSkipped", + detail=f"No model or direction for biological column {name}", + column=name, + ) + elif decision.selectedId == "coefficient:yes": + selected.append(name) + run.actions.append(f"coefficient:{name}") + return selected + + +def _unit_evidence(run: _Run, names: Sequence[str], prefix: str) -> list[EvidenceItem]: + return [ + EvidenceItem( + id=f"{prefix}:{name}", + label=name, + summary=(f"domain={run.domains.get(name, 'unknown')}; {run.summary(name)}"), + ) + for name in names + ] + + +def _observation_unit_candidates( + run: _Run, + coefficient: str, + pool: Sequence[str], +) -> list[str]: + """Design/technical columns that can host a between-unit coefficient. + + A valid observation unit is a metadata fact, not a judgement: the coefficient + must be constant within each level, and the unit must leave more than one + design row while still being coarser than cell-level variation. + """ + return [ + name + for name in pool + if name != coefficient + and name in run.profiles + and column_constant_within( + run.store.cells, + coefficient, + name, + cell_key=run.cell_key, + ) + and 2 <= run.digest(name).nLevels < run.n_rows + ] + + +def _independent_is_coarser( + store: Any, + *, + cell_key: str, + observation: str, + independent: str, +) -> bool: + """True when each observation level maps to one independent level. + + The independent unit must be coarser (or equal), never finer. A finer + independent unit inflates the design table with pseudo-replicated rows. + """ + observation_values = store.cells.fetch(observation, key=cell_key) + independent_values = store.cells.fetch(independent, key=cell_key) + nesting = directional_mapping(observation_values, independent_values).get("nesting") + return nesting in {"leftInRight", "equivalent"} + + +def _resolve_units( + run: _Run, + coefficient: str, + *, + directed: Mapping[str, Any], + design_columns: Sequence[str], + unit_candidates: Sequence[str], +) -> tuple[str | None, str | None]: + unit_map = dict(directed.get(coefficient) or {}) + observation = unit_map.get("observationUnit") + independent = unit_map.get("independentUnit") + valid_observation = _observation_unit_candidates(run, coefficient, unit_candidates) + + if observation is not None: + if observation not in run.profiles: + run.note( + kind="invalidObservationUnit", + detail=f"Directed observation unit {observation!r} is missing", + column=coefficient, + observationUnit=observation, + ) + observation = None + elif not column_constant_within( + run.store.cells, + coefficient, + observation, + cell_key=run.cell_key, + ): + # Keep the directed unit; _characterize_coefficient records withinUnit. + pass + elif observation not in valid_observation: + run.note( + kind="invalidObservationUnit", + detail=( + f"Directed observation unit {observation!r} is vacuous or " + f"otherwise invalid for {coefficient}" + ), + column=coefficient, + observationUnit=observation, + ) + observation = None + elif len(valid_observation) == 1: + observation = valid_observation[0] + run.actions.append(f"observationUnit:{coefficient}->{observation}") + elif len(valid_observation) >= 2: + decision = run.ask( + task="observationUnit", + column=coefficient, + question=( + f"Choose the observation unit for coefficient {coefficient}. " + "Only columns where this coefficient is constant within each " + "level are listed. Each distinct value is one design-table row. " + f"Study context: {run.context or 'none provided'}." + ), + evidence=_unit_evidence(run, valid_observation, "unit"), + ) + if decision is not None: + observation = decision.selectedId.removeprefix("unit:") + if observation not in valid_observation: + run.note( + kind="invalidObservationUnit", + detail=( + f"Model chose {observation!r}, which is not a valid " + f"observation unit for {coefficient}" + ), + column=coefficient, + observationUnit=observation, + ) + observation = None + else: + run.actions.append(f"observationUnit:{coefficient}->{observation}") + else: + run.note( + kind="noValidObservationUnit", + detail=( + f"No design/technical column keeps {coefficient} constant; " + "cannot build a between-unit design table from available metadata" + ), + column=coefficient, + ) + + if observation is None: + return None, None + + valid_independent = [ + name + for name in design_columns + if name not in {coefficient, observation} + and name in run.profiles + and _independent_is_coarser( + run.store, + cell_key=run.cell_key, + observation=observation, + independent=name, + ) + ] + + if independent is not None: + if independent not in run.profiles: + run.note( + kind="invalidIndependentUnit", + detail=f"Directed independent unit {independent!r} is missing", + column=coefficient, + independentUnit=independent, + ) + independent = None + elif not _independent_is_coarser( + run.store, + cell_key=run.cell_key, + observation=observation, + independent=independent, + ): + run.note( + kind="independentUnitFiner", + detail=( + f"Independent unit {independent!r} is finer than observation " + f"unit {observation!r}; dropped to avoid pseudo-replicated " + "design rows" + ), + column=coefficient, + observationUnit=observation, + independentUnit=independent, + ) + independent = None + elif valid_independent: + decision = run.ask( + task="independentUnit", + column=coefficient, + question=( + f"Optional independent unit for coefficient {coefficient} " + f"with observation unit {observation}. Listed columns are " + "coarser than the observation unit (each observation level " + "maps to one independent level)." + ), + evidence=[ + EvidenceItem( + id="independentUnit:none", + label="none", + summary="No separate independent unit or subject column", + ), + *_unit_evidence(run, valid_independent, "independentUnit"), + ], + ) + if decision is not None and decision.selectedId != "independentUnit:none": + chosen = decision.selectedId.removeprefix("independentUnit:") + if chosen not in valid_independent: + run.note( + kind="invalidIndependentUnit", + detail=( + f"Model chose {chosen!r}, which is not coarser than " + f"observation unit {observation!r}" + ), + column=coefficient, + independentUnit=chosen, + ) + else: + independent = chosen + run.actions.append(f"independentUnit:{coefficient}->{independent}") + return observation, independent + + +def _characterize_coefficient( + run: _Run, + coefficient: str, + *, + observation_unit: str | None, + independent_unit: str | None, + technical: Sequence[str], +) -> tuple[dict[str, Any], dict[str, Any] | None]: + record: dict[str, Any] = { + "name": coefficient, + "kind": run.kind(coefficient), + "observationUnit": observation_unit, + "independentUnit": independent_unit, + "scope": "unresolvedUnit", + } + if observation_unit is None or observation_unit not in run.profiles: + run.note( + kind="unresolvedUnit", + detail=f"No usable observation unit for coefficient {coefficient}", + column=coefficient, + ) + return record, None + + if not column_constant_within( + run.store.cells, + coefficient, + observation_unit, + cell_key=run.cell_key, + ): + record["scope"] = "withinUnit" + run.note( + kind="withinUnit", + detail=( + f"{coefficient} varies within {observation_unit}; recorded as " + "composition and skipped for design-table association" + ), + column=coefficient, + observationUnit=observation_unit, + ) + return record, None + + # Technically only columns constant inside the observation unit have a + # well-defined design-table value; drop_duplicates would otherwise keep an + # arbitrary cell row. + unit_constant: list[str] = [] + for name in technical: + if name not in run.profiles: + continue + if column_constant_within( + run.store.cells, + name, + observation_unit, + cell_key=run.cell_key, + ): + unit_constant.append(name) + continue + run.note( + kind="technicalVariesWithinUnit", + detail=( + f"{name} varies within {observation_unit}; " + f"excluded from design-table association for {coefficient}" + ), + column=name, + coefficient=coefficient, + observationUnit=observation_unit, + ) + + group_cols = [observation_unit] + if independent_unit is not None and independent_unit in run.profiles: + if _independent_is_coarser( + run.store, + cell_key=run.cell_key, + observation=observation_unit, + independent=independent_unit, + ): + group_cols.append(independent_unit) + else: + run.note( + kind="independentUnitFiner", + detail=( + f"Independent unit {independent_unit!r} is finer than " + f"observation unit {observation_unit!r}; omitted from " + "design table" + ), + column=coefficient, + observationUnit=observation_unit, + independentUnit=independent_unit, + ) + independent_unit = None + record["independentUnit"] = None + columns = list(dict.fromkeys([*group_cols, coefficient, *unit_constant])) + design = reduce_observation_units( + run.store.cells, + observation_unit, + columns, + cell_key=run.cell_key, + ) + design_rows = int(len(design)) + if not (2 <= design_rows < run.n_rows): + run.note( + kind="invalidObservationUnit", + detail=( + f"Observation unit {observation_unit!r} yields {design_rows} " + f"design rows for {coefficient}; need at least 2 and fewer " + f"than {run.n_rows} active cells" + ), + column=coefficient, + observationUnit=observation_unit, + designRows=design_rows, + ) + record["scope"] = "unresolvedUnit" + return record, None + + record["scope"] = "betweenUnit" + record["designRows"] = design_rows + + report = report_confounding( + design, + coefficient=coefficient, + technicalColumns=unit_constant, + columnKinds={ + coefficient: run.kind(coefficient), + **{name: run.kind(name) for name in unit_constant}, + }, + associationFloor=_ASSOCIATION_FLOOR, + ) + report["observationUnit"] = observation_unit + report["independentUnit"] = independent_unit + run.actions.append(f"confounding:{coefficient}") + return record, report + + +def _characterize_coefficients( + run: _Run, + coefficients: Sequence[str], + *, + unit_directions: Mapping[str, Any], + design_columns: Sequence[str], + technical: Sequence[str], +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + records: list[dict[str, Any]] = [] + reports: list[dict[str, Any]] = [] + for coefficient in coefficients: + observation, independent = _resolve_units( + run, + coefficient, + directed=unit_directions, + design_columns=design_columns, + unit_candidates=[*design_columns, *technical], + ) + record, report = _characterize_coefficient( + run, + coefficient, + observation_unit=observation, + independent_unit=independent, + technical=technical, + ) + records.append(record) + if report is not None: + reports.append(report) + return records, reports + + +def _technical_nesting_reports( + store: Any, + names: Sequence[str], + *, + cell_key: str, +) -> list[dict[str, Any]]: + """Directional nesting among categorical technical columns without bulk fetch.""" + name_list = list(names) + reports: list[dict[str, Any]] = [] + for index, left_name in enumerate(name_list): + left_values = store.cells.fetch(left_name, key=cell_key) + for right_name in name_list[index + 1 :]: + right_values = store.cells.fetch(right_name, key=cell_key) + mapping = directional_mapping(left_values, right_values) + if mapping["nesting"] == "none": + continue + reports.append( + { + "left": left_name, + "right": right_name, + "nesting": mapping["nesting"], + "directionalMapping": mapping, + } + ) + return reports + + +def _column_records( + run: _Run, + candidates: Sequence[str], + *, + aliases: Mapping[str, list[str]], + dropped: Sequence[tuple[str, str]], +) -> list[dict[str, Any]]: + records = [ + { + "name": name, + "kind": run.kind(name), + "domain": run.domains[name], + "summary": run.summary(name), + "aliases": list(aliases.get(name, [])), + } + for name in candidates + ] + records.extend( + { + "name": name, + "kind": "continuous", + "domain": "ignore", + "summary": f"dropped before triage ({_DROP_REASONS[reason]})", + "aliases": [], + } + for name, reason in dropped + ) + return records + + +def characterize_covariates( + store: Any, + *, + studyContext: str | None = None, + model: Any | None = None, + cellKey: str = "I", + directions: Mapping[str, Any] | None = None, +) -> CovariateCharacterization: + """Label cell covariates and record design-level confounding.""" + direction_map = dict(directions or {}) + available = set(store.cells.columns) + if cellKey not in available: + return CovariateCharacterization( + status="failed", + notes=[f"cellKey {cellKey!r} is not present in cell metadata"], + ) + errors = _validate_directions(direction_map, available) + if errors: + return CovariateCharacterization(status="failed", notes=errors) + + candidates, dropped = _triage_columns( + store, + cell_key=cellKey, + exclude=set(direction_map.get("excludeColumns") or []), + ) + reviewed = len(candidates) + len(dropped) + candidates = [name for name in candidates if name in store.cells.columns] + kind_directions = dict(direction_map.get("columnKinds") or {}) + directed_coefficients = set(direction_map.get("coefficientsOfInterest") or []) + + profiles: dict[str, _ColumnProfile] = {} + n_rows = 0 + varying: list[str] = [] + for name in candidates: + directed_kind = kind_directions.get(name) + profile = _profile_column( + store, + name, + cell_key=cellKey, + kind=cast(ColumnKind, directed_kind) if directed_kind in _KINDS else None, + ) + profiles[name] = profile + n_rows = profile.digest.nRows + if profile.digest.nLevels <= 1: + dropped.append((name, "dropConstant")) + continue + varying.append(name) + candidates = varying + candidates, aliases, alias_notes = _collapse_ontology_aliases( + store, + candidates, + profiles, + cell_key=cellKey, + ) + + run = _Run( + store=store, + cell_key=cellKey, + n_rows=n_rows, + context=_bounded_context(studyContext), + model=model, + profiles=profiles, + ) + for name, reason in dropped: + detail = f"Dropped {_DROP_REASONS[reason]} {name}" + if reason == "dropConstant" and name in directed_coefficients: + detail = ( + f"{detail}; also listed in coefficientsOfInterest but has no variation" + ) + run.note( + kind=reason, + detail=detail, + column=name, + ) + run.audit.extend(alias_notes) + + domain_directions = dict(direction_map.get("columnDomains") or {}) + for name in candidates: + run.domains[name] = _assign_domain(run, name, domain_directions) + candidates = _collapse_equivalent_columns(run, candidates, aliases) + + technical = [name for name in candidates if run.domains[name] == "technical"] + design_columns = [name for name in candidates if run.domains[name] == "design"] + records, reports = _characterize_coefficients( + run, + _select_coefficients( + run, + biological=[ + name for name in candidates if run.domains[name] == "biological" + ], + directed=set(direction_map.get("coefficientsOfInterest") or []), + ), + unit_directions=dict(direction_map.get("unitsOfInference") or {}), + design_columns=design_columns, + technical=technical, + ) + + categorical_technical = [ + name for name in technical if run.kind(name) == "categorical" + ] + return CovariateCharacterization( + status="done", + auditLog=run.audit, + actions=run.actions, + notes=[ + f"Reviewed {reviewed} columns; {len(candidates)} triaged after " + "deterministic drops and ontology alias collapse" + ], + decisions=run.decisions, + columns=_column_records(run, candidates, aliases=aliases, dropped=dropped), + coefficients=records, + technicalNesting=( + _technical_nesting_reports( + store, + categorical_technical, + cell_key=cellKey, + ) + if len(categorical_technical) >= 2 + else [] + ), + confounding=reports, + ) diff --git a/scarf/agent/characterize_features.py b/scarf/agent/characterize_features.py new file mode 100644 index 00000000..15af75e6 --- /dev/null +++ b/scarf/agent/characterize_features.py @@ -0,0 +1,554 @@ +"""Characterize feature identity, species, families, and exogenous candidates.""" + +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +from ..assay import RNAassay +from ..features.gene_reference import ( + GeneReference, + default_cache_dir, + ensure_reference, + load_reference, + species_registry, +) +from ..features.identity import ( + audit_feature_identity, + backfill_symbols, + exogenous_candidates, + observe_families, + reference_misses, + resolve_species, +) +from ..quality_control.cell_cycle_genes import ( + g2m_phase_genes, + g2m_phase_genes_mouse, + s_phase_genes, + s_phase_genes_mouse, +) +from ._deps import AGENT_INSTALL_HINT +from .decide import DecisionValidationError, decide +from .types import Decision, EvidenceItem, StageStatus + +try: + from pydantic import BaseModel, Field +except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + +__all__ = [ + "FeatureCharacterization", + "characterize_features", +] + +_DEFAULT_MAX_EXOGENOUS = 25 +_CONTEXT_LIMIT = 1200 +_AUTO_DOWNLOAD_SPECIES = frozenset({"homo_sapiens", "mus_musculus"}) +_CELL_CYCLE = { + "homo_sapiens": {"s": s_phase_genes, "g2m": g2m_phase_genes}, + "mus_musculus": {"s": s_phase_genes_mouse, "g2m": g2m_phase_genes_mouse}, +} +_SEX_COEFFICIENT_TOKENS = frozenset({"sex", "gender", "Sex", "Gender"}) + + +class FeatureCharacterization(BaseModel): + status: StageStatus + auditLog: list[dict[str, Any]] = Field(default_factory=list) + actions: list[str] = Field(default_factory=list) + notes: list[str] = Field(default_factory=list) + decisions: list[dict[str, Any]] = Field(default_factory=list) + assays: list[dict[str, Any]] = Field(default_factory=list) + + +def _bounded_context(study_context: str | None) -> str: + text = (study_context or "").strip() + return text if len(text) <= _CONTEXT_LIMIT else text[: _CONTEXT_LIMIT - 3] + "..." + + +def _audit( + audit_log: list[dict[str, Any]], + *, + kind: str, + detail: str, + **fields: Any, +) -> None: + audit_log.append({"kind": kind, "detail": detail, **fields}) + + +def _ask( + *, + model: Any | None, + question: str, + evidence: Sequence[EvidenceItem], + decisions: list[dict[str, Any]], + audit_log: list[dict[str, Any]], + task: str, + assay: str | None = None, +) -> Decision | None: + if model is None or len(evidence) < 2: + return None + try: + decision = decide(model=model, question=question, evidence=evidence) + except DecisionValidationError as exc: + _audit( + audit_log, + kind="decisionInvalid", + detail=str(exc), + task=task, + assay=assay, + ) + return None + record: dict[str, Any] = { + "task": task, + "selectedId": decision.selectedId, + "rationale": decision.rationale, + "evidenceIds": list(decision.evidenceIds), + } + if assay is not None: + record["assay"] = assay + decisions.append(record) + return decision + + +def _sex_coefficient_note( + covariates: Any | None, +) -> str | None: + if covariates is None: + return None + coefficients = getattr(covariates, "coefficients", None) or [] + for item in coefficients: + name = item.get("name") if isinstance(item, Mapping) else None + if name in _SEX_COEFFICIENT_TOKENS: + return ( + f"Prior covariate characterization marked {name!r} as a coefficient " + "of interest; tracked sex-chromosome genes must not be excluded later" + ) + return None + + +def _load_or_fetch_reference( + species: str, + *, + cache_dir: Path, + allow_download: bool, + audit_log: list[dict[str, Any]], + assay: str, +) -> GeneReference | None: + if species == "unknown": + return None + cached = load_reference(species, cacheDir=cache_dir) + if cached is not None: + return cached + if not allow_download: + _audit( + audit_log, + kind="referenceUnavailable", + detail=f"No cached reference for {species}; download disabled", + assay=assay, + species=species, + ) + return None + try: + reference = ensure_reference(species, cacheDir=cache_dir) + except Exception as exc: + _audit( + audit_log, + kind="referenceDownloadFailed", + detail=f"Failed to download reference for {species}: {exc}", + assay=assay, + species=species, + ) + return None + _audit( + audit_log, + kind="referenceDownloaded", + detail=f"Cached gene reference for {species} release {reference.release}", + assay=assay, + species=species, + release=reference.release, + ) + return reference + + +def _assist_species( + *, + model: Any | None, + unresolved: Mapping[str, Any], + context: str, + decisions: list[dict[str, Any]], + audit_log: list[dict[str, Any]], + assay: str, +) -> str | None: + # Only ask among species that already have overlap evidence. Expanding to the + # full registry would let a guess trigger a non-human/mouse download. + candidates = [ + key for key in (unresolved.get("candidates") or []) if key in species_registry() + ] + if len(candidates) < 2: + return None + evidence = [ + EvidenceItem( + id=f"species:{key}", + label=species_registry()[key].label, + summary=( + f"overlap hits=" + f"{(unresolved.get('overlap') or {}).get('scores', {}).get(key, {}).get('hits', 0)}; " + f"prefix count=" + f"{(unresolved.get('prefixCounts') or {}).get(key, 0)}" + ), + ) + for key in candidates + ] + evidence.append( + EvidenceItem( + id="species:unknown", + label="unknown", + summary="Cannot settle species from available evidence", + ) + ) + decision = _ask( + model=model, + question=( + f"Choose the species for assay {assay}. " + f"Identity notes: {unresolved.get('reason', '')}. " + f"Study context: {context or 'none provided'}." + ), + evidence=evidence, + decisions=decisions, + audit_log=audit_log, + task="species", + assay=assay, + ) + if decision is None: + return None + selected = decision.selectedId.removeprefix("species:") + return selected if selected in species_registry() or selected == "unknown" else None + + +def _classify_exogenous( + *, + model: Any | None, + candidates: Sequence[Mapping[str, Any]], + context: str, + decisions: list[dict[str, Any]], + audit_log: list[dict[str, Any]], + assay: str, + species: str, +) -> list[dict[str, Any]]: + classified: list[dict[str, Any]] = [] + for item in candidates: + label = item.get("name") or item.get("id") or "feature" + evidence = [ + EvidenceItem( + id="exogenous:potentialExogenous", + label="potentialExogenous", + summary="Spike-in, transgene, guide, antibody tag, or other non-endogenous feature", + ), + EvidenceItem( + id="exogenous:unresolved", + label="unresolved", + summary="Not enough evidence to treat as exogenous", + ), + ] + decision = _ask( + model=model, + question=( + f"Classify feature {label!r} (id={item.get('id')!r}) for assay {assay} " + f"under species {species}. Study context: {context or 'none provided'}." + ), + evidence=evidence, + decisions=decisions, + audit_log=audit_log, + task="exogenous", + assay=assay, + ) + record = dict(item) + if decision is None: + record["class"] = "unresolved" + else: + record["class"] = decision.selectedId.removeprefix("exogenous:") + classified.append(record) + return classified + + +def _characterize_assay( + store: Any, + assay_name: str, + *, + model: Any | None, + context: str, + directions: Mapping[str, Any], + cache_dir: Path, + allow_download: bool, + audit_log: list[dict[str, Any]], + actions: list[str], + decisions: list[dict[str, Any]], + sex_note: str | None, +) -> dict[str, Any]: + assay = store.get_assay(assay_name) + ids = [str(value) for value in assay.feats.fetch_all("ids")] + names = [str(value) for value in assay.feats.fetch_all("names")] + identity = audit_feature_identity(ids, names) + record: dict[str, Any] = { + "assay": assay_name, + "assayKind": type(assay).__name__, + "identity": identity, + "species": "unknown", + "speciesMethod": None, + "families": [], + "exogenous": [], + "symbolBackfill": None, + } + + if not isinstance(assay, RNAassay): + record["skipped"] = "familyPlanningNotApplicable" + _audit( + audit_log, + kind="nonRnaAssay", + detail=f"Assay {assay_name} is not RNA; stopped after identity audit", + assay=assay_name, + ) + return record + + species_by_assay = dict(directions.get("speciesByAssay") or {}) + directed_species = species_by_assay.get(assay_name) + resolution = resolve_species( + ids, + names, + directed=directed_species, + cacheDir=cache_dir, + allowDownload=allow_download, + ) + species = resolution["species"] + if species == "unknown" and resolution.get("method") == "inconclusive": + assisted = _assist_species( + model=model, + unresolved=resolution, + context=context, + decisions=decisions, + audit_log=audit_log, + assay=assay_name, + ) + if assisted is not None: + species = assisted + resolution = { + **resolution, + "species": species, + "method": "llmAssist", + "reason": "model choice among inconclusive candidates", + } + actions.append(f"species:{assay_name}->{species}") + + record["species"] = species + record["speciesMethod"] = resolution.get("method") + record["speciesResolution"] = { + key: value + for key, value in resolution.items() + if key not in {"overlap"} or value is not None + } + _audit( + audit_log, + kind="speciesResolved", + detail=resolution.get("reason", f"species={species}"), + assay=assay_name, + species=species, + method=resolution.get("method"), + ) + + if species == "unknown": + record["families"] = observe_families( + species="unknown", + ids=ids, + symbols=names, + reference=None, + ) + _audit( + audit_log, + kind="speciesUnknown", + detail=f"Skipped species-dependent steps for assay {assay_name}", + assay=assay_name, + ) + return record + + # Only human/mouse auto-download; any other species needs an explicit direction. + may_download = allow_download and ( + species in _AUTO_DOWNLOAD_SPECIES or directed_species == species + ) + reference = _load_or_fetch_reference( + species, + cache_dir=cache_dir, + allow_download=may_download, + audit_log=audit_log, + assay=assay_name, + ) + + symbols = list(names) + if reference is not None: + backfill = backfill_symbols(ids, names, reference) + if backfill["nRecovered"]: + symbols = backfill["symbols"] + record["symbolBackfill"] = { + "nRecovered": backfill["nRecovered"], + "joinRate": backfill["joinRate"], + } + actions.append( + f"symbolBackfill:{assay_name}:{backfill['nRecovered']}/{backfill['nFeatures']}" + ) + elif identity.get("idsEqualNames") or identity.get("nEmptyNames", 0) > 0: + _audit( + audit_log, + kind="familiesNotAssessable", + detail=( + f"Names are empty or ID-shaped on {assay_name} and no reference " + "is available to recover symbols" + ), + assay=assay_name, + ) + + record["families"] = observe_families( + species=species, + ids=ids, + symbols=symbols, + reference=reference, + cellCycleGenes=_CELL_CYCLE, + ) + if sex_note is not None: + record["notes"] = [sex_note] + _audit( + audit_log, + kind="sexCoefficientNote", + detail=sex_note, + assay=assay_name, + ) + elif species != "unknown": + _audit( + audit_log, + kind="sexChromosomeTracked", + detail=( + "Sex-chromosome genes are tracked with defaultExclude=false; " + "Phase 3 may exclude them when sex is not a coefficient of interest" + ), + assay=assay_name, + ) + + raw_max = directions.get("maxExogenousCandidates", _DEFAULT_MAX_EXOGENOUS) + try: + max_exogenous = int(raw_max) + except (TypeError, ValueError): + _audit( + audit_log, + kind="invalidDirection", + detail=f"maxExogenousCandidates={raw_max!r}; using {_DEFAULT_MAX_EXOGENOUS}", + assay=assay_name, + ) + max_exogenous = _DEFAULT_MAX_EXOGENOUS + if reference is not None: + misses = reference_misses(ids, symbols, reference) + if misses["count"]: + record["referenceMisses"] = misses + _audit( + audit_log, + kind="referenceMiss", + detail=( + f"{misses['count']} Ensembl-shaped id(s) absent from the " + f"{species} reference (release drift, not exogenous)" + ), + assay=assay_name, + count=misses["count"], + examples=misses["examples"], + ) + candidates = exogenous_candidates( + ids, + symbols, + reference=reference, + maxCandidates=max_exogenous, + ) + if reference is None and not candidates: + _audit( + audit_log, + kind="exogenousUnresolved", + detail=( + f"No reference and no structural exogenous candidates for {assay_name}" + ), + assay=assay_name, + ) + record["exogenous"] = _classify_exogenous( + model=model, + candidates=candidates, + context=context, + decisions=decisions, + audit_log=audit_log, + assay=assay_name, + species=species, + ) + actions.append(f"families:{assay_name}") + return record + + +def characterize_features( + store: Any, + *, + studyContext: str | None = None, + model: Any | None = None, + assays: Sequence[str] | None = None, + directions: Mapping[str, Any] | None = None, + covariates: Any | None = None, + cacheDir: Path | str | None = None, + allowDownload: bool = False, +) -> FeatureCharacterization: + """Label feature identity, species, families, and exogenous candidates.""" + direction_map = dict(directions or {}) + audit_log: list[dict[str, Any]] = [] + actions: list[str] = [] + decisions: list[dict[str, Any]] = [] + notes: list[str] = [] + context = _bounded_context(studyContext) + cache_dir = Path(cacheDir) if cacheDir is not None else default_cache_dir() + + available = list(store.assay_names) + selected = list(assays) if assays is not None else available + unknown = sorted(set(selected) - set(available)) + if unknown: + return FeatureCharacterization( + status="failed", + notes=[f"unknown assays: {unknown}"], + ) + species_by_assay = direction_map.get("speciesByAssay") + if species_by_assay is not None: + if not isinstance(species_by_assay, Mapping): + return FeatureCharacterization( + status="failed", + notes=["speciesByAssay must be a mapping"], + ) + bad = sorted(set(species_by_assay) - set(available)) + if bad: + return FeatureCharacterization( + status="failed", + notes=[f"speciesByAssay cites unknown assays: {bad}"], + ) + + sex_note = _sex_coefficient_note(covariates) + assay_records = [ + _characterize_assay( + store, + assay_name, + model=model, + context=context, + directions=direction_map, + cache_dir=cache_dir, + allow_download=allowDownload, + audit_log=audit_log, + actions=actions, + decisions=decisions, + sex_note=sex_note, + ) + for assay_name in selected + ] + notes.append(f"Characterized {len(assay_records)} assay(s)") + return FeatureCharacterization( + status="done", + auditLog=audit_log, + actions=actions, + notes=notes, + decisions=decisions, + assays=assay_records, + ) diff --git a/scarf/agent/decide.py b/scarf/agent/decide.py new file mode 100644 index 00000000..12ae7886 --- /dev/null +++ b/scarf/agent/decide.py @@ -0,0 +1,131 @@ +"""Grounded structured decisions over evidence IDs.""" + +from collections.abc import Sequence +from typing import Any + +from ._deps import require_pydantic_ai +from .types import Decision, EvidenceItem + +_SYSTEM_PROMPT = ( + "Choose exactly one evidence id that answers the question. " + "selectedId must be copied exactly from the id= values. " + "Do not put labels or summaries in selectedId. " + "evidenceIds must include selectedId and only provided ids. " + "Keep the rationale to one short sentence." +) + + +class DecisionValidationError(ValueError): + """Raised when a model decision cites unknown or invalid evidence.""" + + +def _coerce_evidence_id(evidence_id: str, allowed: set[str]) -> str: + """Map a model-emitted id onto an allowed evidence id when unambiguous. + + Live models often echo prompt scaffolding such as ``id=domain:biological`` + instead of the bare id. Accept that when exactly one allowed id is embedded. + """ + if evidence_id in allowed: + return evidence_id + stripped = evidence_id.strip() + if stripped.startswith("id="): + stripped = stripped[3:].strip() + if stripped in allowed: + return stripped + matches = [allowed_id for allowed_id in allowed if allowed_id in evidence_id] + if len(matches) == 1: + return matches[0] + return evidence_id + + +def validate_decision( + decision: Decision, + evidence: Sequence[EvidenceItem], +) -> Decision: + allowed = {item.id for item in evidence} + if not allowed: + raise DecisionValidationError("evidence must contain at least one item") + selected_id = _coerce_evidence_id(decision.selectedId, allowed) + evidence_ids = [ + _coerce_evidence_id(evidence_id, allowed) + for evidence_id in decision.evidenceIds + ] + if selected_id not in evidence_ids and selected_id in allowed: + evidence_ids = [selected_id, *evidence_ids] + if selected_id != decision.selectedId or evidence_ids != list(decision.evidenceIds): + decision = Decision( + selectedId=selected_id, + rationale=decision.rationale, + evidenceIds=evidence_ids, + ) + if decision.selectedId not in allowed: + raise DecisionValidationError( + f"selectedId {decision.selectedId!r} is not in evidence ids {sorted(allowed)}" + ) + unknown = [ + evidence_id + for evidence_id in decision.evidenceIds + if evidence_id not in allowed + ] + if unknown: + raise DecisionValidationError( + f"evidenceIds cite unknown ids {unknown}; allowed {sorted(allowed)}" + ) + if decision.selectedId not in decision.evidenceIds: + raise DecisionValidationError( + f"evidenceIds must include selectedId {decision.selectedId!r}" + ) + return decision + + +def _format_user_prompt(question: str, evidence: Sequence[EvidenceItem]) -> str: + lines = [ + question.strip(), + "", + "Choose selectedId from these exact id values:", + ] + for item in evidence: + lines.append(f"- id={item.id} | label={item.label} | summary={item.summary}") + lines.append("") + lines.append( + f"Allowed selectedId values: {', '.join(item.id for item in evidence)}" + ) + return "\n".join(lines) + + +def decide( + *, + model: Any, + question: str, + evidence: Sequence[EvidenceItem], + systemPrompt: str = _SYSTEM_PROMPT, +) -> Decision: + """Ask the model to choose among evidence IDs and validate citations.""" + if not question.strip(): + raise ValueError("question must be non-empty") + if not evidence: + raise ValueError("evidence must contain at least one item") + seen: set[str] = set() + duplicates: set[str] = set() + for item in evidence: + if item.id in seen: + duplicates.add(item.id) + else: + seen.add(item.id) + if duplicates: + raise ValueError( + f"evidence ids must be unique; duplicates: {sorted(duplicates)}" + ) + + require_pydantic_ai() + from pydantic_ai import Agent + from pydantic_ai.settings import ModelSettings + + agent = Agent( + model, + output_type=Decision, + system_prompt=systemPrompt, + model_settings=ModelSettings(thinking=False, extra_body={"think": False}), + ) + result = agent.run_sync(_format_user_prompt(question, evidence)) + return validate_decision(result.output, evidence) diff --git a/scarf/agent/ingest/__init__.py b/scarf/agent/ingest/__init__.py new file mode 100644 index 00000000..a90cfbaa --- /dev/null +++ b/scarf/agent/ingest/__init__.py @@ -0,0 +1,135 @@ +"""Ingest Scarf-supported inputs into a typed Zarr store.""" + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from .cellranger import ingest_cellranger +from .common import CONVERT_FORMATS, ensure_convert_destination +from .detect import detect_format +from .h5ad import ingest_h5ad +from .loom import ingest_loom +from .mtx import ingest_mtx +from .result import IngestResult, needs_input +from .seurat import ingest_seurat +from .zarr_store import ingest_zarr + +__all__ = [ + "IngestResult", + "detect_format", + "ingest", +] + + +def ingest( + *, + path: str | Path, + zarrPath: str | Path | None = None, + model: Any | None = None, + directions: Mapping[str, Any] | None = None, +) -> IngestResult: + """Detect, inspect, convert, and open a Scarf store, or return NeedsInput.""" + source = Path(path) + if not source.exists(): + return IngestResult( + status="failed", + notes=[f"Input path does not exist: {source}"], + ) + + direction_map = dict(directions or {}) + notes: list[str] = [] + format_name = str(direction_map.get("format") or detect_format(source)) + notes.append(f"Detected format: {format_name}") + + destination: str | None = None + if format_name in CONVERT_FORMATS: + preflight = ensure_convert_destination( + source, + zarrPath, + direction_map, + format_name=format_name, + ) + if isinstance(preflight, IngestResult): + preflight.notes = [*notes, *preflight.notes] + return preflight + destination = preflight + if direction_map.get("overwrite") is True: + notes.append(f"Overwrite authorized for destination {destination}") + + if format_name == "zarr": + return ingest_zarr( + source, + notes, + default_assay=direction_map.get("defaultAssay"), + ) + if format_name == "h5ad": + assert destination is not None + return ingest_h5ad( + source, + zarrPath=destination, + model=model, + directions=direction_map, + notes=notes, + ) + if format_name == "10x_h5": + assert destination is not None + return ingest_cellranger( + source, + format_name=format_name, + reader_class_name="CrH5Reader", + zarrPath=destination, + model=model, + directions=direction_map, + notes=notes, + ) + if format_name == "10x_dir": + assert destination is not None + return ingest_cellranger( + source, + format_name=format_name, + reader_class_name="CrDirReader", + zarrPath=destination, + model=model, + directions=direction_map, + notes=notes, + ) + if format_name == "mtx": + assert destination is not None + return ingest_mtx( + source, + zarrPath=destination, + directions=direction_map, + notes=notes, + ) + if format_name == "loom": + assert destination is not None + return ingest_loom( + source, + zarrPath=destination, + directions=direction_map, + notes=notes, + ) + if format_name == "seurat": + assert destination is not None + return ingest_seurat( + source, + zarrPath=destination, + directions=direction_map, + notes=notes, + ) + if format_name == "csv": + return needs_input( + format_name="csv", + question=( + "CSV/TSV import needs explicit parameters " + "(assay name, cell/feature id columns). Provide directions to continue." + ), + options=[], + evidence_ids=[], + notes=notes, + ) + return IngestResult( + status="failed", + format=format_name, + notes=[*notes, f"Unsupported input format for path {source}"], + ) diff --git a/scarf/agent/ingest/cellranger.py b/scarf/agent/ingest/cellranger.py new file mode 100644 index 00000000..ff0ad7fb --- /dev/null +++ b/scarf/agent/ingest/cellranger.py @@ -0,0 +1,87 @@ +"""Cell Ranger H5 and directory ingest handlers.""" + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from .common import CONVERSION_DATA_ERRORS, finish, resolve_modality_choice +from .result import IngestResult, failed_from_exception + + +def ingest_cellranger( + path: Path, + *, + format_name: str, + reader_class_name: str, + zarrPath: str | Path, + model: Any | None, + directions: Mapping[str, Any], + notes: list[str], +) -> IngestResult: + from ...readers.cellranger import CrDirReader, CrH5Reader + from ...writers.cellranger import CrToZarr + + zarr_path = str(zarrPath) + overwrite = directions.get("overwrite") is True + + reader_cls = CrH5Reader if reader_class_name == "CrH5Reader" else CrDirReader + reader = None + writer_started = False + decision = None + rename_assays: dict[str, str] = dict(directions.get("renameAssays") or {}) + try: + reader = reader_cls(str(path)) + assay_columns = list(reader.assayFeats.columns) + if "ADT" in assay_columns and "HTO" not in assay_columns: + adt_names = [str(name) for name in reader.feature_names("ADT")] + choice, decision, blocked = resolve_modality_choice( + model=model, + directions=directions, + feature_names=adt_names, + format_name=format_name, + ) + if blocked is not None: + blocked.notes = [*notes, *blocked.notes] + return blocked + if choice == "HTO": + rename_assays.setdefault("ADT", "HTO") + notes.append("Renamed ADT assay to HTO") + + if rename_assays: + reader.rename_assays(rename_assays) + + writer = CrToZarr(reader, zarr_loc=zarr_path) + writer_started = True + writer.dump() + except CONVERSION_DATA_ERRORS as exc: + return failed_from_exception( + format_name=format_name, + operation="convert cellranger", + exc=exc, + zarr_path=zarr_path, + notes=notes, + partial_store=writer_started, + ) + finally: + if reader is not None: + reader.close() + + convert_action: dict[str, Any] = { + "op": "CrToZarr", + "path": str(path), + "zarrPath": zarr_path, + "readerClass": reader_class_name, + "renameAssays": rename_assays or None, + } + if overwrite: + convert_action["overwrite"] = True + + return finish( + format_name=format_name, + zarr_path=zarr_path, + notes=notes, + convert_actions=[convert_action], + action_labels=["convert_cellranger", "open_datastore"], + default_assay=directions.get("defaultAssay"), + decision=decision, + ) diff --git a/scarf/agent/ingest/common.py b/scarf/agent/ingest/common.py new file mode 100644 index 00000000..6c3a8610 --- /dev/null +++ b/scarf/agent/ingest/common.py @@ -0,0 +1,340 @@ +"""Shared helpers for format-specific ingest handlers.""" + +import re +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +from ...storage.profiles import is_local_zarr_path +from ...storage.stores import zarr_location_has_content +from ..decide import DecisionValidationError, decide +from ..types import Decision, EvidenceItem +from .result import ( + IngestResult, + done, + failed, + failed_from_exception, + failure_note, + needs_input, +) + +# `hto(?![a-z])` matches HTO, HTO1, HTO-1; avoids requiring a word boundary +# after HTO (digits are word chars, so `hto\b` misses HTO1). +_HTO_NAME_RE = re.compile( + r"(hashtag|hto(?![a-z])|totalseq[^a-z0-9]*hash|hash[^a-z0-9]*tag)", + re.IGNORECASE, +) + +CONVERT_FORMATS = frozenset({"h5ad", "10x_h5", "10x_dir", "mtx", "loom", "seurat"}) + +# Inspection and summary boundaries expose data, layout, and I/O failures. +# Some readers use RuntimeError for data-dependent conversion failures, so that +# broader set is restricted to reader and writer execution. +DATA_LAYOUT_ERRORS = (OSError, ValueError, KeyError) +CONVERSION_DATA_ERRORS = (*DATA_LAYOUT_ERRORS, RuntimeError) + + +def _local_path(location: str) -> Path: + if location.startswith("file://"): + return Path(location[7:]) + return Path(location) + + +def _paths_overlap(source: Path, destination: Path) -> bool: + try: + source.relative_to(destination) + return True + except ValueError: + pass + try: + destination.relative_to(source) + return True + except ValueError: + return False + + +def ensure_convert_destination( + source: Path, + zarrPath: str | Path | None, + directions: Mapping[str, Any], + *, + format_name: str, +) -> str | IngestResult: + """Validate conversion destination before inspect, reader, or model work.""" + if zarrPath is None: + return failed( + format_name=format_name, + notes=[f"zarrPath is required when converting {format_name} inputs"], + ) + + destination = str(zarrPath) + overwrite = directions.get("overwrite") + if overwrite is not None and type(overwrite) is not bool: + return failed( + format_name=format_name, + zarr_path=destination, + notes=[ + "overwrite must be boolean true; " + f"got {type(overwrite).__name__}: {overwrite!r}" + ], + ) + + if is_local_zarr_path(destination): + try: + source_resolved = source.resolve() + dest_resolved = _local_path(destination).resolve() + except OSError as exc: + return failed_from_exception( + format_name=format_name, + operation="resolve destination paths", + exc=exc, + zarr_path=destination, + notes=[], + ) + if source_resolved == dest_resolved or _paths_overlap( + source_resolved, dest_resolved + ): + return failed( + format_name=format_name, + zarr_path=destination, + notes=[ + "destination must not equal or nest with the source path; " + f"source={source_resolved} destination={dest_resolved}" + ], + ) + + try: + exists = zarr_location_has_content(destination) + except Exception as exc: + return failed_from_exception( + format_name=format_name, + operation="probe destination", + exc=exc, + zarr_path=destination, + notes=[], + extra_notes=[ + "Destination existence could not be verified; refusing to write", + ], + ) + + if exists and overwrite is not True: + return failed( + format_name=format_name, + zarr_path=destination, + notes=[ + f"Destination already exists: {destination}. " + 'Pass directions={"overwrite": true} to replace it.' + ], + ) + return destination + + +def open_summary( + zarr_path: str, + *, + default_assay: str | None = None, +) -> tuple[list[str], str | None, dict[str, Any]]: + """Open a converted store for first-time QC initialization and summary.""" + import zarr + + from ...datastore.datastore import DataStore + from ...storage.stores import load_zarr + + resolved_default = default_assay + if resolved_default is None: + root = load_zarr(zarr_loc=zarr_path, mode="r") + assay_names = [ + name + for name in sorted(dict.fromkeys(root.group_keys())) + if isinstance(root[name], zarr.Group) and "is_assay" in root[name].attrs + ] + if "RNA" in assay_names: + resolved_default = "RNA" + elif assay_names: + resolved_default = assay_names[0] + + ds = DataStore(zarr_path, default_assay=resolved_default) + summary = ds.summary().to_dict() + return list(ds.assay_names), ds._defaultAssay, summary + + +def open_readonly_summary( + zarr_path: str, + *, + default_assay: str | None = None, +) -> tuple[list[str], str | None, dict[str, Any]]: + """Summarize an existing store without constructing a mutable DataStore.""" + from ...datastore.summary import summarize_zarr_readonly + + summary = summarize_zarr_readonly(zarr_path, default_assay=default_assay) + return ( + [assay.name for assay in summary.assays], + summary.default_assay, + summary.to_dict(), + ) + + +def datastore_action( + zarr_path: str, + default_assay: str | None, + *, + zarr_mode: str = "r+", +) -> dict[str, Any]: + action: dict[str, Any] = { + "op": "DataStore", + "zarrPath": zarr_path, + "zarrMode": zarr_mode, + } + if default_assay is not None: + action["defaultAssay"] = default_assay + return action + + +def finish( + *, + format_name: str, + zarr_path: str, + notes: list[str], + convert_actions: list[dict[str, Any]], + action_labels: list[str], + default_assay: str | None = None, + decision: Decision | None = None, + summary_mode: str = "r+", +) -> IngestResult: + """Open the store, append DataStore replay action, and return a done result.""" + try: + if summary_mode == "r": + assay_names, resolved_default, summary = open_readonly_summary( + zarr_path, + default_assay=default_assay, + ) + action = { + "op": "summarizeZarr", + "zarrPath": zarr_path, + "zarrMode": "r", + } + if resolved_default is not None: + action["defaultAssay"] = resolved_default + else: + assay_names, resolved_default, summary = open_summary( + zarr_path, + default_assay=default_assay, + ) + action = datastore_action(zarr_path, resolved_default, zarr_mode="r+") + except DATA_LAYOUT_ERRORS as exc: + read_only = summary_mode == "r" + return failed_from_exception( + format_name=format_name, + operation=( + "summarize existing Zarr" if read_only else "open converted store" + ), + exc=exc, + zarr_path=zarr_path, + notes=notes, + extra_notes=( + () + if read_only + else (f"Destination may contain a converted store at {zarr_path}",) + ), + ) + return done( + format_name=format_name, + zarr_path=zarr_path, + assay_names=assay_names, + summary=summary, + accepted_actions=[*convert_actions, action], + action_labels=action_labels, + notes=notes, + decision=decision, + ) + + +def antibody_names_look_like_hto(names: Sequence[str]) -> bool: + if not names: + return False + hits = sum(1 for name in names if _HTO_NAME_RE.search(str(name))) + return hits >= max(1, (len(names) + 1) // 2) + + +def _modality_needs_input( + *, + format_name: str, + evidence: Sequence[EvidenceItem], + notes: list[str], +) -> IngestResult: + return needs_input( + format_name=format_name, + question=( + "Antibody Capture features look like hashtags. " + "Should they be imported as ADT or HTO?" + ), + options=["ADT", "HTO"], + evidence_ids=[item.id for item in evidence], + notes=notes, + ) + + +def resolve_modality_choice( + *, + model: Any | None, + directions: Mapping[str, Any], + feature_names: Sequence[str], + format_name: str, +) -> tuple[str | None, Decision | None, IngestResult | None]: + forced = directions.get("modalityChoice") + if forced in {"ADT", "HTO"}: + return str(forced), None, None + if not antibody_names_look_like_hto(feature_names): + return "ADT", None, None + + evidence = [ + EvidenceItem( + id="modality:ADT", + label="ADT", + summary="Treat Antibody Capture features as surface protein ADT assay", + ), + EvidenceItem( + id="modality:HTO", + label="HTO", + summary=( + "Treat Antibody Capture features as multiplexing HTO assay; " + f"names include {[str(name) for name in feature_names[:8]]}" + ), + ), + ] + if model is None: + return ( + None, + None, + _modality_needs_input( + format_name=format_name, + evidence=evidence, + notes=[ + "Hashtag-like Antibody Capture features require an explicit choice" + ], + ), + ) + try: + decision = decide( + model=model, + question=( + "Antibody Capture features look like hashtags. " + "Choose modality:ADT or modality:HTO." + ), + evidence=evidence, + ) + except DecisionValidationError as exc: + return ( + None, + None, + _modality_needs_input( + format_name=format_name, + evidence=evidence, + notes=[ + "Hashtag-like Antibody Capture features require an explicit choice", + failure_note("modalityChoice", exc), + ], + ), + ) + choice = "HTO" if decision.selectedId.endswith("HTO") else "ADT" + return choice, decision, None diff --git a/scarf/agent/ingest/detect.py b/scarf/agent/ingest/detect.py new file mode 100644 index 00000000..4e42d705 --- /dev/null +++ b/scarf/agent/ingest/detect.py @@ -0,0 +1,42 @@ +"""Input format detection for ingest.""" + +from pathlib import Path + + +def detect_format(path: str | Path) -> str: + """Detect a Scarf-supported input family from path layout.""" + source = Path(path) + suffix = "".join(source.suffixes).lower() if source.suffixes else "" + name = source.name.lower() + + if name.endswith(".h5ad") or suffix.endswith(".h5ad"): + return "h5ad" + if name.endswith(".loom") or suffix.endswith(".loom"): + return "loom" + if name.endswith(".rds") or name.endswith(".h5seurat") or suffix.endswith(".rds"): + return "seurat" + if name.endswith(".csv") or name.endswith(".tsv") or name.endswith(".txt"): + return "csv" + if name.endswith(".zarr") or (source.is_dir() and _looks_like_zarr(source)): + return "zarr" + if source.is_file() and name.endswith((".h5", ".hdf5")): + return "10x_h5" + if source.is_dir() and _looks_like_10x_dir(source): + return "10x_dir" + if source.is_dir() or name.endswith((".mtx", ".mtx.gz")): + return "mtx" + return "unknown" + + +def _looks_like_zarr(path: Path) -> bool: + return (path / "zarr.json").exists() or (path / "cellData").exists() + + +def _looks_like_10x_dir(path: Path) -> bool: + names = {child.name.lower() for child in path.iterdir()} if path.exists() else set() + has_matrix = any("matrix.mtx" in name for name in names) + has_barcodes = any("barcode" in name for name in names) + has_features = any( + name.startswith("features") or name.startswith("genes") for name in names + ) + return has_matrix and has_barcodes and has_features diff --git a/scarf/agent/ingest/h5ad.py b/scarf/agent/ingest/h5ad.py new file mode 100644 index 00000000..802a3d11 --- /dev/null +++ b/scarf/agent/ingest/h5ad.py @@ -0,0 +1,175 @@ +"""H5AD ingest handler.""" + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from .common import ( + CONVERSION_DATA_ERRORS, + DATA_LAYOUT_ERRORS, + finish, + resolve_modality_choice, +) +from .result import IngestResult, failed_from_exception, needs_input + + +def _antibody_names(inspection: Any) -> list[str]: + if "ADT" not in inspection.suggestedAssays: + return [] + import h5py + + from ...readers._h5ad_inspect import _as_text, _read_column + + with h5py.File(inspection.h5adFn, mode="r") as h5: + feature_node = h5.get(inspection.featureAttrsKey) + if feature_node is None: + return [] + types = _read_column(feature_node, inspection.assaySplitKey) + names = _read_column(feature_node, inspection.featureNameKey) + if types is None or names is None: + return [] + return [ + _as_text(name) + for feature_type, name in zip(types, names, strict=True) + if _as_text(feature_type) == "Antibody Capture" + ] + + +def ingest_h5ad( + path: Path, + *, + zarrPath: str | Path, + model: Any | None, + directions: Mapping[str, Any], + notes: list[str], +) -> IngestResult: + from ...readers._h5ad_inspect import inspect_h5ad + from ...readers.h5ad import H5adReader + from ...writers.h5ad import H5adToZarr + + zarr_path = str(zarrPath) + overwrite = directions.get("overwrite") is True + + forced_matrix = directions.get("matrixKey") + try: + if forced_matrix is not None: + forced_matrix = str(forced_matrix) + inspection = inspect_h5ad(str(path), matrix_key=forced_matrix) + notes.append( + f"Forced matrix {inspection.matrixKey} via directions " + f"(integerLike={inspection.integerLike})" + ) + else: + inspection = inspect_h5ad(str(path)) + notes.append( + f"Selected matrix {inspection.matrixKey} " + f"(integerLike={inspection.integerLike})" + ) + if not inspection.integerLike: + return needs_input( + format_name="h5ad", + question=( + "No integer-like count matrix matched this H5AD layout. " + "Provide a raw counts matrix, or retry with " + 'directions={"matrixKey": ""} to force a candidate.' + ), + options=list(inspection.matrixCandidates), + evidence_ids=[ + f"matrix:{key}" for key in inspection.matrixCandidates + ], + notes=[ + *notes, + "Prenormalized-only H5AD inputs are not imported silently", + ], + ) + except DATA_LAYOUT_ERRORS as exc: + return failed_from_exception( + format_name="h5ad", + operation="inspect_h5ad", + exc=exc, + zarr_path=zarr_path, + notes=notes, + ) + + assay_name_map = dict(directions.get("assayNameMap") or {}) + decision = None + if inspection.assaySplitKey and "ADT" in inspection.suggestedAssays: + try: + antibody_names = _antibody_names(inspection) + except DATA_LAYOUT_ERRORS as exc: + return failed_from_exception( + format_name="h5ad", + operation="read antibody names", + exc=exc, + zarr_path=zarr_path, + notes=notes, + ) + choice, decision, blocked = resolve_modality_choice( + model=model, + directions=directions, + feature_names=antibody_names, + format_name="h5ad", + ) + if blocked is not None: + blocked.notes = [*notes, *blocked.notes] + return blocked + if choice == "HTO": + assay_name_map.setdefault("Antibody Capture", "HTO") + notes.append("Mapped Antibody Capture to HTO") + + writer_started = False + writer_kwargs: dict[str, Any] = {"zarr_loc": zarr_path} + reader = None + try: + reader = H5adReader.from_inspect(inspection) + if inspection.assaySplitKey is not None: + writer_kwargs["assay_split_key"] = inspection.assaySplitKey + if assay_name_map: + writer_kwargs["assay_name_map"] = assay_name_map + else: + writer_kwargs["assay_name"] = directions.get("assayName") or "RNA" + writer = H5adToZarr(reader, **writer_kwargs) + writer_started = True + writer.dump() + except CONVERSION_DATA_ERRORS as exc: + return failed_from_exception( + format_name="h5ad", + operation="convert h5ad", + exc=exc, + zarr_path=zarr_path, + notes=notes, + partial_store=writer_started, + ) + finally: + if reader is not None: + reader.h5.close() + + convert_action: dict[str, Any] = { + "op": "H5adToZarr", + "path": str(path), + "zarrPath": zarr_path, + "assaySplitKey": inspection.assaySplitKey, + "assayNameMap": assay_name_map or None, + "assayName": None + if inspection.assaySplitKey is not None + else writer_kwargs.get("assay_name"), + } + if overwrite: + convert_action["overwrite"] = True + + return finish( + format_name="h5ad", + zarr_path=zarr_path, + notes=notes, + convert_actions=[ + { + "op": "inspect_h5ad", + "path": str(path), + "matrixKey": inspection.matrixKey, + }, + convert_action, + ], + action_labels=["inspect_h5ad", "convert_h5ad", "open_datastore"], + default_assay=directions.get("defaultAssay"), + decision=decision, + ) diff --git a/scarf/agent/ingest/loom.py b/scarf/agent/ingest/loom.py new file mode 100644 index 00000000..89761ca3 --- /dev/null +++ b/scarf/agent/ingest/loom.py @@ -0,0 +1,69 @@ +"""Loom ingest handler.""" + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from .common import CONVERSION_DATA_ERRORS, finish +from .result import IngestResult, failed_from_exception + + +def ingest_loom( + path: Path, + *, + zarrPath: str | Path, + directions: Mapping[str, Any], + notes: list[str], +) -> IngestResult: + from ...readers.loom import LoomReader + from ...writers.loom import LoomToZarr + + zarr_path = str(zarrPath) + overwrite = directions.get("overwrite") is True + + reader_kwargs = {} + if directions.get("cellNamesKey"): + reader_kwargs["cell_names_key"] = directions["cellNamesKey"] + if directions.get("featureNamesKey"): + reader_kwargs["feature_names_key"] = directions["featureNamesKey"] + + reader = None + writer_started = False + try: + reader = LoomReader(str(path), **reader_kwargs) + writer = LoomToZarr( + reader, + zarr_loc=zarr_path, + assay_name=directions.get("assayName") or "RNA", + ) + writer_started = True + writer.dump() + except CONVERSION_DATA_ERRORS as exc: + return failed_from_exception( + format_name="loom", + operation="convert loom", + exc=exc, + zarr_path=zarr_path, + notes=notes, + partial_store=writer_started, + ) + finally: + if reader is not None: + reader.h5.close() + + convert_action: dict[str, Any] = { + "op": "LoomToZarr", + "path": str(path), + "zarrPath": zarr_path, + } + if overwrite: + convert_action["overwrite"] = True + + return finish( + format_name="loom", + zarr_path=zarr_path, + notes=notes, + convert_actions=[convert_action], + action_labels=["convert_loom", "open_datastore"], + default_assay=directions.get("defaultAssay"), + ) diff --git a/scarf/agent/ingest/mtx.py b/scarf/agent/ingest/mtx.py new file mode 100644 index 00000000..a751436c --- /dev/null +++ b/scarf/agent/ingest/mtx.py @@ -0,0 +1,125 @@ +"""Matrix Market ingest handler.""" + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from .common import CONVERSION_DATA_ERRORS, DATA_LAYOUT_ERRORS, finish +from .result import IngestResult, failed, failed_from_exception, needs_input + + +def _resolve_mtx_index( + raw: Any, + *, + n_candidates: int, +) -> int | IngestResult: + if raw is None: + return 0 + if type(raw) is bool or not isinstance(raw, int | float | str): + return failed( + format_name="mtx", + notes=[f"mtxIndex must be an integer index; got {raw!r}"], + ) + try: + if isinstance(raw, float) and not raw.is_integer(): + raise ValueError("non-integral float") + index = int(raw) + except (TypeError, ValueError): + return failed( + format_name="mtx", + notes=[f"mtxIndex must be an integer index; got {raw!r}"], + ) + if index < 0 or index >= n_candidates: + return failed( + format_name="mtx", + notes=[ + f"mtxIndex {index} is out of range for {n_candidates} MTX candidates" + ], + ) + return index + + +def ingest_mtx( + path: Path, + *, + zarrPath: str | Path, + directions: Mapping[str, Any], + notes: list[str], +) -> IngestResult: + from ...readers.mtx import MtxReader, inspect_mtx + from ...writers.cellranger import MtxToZarr + + zarr_path = str(zarrPath) + overwrite = directions.get("overwrite") is True + + try: + candidates = inspect_mtx(path) + except DATA_LAYOUT_ERRORS as exc: + return failed_from_exception( + format_name="mtx", + operation="inspect_mtx", + exc=exc, + zarr_path=zarr_path, + notes=notes, + ) + if not candidates: + return failed( + format_name="mtx", + zarr_path=zarr_path, + notes=[*notes, f"No MTX matrix candidates found under {path}"], + ) + if len(candidates) > 1 and directions.get("mtxIndex") is None: + return needs_input( + format_name="mtx", + question="Multiple MTX layouts found. Which candidate index should be used?", + options=[str(index) for index in range(len(candidates))], + evidence_ids=[f"mtx:{index}" for index in range(len(candidates))], + notes=[*notes, f"Found {len(candidates)} MTX candidates"], + ) + + resolved = _resolve_mtx_index( + directions.get("mtxIndex"), + n_candidates=len(candidates), + ) + if isinstance(resolved, IngestResult): + resolved.notes = [*notes, *resolved.notes] + resolved.zarrPath = zarr_path + return resolved + + reader = None + writer_started = False + try: + reader = MtxReader(candidates[resolved]) + writer = MtxToZarr(reader, zarr_loc=zarr_path) + writer_started = True + writer.dump() + except CONVERSION_DATA_ERRORS as exc: + return failed_from_exception( + format_name="mtx", + operation="convert mtx", + exc=exc, + zarr_path=zarr_path, + notes=notes, + partial_store=writer_started, + ) + finally: + if reader is not None: + reader.close() + + convert_action: dict[str, Any] = { + "op": "MtxToZarr", + "path": str(path), + "zarrPath": zarr_path, + "mtxIndex": resolved, + } + if overwrite: + convert_action["overwrite"] = True + + return finish( + format_name="mtx", + zarr_path=zarr_path, + notes=notes, + convert_actions=[convert_action], + action_labels=["convert_mtx", "open_datastore"], + default_assay=directions.get("defaultAssay"), + ) diff --git a/scarf/agent/ingest/result.py b/scarf/agent/ingest/result.py new file mode 100644 index 00000000..3f38d3b2 --- /dev/null +++ b/scarf/agent/ingest/result.py @@ -0,0 +1,114 @@ +"""Ingest result types and stage helpers.""" + +from collections.abc import Sequence +from typing import Any + +from .._deps import AGENT_INSTALL_HINT +from ..types import Decision, NeedsInput, StageStatus + +try: + from pydantic import BaseModel, Field +except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + + +class IngestResult(BaseModel): + status: StageStatus + format: str | None = None + zarrPath: str | None = None + assayNames: list[str] = Field(default_factory=list) + summary: dict[str, Any] | None = None + decision: Decision | None = None + needsInput: NeedsInput | None = None + actions: list[str] = Field(default_factory=list) + acceptedActions: list[dict[str, Any]] = Field(default_factory=list) + notes: list[str] = Field(default_factory=list) + + +def done( + *, + format_name: str, + zarr_path: str, + assay_names: list[str], + summary: dict[str, Any], + accepted_actions: list[dict[str, Any]], + action_labels: list[str], + notes: list[str], + decision: Decision | None = None, +) -> IngestResult: + return IngestResult( + status="done", + format=format_name, + zarrPath=zarr_path, + assayNames=assay_names, + summary=summary, + decision=decision, + actions=action_labels, + acceptedActions=accepted_actions, + notes=notes, + ) + + +def needs_input( + *, + format_name: str, + question: str, + options: list[str], + evidence_ids: list[str], + notes: list[str] | None = None, +) -> IngestResult: + return IngestResult( + status="needsInput", + format=format_name, + needsInput=NeedsInput( + question=question, + options=options, + evidenceIds=evidence_ids, + ), + notes=notes or [], + ) + + +def failed( + *, + format_name: str | None = None, + notes: list[str], + zarr_path: str | None = None, +) -> IngestResult: + return IngestResult( + status="failed", + format=format_name, + zarrPath=zarr_path, + notes=notes, + ) + + +def failure_note(operation: str, exc: BaseException) -> str: + return f"{operation} failed: {type(exc).__name__}: {exc}" + + +def failed_from_exception( + *, + format_name: str, + operation: str, + exc: BaseException, + notes: Sequence[str], + zarr_path: str | None = None, + extra_notes: Sequence[str] = (), + partial_store: bool = False, +) -> IngestResult: + partial_notes = ( + [f"Destination may contain a partial store at {zarr_path}"] + if partial_store and zarr_path is not None + else [] + ) + return failed( + format_name=format_name, + zarr_path=zarr_path, + notes=[ + *notes, + failure_note(operation, exc), + *partial_notes, + *extra_notes, + ], + ) diff --git a/scarf/agent/ingest/seurat.py b/scarf/agent/ingest/seurat.py new file mode 100644 index 00000000..e3976d0a --- /dev/null +++ b/scarf/agent/ingest/seurat.py @@ -0,0 +1,59 @@ +"""Seurat ingest handler.""" + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from .common import CONVERSION_DATA_ERRORS, finish +from .result import IngestResult, failed_from_exception + + +def ingest_seurat( + path: Path, + *, + zarrPath: str | Path, + directions: Mapping[str, Any], + notes: list[str], +) -> IngestResult: + from ...readers.seurat import SeuratReader + from ...writers.seurat import SeuratToZarr + + zarr_path = str(zarrPath) + overwrite = directions.get("overwrite") is True + + reader = None + writer_started = False + try: + reader = SeuratReader(str(path)) + writer = SeuratToZarr(reader, zarr_loc=zarr_path) + writer_started = True + writer.dump() + except CONVERSION_DATA_ERRORS as exc: + return failed_from_exception( + format_name="seurat", + operation="convert seurat", + exc=exc, + zarr_path=zarr_path, + notes=notes, + partial_store=writer_started, + ) + finally: + if reader is not None: + reader.close() + + convert_action: dict[str, Any] = { + "op": "SeuratToZarr", + "path": str(path), + "zarrPath": zarr_path, + } + if overwrite: + convert_action["overwrite"] = True + + return finish( + format_name="seurat", + zarr_path=zarr_path, + notes=notes, + convert_actions=[convert_action], + action_labels=["convert_seurat", "open_datastore"], + default_assay=directions.get("defaultAssay"), + ) diff --git a/scarf/agent/ingest/zarr_store.py b/scarf/agent/ingest/zarr_store.py new file mode 100644 index 00000000..104e8c5a --- /dev/null +++ b/scarf/agent/ingest/zarr_store.py @@ -0,0 +1,23 @@ +"""Open an existing Scarf Zarr store.""" + +from pathlib import Path + +from .common import finish +from .result import IngestResult + + +def ingest_zarr( + path: Path, + notes: list[str], + *, + default_assay: str | None, +) -> IngestResult: + return finish( + format_name="zarr", + zarr_path=str(path), + notes=notes, + convert_actions=[], + action_labels=["summarize_zarr"], + default_assay=default_assay, + summary_mode="r", + ) diff --git a/scarf/agent/runtime.py b/scarf/agent/runtime.py new file mode 100644 index 00000000..0f5d8616 --- /dev/null +++ b/scarf/agent/runtime.py @@ -0,0 +1,116 @@ +"""Runtime checks and local env loading for LLM endpoints.""" + +import json +import os +import urllib.error +import urllib.request +from pathlib import Path +from typing import Any + +_ENV_PATH = Path(__file__).resolve().parent / ".env" + + +def load_env(path: Path | None = None) -> Path | None: + """Load scarf/agent/.env into os.environ without overriding existing vars. + + Returns the path loaded, or None if no file was found. + """ + env_path = path or _ENV_PATH + if not env_path.is_file(): + return None + for line in env_path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, value = line.split("=", 1) + key = key.strip() + value = value.strip().strip('"').strip("'") + if key and key not in os.environ: + os.environ[key] = value + return env_path + + +def _unreachable_message(baseUrl: str, model: str) -> str: + return ( + f"OpenAI-compatible endpoint unreachable at {baseUrl!r} " + f"(requested model {model!r})." + ) + + +def _missing_model_message(baseUrl: str, model: str, available: set[str]) -> str: + listed = ", ".join(sorted(available)) if available else "(none)" + return f"Model {model!r} not found at {baseUrl!r}. Available models: {listed}." + + +def _missing_config_message() -> str: + example = _ENV_PATH.with_name(".env.example") + return ( + "baseUrl and model are required. Pass them as arguments or set " + "OLLAMA_BASE_URL and OLLAMA_MODEL in the environment or in " + f"{_ENV_PATH} (see {example})." + ) + + +def _request_json( + url: str, + *, + timeout: float, + api_key: str | None, +) -> Any: + headers = {"Accept": "application/json"} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + request = urllib.request.Request(url, headers=headers, method="GET") + with urllib.request.urlopen(request, timeout=timeout) as response: + return json.loads(response.read().decode("utf-8")) + + +def _listed_model_ids(payload: Any) -> set[str]: + ids: set[str] = set() + if not isinstance(payload, dict): + return ids + data = payload.get("data") + if isinstance(data, list): + for item in data: + if isinstance(item, dict) and isinstance(item.get("id"), str): + ids.add(item["id"]) + models = payload.get("models") + if isinstance(models, list): + for item in models: + if not isinstance(item, dict): + continue + name = item.get("name") or item.get("model") + if isinstance(name, str): + ids.add(name) + return ids + + +def check_runtime( + *, + baseUrl: str | None = None, + model: str | None = None, + timeout: float = 5.0, +) -> None: + """Fail fast if the OpenAI-compatible endpoint or model is unavailable. + + Loads `scarf/agent/.env` first (without overriding existing env vars). + Uses `OLLAMA_BASE_URL`, `OLLAMA_MODEL`, and `OLLAMA_API_KEY` when args are omitted. + """ + load_env() + resolved_base = baseUrl or os.environ.get("OLLAMA_BASE_URL") + resolved_model = model or os.environ.get("OLLAMA_MODEL") + if not resolved_base or not resolved_model: + raise ValueError(_missing_config_message()) + + api_key = os.environ.get("OLLAMA_API_KEY") or None + models_url = f"{resolved_base.rstrip('/')}/models" + try: + payload = _request_json(models_url, timeout=timeout, api_key=api_key) + except (urllib.error.URLError, TimeoutError, json.JSONDecodeError) as exc: + raise RuntimeError(_unreachable_message(resolved_base, resolved_model)) from exc + + available = _listed_model_ids(payload) + if resolved_model not in available: + raise RuntimeError( + _missing_model_message(resolved_base, resolved_model, available) + ) diff --git a/scarf/agent/types.py b/scarf/agent/types.py new file mode 100644 index 00000000..074ec7b2 --- /dev/null +++ b/scarf/agent/types.py @@ -0,0 +1,44 @@ +"""Shared result types for scarf.agent stages and decide().""" + +from typing import Literal + +from ._deps import AGENT_INSTALL_HINT + +try: + from pydantic import BaseModel, Field +except ImportError as exc: + raise ImportError(AGENT_INSTALL_HINT) from exc + + +type StageStatus = Literal["done", "needsInput", "failed"] + + +class EvidenceItem(BaseModel): + id: str + label: str + summary: str + + +class Decision(BaseModel): + selectedId: str = Field( + description="Exact evidence id string from the provided list, nothing else" + ) + rationale: str = Field(description="Short reason for the choice") + evidenceIds: list[str] = Field( + default_factory=list, + description="Evidence ids used; must include selectedId and only provided ids", + ) + + +class NeedsInput(BaseModel): + question: str + options: list[str] = Field(default_factory=list) + evidenceIds: list[str] = Field(default_factory=list) + + +class StageResult(BaseModel): + status: StageStatus + decision: Decision | None = None + needsInput: NeedsInput | None = None + actions: list[str] = Field(default_factory=list) + notes: list[str] = Field(default_factory=list) diff --git a/scarf/datastore/summary.py b/scarf/datastore/summary.py index 7029cf45..dad48d8b 100644 --- a/scarf/datastore/summary.py +++ b/scarf/datastore/summary.py @@ -5,13 +5,19 @@ import zarr from ..assay import Assay -from ..graph.state import AssayState +from ..graph.state import AssayState, read_assay_state from ..metadata import MetaData -from ..storage.artifacts import ArtifactStatus -from ..storage.budget import ResourceBudget -from ..storage.profiles import StorageProfile +from ..storage.artifacts import ( + ArtifactStatus, + inspect_artifact, + list_artifacts as list_artifact_refs, +) +from ..storage.budget import ResourceBudget, resolve_budget +from ..storage.profiles import StorageProfile, resolve_storage_profile from ..storage.refs import ArtifactRef, ArtifactScope -from ..storage.types import ZarrMode +from ..storage.schema import validate_assay_name +from ..storage.stores import load_zarr +from ..storage.types import ZarrMode, as_zarr_group @dataclass(frozen=True, slots=True) @@ -95,6 +101,11 @@ def to_dict(self) -> dict[str, Any]: } +class _AssaySummaryView(Protocol): + feats: MetaData + attrs: Mapping[str, Any] + + class _SummaryStore(Protocol): zarr_mode: ZarrMode workspace: str | None @@ -109,7 +120,7 @@ def assay_names(self) -> list[str]: ... @property def zw(self) -> zarr.Group: ... - def _get_assay(self, from_assay: str | None) -> Assay: ... + def _get_assay(self, from_assay: str | None) -> Assay | _AssaySummaryView: ... def get_assay_state(self, from_assay: str | None = None) -> AssayState | None: ... @@ -125,6 +136,117 @@ def list_artifacts( def inspect_artifact(self, ref: ArtifactRef) -> ArtifactStatus: ... +@dataclass(slots=True) +class _ReadOnlyAssayView: + name: str + feats: MetaData + attrs: Mapping[str, Any] + + +class _ReadOnlySummaryStore: + """Read-only adapter that summarizes a Scarf Zarr without DataStore init.""" + + zarr_mode: ZarrMode = "r" + + def __init__( + self, + zarr_path: str, + *, + default_assay: str | None = None, + workspace: str | None = None, + storage_options: dict[str, Any] | None = None, + ) -> None: + self.workspace = workspace + self.resources = resolve_budget() + self.storageProfile = resolve_storage_profile(zarr_path) + self._root = load_zarr(zarr_path, mode="r", storage_options=storage_options) + self.cells = MetaData(as_zarr_group(self.zw["cellData"], name="cellData")) + names = self.assay_names + if not names: + raise ValueError(f"No assays found in Zarr store at {zarr_path}") + self._defaultAssay = self._resolve_default_assay(default_assay, names) + + @property + def zw(self) -> zarr.Group: + if self.workspace is None: + return self._root + return as_zarr_group(self._root[self.workspace], name=self.workspace) + + @property + def assay_names(self) -> list[str]: + names: list[str] = [] + for name in sorted(dict.fromkeys(self.zw.group_keys())): + node = self.zw[name] + if isinstance(node, zarr.Group) and "is_assay" in node.attrs: + validate_assay_name(name) + names.append(name) + return names + + def _resolve_default_assay( + self, + requested: str | None, + assay_names: list[str], + ) -> str: + if requested is not None: + if requested not in assay_names: + raise ValueError( + f"Default assay {requested!r} was not found. " + f"Choose one from: {' '.join(assay_names)}" + ) + return requested + stored = self.zw.attrs.get("defaultAssay") + if isinstance(stored, str) and stored in assay_names: + return stored + if "RNA" in assay_names: + return "RNA" + return assay_names[0] + + def _get_assay(self, from_assay: str | None) -> Assay | _AssaySummaryView: + assay = from_assay or self._defaultAssay + if assay not in self.assay_names: + raise ValueError(f"Assay {assay!r} not found in the Zarr file") + feature_path = f"{assay}/featureData" + display_path = ( + feature_path + if self.workspace is None + else f"{self.workspace}/{feature_path}" + ) + assay_group = as_zarr_group( + self.zw[assay], + name=assay if self.workspace is None else f"{self.workspace}/{assay}", + ) + return _ReadOnlyAssayView( + name=assay, + feats=MetaData(as_zarr_group(self.zw[feature_path], name=display_path)), + attrs=assay_group.attrs, + ) + + def get_assay_state(self, from_assay: str | None = None) -> AssayState | None: + assay = from_assay or self._defaultAssay + return read_assay_state(self.zw, assay) + + def list_artifacts( + self, + *, + kind: str | None = None, + from_assay: str | None = None, + scope: ArtifactScope = "assay", + complete_only: bool = False, + ) -> list[ArtifactRef]: + if scope == "assay" and from_assay is None: + from_assay = self._defaultAssay + return list_artifact_refs( + self.zw, + scope=scope, + assay=from_assay, + kind=kind, + complete_only=complete_only, + ) + + def inspect_artifact(self, ref: ArtifactRef) -> ArtifactStatus: + return inspect_artifact(self.zw, ref) + + def _count_active(metadata: MetaData) -> int: if metadata.N == 0: return 0 @@ -211,3 +333,22 @@ def build_datastore_summary( store.list_artifacts(scope="datastore"), ), ) + + +def summarize_zarr_readonly( + zarr_path: str, + *, + default_assay: str | None = None, + workspace: str | None = None, + storage_options: dict[str, Any] | None = None, +) -> DataStoreSummary: + """Summarize an existing Scarf Zarr without mutating it.""" + from .. import __version__ + + store = _ReadOnlySummaryStore( + zarr_path, + default_assay=default_assay, + workspace=workspace, + storage_options=storage_options, + ) + return build_datastore_summary(store, scarf_version=__version__) diff --git a/scarf/features/gene_reference.py b/scarf/features/gene_reference.py new file mode 100644 index 00000000..b20682df --- /dev/null +++ b/scarf/features/gene_reference.py @@ -0,0 +1,374 @@ +"""Per-species Ensembl gene reference download and local lookup.""" + +import gzip +import os +import re +import urllib.request +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any, TextIO + +__all__ = [ + "GeneReference", + "SpeciesSpec", + "cached_species", + "default_cache_dir", + "ensure_reference", + "load_reference", + "parse_gff3_genes", + "prefix_species", + "reference_summary", + "species_registry", + "write_reference_fixture", +] + +_ENSEMBL_GFF3 = "https://ftp.ensembl.org/pub/current_gff3/{species}/" +_ENSEMBL_GENOMES_GFF3 = ( + "https://ftp.ensemblgenomes.ebi.ac.uk/pub/{division}/current/gff3/{species}/" +) +_GENE_TYPES = frozenset({"gene", "ncRNA_gene", "pseudogene"}) +_GFF_NAME = re.compile( + r"^(?P