From 2bec3dafc5dcb2c477697dd95e3027e88430357c Mon Sep 17 00:00:00 2001 From: David Kelley Date: Thu, 8 Oct 2026 10:47:18 -0700 Subject: [PATCH 1/6] Replace fast and qnorm spec with one-pass qmap-partial spec hound_data writes exact whole-genome per-track histograms of the stored fp16 targets (strand sums) as target_hist; SpecPearsonCorrCoef builds qnorm's tie-averaged quantile maps from them, maps preds and targets through each track's table, regresses out the group mean per track, and streams per-track moments. hound_eval reports it in its single pass. Removes hound_eval_spec, spec_norm, combine_pairs, and hound_eval_folds --spec/--band/--seq_chunk/--ram/--scratch_dir. Co-Authored-By: Claude Opus 5.5 --- docs/gcprunner.md | 4 - docs/train.md | 4 +- docs/train_folds.md | 18 +- pyproject.toml | 1 - src/baskerville/dataset.py | 72 ++-- src/baskerville/metrics.py | 123 +++++-- src/baskerville/scripts/hound_data.py | 4 +- src/baskerville/scripts/hound_eval_folds.py | 147 +------- src/baskerville/scripts/hound_eval_spec.py | 337 ------------------ .../utils/hound_data_subset_zarr_targets.py | 12 + src/baskerville/seqnn.py | 93 +---- src/baskerville/spec_norm.py | 255 ------------- src/baskerville/trainer.py | 2 +- src/tests/test_eval.py | 75 +--- src/tests/test_eval_folds_gcp.py | 32 +- src/tests/test_spec_metrics.py | 148 ++++++-- src/tests/test_spec_norm.py | 221 ------------ 17 files changed, 312 insertions(+), 1236 deletions(-) delete mode 100755 src/baskerville/scripts/hound_eval_spec.py delete mode 100644 src/baskerville/spec_norm.py delete mode 100644 src/tests/test_spec_norm.py diff --git a/docs/gcprunner.md b/docs/gcprunner.md index 67da5e4..039773c 100644 --- a/docs/gcprunner.md +++ b/docs/gcprunner.md @@ -852,10 +852,6 @@ mounted at `/workspace/cache`. Use this when the weights are actually on disk. - `--test` / `--valid` evaluate only the held-out test / validation fold per model (from `folds.json`); the default evaluates every fold. - `--save` / `--aggregate_genes` need ~60 GB RAM — use `-q l4-large`. -- `--spec` (specificity) is **auto-pinned to `-q l4-large`** (g2-standard-32, - 32 vCPU, ~122 GB) regardless of `-q`, since quantile normalization is RAM- and - CPU-hungry; it runs with `--ncpus 32`. Your `-q` still governs the lighter - coverage-eval jobs. - On GCP, `mem` is the VM's hard RAM (the task owns the whole VM), not a soft Slurm request — size it via `-q`, not a mem flag. Eval warns if the chosen profile is too small. diff --git a/docs/train.md b/docs/train.md index 60eed51..10f8109 100644 --- a/docs/train.md +++ b/docs/train.md @@ -126,9 +126,9 @@ Each epoch, `log.txt` gets a summary line per dataset (`loss`, `r`, `r2`, `spec` Tracks are grouped by the targets `group` column (else the `description` prefix). Each group with at least `train.spec_group_min` tracks (default 20, after merging strand pairs) gets: - `r/`, `r2/`: mean per-track Pearson r and R² over the group's tracks. -- `spec/`: specificity, the mean per-track Pearson r after dividing every track by its dataset-wide mean and subtracting the group mean at each position. It rewards predicting how a track differs from its group, not the shared signal. +- `spec/`: specificity, the mean per-track Pearson r of what remains after removing the group's shared signal. Each track's values (strand pairs summed) are quantile-normalized to the group with a fixed per-track lookup table; predictions use the same table as their targets. The group mean at each position is then regressed out of each track, with the track's own slope and intercept. It rewards predicting how a track differs from its group, not the shared signal or the track's depth and signal-to-noise. -`spec` is the unweighted mean of `spec/` over groups. It needs the per-track means that `hound_data` stores in each `examples/*.zarr`; for an older dataset, add them with `python -c "from baskerville import dataset; dataset.write_target_means('data_dir')"`. +`spec` is the unweighted mean of `spec/` over groups. `hound_eval` reports the same metric per track (`spec` column of `acc.txt`), in its single pass. The tables are built from the exact whole-genome histograms of each track's stored fp16 values, which `hound_data` writes to each `examples/*.zarr` as `target_hist`. For an older dataset, add them with `python -c "from baskerville import dataset; dataset.write_target_hist('data_dir')"`. `train.stop_stat` chooses the statistic that selects `model_best.pth` and drives `patience`. It is a `{metric: weight}` dict over the keys above plus `loss`, `r_gene`, and `r2_gene`, and the weighted sum is maximized. The presets `"weighted_r_r2"` (default, `{"r": 1, "r2": 0.25}`), `"r"`, `"r2"`, and `"loss"` (`{"loss": -1}`) remain available as strings; these presets score gene-only datasets by −loss. A dict applies only its own keys, e.g. `{"r_gene": 1}` or `{"loss": -1}` for gene-only datasets. diff --git a/docs/train_folds.md b/docs/train_folds.md index 7060358..5c216ce 100644 --- a/docs/train_folds.md +++ b/docs/train_folds.md @@ -258,8 +258,6 @@ hound_eval_folds --backend local --rank -o train_folds params.json data_folds # Evaluate only test sets (faster) hound_eval_folds --backend local --test -o train_folds params.json data_folds -# Include specificity analysis -hound_eval_folds --backend local --spec -o train_folds params.json data_folds ``` ### Understanding Evaluation Output @@ -269,20 +267,18 @@ After evaluation, each fold directory contains: ``` train_folds/f0c0/ ├── train/ # Training artifacts -├── eval/ # Evaluation results -│ ├── fold0/ # Model evaluated on fold 0 -│ ├── fold1/ # Model evaluated on fold 1 -│ ├── ... -│ ├── test -> fold0 # Symlink to test fold -│ └── test.out -└── spec/ # Specificity analysis (if --spec used) +└── eval/ # Evaluation results + ├── fold0/ # Model evaluated on fold 0 + ├── fold1/ # Model evaluated on fold 1 + ├── ... + ├── test -> fold0 # Symlink to test fold + └── test.out ``` Key files: - `eval/fold{X}/acc.txt`: Accuracy metrics for each target -- `eval/test/acc.txt`: Test set performance (most important) -- `spec/acc.txt`: Specificity analysis results +- `eval/test/acc.txt`: Test set performance (most important), with per-track `spec` (see [train.md](train.md)) ## Analyzing Results diff --git a/pyproject.toml b/pyproject.toml index 3445064..7d279a1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,7 +80,6 @@ hound_eval = "baskerville.scripts.hound_eval:main" hound_eval_folds = "baskerville.scripts.hound_eval_folds:main" hound_eval_genes = "baskerville.scripts.hound_eval_genes:main" hound_eval_genes_folds = "baskerville.scripts.hound_eval_genes_folds:main" -hound_eval_spec = "baskerville.scripts.hound_eval_spec:main" hound_grad = "baskerville.scripts.hound_grad:main" hound_grad_folds = "baskerville.scripts.hound_grad_folds:main" hound_ism_bed = "baskerville.scripts.hound_ism_bed:main" diff --git a/src/baskerville/dataset.py b/src/baskerville/dataset.py index 6421ecf..0b68997 100644 --- a/src/baskerville/dataset.py +++ b/src/baskerville/dataset.py @@ -187,10 +187,10 @@ def __init__( # don't persist handles (not fork-safe with DataLoader workers) zarr_data = [zarr.open(zf, mode="r") for zf in self.zarr_files] - # per-track means of the stored targets (write_target_means) - self.target_means = None - if self.has_coverage and "mean" in zarr_data[0]["target"].attrs: - self.target_means = np.array(zarr_data[0]["target"].attrs["mean"]) + # per-track histograms of the stored targets (write_target_hist) + self.target_hist = None + if self.has_coverage and "target_hist" in zarr_data[0]: + self.target_hist = zarr_data[0]["target_hist"][:] # count sequences zarr_seqs = [zd["sequence"].shape[0] for zd in zarr_data] @@ -864,47 +864,69 @@ def target_groups(targets_df): return np.array(groups) -def _target_block_sum(args): - """Float64 per-track sum of the stored targets over sequences [start, end).""" - zarr_file, start, end = args +# non-negative finite fp16 bit patterns: 0x0000 (0.0) to 0x7BFF (65504) +NUM_HIST_BINS = 0x7C00 + + +def _target_block_hist(args): + """Per-track counts of strand-summed fp16 target bit patterns over + sequences [start, end).""" + zarr_file, start, end, pair = args targets = zarr.open(zarr_file, mode="r")["target"] - block_sum = np.zeros(targets.shape[1]) + num_targets = targets.shape[1] + offsets = np.arange(num_targets)[:, None] * NUM_HIST_BINS + hist = np.zeros(num_targets * NUM_HIST_BINS, dtype=np.int64) for si in range(start, end): - block_sum += targets[si].astype("float64").sum(axis=-1) - return block_sum + y = targets[si].astype(np.float32) + # + 0.0 turns -0.0 into 0.0 + bits = (y + y[pair] + 0.0).astype(np.float16).view(np.uint16) + if bits.max() >= NUM_HIST_BINS: + raise ValueError(f"{zarr_file} seq {si}: negative or non-finite targets") + hist += np.bincount((bits + offsets).ravel(), minlength=hist.size) + return hist.reshape(num_targets, NUM_HIST_BINS) -def write_target_means(data_dir, processes=16, block_seqs=64): - """Store exact per-track means of the stored targets in each zarr. +def write_target_hist(data_dir, processes=16, block_seqs=256): + """Store exact per-track histograms of the stored targets in each zarr. - The mean is over every position of every sequence in all examples/*.zarr, - in the stored (transformed) space, and is written to each zarr's - ``target.attrs["mean"]``. + Counts are over every position of every sequence in all examples/*.zarr, + of the fp16 value y_t + y_pair(t) (strand_pair from targets.txt; unstranded + tracks are their own pair, so 2 * y_t), binned by fp16 bit pattern. They are + written to each zarr as ``target_hist`` (num_targets, NUM_HIST_BINS) int64, + from which SpecPearsonCorrCoef builds its quantile maps. Args: - data_dir: dataset directory with examples/*.zarr. + data_dir: dataset directory with targets.txt and examples/*.zarr. processes: parallel readers. block_seqs: sequences per read task. Returns: - np.ndarray of per-track means. + np.ndarray of per-track histograms. """ + targets_df = pd.read_csv(f"{data_dir}/targets.txt", sep="\t", index_col=0) + pair = strand_pair_indices(targets_df) zarr_files = natsorted(glob.glob(f"{data_dir}/examples/*.zarr")) tasks = [] - num_positions = 0 for zarr_file in zarr_files: - num_seqs, _, length = zarr.open(zarr_file, mode="r")["target"].shape - num_positions += num_seqs * length + num_seqs = zarr.open(zarr_file, mode="r")["target"].shape[0] for start in range(0, num_seqs, block_seqs): - tasks.append((zarr_file, start, min(start + block_seqs, num_seqs))) + tasks.append((zarr_file, start, min(start + block_seqs, num_seqs), pair)) + hist = 0 with multiprocessing.get_context("spawn").Pool(processes) as pool: - target_sum = np.sum(pool.map(_target_block_sum, tasks), axis=0) - target_means = target_sum / num_positions + for block_hist in pool.imap_unordered(_target_block_hist, tasks): + hist = hist + block_hist for zarr_file in zarr_files: - zarr.open(zarr_file, mode="r+")["target"].attrs["mean"] = target_means.tolist() - return target_means + root = zarr.open_group(zarr_file, mode="r+") + root.create_array( + "target_hist", + shape=hist.shape, + dtype="int64", + chunks=(1, NUM_HIST_BINS), + overwrite=True, + )[:] = hist + return hist def targets_prep_strand(targets_df): diff --git a/src/baskerville/metrics.py b/src/baskerville/metrics.py index c5ecb42..0cec0cd 100644 --- a/src/baskerville/metrics.py +++ b/src/baskerville/metrics.py @@ -105,25 +105,69 @@ def reset(self): self.r2score.reset() +def qmap_tables(hists): + """Quantile-normalization lookup tables from per-track value histograms. + + Exactly qnorm's map for the histogrammed data: the group reference is the + average of the tracks' sorted columns, and a value occupying ranks [s, e) + in its track maps to the mean reference over [s, e) (ties averaged). Values + absent from a track interpolate between its present values and clamp + outside them, so any fp16 prediction has an entry. + + Args: + hists: (T, dataset.NUM_HIST_BINS) counts per fp16 bit pattern, with + equal totals. + + Returns: + np.ndarray (T, NUM_HIST_BINS) float32 tables, indexed by bit pattern. + """ + hists = np.asarray(hists, dtype=np.int64) + totals = hists.sum(axis=1) + if totals.min() == 0 or totals.min() != totals.max(): + raise ValueError("target histograms must have equal nonzero totals") + values = np.arange(dataset.NUM_HIST_BINS, dtype=np.uint16).view(np.float16) + values = values.astype(np.float64) + + # per track: present values, rank bounds of their tie blocks, and the + # running sum of the sorted column at those bounds (linear within blocks) + knots = [] + for h in hists: + nz = np.flatnonzero(h) + bounds = np.concatenate([[0], np.cumsum(h[nz])]) + sums = np.concatenate([[0.0], np.cumsum(h[nz] * values[nz])]) + knots.append((nz, bounds, sums)) + + tables = np.empty(hists.shape, dtype=np.float32) + for t, (nz, bounds, _) in enumerate(knots): + # running sum of the reference = mean of the tracks' running sums + ref_sum = np.mean([np.interp(bounds, b, s) for _, b, s in knots], axis=0) + tables[t] = np.interp(values, values[nz], np.diff(ref_sum) / np.diff(bounds)) + return tables + + class SpecPearsonCorrCoef: """Per-track specificity Pearson within target groups. For each group (dataset.target_groups) with at least ``group_min`` tracks - after strand collapse: sum strand pairs, divide preds and targets by the - dataset's track means (SeqDataset.target_means), subtract the group mean - at each position, and correlate each track's residuals. + after strand collapse: sum strand pairs, map preds and targets through each + track's quantile-normalization table (qmap_tables on the dataset's + whole-genome target histograms, SeqDataset.target_hist), regress the group + mean at each position out of each track (its own slope and intercept), and + correlate each track's pred and target residuals. Streaming state is the + per-track first and second moments of (pred, target, group-mean pred, + group-mean target). Args: targets_df: targets table aligned with the output tracks. - target_means: per-track means of the stored targets, or None. + target_hist: per-track histograms of the stored targets, or None. group_min: minimum tracks per group. ``groups`` maps group name to (rep, pair) track positions. compute() returns a per-track tensor with values at rep positions, NaN elsewhere. - Without target_means, ``enabled`` is False and nothing is scored. + Without target_hist, ``enabled`` is False and nothing is scored. """ - def __init__(self, targets_df, target_means, group_min: int = 20): + def __init__(self, targets_df, target_hist, group_min: int = 20): self.num_targets = len(targets_df) rep_pos, pair_pos = dataset.strand_collapse(targets_df) rep_groups = dataset.target_groups(targets_df)[rep_pos] @@ -133,51 +177,78 @@ def __init__(self, targets_df, target_means, group_min: int = 20): if gi.sum() >= group_min: self.groups[g] = (rep_pos[gi], pair_pos[gi]) - self.enabled = target_means is not None + self.enabled = target_hist is not None if not self.enabled: - print("Warning: data lack target means; spec metric disabled.") + print("Warning: data lack target histograms; spec metric disabled.") return - means = np.asarray(target_means) - self.index, self.b, self.pearson = {}, {}, {} + hist = np.asarray(target_hist) + self.index, self.table = {}, {} for g, (rep, pair) in self.groups.items(): self.index[g] = (torch.tensor(rep), torch.tensor(pair)) - # unstranded tracks have pair == rep, so the doubling cancels - b = np.maximum(means[rep] + means[pair], 1e-6) - self.b[g] = torch.tensor(b, dtype=torch.float32).view(1, -1, 1) - self.pearson[g] = PearsonCorrCoef(len(rep), average=False) + # histograms are of the strand sum, stored at both partners + self.table[g] = torch.from_numpy(qmap_tables(hist[rep])) + self.reset() def to(self, device): if self.enabled: for g in self.groups: self.index[g] = tuple(ix.to(device) for ix in self.index[g]) - self.b[g] = self.b[g].to(device) - self.pearson[g].to(device) + self.table[g] = self.table[g].to(device) + self.sum1[g] = self.sum1[g].to(device) + self.sum2[g] = self.sum2[g].to(device) return self - def _center(self, x, g): + def _map(self, x, g): + """(N, T, L) -> (tracks, N*L) table values of the group's strand sums.""" rep, pair = self.index[g] - x = (x[:, rep].float() + x[:, pair].float()) / self.b[g] - return x - x.mean(dim=1, keepdim=True) + s = x[:, rep].float() + x[:, pair].float() + bits = s.clamp(0, 65504).half().view(torch.int16).long() + return torch.gather(self.table[g], 1, bits.transpose(0, 1).flatten(1)) @torch.no_grad() def update(self, preds: torch.Tensor, target: torch.Tensor) -> None: """preds, target: (N, T, L).""" if self.enabled: for g in self.groups: - self.pearson[g].update(self._center(preds, g), self._center(target, g)) + zp, zy = self._map(preds, g), self._map(target, g) + mp = zp.mean(dim=0, keepdim=True).expand_as(zp) + my = zy.mean(dim=0, keepdim=True).expand_as(zy) + z = torch.stack([zp, zy, mp, my], dim=1).double() # (tracks, 4, N*L) + self.sum1[g] += z.sum(dim=-1) + self.sum2[g] += torch.einsum("tap,tbp->tab", z, z) + self.count[g] += zp.shape[1] def compute(self) -> torch.Tensor: spec = torch.full((self.num_targets,), float("nan")) if self.enabled: for g in self.groups: - spec[self.index[g][0].cpu()] = self.pearson[g].compute().float().cpu() + n = self.count[g] + s1, s2 = self.sum1[g].cpu(), self.sum2[g].cpu() + cov = s2 / n - s1[:, :, None] * s1[:, None, :] / n**2 + # residuals of pred/target after regressing out their group mean + u = torch.zeros(len(cov), 4, dtype=torch.float64) + w = torch.zeros_like(u) + u[:, 0], u[:, 2] = 1, -cov[:, 0, 2] / cov[:, 2, 2] + w[:, 1], w[:, 3] = 1, -cov[:, 1, 3] / cov[:, 3, 3] + cuw = torch.einsum("ta,tab,tb->t", u, cov, w) + cuu = torch.einsum("ta,tab,tb->t", u, cov, u) + cww = torch.einsum("ta,tab,tb->t", w, cov, w) + spec[self.index[g][0].cpu()] = (cuw / torch.sqrt(cuu * cww)).float() return spec def reset(self): if self.enabled: - for p in self.pearson.values(): - p.reset() + device = {g: t.device for g, t in self.table.items()} + self.sum1 = { + g: torch.zeros(len(r), 4, dtype=torch.float64, device=device[g]) + for g, (r, _) in self.groups.items() + } + self.sum2 = { + g: torch.zeros(len(r), 4, 4, dtype=torch.float64, device=device[g]) + for g, (r, _) in self.groups.items() + } + self.count = {g: 0 for g in self.groups} class _MaskedGeneMetric(torchmetrics.Metric): @@ -378,7 +449,7 @@ class DatasetMetrics: device: Device to place metrics on. targets_df: Coverage targets table; enables per-group metrics when it matches num_targets. - target_means: Per-track means of the stored targets; enables spec. + target_hist: Per-track histograms of the stored targets; enables spec. spec_group_min: Minimum tracks for a group to be scored. """ @@ -390,7 +461,7 @@ def __init__( num_gene_targets: int | None, device: str, targets_df=None, - target_means=None, + target_hist=None, spec_group_min: int = 20, ): self.has_coverage = has_coverage @@ -406,7 +477,7 @@ def __init__( self.r2 = R2Score(average=False).to(device) if targets_df is not None and len(targets_df) == num_targets: self.spec = SpecPearsonCorrCoef( - targets_df, target_means, spec_group_min + targets_df, target_hist, spec_group_min ).to(device) else: self.r = None diff --git a/src/baskerville/scripts/hound_data.py b/src/baskerville/scripts/hound_data.py index d9d5b90..9e4ac5e 100755 --- a/src/baskerville/scripts/hound_data.py +++ b/src/baskerville/scripts/hound_data.py @@ -899,9 +899,9 @@ def main(): update_sleep=3, ) - # per-track means, for depth-normalized specificity metrics + # per-track value histograms, for the specificity metric's quantile maps if num_targets > 0: - dataset.write_target_means(args.out_dir, processes=args.processes) + dataset.write_target_hist(args.out_dir, processes=args.processes) ################################################################ # stats diff --git a/src/baskerville/scripts/hound_eval_folds.py b/src/baskerville/scripts/hound_eval_folds.py index ca3dff0..8822804 100755 --- a/src/baskerville/scripts/hound_eval_folds.py +++ b/src/baskerville/scripts/hound_eval_folds.py @@ -63,19 +63,12 @@ def main(): description="Evaluate baskerville model replicates on cross folds using given parameters and data." ) - # eval/spec options + # eval options parser.add_argument( "--aggregate_genes", action="store_true", help="Aggregate predictions per unique gene across sequences [Default: %(default)s]", ) - parser.add_argument( - "--band", - default=None, - type=int, - help="Tracks read per band during spec streaming normalization; lower " - "to reduce peak RAM [Default: hound_eval_spec's default]", - ) parser.add_argument( "-o", "--out_dir", @@ -101,12 +94,6 @@ def main(): action="store_true", help="Save targets and predictions numpy arrays [Default: %(default)s]", ) - parser.add_argument( - "--seq_chunk", - default=None, - type=int, - help="Zarr seq-axis chunk for preds/targets store; [Default: batch_size]", - ) parser.add_argument( "--shifts", default="0", @@ -117,7 +104,7 @@ def main(): "--step", default=1, type=int, - help="Spatial step for specificity/spearmanr [Default: %(default)s]", + help="Spatial step for saved/ranked preds [Default: %(default)s]", ) parser.add_argument( "--test", @@ -187,26 +174,6 @@ def main(): action="store_true", help="Restart evaluation [Default: %(default)s]", ) - parser.add_argument( - "--ram", - default=False, - action="store_true", - help="Hold spec preds/targets in RAM instead of streaming to a Zarr " - "store on disk (big-memory nodes only) [Default: %(default)s]", - ) - parser.add_argument( - "--scratch_dir", - default=None, - help="Directory for the spec preds/targets Zarr store; must be a real " - "disk with room [Default: the job output directory]", - ) - parser.add_argument( - "--spec", - default=False, - action="store_true", - help="Specificity evaluation [Default: %(default)s]", - ) - parser.add_argument("params_file", help="JSON file with model parameters") parser.add_argument( "data_dirs", nargs="+", help="Train/valid/test data directorie(s)" @@ -300,10 +267,8 @@ def eval_folds(args): models_sha or hashlib.sha256(gcs_models_dir.encode()).hexdigest() ) # output-changing options too, so changing them starts fresh instead of - # resuming; spec's are keyed even without --spec so adding it resumes - opts_sha = stage_cache.hash_text( - eval_options(args, container_targets_gene) + spec_options(args) - ) + # resuming + opts_sha = stage_cache.hash_text(eval_options(args, container_targets_gene)) run_id = stage_cache.build_run_id( models_key, params_sha, data_uri_sha, opts_sha, deterministic=True ) @@ -369,7 +334,6 @@ def eval_folds(args): # need conda activation and a $HOSTNAME echo. if gcp_backend: cmd_base = "hound_eval" - spec_cmd_base = "hound_eval_spec" data_mounts = [ DataMount(args.gcp_data_dir, args.gcp_data_local, mode="fuse"), DataMount( @@ -394,7 +358,6 @@ def eval_folds(args): else: env_base = utils.conda_activate(args.conda_env) + "echo $HOSTNAME;" cmd_base = f"{env_base} hound_eval" - spec_cmd_base = f"{env_base} hound_eval_spec" gcp_extra = {} jobs = [] @@ -507,98 +470,6 @@ def eval_folds(args): ) jobs.append(job) - ####################################################### - # evaluate test specificity - - if args.spec: - for ci in range(args.crosses): - for fi in fold_index: - fold_cross = f"f{fi}c{ci}" - it_dir = f"{args.out_dir}/{fold_cross}" - train_dir = f"{it_dir}/train" - - if not _model_present( - gcp_backend, gcs_models_dir, train_dir, fold_cross - ): - continue - - # use the test fold training actually held out (not just fi) - test_fold, _ = _read_fold_splits(train_dir, num_folds, fi, ci) - - model_file = _resolve_model_file( - gcp_backend, - gcs_models_dir, - container_models_dir, - train_dir, - fold_cross, - ) - - for di in range(num_data): - spec_sub = "spec" if num_data == 1 else f"spec{di}" - rel = f"{fold_cross}/{spec_sub}" - - # check if done (GCS for gcp, local otherwise) - if gcp_backend: - acc_loc = f"{args.gcp_output_dir}/{rel}/acc.txt" - already_done = gcs_file_exist(acc_loc) - out_dir = f"/workspace/out/{rel}" - else: - out_dir = f"{it_dir}/{spec_sub}" - acc_loc = f"{out_dir}/acc.txt" - already_done = os.path.isfile(acc_loc) - if already_done: - print(f"{acc_loc} already generated.") - continue - - cmd = spec_cmd_base - cmd += f" --dataset {di}" - cmd += f" -o {out_dir}" - cmd += f" --split fold{test_fold}" - cmd += spec_options(args) - cmd += f" --ncpus {num_cpu}" - if args.band is not None: - cmd += f" --band {args.band}" - if args.seq_chunk is not None: - cmd += f" --seq_chunk {args.seq_chunk}" - if args.ram: - cmd += " --ram" - if args.scratch_dir is not None: - cmd += f" --scratch_dir {args.scratch_dir}" - cmd += f" {params_file}" - cmd += f" {model_file}" - cmd += f" {data_dirs[di]}" - - if slurmrunner is None: - # Run locally - jobs.append(cmd) - else: - # Submit to SLURM / GCP Batch - name = f"{args.name}-spec-{fold_cross}" - spec_extra = dict(gcp_extra) - if gcp_backend: - out_file = f"{args.gcp_output_dir}/{rel}.out" - err_file = f"{args.gcp_output_dir}/{rel}.err" - spec_extra["boot_disk_gb"] = 200 - spec_extra["labels"] = run_identity.identity_labels( - run_id, "eval", fold_cross - ) - else: - out_file = f"{out_dir}.out" - err_file = f"{out_dir}.err" - job = slurmrunner.Job( - cmd, - name=name, - out_file=out_file, - err_file=err_file, - queue=args.queue, - cpu=num_cpu, - gpu=num_gpu, - mem=None if gcp_backend else 64000, - time=f"{3 * time_base}:00:00", - **spec_extra, - ) - jobs.append(job) - # dir exists but had no replicate weights (vs. all jobs already complete) if not found_model: where = gcs_models_dir if gcs_models_dir is not None else args.out_dir @@ -667,16 +538,6 @@ def eval_options(args, targets_gene): return opts -def spec_options(args): - """hound_eval_spec flags that change its output (not memory/IO knobs).""" - opts = f" --step {args.step}" - if args.rc: - opts += " --rc" - if args.shifts: - opts += f" --shifts {args.shifts}" - return opts - - def _link_test_metrics(eval_dir, test_fold): """Symlink test -> fold{test_fold} and test.out -> fold{test_fold}.out. diff --git a/src/baskerville/scripts/hound_eval_spec.py b/src/baskerville/scripts/hound_eval_spec.py deleted file mode 100755 index c5bba39..0000000 --- a/src/baskerville/scripts/hound_eval_spec.py +++ /dev/null @@ -1,337 +0,0 @@ -#!/usr/bin/env python -# Copyright 2020 Calico Life Sciences LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ========================================================================= -import argparse -import json -import os -import shutil -import tempfile -import time - -import numpy as np -import pandas as pd -import torch - -from baskerville import dataset -from baskerville import seqnn -from baskerville import spec_norm -from baskerville.hardware import check_mixed_precision - -""" -hound_eval_spec - -Test the accuracy of a trained model on targets/predictions normalized across targets. -""" - - -def main(): - parser = argparse.ArgumentParser(description="Evaluate a trained model.") - parser.add_argument( - "--band", - default=spec_norm.DEFAULT_BAND_SIZE, - type=int, - help="Tracks read per band during streaming normalization [Default: %(default)s]", - ) - parser.add_argument( - "-di", - "--dataset", - dest="di", - default=0, - type=int, - help="Dataset index to evaluate [Default: %(default)s]", - ) - parser.add_argument( - "-m", - "--group_min", - default=20, - type=int, - help="Minimum target group size to consider [Default: %(default)s]", - ) - parser.add_argument( - "--ncpus", - default=2, - type=int, - help="Threads for column sort/normalization [Default: %(default)s]", - ) - parser.add_argument( - "-o", - "--out_dir", - default="spec_out", - help="Output directory for evaluation statistics [Default: %(default)s]", - ) - parser.add_argument( - "--ram", - default=False, - action="store_true", - help="Hold preds/targets in RAM instead of streaming to a Zarr store on disk " - "(default streams to disk; only use --ram on big-memory nodes)", - ) - parser.add_argument( - "--rc", - default=False, - action="store_true", - help="Average the fwd and rc predictions [Default: %(default)s]", - ) - parser.add_argument( - "--scratch_dir", - default=None, - help="Directory for the temporary preds/targets Zarr store. Must be a " - "real disk with room (NOT a tmpfs like many /tmp mounts) " - "[Default: the output directory]", - ) - parser.add_argument( - "--seq_chunk", - default=None, - type=int, - help="Zarr seq-axis chunk for the preds/targets store; wider chunks cut the " - "file count and per-column read ops (writes are buffered a chunk at a time; " - "RAM ≈ 2 * seq_chunk * num_targets * bins * 2 bytes for preds+targets). " - "Default: None (uses batch_size) [Default: %(default)s]", - ) - parser.add_argument( - "--shifts", - default="0", - help="Ensemble prediction shifts [Default: %(default)s]", - ) - parser.add_argument( - "--split", - default="test", - help="Dataset split label for eg TFR pattern [Default: %(default)s]", - ) - parser.add_argument( - "--step", - default=1, - type=int, - help="Step across positions [Default: %(default)s]", - ) - parser.add_argument( - "-t", - "--targets_file", - default=None, - help="File specifying target indexes and labels in table format", - ) - parser.add_argument( - "--target_groups", - default=None, - type=str, - help="Comma separated string of target groups", - ) - parser.add_argument( - "--var_pct", - default=1.0, - type=float, - help="Highly variable site proportion to take [Default: %(default)s]", - ) - parser.add_argument("params_file", help="JSON file with model parameters") - parser.add_argument("model_file", help="Trained model file.") - parser.add_argument("data_dir", help="Train/valid/test data directory") - args = parser.parse_args() - - os.makedirs(args.out_dir, exist_ok=True) - - # parse shifts to integers - args.shifts = [int(shift) for shift in args.shifts.split(",")] - - ####################################################### - # targets - - # read table - if args.targets_file is None: - args.targets_file = f"{args.data_dir}/targets.txt" - targets_df = pd.read_csv(args.targets_file, index_col=0, sep="\t") - target_slice = targets_df.index.values - num_targets = targets_df.shape[0] - - # set target groups - targets_df["group"] = dataset.target_groups(targets_df) - - if args.target_groups is None: - args.target_groups = sorted(set(targets_df.group)) - else: - args.target_groups = args.target_groups.split(",") - - ####################################################### - # model - - # read parameters - with open(args.params_file) as params_open: - params = json.load(params_open) - params_model = params["model"] - params_train = params.get("train", {}) - - # update strand pairs for new indexing (positional partner per target) - strand_pair_pos = dataset.strand_pair_indices(targets_df) - params_model["strand_pair"] = strand_pair_pos - - # Collapse stranded pairs to their plus representative: one column per - # experiment, summed at eval write-time (combine_pairs below) -- store is - # ~half size for stranded assays. - rep_pos, pair_pos = dataset.strand_collapse(targets_df) - targets_strand_df = targets_df.iloc[rep_pos] - store_columns = [ - (int(r),) if p == r else (int(r), int(p)) for r, p in zip(rep_pos, pair_pos) - ] - num_store_targets = len(store_columns) - - # construct eval data - eval_data = dataset.SeqDataset( - args.data_dir, - split_label=args.split, - mode="eval", - targets_file=args.targets_file, - ) - - # initialize model - seqnn_model = seqnn.SeqNN(params_model, output_slice=target_slice) - seqnn_model.restore(args.model_file) - seqnn_model.ensemble_rc = args.rc - seqnn_model.ensemble_shifts = args.shifts - mix_dtype = params_train.get("mix_dtype", "float32") - if mix_dtype != "float32": - if check_mixed_precision(): - if mix_dtype == "float16": - seqnn_model.mix_dtype = torch.float16 - elif mix_dtype == "bfloat16": - seqnn_model.mix_dtype = torch.bfloat16 - else: - print(f"Warning: Unrecognized mixed precision dtype {mix_dtype}") - else: - print("Warning: Mixed precision training not supported on this GPU.") - - ####################################################### - # targets/predictions - - # By default stream preds/targets to a Zarr store on disk so the full - # (num_seqs, num_targets, bins) arrays never live in RAM (--ram keeps them in - # memory, only viable on big-memory nodes). Scratch is --scratch_dir, else the - # output dir -- deliberately NOT $TMPDIR or /tmp: SLURM sets $TMPDIR=/tmp on the - # compute nodes, a small mount shared across jobs where big stores ENOSPC-crash. - # The output dir is per-job on real disk (NFS here, ~TB free); point - # --scratch_dir at node-local scratch only if it is a real disk with room. - tmp_root = None - zarr_path = None - if args.ram and args.seq_chunk: - print( - "Warning: --seq_chunk is ignored with --ram (no on-disk store)", flush=True - ) - if not args.ram: - scratch_root = args.scratch_dir or args.out_dir - tmp_root = tempfile.mkdtemp(prefix="hound_spec_", dir=scratch_root) - zarr_path = os.path.join(tmp_root, "preds_targets.zarr") - free_gb = shutil.disk_usage(tmp_root).free / 1e9 - # uncompressed upper bound for preds + targets (float16); on-disk is - # smaller after compression, but fail fast if there is clearly no room - est_gb = ( - 2 - * len(eval_data) - * num_store_targets - * (eval_data.target_length // args.step) - * 2 - ) / 1e9 - print( - "Streaming preds/targets to %s (%.0f GB free, ~%.0f GB uncompressed)" - % (zarr_path, free_gb, est_gb), - flush=True, - ) - if free_gb < 0.4 * est_gb: - shutil.rmtree(tmp_root, ignore_errors=True) - raise OSError( - "Not enough scratch space at %s (%.0f GB free, ~%.0f GB needed). " - "Set --scratch_dir to a disk with room, raise --step, or use --ram." - % (scratch_root, free_gb, est_gb) - ) - - try: - # target_chunk chunks the store along targets so each band read (below) - # decompresses only its own tracks; combine_pairs sums each stranded pair - # at write-time -> one column per experiment, in targets_strand_df order. - # None when nothing collapses (identity) to skip a per-batch gather. - combine_pairs = store_columns if num_store_targets < num_targets else None - results = seqnn_model.eval( - eval_data, - hi=args.di, - batch_size=params_train.get("batch_size", 1), - return_values=True, - step=args.step, - zarr_store=zarr_path, - target_chunk=args.band, - seq_chunk=args.seq_chunk, - combine_pairs=combine_pairs, - ) - - if "coverage" not in results: - raise ValueError( - "Expected coverage evaluation results, but coverage head is unavailable." - ) - - cov_results = results["coverage"] - eval_preds = cov_results["preds"] - eval_targets = cov_results["targets"] - - ################################################### - # process groups (collapsed / plus-represented frame) - - group_values = targets_strand_df.group.values - targets_spec = np.full(num_store_targets, np.nan) - - for tg in args.target_groups: - # column indices into the store (== positions in targets_strand_df) - col_idx = np.where(group_values == tg)[0] - num_targets_group = len(col_idx) - print("%-15s %4d" % (tg, num_targets_group), flush=True) - - # group_min counts experiments (columns), the cross-track dimension - if num_targets_group < args.group_min: - continue - - # streaming quantile-normalized cross-track specificity - t0 = time.time() - print(" Quantile normalize + correlate...", flush=True, end="") - col_pearsonr = spec_norm.group_specificity_pearson( - eval_preds, - eval_targets, - col_idx, - var_pct=args.var_pct, - band_size=args.band, - ncpus=args.ncpus, - ) - print("DONE in %ds" % (time.time() - t0)) - - targets_spec[col_idx] = col_pearsonr - - print(" PearsonR %.4f" % np.nanmean(targets_spec[col_idx]), flush=True) - - # write target-level statistics (one row per experiment, plus-represented) - targets_acc_df = pd.DataFrame( - { - "index": targets_strand_df.index, - "pearsonr": targets_spec, - "identifier": targets_strand_df.identifier, - "description": targets_strand_df.description, - } - ) - targets_acc_df.to_csv( - f"{args.out_dir}/acc.txt", sep="\t", index=False, float_format="%.5f" - ) - finally: - if tmp_root is not None: - shutil.rmtree(tmp_root, ignore_errors=True) - - -################################################################################ -# __main__ -################################################################################ -if __name__ == "__main__": - main() diff --git a/src/baskerville/scripts/utils/hound_data_subset_zarr_targets.py b/src/baskerville/scripts/utils/hound_data_subset_zarr_targets.py index a956f03..45c3a2d 100755 --- a/src/baskerville/scripts/utils/hound_data_subset_zarr_targets.py +++ b/src/baskerville/scripts/utils/hound_data_subset_zarr_targets.py @@ -107,6 +107,18 @@ def subset_zarr_for_fold( compressors=compressors, ) + # ** keep the kept tracks' histograms (strand sums: keep pairs together) ** + old_root = zarr.open_group(f"{og_data_dir}/examples/fold{fi}.zarr", mode="r") + if "target_hist" in old_root: + hist = old_root["target_hist"][:][old_indices] + fold_zarr_root.create_array( + "target_hist", + shape=hist.shape, + dtype=hist.dtype, + chunks=(1, hist.shape[1]), + overwrite=True, + )[:] = hist + # --- CREATE DATALOADER --- dataset_obj = IndexedSeqDataset(og_data_dir, split_label=f"fold{fi}", mode="eval") dataloader = torch.utils.data.DataLoader( diff --git a/src/baskerville/seqnn.py b/src/baskerville/seqnn.py index bbbc194..a33994d 100644 --- a/src/baskerville/seqnn.py +++ b/src/baskerville/seqnn.py @@ -222,7 +222,6 @@ def eval( zarr_store: str | None = None, target_chunk: int | None = 128, seq_chunk: int | None = None, - combine_pairs=None, ): """Evaluate model on data. @@ -235,19 +234,12 @@ def eval( zarr_store: Optional zarr store path for predictions. target_chunk: zarr chunk size along the target axis seq_chunk: zarr chunk size along the seq axis for the coverage store - combine_pairs: optional list of coverage column specs (tuples of - target indices in the output_slice frame) collapsing the stored - preds/targets to one column per spec -- ``(p,)`` keeps a track, - ``(p, q)`` sums a strand pair. Reduces the stored/returned - coverage arrays' target axis; coverage 'r'/'r2' are returned as - None when set (per-track metrics no longer align with the - collapsed preds). Default None leaves every track individual. Returns: dict with keys: - 'coverage': dict with 'r', 'r2', 'spec', 'preds', 'targets' (if coverage head); 'spec' is per-track specificity Pearson (NaN - outside scored groups), or None without data.target_means + outside scored groups), or None without data.target_hist - 'gene': dict with 'r', 'r2', 'preds', 'targets', 'masks' (if gene head) """ self.model.eval() @@ -290,34 +282,21 @@ def eval( eval_metric_r2 = metrics.R2Score(average=False) eval_metric_r2.to(self.device) - # specificity, given the dataset's target means + # specificity, given the dataset's target histograms eval_metric_spec = None - target_means = getattr(data, "target_means", None) - if target_means is not None: + target_hist = getattr(data, "target_hist", None) + if target_hist is not None: targets_df = data.targets_df if self.output_slice is not None: targets_df = targets_df.loc[self.output_slice] - target_means = target_means[self.output_slice] + target_hist = target_hist[self.output_slice] # metadata must describe every scored track if len(targets_df) == num_targets: eval_metric_spec = metrics.SpecPearsonCorrCoef( - targets_df, target_means + targets_df, target_hist ) eval_metric_spec.to(self.device) - # optional strand-pair reduction of the stored coverage arrays (metrics - # above stay on the full, un-reduced tracks) - combine_np = combine_torch = None - num_store_targets = num_targets if has_coverage else None - if has_coverage and combine_pairs is not None: - a_idx, pair_cols, pair_b = _strand_column_index(combine_pairs) - num_store_targets = len(combine_pairs) - combine_np = (a_idx, pair_cols, pair_b) - combine_torch = tuple( - torch.as_tensor(ix, device=self.device) - for ix in (a_idx, pair_cols, pair_b) - ) - # prepare gene metrics eval_metric_r_gene = eval_metric_r2_gene = None if has_gene: @@ -364,11 +343,11 @@ def _create_array(name, shape, chunks, dtype): f"step ({step}) must evenly divide target_length ({data.target_length})" ) targets_length = data.target_length // step - vshape = (num_seqs, num_store_targets, targets_length) + vshape = (num_seqs, num_targets, targets_length) tchunk = ( - num_store_targets + num_targets if target_chunk is None - else min(target_chunk, num_store_targets) + else min(target_chunk, num_targets) ) cshape = (cov_seq_chunk or batch_size, tchunk, targets_length) preds = _create_array("preds", vshape, cshape, "float16") @@ -409,8 +388,6 @@ def _create_array(name, shape, chunks, dtype): y_np = y.detach().numpy() if step != 1: y_np = y_np[:, :, ::step] - if combine_np is not None: - y_np = _combine_strand_columns(y_np, *combine_np) w_targets.add(si, y_np) y = y.to(self.device).float() @@ -441,9 +418,6 @@ def _create_array(name, shape, chunks, dtype): eval_metric_spec.update(yh.coverage, y) if return_values: yhs = yh.coverage if step == 1 else yh.coverage[:, :, ::step] - if combine_torch is not None: - # sum strand pairs before the fp16 store cast - yhs = _combine_strand_columns(yhs, *combine_torch) w_preds.add(si, yhs.to(torch.float16).detach().cpu().numpy()) # update gene metrics @@ -466,15 +440,11 @@ def _create_array(name, shape, chunks, dtype): results = {} if has_coverage: - if combine_pairs is not None: - # per-track metrics don't align with collapsed preds -- omit - eval_r = eval_r2 = eval_spec = None - else: - eval_r = eval_metric_r.compute().cpu().detach().numpy() - eval_r2 = eval_metric_r2.compute().cpu().detach().numpy() - eval_spec = None - if eval_metric_spec is not None: - eval_spec = eval_metric_spec.compute().numpy() + eval_r = eval_metric_r.compute().cpu().detach().numpy() + eval_r2 = eval_metric_r2.compute().cpu().detach().numpy() + eval_spec = None + if eval_metric_spec is not None: + eval_spec = eval_metric_spec.compute().numpy() results["coverage"] = { "r": eval_r, "r2": eval_r2, @@ -1082,38 +1052,3 @@ def flush(self): self.arr[self.base : self.base + self.fill] = self.buf[: self.fill] self.base += self.fill self.fill = 0 - - -def _strand_column_index(columns): - """Index arrays for reducing tracks to strand-summed columns along axis 1. - - Args: - columns: list of tuples of source track indices -- ``(p,)`` for a single - (unstranded) track, ``(p, q)`` for a strand pair to be summed. - - Returns: - ``(a_idx, pair_cols, pair_b)``: ``a_idx`` (len num_columns) is the first - source track of each output column; ``pair_cols`` are the output columns - that are strand pairs; ``pair_b`` are those pairs' second source tracks. - """ - for c in columns: - if len(c) not in (1, 2): - raise ValueError(f"combine_pairs spec must have 1 or 2 indices, got {c!r}") - a_idx = np.array([int(c[0]) for c in columns], dtype=np.int64) - is_pair = np.array([len(c) == 2 for c in columns], dtype=bool) - pair_cols = np.where(is_pair)[0].astype(np.int64) - pair_b = np.array([int(columns[c][1]) for c in pair_cols], dtype=np.int64) - return a_idx, pair_cols, pair_b - - -def _combine_strand_columns(arr, a_idx, pair_cols, pair_b): - """Reduce ``arr`` (..., num_tracks, bins) to strand-summed columns along axis 1. - - ``out[:, j, :] = arr[:, a_idx[j], :]`` plus, for paired columns, the partner - track. Works for both numpy arrays and torch tensors (``a_idx`` etc. must be - the matching index type). ``a_idx`` is unique, so the in-place add is safe. - """ - out = arr[:, a_idx, :] # advanced index -> fresh array/tensor (a copy) - if len(pair_cols): - out[:, pair_cols, :] += arr[:, pair_b, :] - return out diff --git a/src/baskerville/spec_norm.py b/src/baskerville/spec_norm.py deleted file mode 100644 index bf045cd..0000000 --- a/src/baskerville/spec_norm.py +++ /dev/null @@ -1,255 +0,0 @@ -"""Memory-efficient cross-track specificity normalization. - -`hound_eval_spec` measures how well a model predicts the *relative* differences -between same-assay tracks at each genomic position (rather than just the shared -mean). Historically this quantile-normalized a dense ``(num_seqs*bins, num_tracks)`` -matrix per target group with the ``qnorm`` package, whose internal int64 argsort -temporaries blow up memory for large track collections. - -This module reproduces the exact quantile-normalization semantics but computes them -**one column at a time**, so peak memory is ``O(num_positions)`` and independent of -the number of tracks. A "column" is one stored track's flattened ``eval[:, ti, :]`` -vector. Stranded experiments are already collapsed to a single total-signal column -upstream (``seqnn.eval(combine_pairs=...)`` sums each plus/minus pair before writing -the store), so this module only ever sees individual columns. Quantile normalization, -the per-position cross-track mean subtraction, the optional high-variance-position -filter, and the per-track Pearson correlation are all computed in a few streaming -passes with float64 accumulators. -""" - -from concurrent.futures import ThreadPoolExecutor - -import numpy as np -from scipy.stats import pearsonr - -# Default columns (tracks) per band. Band memory is ~2 * band_size * R * 4 -# bytes (preds + targets, float32). When the eval store is chunked along the -# target axis (hound_eval_spec passes target_chunk=band), a band read only -# touches its own chunks, so this doubles as the natural target-chunk size. -DEFAULT_BAND_SIZE = 16 - - -def apply_quantile_column(col, reference): - """Map a column onto a reference distribution, exactly like ``qnorm``. - - Each value is replaced by the reference value at its rank; tied values are - all assigned the *average* of the reference over the tied block (matching - ``qnorm.quantile_normalize``'s handling, which a rank-midpoint interpolation - would not reproduce on the large block of zeros typical of coverage data). - - Args: - col: 1D array of length R. - reference: 1D ascending array of length R (the mean of the sorted - columns across the group). - - Returns: - float64 1D array of length R (the normalized column). - """ - col = np.asarray(col) - n = col.shape[0] - out = np.empty(n, dtype=np.float64) - if n == 0: - return out - - # ties are block-averaged below, so intra-tie order is irrelevant: the - # default (faster) quicksort gives identical results to a stable sort. - order = np.argsort(col) - sorted_col = col[order] - - # group id per sorted position: increments at each new (untied) value - change = np.empty(n, dtype=bool) - change[0] = True - np.not_equal(sorted_col[1:], sorted_col[:-1], out=change[1:]) - group_id = np.cumsum(change) - 1 - - # average the reference within each tie block - ref = np.asarray(reference, dtype=np.float64) - block_sum = np.bincount(group_id, weights=ref) - block_cnt = np.bincount(group_id) - block_mean = block_sum / block_cnt - - out[order] = block_mean[group_id] - return out - - -def _read_tracks(eval_arr, track_indices): - """Read ``eval_arr[:, track_indices, :]`` (numpy or zarr) as finite float32. - - Non-finite values raise (QN has no NaN/inf handling). - """ - if isinstance(eval_arr, np.ndarray): - sub = eval_arr[:, track_indices, :] - else: - # zarr: orthogonal indexing along the target axis - sub = eval_arr.oindex[:, list(track_indices), :] - arr = np.asarray(sub, dtype=np.float32) - if not np.isfinite(arr).all(): - raise ValueError("spec_norm requires finite preds/targets (found NaN/inf)") - return arr - - -def _iter_bands(eval_preds, eval_targets, column_indices, band_size): - """Yield, for each band of columns, a list of ``(col_p, col_t)`` pairs. - - A band reads ``band_size`` columns from ``eval_preds``/``eval_targets`` in a - single (chunk-friendly) read and flattens each into its ``num_seqs*bins`` - vector. Re-invoking this generator re-reads the data (one read per pass). - """ - num_cols = len(column_indices) - for start in range(0, num_cols, band_size): - idx = [int(c) for c in column_indices[start : start + band_size]] - band_p = _read_tracks(eval_preds, idx) - band_t = _read_tracks(eval_targets, idx) - band = [ - # reshape already copies (non-contiguous view) - (band_p[:, k, :].reshape(-1), band_t[:, k, :].reshape(-1)) - for k in range(len(idx)) - ] - yield band - - -def _pmap(func, items, executor): - """Map ``func`` over ``items``, on ``executor``'s threads if one is given. - - numpy sort/apply release the GIL, so threads give real parallelism; the - caller reduces the results serially to avoid races on the accumulators. - """ - if executor is not None and len(items) > 1: - return list(executor.map(func, items)) - return [func(it) for it in items] - - -def _stream_columns( - eval_preds, eval_targets, column_indices, band_size, per_column, executor -): - """Yield ``per_column((col_p, col_t))`` for every column, in column order. - - Wraps the band read + optional threaded map so each streaming pass is just a - per-column function plus a serial reduction of the yielded results. - """ - for band in _iter_bands(eval_preds, eval_targets, column_indices, band_size): - yield from _pmap(per_column, band, executor) - - -def _pearson(xp, xt): - """Pearson r, returning NaN for degenerate (constant / <2 point) inputs.""" - if xp.size < 2: - return np.nan - sp = xp.std() - st = xt.std() - if sp == 0 or st == 0 or not np.isfinite(sp) or not np.isfinite(st): - return np.nan - return pearsonr(xp, xt)[0] - - -def group_specificity_pearson( - eval_preds, - eval_targets, - column_indices, - *, - var_pct=1.0, - band_size=DEFAULT_BAND_SIZE, - ncpus=1, -): - """Cross-track specificity Pearson r per column, computed by streaming. - - For a group of columns, quantile-normalize every column to a common marginal, - subtract the per-position cross-track mean, optionally restrict to the most - variable positions, and correlate normalized predictions vs targets per - column. Memory is ``O(band_size * num_positions)`` regardless of the number - of columns; the data is read one band at a time (three passes). - - For a zarr input, reads are only cheap if the store is chunked along the - target axis (see ``seqnn.eval(target_chunk=...)``); otherwise every band read - decompresses the whole store. - - Args: - eval_preds: array ``(num_seqs, num_targets, bins)`` (numpy in RAM or a - zarr array on disk), float16/float32. - eval_targets: same shape/type as ``eval_preds``. - column_indices: 1D sequence of track indices (into the store's target - axis) forming the group; stranded pairs are collapsed upstream. - var_pct: proportion of highest-variance positions to keep (1.0 = all). - band_size: number of columns read per band (memory/I-O knob). - ncpus: threads for the per-column sort/apply (default 1); one executor is - shared across all bands and passes. - - Returns: - float64 array of length ``len(column_indices)`` with the per-column - Pearson r (NaN where degenerate), aligned with ``column_indices``. - """ - column_indices = list(column_indices) - num_cols = len(column_indices) - if num_cols == 0: - return np.zeros(0, dtype=np.float64) - - if band_size < 1: - raise ValueError("band_size must be >= 1") - if not (0 < var_pct <= 1): - raise ValueError("var_pct must be in (0, 1]") - - num_seqs, _, bins = eval_preds.shape - num_rows = num_seqs * bins - - executor = ThreadPoolExecutor(max_workers=ncpus) if (ncpus and ncpus > 1) else None - try: - - def stream(per_column): - return _stream_columns( - eval_preds, - eval_targets, - column_indices, - band_size, - per_column, - executor, - ) - - # Pass 1: quantile-normalization references (mean of sorted columns). - ref_p = np.zeros(num_rows, dtype=np.float64) - ref_t = np.zeros(num_rows, dtype=np.float64) - for sorted_p, sorted_t in stream( - lambda pair: (np.sort(pair[0]), np.sort(pair[1])) - ): - ref_p += sorted_p - ref_t += sorted_t - ref_p /= num_cols - ref_t /= num_cols - - # Pass 2: cross-track (per-position) means + target variance for the mask. - rowsum_p = np.zeros(num_rows, dtype=np.float64) - rowsum_t = np.zeros(num_rows, dtype=np.float64) - rowsumsq_t = np.zeros(num_rows, dtype=np.float64) - - def _normalize(pair): - cp, ct = pair - return apply_quantile_column(cp, ref_p), apply_quantile_column(ct, ref_t) - - for npv, ntv in stream(_normalize): - rowsum_p += npv - rowsum_t += ntv - rowsumsq_t += ntv * ntv - rowmean_p = rowsum_p / num_cols - rowmean_t = rowsum_t / num_cols - - mask = None - if var_pct < 1: - var_t = rowsumsq_t / num_cols - rowmean_t**2 - np.maximum(var_t, 0, out=var_t) - thresh = np.percentile(var_t, 100 * (1 - var_pct)) - mask = var_t >= thresh - - # Pass 3: per-column specificity Pearson r. - def _corr(pair): - cp, ct = pair - xp = apply_quantile_column(cp, ref_p) - rowmean_p - xt = apply_quantile_column(ct, ref_t) - rowmean_t - if mask is not None: - xp = xp[mask] - xt = xt[mask] - return _pearson(xp, xt) - - corr = np.fromiter(stream(_corr), dtype=np.float64, count=num_cols) - finally: - if executor is not None: - executor.shutdown() - return corr diff --git a/src/baskerville/trainer.py b/src/baskerville/trainer.py index 594a321..52b41a2 100644 --- a/src/baskerville/trainer.py +++ b/src/baskerville/trainer.py @@ -583,7 +583,7 @@ def _init_metrics(self): num_gene_targets, self.device, targets_df=getattr(self.train_data[di], "targets_df", None), - target_means=getattr(self.train_data[di], "target_means", None), + target_hist=getattr(self.train_data[di], "target_hist", None), spec_group_min=self.spec_group_min, ) ) diff --git a/src/tests/test_eval.py b/src/tests/test_eval.py index cae832e..9626c9e 100644 --- a/src/tests/test_eval.py +++ b/src/tests/test_eval.py @@ -14,7 +14,7 @@ import zarr from baskerville import metrics -from baskerville.dataset import SeqDataset +from baskerville.dataset import NUM_HIST_BINS, SeqDataset from baskerville.seqnn import SeqNN @@ -324,7 +324,10 @@ def test_eval_spec_target_subset( ): target_slice = np.array([2, 0]) data = SeqDataset(coverage_only_dataset, split_label="test") - data.target_means = np.array([1.0, 2.0, 4.0]) + y = np.random.default_rng(0).gamma(2, size=(3, 1000)).astype(np.float16) + data.target_hist = np.stack( + [np.bincount(v.view(np.uint16), minlength=NUM_HIST_BINS) for v in y] + ) if subset_metadata: data.targets_df = data.targets_df.loc[[0, 2]] @@ -337,7 +340,7 @@ def test_eval_spec_target_subset( cov = model.eval(data, batch_size=2, return_values=True)["coverage"] expected = spec_metric( data.targets_df.loc[target_slice], - data.target_means[target_slice], + data.target_hist[target_slice], group_min=2, ) expected.update( @@ -598,69 +601,3 @@ def test_eval_coverage_zarr_only(self, coverage_only_dataset): assert "targets" in root assert "gene_preds" not in root assert "gene_targets" not in root - - def test_eval_combine_pairs(self, coverage_only_dataset): - """combine_pairs collapses the coverage store to summed columns.""" - model_def = self._create_model_def(has_coverage=True, has_gene=False) - model = SeqNN(model_def) # one instance -> deterministic preds across runs - dataset = SeqDataset( - data_dir=coverage_only_dataset, split_label="test", mode="eval" - ) - - # full (no combine): individual tracks, shape num_targets - full = model.eval(dataset, batch_size=2, return_values=True) - fp = np.asarray(full["coverage"]["preds"], dtype=np.float32) - ft = np.asarray(full["coverage"]["targets"], dtype=np.float32) - assert fp.shape[1] == self.num_targets - assert full["coverage"]["r"] is not None - assert full["coverage"]["r2"] is not None - - # combine: sum pair (0,1), keep 2 -> 2 columns in spec order - columns = [(0, 1), (2,)] - comb = model.eval( - dataset, batch_size=2, return_values=True, combine_pairs=columns - ) - cp = np.asarray(comb["coverage"]["preds"], dtype=np.float32) - ct = np.asarray(comb["coverage"]["targets"], dtype=np.float32) - - assert cp.shape == (self.num_seqs, len(columns), self.target_length) - assert comb["coverage"]["r"] is None - assert comb["coverage"]["r2"] is None - # column 0 == track0 + track1, column 1 == track2 (fp16-store tolerance) - np.testing.assert_allclose(cp[:, 0], fp[:, 0] + fp[:, 1], atol=1e-2) - np.testing.assert_allclose(cp[:, 1], fp[:, 2], atol=1e-3) - np.testing.assert_allclose(ct[:, 0], ft[:, 0] + ft[:, 1], atol=1e-2) - np.testing.assert_allclose(ct[:, 1], ft[:, 2], atol=1e-3) - - -def test_strand_column_helpers(): - """_strand_column_index / _combine_strand_columns sum pairs along axis 1.""" - from baskerville.seqnn import ( - _combine_strand_columns, - _strand_column_index, - ) - - columns = [(0, 1), (2,), (3, 4)] - a_idx, pair_cols, pair_b = _strand_column_index(columns) - np.testing.assert_array_equal(a_idx, [0, 2, 3]) - np.testing.assert_array_equal(pair_cols, [0, 2]) - np.testing.assert_array_equal(pair_b, [1, 4]) - - arr = np.arange(2 * 5 * 3).reshape(2, 5, 3).astype(np.float32) - out = _combine_strand_columns(arr, a_idx, pair_cols, pair_b) - expected = np.stack( - [arr[:, 0] + arr[:, 1], arr[:, 2], arr[:, 3] + arr[:, 4]], axis=1 - ) - np.testing.assert_array_equal(out, expected) - # source array is not mutated (advanced indexing copies) - np.testing.assert_array_equal(arr, np.arange(2 * 5 * 3).reshape(2, 5, 3)) - - # no-pairs case is an identity gather - a2, pc2, pb2 = _strand_column_index([(0,), (1,)]) - out2 = _combine_strand_columns(arr[:, :2], a2, pc2, pb2) - np.testing.assert_array_equal(out2, arr[:, :2]) - - # malformed specs (empty / too many indices) raise - for bad in [(), (0, 1, 2)]: - with pytest.raises(ValueError): - _strand_column_index([bad]) diff --git a/src/tests/test_eval_folds_gcp.py b/src/tests/test_eval_folds_gcp.py index 881ebb8..8edffb4 100644 --- a/src/tests/test_eval_folds_gcp.py +++ b/src/tests/test_eval_folds_gcp.py @@ -29,11 +29,6 @@ def __init__(self, **kwargs): self.test_only = False self.valid_only = False self.targets_gene_file = None - # spec streaming knobs (hound_eval_spec pass-through) - self.band = None - self.seq_chunk = None - self.ram = False - self.scratch_dir = None # replication options self.crosses = 1 self.conda_env = "torch2.6" @@ -42,7 +37,6 @@ def __init__(self, **kwargs): self.processes = 1 self.queue = "l4" self.restart = False - self.spec = False self.params_file = "params.json" self.data_dirs = ["/data/hg38"] # runner backend (add_argparse_group) @@ -233,34 +227,19 @@ def test_default_evaluates_all_folds(tmp_path, monkeypatch): assert splits == ["0", "1", "2"] -def test_spec_gcp_uses_queue(tmp_path, monkeypatch, capsys): +def test_eval_gcp_uses_queue(tmp_path, monkeypatch): out_dir = _make_models_dir(tmp_path, num_folds=2, test_fold=0, valid_fold=1) captured = _patch_gcp(monkeypatch, num_folds=2) - args = MockArgs( - out_dir=out_dir, fold_subset=1, spec=True, test_only=True, queue="l4" + hef.eval_folds( + MockArgs(out_dir=out_dir, fold_subset=1, test_only=True, queue="l4") ) - hef.eval_folds(args) - spec_jobs = [j for j in captured if j.cmd.startswith("hound_eval_spec")] - assert spec_jobs - for j in spec_jobs: - # spec now streams (a few GB), so it honors --queue like coverage eval - assert j.kwargs["queue"] == "l4" - assert j.kwargs["cpu"] == 8 - assert " --ncpus 8" in j.cmd - # bumped boot disk: the streamed store is tight on the 100 GiB default - assert j.kwargs["boot_disk_gb"] == 200 - - # coverage eval jobs also honor the user's --queue eval_jobs = [j for j in captured if j.cmd.startswith("hound_eval ")] assert eval_jobs for j in eval_jobs: assert j.kwargs["queue"] == "l4" - # no auto-pin note is emitted anymore - assert "pinned to --queue" not in capsys.readouterr().out - def test_read_fold_splits_from_json(tmp_path): train_dir = tmp_path / "f0c1" / "train" @@ -400,11 +379,8 @@ def _gcp_output_dir(tmp_path, monkeypatch, **kwargs): def test_gcp_output_dir_keyed_on_options(tmp_path, monkeypatch): - """Output-changing options start fresh; memory/IO knobs still resume.""" + """Output-changing options start fresh; identical reruns resume.""" base = _gcp_output_dir(tmp_path, monkeypatch) assert _gcp_output_dir(tmp_path, monkeypatch) == base assert _gcp_output_dir(tmp_path, monkeypatch, rc=True) != base assert _gcp_output_dir(tmp_path, monkeypatch, shifts="0,1") != base - assert _gcp_output_dir(tmp_path, monkeypatch, band=8, ram=True) == base - # adding --spec to a finished eval run resumes its eval jobs - assert _gcp_output_dir(tmp_path, monkeypatch, spec=True) == base diff --git a/src/tests/test_spec_metrics.py b/src/tests/test_spec_metrics.py index 24812ec..8dbd3bd 100644 --- a/src/tests/test_spec_metrics.py +++ b/src/tests/test_spec_metrics.py @@ -28,7 +28,20 @@ def targets_df(): return pd.DataFrame(rows, columns=["identifier", "description", "strand_pair"]) -MEANS = np.linspace(0.5, 2.0, 9) +def _hist_reference(y, pair): + """Per-track fp16 bit-pattern counts of y_t + y_pair(t); y (N, T, L).""" + s = (y.astype(np.float32) + y[:, pair].astype(np.float32)).astype(np.float16) + bits = s.view(np.uint16).transpose(1, 0, 2).reshape(y.shape[1], -1) + return np.stack([np.bincount(b, minlength=dataset.NUM_HIST_BINS) for b in bits]) + + +PAIR = np.array([0, 2, 1, 3, 5, 4, 6, 7, 8]) + + +def _hist(seed=9): + """Histograms of a random dataset for the 9-track fixture.""" + y = np.random.default_rng(seed).gamma(2, size=(2, 9, 64)).astype("float16") + return _hist_reference(y, PAIR) def test_groups_and_collapse(targets_df): @@ -39,39 +52,97 @@ def test_groups_and_collapse(targets_df): assert list(pair) == [0, 2, 3, 5, 6, 7, 8] -def _spec_numpy(preds, targets, rep, pair, means): - """Direct reference: pair-sum, mean-normalize, center, per-track Pearson.""" - b = (means[rep] + means[pair])[None, :, None] +def _qnorm_reference(cols): + """Sort-based quantile normalization with tie blocks averaged; cols (T, n).""" + ref = np.sort(cols.astype(np.float64), axis=1).mean(axis=0) + out = np.empty(cols.shape) + for t, c in enumerate(cols): + order = np.argsort(c, kind="stable") + _, start, counts = np.unique(c[order], return_index=True, return_counts=True) + out[t, order] = np.repeat(np.add.reduceat(ref, start) / counts, counts) + return out + + +def test_qmap_tables_match_qnorm(): + rng = np.random.default_rng(3) + y = rng.gamma(0.3, 4, size=(5, 4, 50)).astype("float16") + y[y < 0.2] = 0 # a large tie block at zero + cols = y.transpose(1, 0, 2).reshape(4, -1) + tables = metrics.qmap_tables(_hist_reference(y, np.arange(4)) ) + # identity pair doubles values; doubling is exact in fp16 + cols2 = (2 * cols.astype(np.float32)).astype(np.float16) + mapped = np.take_along_axis(tables, cols2.view(np.uint16).astype(int), axis=1) + np.testing.assert_allclose(mapped, _qnorm_reference(cols2), rtol=1e-5) + # absent values interpolate between neighbors and clamp outside + assert np.all(np.diff(tables, axis=1) >= 0) + - def center(x): - x = (x[:, rep] + x[:, pair]) / b - return x - x.mean(axis=1, keepdims=True) +def _spec_numpy(preds, targets, rep, pair, tables): + """Direct reference: pair-sum, table-map, regress out the group mean, Pearson.""" - p, t = center(preds), center(targets) - p = p.transpose(0, 2, 1).reshape(-1, len(rep)) - t = t.transpose(0, 2, 1).reshape(-1, len(rep)) - return np.array([np.corrcoef(p[:, i], t[:, i])[0, 1] for i in range(len(rep))]) + def mapped(x): + s = (x[:, rep].astype(np.float32) + x[:, pair]).astype(np.float16) + bits = s.view(np.uint16).transpose(1, 0, 2).reshape(len(rep), -1) + return np.take_along_axis(tables, bits.astype(int), axis=1).astype(np.float64) + + def resid(z): + m = z.mean(axis=0) + return np.array([zi - np.polyval(np.polyfit(m, zi, 1), m) for zi in z]) + + rp, rt = resid(mapped(preds)), resid(mapped(targets)) + return np.array([np.corrcoef(a, b)[0, 1] for a, b in zip(rp, rt)]) def test_spec_matches_numpy(targets_df): rng = np.random.default_rng(0) preds = rng.gamma(2, size=(6, 9, 40)).astype("float32") - targets = (preds + rng.gamma(1, size=preds.shape)).astype("float32") + targets = (preds + rng.gamma(1, size=preds.shape)).astype("float16") + hist = _hist_reference(targets, PAIR) - spec = metrics.SpecPearsonCorrCoef(targets_df, MEANS, group_min=3) + spec = metrics.SpecPearsonCorrCoef(targets_df, hist, group_min=3) assert list(spec.groups) == ["RNA"] for bi in [slice(0, 4), slice(4, 6)]: spec.update(torch.tensor(preds[bi]), torch.tensor(targets[bi])) spec_t = spec.compute().numpy() rep, pair = spec.groups["RNA"] - expected = _spec_numpy(preds, targets, rep, pair, MEANS) + tables = metrics.qmap_tables(hist[rep]) + expected = _spec_numpy(preds, targets, rep, pair, tables) np.testing.assert_allclose(spec_t[rep], expected, rtol=1e-4, atol=1e-5) assert np.isnan(np.delete(spec_t, rep)).all() +def test_spec_gain_and_specific(): + """Depth gain on a shared signal scores ~0; track-specific signal scores.""" + num_tracks = 8 + targets_df = pd.DataFrame( + {"description": [f"CHIP:K4:t{i}" for i in range(num_tracks)]} + ) + rng = np.random.default_rng(4) + shape = (16, num_tracks, 512) + shared = rng.gamma(0.3, 3, size=(16, 1, 512)) + gain = rng.uniform(0.3, 3, size=(1, num_tracks, 1)) + noise = rng.gamma(1, 0.3, size=shape) + jitter = 1 + 0.01 * rng.standard_normal(shape) + specific = rng.gamma(0.5, 6, size=shape) * (rng.random(shape) < 0.05) + + def score(preds, targets): + targets = targets.astype("float16") + hist = _hist_reference(targets, np.arange(num_tracks)) + spec = metrics.SpecPearsonCorrCoef(targets_df, hist, group_min=2) + spec.update(torch.tensor(preds, dtype=torch.float32), torch.tensor(targets)) + return spec.compute().numpy().mean() + + # depth: tracks are scaled copies, so the tables remove gain exactly + assert abs(score(gain * shared * jitter, gain * (shared + noise))) < 0.05 + # unequal SNR: the tables warp noise-free preds nonlinearly per track, so + # some credit remains (qnorm, re-ranking preds, gives 0; fast spec 0.6) + assert score(gain * shared, gain * shared + noise) < 0.3 + assert score(gain * shared + specific, gain * shared + noise + specific) > 0.9 + + def test_dataset_metrics_groups(targets_df): - dm = metrics.DatasetMetrics(True, False, 9, None, "cpu", targets_df, MEANS, 2) + dm = metrics.DatasetMetrics(True, False, 9, None, "cpu", targets_df, _hist(), 2) rng = np.random.default_rng(1) y = torch.tensor(rng.gamma(2, size=(4, 9, 32)), dtype=torch.float32) yh = SimpleNamespace(has_coverage=True, has_gene=False, coverage=y * 0.9 + 0.1) @@ -83,7 +154,7 @@ def test_dataset_metrics_groups(targets_df): assert res["spec"] == pytest.approx((res["spec/RNA"] + res["spec/CHIP/K4"]) / 2) assert "valid_spec" in dm.format_log("valid", res) - # no means: groups still get r, but no spec + # no histograms: groups still get r, but no spec dm = metrics.DatasetMetrics(True, False, 9, None, "cpu", targets_df, None, 2) dm.update(yh, y, None, None, torch.tensor(1.0), 4) res = dm.compute() @@ -146,7 +217,7 @@ def test_spec_stop_validation( has_genes=False, num_targets=len(df), targets_df=df, - target_means=MEANS if means else None, + target_hist=_hist() if means else None, ) for df in [targets_df.assign(group="other"), targets_df] ] @@ -169,29 +240,42 @@ def test_spec_stop_validation( Trainer._init_metrics(trainer) -def test_write_target_means(tmp_path): +def test_write_target_hist(tmp_path): + # tracks 0/1 are a stranded pair, track 2 unstranded + pd.DataFrame( + {"identifier": ["a+", "a-", "b"], "strand_pair": [1, 0, 2]} + ).to_csv(tmp_path / "targets.txt", sep="\t") rng = np.random.default_rng(2) values = [] for fold, num_seqs in enumerate([5, 3]): - x = rng.gamma(2, size=(num_seqs, 9, 16)).astype("float16") + x = rng.gamma(0.5, 3, size=(num_seqs, 3, 16)).astype("float16") + x[x < 0.5] = 0 root = zarr.open_group(str(tmp_path / f"examples/fold{fold}.zarr"), mode="w") root.create_array("target", shape=x.shape, dtype="float16")[:] = x values.append(x) - means = dataset.write_target_means(str(tmp_path), processes=2, block_seqs=2) - expected = np.concatenate(values).astype("float64").mean(axis=(0, 2)) - np.testing.assert_allclose(means, expected, rtol=1e-12) + hist = dataset.write_target_hist(str(tmp_path), processes=2, block_seqs=2) + expected = _hist_reference(np.concatenate(values), np.array([1, 0, 2])) + np.testing.assert_array_equal(hist, expected) + assert hist.sum(axis=1).tolist() == [8 * 16] * 3 for fold in range(2): - target = zarr.open(str(tmp_path / f"examples/fold{fold}.zarr"), mode="r")[ - "target" - ] - np.testing.assert_allclose(target.attrs["mean"], expected, rtol=1e-12) + stored = zarr.open(str(tmp_path / f"examples/fold{fold}.zarr"), mode="r") + np.testing.assert_array_equal(stored["target_hist"][:], expected) + + +def test_write_target_hist_rejects_negative(tmp_path): + pd.DataFrame({"identifier": ["a"]}).to_csv(tmp_path / "targets.txt", sep="\t") + root = zarr.open_group(str(tmp_path / "examples/fold0.zarr"), mode="w") + root.create_array("target", shape=(1, 1, 4), dtype="float16")[:] = -1 + with pytest.raises(ValueError, match="negative or non-finite"): + dataset.write_target_hist(str(tmp_path), processes=1) -def test_hound_data_writes_means(data_me_dir): +def test_hound_data_writes_hist(data_me_dir): + targets_df = pd.read_csv(f"{data_me_dir}/targets.txt", sep="\t", index_col=0) + pair = dataset.strand_pair_indices(targets_df) zarr_files = sorted(glob.glob(f"{data_me_dir}/examples/*.zarr")) - targets = [zarr.open(zf, mode="r")["target"] for zf in zarr_files] - expected = np.concatenate([t[:].astype("float64") for t in targets]) - expected = expected.mean(axis=(0, 2)) - for t in targets: - np.testing.assert_allclose(t.attrs["mean"], expected, rtol=1e-12) + roots = [zarr.open(zf, mode="r") for zf in zarr_files] + expected = _hist_reference(np.concatenate([r["target"][:] for r in roots]), pair) + for r in roots: + np.testing.assert_array_equal(r["target_hist"][:], expected) diff --git a/src/tests/test_spec_norm.py b/src/tests/test_spec_norm.py deleted file mode 100644 index f5d3848..0000000 --- a/src/tests/test_spec_norm.py +++ /dev/null @@ -1,221 +0,0 @@ -"""Tests for streaming cross-track specificity normalization (spec_norm). - -These are CPU-only: they compare the streaming implementation against the -original ``qnorm``-based reshape path on synthetic arrays. Strand pairs are -collapsed to one column upstream (``seqnn.eval(combine_pairs=...)``), so -spec_norm here only ever handles individual columns. -""" - -import numpy as np -import pandas as pd -import pytest -from qnorm import quantile_normalize -from scipy.stats import pearsonr -import zarr - -from baskerville import spec_norm as sn -from baskerville.dataset import strand_pair_indices - - -# --------------------------------------------------------------------------- # -# reference implementation (mirrors the original qnorm-based hound_eval_spec) -# --------------------------------------------------------------------------- # -def _reference_spec(eval_preds, eval_targets, column_indices, var_pct=1.0): - """Original path: build the (rows, cols) matrix, qnorm, mean-subtract, corr. - - Each column is one track's flattened vector, matching the streaming column - definition up to a row permutation (which the metric is invariant to). - """ - cols_p = [ - eval_preds[:, c, :].astype(np.float32).reshape(-1) for c in column_indices - ] - cols_t = [ - eval_targets[:, c, :].astype(np.float32).reshape(-1) for c in column_indices - ] - ep = np.stack(cols_p, axis=1) - et = np.stack(cols_t, axis=1) - - epn = quantile_normalize(ep, axis=1) - etn = quantile_normalize(et, axis=1) - - if var_pct < 1: - v = etn.var(axis=1) - thr = np.percentile(v, 100 * (1 - var_pct)) - mask = v >= thr - epn, etn = epn[mask], etn[mask] - - epn -= epn.mean(axis=-1, keepdims=True) - etn -= etn.mean(axis=-1, keepdims=True) - return np.array( - [pearsonr(epn[:, j], etn[:, j])[0] for j in range(len(column_indices))] - ) - - -def _synthetic(num_seqs, num_targets, bins, seed=0, ties=False): - rng = np.random.default_rng(seed) - preds = rng.normal(1.0, 1.0, (num_seqs, num_targets, bins)) - targets = preds + rng.normal(0.0, 0.5, (num_seqs, num_targets, bins)) - if ties: - # inject a large block of zeros / repeated values (like coverage data) - preds = np.maximum(preds - 1.0, 0.0) - targets = np.maximum(targets - 1.0, 0.0) - preds = np.round(preds * 2) / 2 - targets = np.round(targets * 2) / 2 - return preds.astype(np.float16), targets.astype(np.float16) - - -# --------------------------------------------------------------------------- # -# apply_quantile_column -# --------------------------------------------------------------------------- # -def test_apply_quantile_column_matches_qnorm_ties(): - rng = np.random.default_rng(3) - rows, cols = 400, 5 - # ties-heavy integer matrix + a few jittered values - M = rng.integers(0, 3, size=(rows, cols)).astype(np.float32) - jitter = (rng.random((rows, cols)) > 0.8) * rng.normal(0, 0.01, (rows, cols)) - M += jitter.astype(np.float32) - - ref_qn = quantile_normalize(M.copy(), axis=1) - reference = np.sort(M, axis=0).mean(axis=1) - mine = np.stack( - [sn.apply_quantile_column(M[:, j], reference) for j in range(cols)], axis=1 - ) - np.testing.assert_allclose(mine, ref_qn, atol=1e-4) - - -# --------------------------------------------------------------------------- # -# group_specificity_pearson vs the original path -# --------------------------------------------------------------------------- # -def test_group_specificity_unstranded(): - preds, targets = _synthetic(8, 6, 40, seed=10) - cols = list(range(6)) - mine = sn.group_specificity_pearson(preds, targets, cols, band_size=4) - ref = _reference_spec(preds, targets, cols) - np.testing.assert_allclose(mine, ref, atol=1e-4) - - -def test_group_specificity_ties_heavy(): - preds, targets = _synthetic(10, 8, 50, seed=11, ties=True) - cols = list(range(8)) - mine = sn.group_specificity_pearson(preds, targets, cols, band_size=3) - ref = _reference_spec(preds, targets, cols) - np.testing.assert_allclose(mine, ref, atol=1e-4) - - -def test_group_specificity_var_pct(): - preds, targets = _synthetic(8, 6, 40, seed=12) - cols = list(range(6)) - mine = sn.group_specificity_pearson(preds, targets, cols, band_size=2, var_pct=0.5) - ref = _reference_spec(preds, targets, cols, var_pct=0.5) - np.testing.assert_allclose(mine, ref, atol=1e-4) - - -def test_group_specificity_subset_columns(): - # a group need not be every track: pass an out-of-order index subset - preds, targets = _synthetic(8, 8, 40, seed=19) - cols = [5, 1, 6, 2] - mine = sn.group_specificity_pearson(preds, targets, cols, band_size=2) - ref = _reference_spec(preds, targets, cols) - np.testing.assert_allclose(mine, ref, atol=1e-4) - - -def test_group_specificity_zarr_matches_numpy(): - preds, targets = _synthetic(8, 6, 40, seed=13) - # target-axis chunked store (chunk=3 tracks), as seqnn.eval(target_chunk=...) - # would write it — the band reader touches only its own target-chunks - zp = zarr.create_array( - store=zarr.storage.MemoryStore(), - shape=preds.shape, - chunks=(2, 3, 40), - dtype="float16", - ) - zt = zarr.create_array( - store=zarr.storage.MemoryStore(), - shape=targets.shape, - chunks=(2, 3, 40), - dtype="float16", - ) - zp[:] = preds - zt[:] = targets - cols = list(range(6)) - r_np = sn.group_specificity_pearson(preds, targets, cols, band_size=3) - r_zr = sn.group_specificity_pearson(zp, zt, cols, band_size=3) - np.testing.assert_allclose(r_np, r_zr, atol=1e-6) - - -def test_group_specificity_ncpus_matches_serial(): - preds, targets = _synthetic(8, 6, 40, seed=14) - cols = list(range(6)) - r1 = sn.group_specificity_pearson(preds, targets, cols, band_size=2, ncpus=1) - r4 = sn.group_specificity_pearson(preds, targets, cols, band_size=2, ncpus=4) - np.testing.assert_allclose(r1, r4, atol=1e-6) - - -def test_constant_track_matches_reference(): - # A constant raw track becomes constant after quantile norm, but the - # per-position cross-track mean subtraction still makes it vary, so it is not - # degenerate. The key property is that streaming matches the qnorm path. - preds, targets = _synthetic(6, 4, 20, seed=16) - targets = targets.astype(np.float32) - targets[:, 0, :] = 5.0 - targets = targets.astype(np.float16) - cols = list(range(4)) - mine = sn.group_specificity_pearson(preds, targets, cols, band_size=4) - ref = _reference_spec(preds, targets, cols) - np.testing.assert_allclose(mine, ref, atol=1e-4) - - -def test_non_finite_input_raises(): - preds, targets = _synthetic(6, 4, 20, seed=18) - bad = preds.astype(np.float32) - bad[0, 0, 0] = np.nan - with pytest.raises(ValueError): - sn.group_specificity_pearson(bad.astype(np.float16), targets, list(range(4))) - - -def test_single_column_group(): - preds, targets = _synthetic(6, 1, 20, seed=17) - # one column: after cross-track mean subtraction everything is 0 -> NaN - mine = sn.group_specificity_pearson(preds, targets, [0], band_size=1) - assert mine.shape == (1,) - assert np.isnan(mine[0]) - - -def test_collapsed_group_loop(): - """Mirror hound_eval_spec's per-group loop over a pre-summed (collapsed) store.""" - num_seqs, bins = 8, 30 - # one column per experiment: RNA (4), DNASE (5), SMALL (2, below group_min). - group = ["RNA"] * 4 + ["DNASE"] * 5 + ["SMALL"] * 2 - num_cols = len(group) - preds, targets = _synthetic(num_seqs, num_cols, bins, seed=42) - group_values = np.array(group) - - group_min = 4 - targets_spec = np.full(num_cols, np.nan) - for tg in sorted(set(group)): - col_idx = np.where(group_values == tg)[0] - if len(col_idx) < group_min: - continue - r = sn.group_specificity_pearson(preds, targets, col_idx, band_size=2) - ref = _reference_spec(preds, targets, col_idx) - np.testing.assert_allclose(r, ref, atol=1e-4) - targets_spec[col_idx] = r - - assert np.all(np.isnan(targets_spec[9:11])) # SMALL skipped - assert np.all(np.isfinite(targets_spec[:9])) - - -# --------------------------------------------------------------------------- # -# strand_pair_indices helper (used by hound_eval_spec to build combine_pairs) -# --------------------------------------------------------------------------- # -def test_strand_pair_indices_identity_without_column(): - df = pd.DataFrame({"identifier": ["x", "y"]}, index=[5, 6]) - np.testing.assert_array_equal(strand_pair_indices(df), [0, 1]) - - -def test_strand_pair_indices_remaps_labels(): - df = pd.DataFrame( - {"strand_pair": [2, 3, 0, 1]}, # labels happen to be 0..N-1 here - index=[0, 1, 2, 3], - ) - np.testing.assert_array_equal(strand_pair_indices(df), [2, 3, 0, 1]) From a90ad41df47c26551bb1700550bb6f61e29505af Mon Sep 17 00:00:00 2001 From: David Kelley Date: Thu, 8 Oct 2026 11:09:20 -0700 Subject: [PATCH 2/6] Fix fp16 boundary handling in specificity metrics --- src/baskerville/dataset.py | 12 +++++----- src/baskerville/metrics.py | 2 +- src/tests/test_spec_metrics.py | 40 ++++++++++++++++++++++++++++++---- 3 files changed, 44 insertions(+), 10 deletions(-) diff --git a/src/baskerville/dataset.py b/src/baskerville/dataset.py index 0b68997..798b37c 100644 --- a/src/baskerville/dataset.py +++ b/src/baskerville/dataset.py @@ -878,10 +878,11 @@ def _target_block_hist(args): hist = np.zeros(num_targets * NUM_HIST_BINS, dtype=np.int64) for si in range(start, end): y = targets[si].astype(np.float32) - # + 0.0 turns -0.0 into 0.0 - bits = (y + y[pair] + 0.0).astype(np.float16).view(np.uint16) - if bits.max() >= NUM_HIST_BINS: + if not np.isfinite(y).all() or (y < 0).any(): raise ValueError(f"{zarr_file} seq {si}: negative or non-finite targets") + # + 0.0 turns -0.0 into 0.0 + summed = np.minimum(y + y[pair], 65504) + 0.0 + bits = summed.astype(np.float16).view(np.uint16) hist += np.bincount((bits + offsets).ravel(), minlength=hist.size) return hist.reshape(num_targets, NUM_HIST_BINS) @@ -891,8 +892,9 @@ def write_target_hist(data_dir, processes=16, block_seqs=256): Counts are over every position of every sequence in all examples/*.zarr, of the fp16 value y_t + y_pair(t) (strand_pair from targets.txt; unstranded - tracks are their own pair, so 2 * y_t), binned by fp16 bit pattern. They are - written to each zarr as ``target_hist`` (num_targets, NUM_HIST_BINS) int64, + tracks are their own pair, so 2 * y_t), clamped to 65504 and binned by fp16 + bit pattern. Counts are written to each zarr as ``target_hist`` + (num_targets, NUM_HIST_BINS) int64, from which SpecPearsonCorrCoef builds its quantile maps. Args: diff --git a/src/baskerville/metrics.py b/src/baskerville/metrics.py index 0cec0cd..95e1997 100644 --- a/src/baskerville/metrics.py +++ b/src/baskerville/metrics.py @@ -203,7 +203,7 @@ def _map(self, x, g): """(N, T, L) -> (tracks, N*L) table values of the group's strand sums.""" rep, pair = self.index[g] s = x[:, rep].float() + x[:, pair].float() - bits = s.clamp(0, 65504).half().view(torch.int16).long() + bits = (s.clamp(0, 65504) + 0.0).half().view(torch.int16).long() return torch.gather(self.table[g], 1, bits.transpose(0, 1).flatten(1)) @torch.no_grad() diff --git a/src/tests/test_spec_metrics.py b/src/tests/test_spec_metrics.py index 8dbd3bd..0a0dd60 100644 --- a/src/tests/test_spec_metrics.py +++ b/src/tests/test_spec_metrics.py @@ -68,7 +68,7 @@ def test_qmap_tables_match_qnorm(): y = rng.gamma(0.3, 4, size=(5, 4, 50)).astype("float16") y[y < 0.2] = 0 # a large tie block at zero cols = y.transpose(1, 0, 2).reshape(4, -1) - tables = metrics.qmap_tables(_hist_reference(y, np.arange(4)) ) + tables = metrics.qmap_tables(_hist_reference(y, np.arange(4))) # identity pair doubles values; doubling is exact in fp16 cols2 = (2 * cols.astype(np.float32)).astype(np.float16) mapped = np.take_along_axis(tables, cols2.view(np.uint16).astype(int), axis=1) @@ -240,11 +240,43 @@ def test_spec_stop_validation( Trainer._init_metrics(trainer) +def test_target_hist_and_map_fp16_boundaries(monkeypatch): + y = np.tile(np.array([-0.0, 1, 40000, 65504], dtype=np.float16), (1, 3, 1)) + monkeypatch.setattr(dataset.zarr, "open", lambda *a, **kw: {"target": y}) + hist = dataset._target_block_hist(("unused", 0, 1, np.array([1, 0, 2]))) + expected = np.zeros_like(hist) + expected[:, 0] = 1 + expected[:, 0x4000] = 1 # 2.0 + expected[:, 0x7BFF] = 2 # 65504.0 + np.testing.assert_array_equal(hist, expected) + + df = pd.DataFrame( + { + "identifier": ["a+", "a-", "b"], + "description": ["RNA:a", "RNA:a", "RNA:b"], + "strand_pair": [1, 0, 2], + } + ) + spec = metrics.SpecPearsonCorrCoef(df, hist, group_min=2) + mapped = spec._map(torch.from_numpy(y), "RNA") + torch.testing.assert_close( + mapped, torch.tensor([[0.0, 2, 65504, 65504]]).expand(2, -1) + ) + + +@pytest.mark.parametrize("value", [-1, np.inf, -np.inf, np.nan]) +def test_target_hist_rejects_invalid_before_clamping(monkeypatch, value): + y = np.array([[[value]]], dtype=np.float16) + monkeypatch.setattr(dataset.zarr, "open", lambda *a, **kw: {"target": y}) + with pytest.raises(ValueError, match="negative or non-finite"): + dataset._target_block_hist(("unused", 0, 1, np.array([0]))) + + def test_write_target_hist(tmp_path): # tracks 0/1 are a stranded pair, track 2 unstranded - pd.DataFrame( - {"identifier": ["a+", "a-", "b"], "strand_pair": [1, 0, 2]} - ).to_csv(tmp_path / "targets.txt", sep="\t") + pd.DataFrame({"identifier": ["a+", "a-", "b"], "strand_pair": [1, 0, 2]}).to_csv( + tmp_path / "targets.txt", sep="\t" + ) rng = np.random.default_rng(2) values = [] for fold, num_seqs in enumerate([5, 3]): From 7588bf3cd7696d8816150a5f809a7e0c85ea6b60 Mon Sep 17 00:00:00 2001 From: David Kelley Date: Thu, 8 Oct 2026 11:48:53 -0700 Subject: [PATCH 3/6] Spec metric: NaN-safe lookup, O(TK log TK) tables, lean moments Map NaN predictions as 0 instead of indexing outside the tables; build the group reference once by sweeping all tracks' tie-block bounds; accumulate group-mean moments per group and only cross terms per track; share tables between train and valid metrics; count histograms per track in int32; reject strand-pair splits when subsetting target_hist. Co-Authored-By: Claude Opus 5.5 --- docs/train.md | 2 +- src/baskerville/dataset.py | 12 +-- src/baskerville/metrics.py | 100 +++++++++++++----- .../utils/hound_data_subset_zarr_targets.py | 8 +- src/baskerville/trainer.py | 35 +++--- src/tests/test_spec_metrics.py | 17 ++- 6 files changed, 124 insertions(+), 50 deletions(-) diff --git a/docs/train.md b/docs/train.md index 10f8109..69eeba6 100644 --- a/docs/train.md +++ b/docs/train.md @@ -128,7 +128,7 @@ Tracks are grouped by the targets `group` column (else the `description` prefix) - `r/`, `r2/`: mean per-track Pearson r and R² over the group's tracks. - `spec/`: specificity, the mean per-track Pearson r of what remains after removing the group's shared signal. Each track's values (strand pairs summed) are quantile-normalized to the group with a fixed per-track lookup table; predictions use the same table as their targets. The group mean at each position is then regressed out of each track, with the track's own slope and intercept. It rewards predicting how a track differs from its group, not the shared signal or the track's depth and signal-to-noise. -`spec` is the unweighted mean of `spec/` over groups. `hound_eval` reports the same metric per track (`spec` column of `acc.txt`), in its single pass. The tables are built from the exact whole-genome histograms of each track's stored fp16 values, which `hound_data` writes to each `examples/*.zarr` as `target_hist`. For an older dataset, add them with `python -c "from baskerville import dataset; dataset.write_target_hist('data_dir')"`. +`spec` is the unweighted mean of `spec/` over groups. `hound_eval` reports the same metric per track (`spec` column of `acc.txt`). The tables are built from the exact whole-genome histograms of each track's stored fp16 values, which `hound_data` writes to each `examples/*.zarr` as `target_hist`. For an older dataset, add them with `python -c "from baskerville import dataset; dataset.write_target_hist('data_dir')"`. `train.stop_stat` chooses the statistic that selects `model_best.pth` and drives `patience`. It is a `{metric: weight}` dict over the keys above plus `loss`, `r_gene`, and `r2_gene`, and the weighted sum is maximized. The presets `"weighted_r_r2"` (default, `{"r": 1, "r2": 0.25}`), `"r"`, `"r2"`, and `"loss"` (`{"loss": -1}`) remain available as strings; these presets score gene-only datasets by −loss. A dict applies only its own keys, e.g. `{"r_gene": 1}` or `{"loss": -1}` for gene-only datasets. diff --git a/src/baskerville/dataset.py b/src/baskerville/dataset.py index 798b37c..316e97d 100644 --- a/src/baskerville/dataset.py +++ b/src/baskerville/dataset.py @@ -873,9 +873,8 @@ def _target_block_hist(args): sequences [start, end).""" zarr_file, start, end, pair = args targets = zarr.open(zarr_file, mode="r")["target"] - num_targets = targets.shape[1] - offsets = np.arange(num_targets)[:, None] * NUM_HIST_BINS - hist = np.zeros(num_targets * NUM_HIST_BINS, dtype=np.int64) + # int32: a block's counts per bin are at most (end - start) * length + hist = np.zeros((targets.shape[1], NUM_HIST_BINS), dtype=np.int32) for si in range(start, end): y = targets[si].astype(np.float32) if not np.isfinite(y).all() or (y < 0).any(): @@ -883,8 +882,9 @@ def _target_block_hist(args): # + 0.0 turns -0.0 into 0.0 summed = np.minimum(y + y[pair], 65504) + 0.0 bits = summed.astype(np.float16).view(np.uint16) - hist += np.bincount((bits + offsets).ravel(), minlength=hist.size) - return hist.reshape(num_targets, NUM_HIST_BINS) + for ti, track_bits in enumerate(bits): + hist[ti] += np.bincount(track_bits, minlength=NUM_HIST_BINS) + return hist def write_target_hist(data_dir, processes=16, block_seqs=256): @@ -917,7 +917,7 @@ def write_target_hist(data_dir, processes=16, block_seqs=256): hist = 0 with multiprocessing.get_context("spawn").Pool(processes) as pool: for block_hist in pool.imap_unordered(_target_block_hist, tasks): - hist = hist + block_hist + hist = hist + block_hist.astype(np.int64) for zarr_file in zarr_files: root = zarr.open_group(zarr_file, mode="r+") diff --git a/src/baskerville/metrics.py b/src/baskerville/metrics.py index 95e1997..3ff888f 100644 --- a/src/baskerville/metrics.py +++ b/src/baskerville/metrics.py @@ -1,3 +1,5 @@ +import copy + import numpy as np import torch import torch.nn as nn @@ -128,20 +130,28 @@ def qmap_tables(hists): values = np.arange(dataset.NUM_HIST_BINS, dtype=np.uint16).view(np.float16) values = values.astype(np.float64) - # per track: present values, rank bounds of their tie blocks, and the - # running sum of the sorted column at those bounds (linear within blocks) + # per track: present values and the rank bounds of their tie blocks knots = [] for h in hists: nz = np.flatnonzero(h) - bounds = np.concatenate([[0], np.cumsum(h[nz])]) - sums = np.concatenate([[0.0], np.cumsum(h[nz] * values[nz])]) - knots.append((nz, bounds, sums)) + knots.append((nz, np.concatenate([[0], np.cumsum(h[nz])]))) + + # the reference (mean sorted column) is a step function of rank: it starts + # at the mean minimum and steps by (value change) / T at every track's + # block bound; sweep all bounds in order for its running sum + num_tracks, num_ranks = len(hists), totals[0] + pos = np.concatenate([b[1:-1] for _, b in knots]) + step = np.concatenate([np.diff(values[nz]) for nz, _ in knots]) + order = np.argsort(pos, kind="stable") + ref_pos = np.concatenate([[0], pos[order], [num_ranks]]) + ref_level = sum(values[nz[0]] for nz, _ in knots) + ref_level = (ref_level + np.concatenate([[0.0], np.cumsum(step[order])])) / num_tracks + ref_sum = np.concatenate([[0.0], np.cumsum(ref_level * np.diff(ref_pos))]) tables = np.empty(hists.shape, dtype=np.float32) - for t, (nz, bounds, _) in enumerate(knots): - # running sum of the reference = mean of the tracks' running sums - ref_sum = np.mean([np.interp(bounds, b, s) for _, b, s in knots], axis=0) - tables[t] = np.interp(values, values[nz], np.diff(ref_sum) / np.diff(bounds)) + for t, (nz, bounds) in enumerate(knots): + block_sum = np.diff(np.interp(bounds, ref_pos, ref_sum)) + tables[t] = np.interp(values, values[nz], block_sum / np.diff(bounds)) return tables @@ -154,8 +164,8 @@ class SpecPearsonCorrCoef: whole-genome target histograms, SeqDataset.target_hist), regress the group mean at each position out of each track (its own slope and intercept), and correlate each track's pred and target residuals. Streaming state is the - per-track first and second moments of (pred, target, group-mean pred, - group-mean target). + first and second moments of (pred, target, group-mean pred, group-mean + target): per track for terms involving the track, per group otherwise. Args: targets_df: targets table aligned with the output tracks. @@ -190,19 +200,27 @@ def __init__(self, targets_df, target_hist, group_min: int = 20): self.table[g] = torch.from_numpy(qmap_tables(hist[rep])) self.reset() + def fresh(self): + """A new metric with zeroed state, sharing this one's groups and tables.""" + new = copy.copy(self) + if self.enabled: + new.index, new.table = dict(self.index), dict(self.table) + new.reset() + return new + def to(self, device): if self.enabled: for g in self.groups: self.index[g] = tuple(ix.to(device) for ix in self.index[g]) self.table[g] = self.table[g].to(device) - self.sum1[g] = self.sum1[g].to(device) - self.sum2[g] = self.sum2[g].to(device) + self.track_sums[g] = self.track_sums[g].to(device) + self.mean_sums[g] = self.mean_sums[g].to(device) return self def _map(self, x, g): """(N, T, L) -> (tracks, N*L) table values of the group's strand sums.""" rep, pair = self.index[g] - s = x[:, rep].float() + x[:, pair].float() + s = torch.nan_to_num(x[:, rep].float() + x[:, pair].float(), nan=0.0) bits = (s.clamp(0, 65504) + 0.0).half().view(torch.int16).long() return torch.gather(self.table[g], 1, bits.transpose(0, 1).flatten(1)) @@ -211,12 +229,22 @@ def update(self, preds: torch.Tensor, target: torch.Tensor) -> None: """preds, target: (N, T, L).""" if self.enabled: for g in self.groups: - zp, zy = self._map(preds, g), self._map(target, g) - mp = zp.mean(dim=0, keepdim=True).expand_as(zp) - my = zy.mean(dim=0, keepdim=True).expand_as(zy) - z = torch.stack([zp, zy, mp, my], dim=1).double() # (tracks, 4, N*L) - self.sum1[g] += z.sum(dim=-1) - self.sum2[g] += torch.einsum("tap,tbp->tab", z, z) + zp, zy = self._map(preds, g), self._map(target, g) # (tracks, P) + mp, my = zp.mean(dim=0), zy.mean(dim=0) # (P,) + + def total(x): + return x.sum(dim=-1, dtype=torch.float64) + + self.track_sums[g] += torch.stack( + [total(zp), total(zy), total(zp * zp), total(zy * zy)] + + [total(zp * zy), total(zp * mp), total(zp * my)] + + [total(zy * mp), total(zy * my)], + dim=-1, + ) + self.mean_sums[g] += torch.stack( + [total(mp), total(my), total(mp * mp), total(my * my)] + + [total(mp * my)] + ) self.count[g] += zp.shape[1] def compute(self) -> torch.Tensor: @@ -224,7 +252,20 @@ def compute(self) -> torch.Tensor: if self.enabled: for g in self.groups: n = self.count[g] - s1, s2 = self.sum1[g].cpu(), self.sum2[g].cpu() + ts, ms = self.track_sums[g].cpu(), self.mean_sums[g].cpu() + # assemble per-track sums over (pred, target, mean pred, mean target) + s1 = torch.stack( + [ts[:, 0], ts[:, 1], ms[0].expand(len(ts)), ms[1].expand(len(ts))], + dim=-1, + ) + s2 = torch.empty(len(ts), 4, 4, dtype=torch.float64) + pairs = { + (0, 0): ts[:, 2], (1, 1): ts[:, 3], (0, 1): ts[:, 4], + (0, 2): ts[:, 5], (0, 3): ts[:, 6], (1, 2): ts[:, 7], + (1, 3): ts[:, 8], (2, 2): ms[2], (3, 3): ms[3], (2, 3): ms[4], + } + for (a, b), v in pairs.items(): + s2[:, a, b] = s2[:, b, a] = v cov = s2 / n - s1[:, :, None] * s1[:, None, :] / n**2 # residuals of pred/target after regressing out their group mean u = torch.zeros(len(cov), 4, dtype=torch.float64) @@ -240,13 +281,13 @@ def compute(self) -> torch.Tensor: def reset(self): if self.enabled: device = {g: t.device for g, t in self.table.items()} - self.sum1 = { - g: torch.zeros(len(r), 4, dtype=torch.float64, device=device[g]) + self.track_sums = { + g: torch.zeros(len(r), 9, dtype=torch.float64, device=device[g]) for g, (r, _) in self.groups.items() } - self.sum2 = { - g: torch.zeros(len(r), 4, 4, dtype=torch.float64, device=device[g]) - for g, (r, _) in self.groups.items() + self.mean_sums = { + g: torch.zeros(5, dtype=torch.float64, device=device[g]) + for g in self.groups } self.count = {g: 0 for g in self.groups} @@ -451,6 +492,8 @@ class DatasetMetrics: matches num_targets. target_hist: Per-track histograms of the stored targets; enables spec. spec_group_min: Minimum tracks for a group to be scored. + spec_like: Optional SpecPearsonCorrCoef whose groups and tables to share + (replaces building them from targets_df/target_hist). """ def __init__( @@ -463,6 +506,7 @@ def __init__( targets_df=None, target_hist=None, spec_group_min: int = 20, + spec_like=None, ): self.has_coverage = has_coverage self.has_gene = has_gene @@ -475,7 +519,9 @@ def __init__( if has_coverage and num_targets is not None: self.r = PearsonCorrCoef(num_targets, average=False).to(device) self.r2 = R2Score(average=False).to(device) - if targets_df is not None and len(targets_df) == num_targets: + if spec_like is not None: + self.spec = spec_like.fresh() + elif targets_df is not None and len(targets_df) == num_targets: self.spec = SpecPearsonCorrCoef( targets_df, target_hist, spec_group_min ).to(device) diff --git a/src/baskerville/scripts/utils/hound_data_subset_zarr_targets.py b/src/baskerville/scripts/utils/hound_data_subset_zarr_targets.py index 45c3a2d..35f2350 100755 --- a/src/baskerville/scripts/utils/hound_data_subset_zarr_targets.py +++ b/src/baskerville/scripts/utils/hound_data_subset_zarr_targets.py @@ -107,10 +107,14 @@ def subset_zarr_for_fold( compressors=compressors, ) - # ** keep the kept tracks' histograms (strand sums: keep pairs together) ** + # ** keep the kept tracks' histograms (of strand sums, so pairs must stay) ** old_root = zarr.open_group(f"{og_data_dir}/examples/fold{fi}.zarr", mode="r") if "target_hist" in old_root: - hist = old_root["target_hist"][:][old_indices] + old_targets = pd.read_csv(f"{og_data_dir}/targets.txt", sep="\t", index_col=0) + pair = dataset.strand_pair_indices(old_targets) + if not np.isin(pair[old_indices], old_indices).all(): + raise ValueError("old_indices split a strand pair; target_hist would be wrong") + hist = old_root["target_hist"].oindex[old_indices, :] fold_zarr_root.create_array( "target_hist", shape=hist.shape, diff --git a/src/baskerville/trainer.py b/src/baskerville/trainer.py index 52b41a2..2280745 100644 --- a/src/baskerville/trainer.py +++ b/src/baskerville/trainer.py @@ -574,19 +574,30 @@ def _init_metrics(self): self.train_data[di].num_gene_targets if has_gene else None ) - for metrics_list in [self.train_metrics, self.valid_metrics]: - metrics_list.append( - metrics.DatasetMetrics( - has_coverage, - has_gene, - num_targets, - num_gene_targets, - self.device, - targets_df=getattr(self.train_data[di], "targets_df", None), - target_hist=getattr(self.train_data[di], "target_hist", None), - spec_group_min=self.spec_group_min, - ) + train_metrics = metrics.DatasetMetrics( + has_coverage, + has_gene, + num_targets, + num_gene_targets, + self.device, + targets_df=getattr(self.train_data[di], "targets_df", None), + target_hist=getattr(self.train_data[di], "target_hist", None), + spec_group_min=self.spec_group_min, + ) + self.train_metrics.append(train_metrics) + # valid shares train's spec tables (same targets) + self.valid_metrics.append( + metrics.DatasetMetrics( + has_coverage, + has_gene, + num_targets, + num_gene_targets, + self.device, + targets_df=getattr(self.train_data[di], "targets_df", None), + spec_group_min=self.spec_group_min, + spec_like=train_metrics.spec, ) + ) # Validate specificity keys against the datasets used for stopping. stop_metrics = ( diff --git a/src/tests/test_spec_metrics.py b/src/tests/test_spec_metrics.py index 0a0dd60..9a1f5a9 100644 --- a/src/tests/test_spec_metrics.py +++ b/src/tests/test_spec_metrics.py @@ -135,12 +135,25 @@ def score(preds, targets): # depth: tracks are scaled copies, so the tables remove gain exactly assert abs(score(gain * shared * jitter, gain * (shared + noise))) < 0.05 - # unequal SNR: the tables warp noise-free preds nonlinearly per track, so - # some credit remains (qnorm, re-ranking preds, gives 0; fast spec 0.6) + # unequal SNR: the target tables warp noise-free preds nonlinearly per + # track, so some credit remains assert score(gain * shared, gain * shared + noise) < 0.3 assert score(gain * shared + specific, gain * shared + noise + specific) > 0.9 +def test_spec_nan_preds_and_fresh(targets_df): + spec = metrics.SpecPearsonCorrCoef(targets_df, _hist(), group_min=2) + y = torch.rand(2, 9, 16) + preds = y.clone() + preds[0, 0, 0] = float("nan") + spec.update(preds, y) # NaN maps like 0 instead of indexing out of bounds + assert np.isfinite(spec.compute().numpy()[spec.groups["RNA"][0]]).all() + + other = spec.fresh() + assert other.table["RNA"] is spec.table["RNA"] + assert other.count["RNA"] == 0 and spec.count["RNA"] == 32 + + def test_dataset_metrics_groups(targets_df): dm = metrics.DatasetMetrics(True, False, 9, None, "cpu", targets_df, _hist(), 2) rng = np.random.default_rng(1) From 6274757e65378d2b6319f6d5b2d73d011be24859 Mon Sep 17 00:00:00 2001 From: David Kelley Date: Thu, 8 Oct 2026 11:58:36 -0700 Subject: [PATCH 4/6] Fix specificity metric precision and constant group means --- src/baskerville/metrics.py | 14 ++++++++++--- src/tests/test_spec_metrics.py | 36 ++++++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 3 deletions(-) diff --git a/src/baskerville/metrics.py b/src/baskerville/metrics.py index 3ff888f..828c4f6 100644 --- a/src/baskerville/metrics.py +++ b/src/baskerville/metrics.py @@ -229,7 +229,9 @@ def update(self, preds: torch.Tensor, target: torch.Tensor) -> None: """preds, target: (N, T, L).""" if self.enabled: for g in self.groups: - zp, zy = self._map(preds, g), self._map(target, g) # (tracks, P) + # Residual covariance subtracts large, nearly equal moments. + zp = self._map(preds, g).double() # (tracks, P) + zy = self._map(target, g).double() mp, my = zp.mean(dim=0), zy.mean(dim=0) # (P,) def total(x): @@ -270,8 +272,14 @@ def compute(self) -> torch.Tensor: # residuals of pred/target after regressing out their group mean u = torch.zeros(len(cov), 4, dtype=torch.float64) w = torch.zeros_like(u) - u[:, 0], u[:, 2] = 1, -cov[:, 0, 2] / cov[:, 2, 2] - w[:, 1], w[:, 3] = 1, -cov[:, 1, 3] / cov[:, 3, 3] + u[:, 0] = w[:, 1] = 1 + # A constant group mean removes only the intercept. + u[:, 2] = torch.where( + cov[:, 2, 2] > 0, -cov[:, 0, 2] / cov[:, 2, 2], 0 + ) + w[:, 3] = torch.where( + cov[:, 3, 3] > 0, -cov[:, 1, 3] / cov[:, 3, 3], 0 + ) cuw = torch.einsum("ta,tab,tb->t", u, cov, w) cuu = torch.einsum("ta,tab,tb->t", u, cov, u) cww = torch.einsum("ta,tab,tb->t", w, cov, w) diff --git a/src/tests/test_spec_metrics.py b/src/tests/test_spec_metrics.py index 9a1f5a9..2df2561 100644 --- a/src/tests/test_spec_metrics.py +++ b/src/tests/test_spec_metrics.py @@ -112,6 +112,42 @@ def test_spec_matches_numpy(targets_df): assert np.isnan(np.delete(spec_t, rep)).all() +def test_spec_nearly_identical_tracks(): + rng = np.random.default_rng(1) + num_tracks = 20 + shape = (1, num_tracks, 10000) + shared = rng.uniform(100, 200, (1, 1, shape[-1])) + preds = (shared + rng.normal(0, 0.001, shape)).astype(np.float32) + targets = (shared + rng.normal(0, 0.001, shape)).astype(np.float16) + df = pd.DataFrame({"description": ["RNA:a"] * num_tracks}) + hist = _hist_reference(targets, np.arange(num_tracks)) + spec = metrics.SpecPearsonCorrCoef(df, hist) + for start in range(0, shape[-1], 2000): + spec.update( + torch.from_numpy(preds[:, :, start : start + 2000]), + torch.from_numpy(targets[:, :, start : start + 2000]), + ) + + rep, pair = spec.groups["RNA"] + expected = _spec_numpy(preds, targets, rep, pair, spec.table["RNA"].numpy()) + np.testing.assert_allclose(spec.compute().numpy(), expected, atol=1e-5) + + +def test_spec_constant_group_mean(): + targets = np.tile([0, 1], (1, 20, 20)).astype(np.float16) + targets[:, 10:] = 1 - targets[:, 10:] + df = pd.DataFrame({"description": ["RNA:a"] * 20}) + spec = metrics.SpecPearsonCorrCoef(df, _hist_reference(targets, np.arange(20))) + y = torch.from_numpy(targets) + spec.update(y, y) + torch.testing.assert_close(spec.compute(), torch.ones(20)) + + # Constant tracks still have undefined residual correlation. + spec.reset() + spec.update(torch.zeros_like(y), y) + assert spec.compute().isnan().all() + + def test_spec_gain_and_specific(): """Depth gain on a shared signal scores ~0; track-specific signal scores.""" num_tracks = 8 From 8b5ac1a0caa2562917c89de81edea8f0a66282a4 Mon Sep 17 00:00:00 2001 From: David Kelley Date: Thu, 8 Oct 2026 12:11:44 -0700 Subject: [PATCH 5/6] Format with ruff 0.15 Co-Authored-By: Claude Opus 5.5 --- src/baskerville/metrics.py | 25 +++++++++++-------- .../utils/hound_data_subset_zarr_targets.py | 4 ++- src/tests/test_eval_folds_gcp.py | 4 +-- 3 files changed, 19 insertions(+), 14 deletions(-) diff --git a/src/baskerville/metrics.py b/src/baskerville/metrics.py index 828c4f6..d4ad262 100644 --- a/src/baskerville/metrics.py +++ b/src/baskerville/metrics.py @@ -145,7 +145,9 @@ def qmap_tables(hists): order = np.argsort(pos, kind="stable") ref_pos = np.concatenate([[0], pos[order], [num_ranks]]) ref_level = sum(values[nz[0]] for nz, _ in knots) - ref_level = (ref_level + np.concatenate([[0.0], np.cumsum(step[order])])) / num_tracks + ref_level = ( + ref_level + np.concatenate([[0.0], np.cumsum(step[order])]) + ) / num_tracks ref_sum = np.concatenate([[0.0], np.cumsum(ref_level * np.diff(ref_pos))]) tables = np.empty(hists.shape, dtype=np.float32) @@ -262,9 +264,16 @@ def compute(self) -> torch.Tensor: ) s2 = torch.empty(len(ts), 4, 4, dtype=torch.float64) pairs = { - (0, 0): ts[:, 2], (1, 1): ts[:, 3], (0, 1): ts[:, 4], - (0, 2): ts[:, 5], (0, 3): ts[:, 6], (1, 2): ts[:, 7], - (1, 3): ts[:, 8], (2, 2): ms[2], (3, 3): ms[3], (2, 3): ms[4], + (0, 0): ts[:, 2], + (1, 1): ts[:, 3], + (0, 1): ts[:, 4], + (0, 2): ts[:, 5], + (0, 3): ts[:, 6], + (1, 2): ts[:, 7], + (1, 3): ts[:, 8], + (2, 2): ms[2], + (3, 3): ms[3], + (2, 3): ms[4], } for (a, b), v in pairs.items(): s2[:, a, b] = s2[:, b, a] = v @@ -274,12 +283,8 @@ def compute(self) -> torch.Tensor: w = torch.zeros_like(u) u[:, 0] = w[:, 1] = 1 # A constant group mean removes only the intercept. - u[:, 2] = torch.where( - cov[:, 2, 2] > 0, -cov[:, 0, 2] / cov[:, 2, 2], 0 - ) - w[:, 3] = torch.where( - cov[:, 3, 3] > 0, -cov[:, 1, 3] / cov[:, 3, 3], 0 - ) + u[:, 2] = torch.where(cov[:, 2, 2] > 0, -cov[:, 0, 2] / cov[:, 2, 2], 0) + w[:, 3] = torch.where(cov[:, 3, 3] > 0, -cov[:, 1, 3] / cov[:, 3, 3], 0) cuw = torch.einsum("ta,tab,tb->t", u, cov, w) cuu = torch.einsum("ta,tab,tb->t", u, cov, u) cww = torch.einsum("ta,tab,tb->t", w, cov, w) diff --git a/src/baskerville/scripts/utils/hound_data_subset_zarr_targets.py b/src/baskerville/scripts/utils/hound_data_subset_zarr_targets.py index 35f2350..68696f6 100755 --- a/src/baskerville/scripts/utils/hound_data_subset_zarr_targets.py +++ b/src/baskerville/scripts/utils/hound_data_subset_zarr_targets.py @@ -113,7 +113,9 @@ def subset_zarr_for_fold( old_targets = pd.read_csv(f"{og_data_dir}/targets.txt", sep="\t", index_col=0) pair = dataset.strand_pair_indices(old_targets) if not np.isin(pair[old_indices], old_indices).all(): - raise ValueError("old_indices split a strand pair; target_hist would be wrong") + raise ValueError( + "old_indices split a strand pair; target_hist would be wrong" + ) hist = old_root["target_hist"].oindex[old_indices, :] fold_zarr_root.create_array( "target_hist", diff --git a/src/tests/test_eval_folds_gcp.py b/src/tests/test_eval_folds_gcp.py index 8edffb4..90f5196 100644 --- a/src/tests/test_eval_folds_gcp.py +++ b/src/tests/test_eval_folds_gcp.py @@ -231,9 +231,7 @@ def test_eval_gcp_uses_queue(tmp_path, monkeypatch): out_dir = _make_models_dir(tmp_path, num_folds=2, test_fold=0, valid_fold=1) captured = _patch_gcp(monkeypatch, num_folds=2) - hef.eval_folds( - MockArgs(out_dir=out_dir, fold_subset=1, test_only=True, queue="l4") - ) + hef.eval_folds(MockArgs(out_dir=out_dir, fold_subset=1, test_only=True, queue="l4")) eval_jobs = [j for j in captured if j.cmd.startswith("hound_eval ")] assert eval_jobs From 77f1ac6bb055a4a82a58658671717bc015112356 Mon Sep 17 00:00:00 2001 From: David Kelley Date: Thu, 15 Oct 2026 17:14:35 -0700 Subject: [PATCH 6/6] hound_train_folds --eval: score the test fold at the end of each training job With one-pass spec, the test eval is cheap enough to run on the training GPU right after training. Each job chains hound_eval per dataset on its held-out fold, writing /eval*/fold as hound_eval_folds does; the GCP run links test -> fold after the final pull. A fold with COMPLETE but no test acc.txt is relaunched (the trainer resumes finished and exits; the eval reruns). --rc passes through. link_test_metrics moves to helpers/fold_utils. Co-Authored-By: Claude Opus 5.5 --- src/baskerville/helpers/fold_utils.py | 29 +++++- src/baskerville/scripts/hound_eval_folds.py | 31 ++----- src/baskerville/scripts/hound_train_folds.py | 93 ++++++++++++++++++-- src/tests/test_train_folds_gcp.py | 82 +++++++++++++++++ 4 files changed, 199 insertions(+), 36 deletions(-) diff --git a/src/baskerville/helpers/fold_utils.py b/src/baskerville/helpers/fold_utils.py index f6e7a8e..c57a18a 100644 --- a/src/baskerville/helpers/fold_utils.py +++ b/src/baskerville/helpers/fold_utils.py @@ -3,8 +3,10 @@ Both ``hound_eval_folds`` and ``hound_eval_genes_folds`` need the same logic to (a) read the authoritative fold assignment a replicate was trained with, (b) decide whether a replicate's weights exist (locally or in GCS marker mode), -and (c) resolve the ``model_best.pth`` path to pass to the eval command. They're -collected here so the fold scripts share one implementation. +(c) resolve the ``model_best.pth`` path to pass to the eval command, and (d) +link ``test`` to the held-out fold's eval dir. They're collected here so the +fold scripts (including ``hound_train_folds --eval``) share one +implementation. """ from __future__ import annotations @@ -60,3 +62,26 @@ def resolve_model_file( if gcs_models_dir is not None: return f"{GCP_MODELS_MOUNT}/{fold_cross}/train/model_best.pth" return f"{container_models_dir}/{fold_cross}/train/model_best.pth" + + +def link_test_metrics(eval_dir, test_fold): + """Symlink test -> fold{test_fold} and test.out -> fold{test_fold}.out. + + Relative targets so the links resolve wherever eval_dir lives. Targets may + be dangling until eval completes. A stale link (e.g. left by the pre-GCP + script that linked test -> fold{fi}, wrong when test_fold != fi for + cross > 0) is repointed; a real file/dir at the link path is left alone. + """ + os.makedirs(eval_dir, exist_ok=True) + for name, target in ( + ("test", f"fold{test_fold}"), + ("test.out", f"fold{test_fold}.out"), + ): + link = f"{eval_dir}/{name}" + if os.path.islink(link): + if os.readlink(link) == target: + continue # already correct + os.unlink(link) # stale link → repoint below + elif os.path.lexists(link): + continue # real file/dir squatting the name → don't clobber + os.symlink(target, link) diff --git a/src/baskerville/scripts/hound_eval_folds.py b/src/baskerville/scripts/hound_eval_folds.py index 8822804..0dd9fa6 100755 --- a/src/baskerville/scripts/hound_eval_folds.py +++ b/src/baskerville/scripts/hound_eval_folds.py @@ -38,6 +38,7 @@ from gcprunner.batch_spec import DataMount from baskerville.helpers import fold_utils, stage_cache from baskerville.helpers.fold_utils import ( + link_test_metrics as _link_test_metrics, model_present as _model_present, read_fold_splits as _read_fold_splits, resolve_model_file as _resolve_model_file, @@ -513,9 +514,10 @@ def eval_folds(args): ################################################################################ # helpers ################################################################################ -# The shared fold helpers (_read_fold_splits, _model_present, _resolve_model_file) -# now live in baskerville.helpers.fold_utils and are imported above under -# their original private names so existing call sites and tests are unchanged. +# The shared fold helpers (_read_fold_splits, _model_present, _resolve_model_file, +# _link_test_metrics) now live in baskerville.helpers.fold_utils and are imported +# above under their original private names so existing call sites and tests are +# unchanged. def eval_options(args, targets_gene): @@ -538,29 +540,6 @@ def eval_options(args, targets_gene): return opts -def _link_test_metrics(eval_dir, test_fold): - """Symlink test -> fold{test_fold} and test.out -> fold{test_fold}.out. - - Relative targets so the links resolve wherever eval_dir lives. Targets may - be dangling until eval completes. A stale link (e.g. left by the pre-GCP - script that linked test -> fold{fi}, wrong when test_fold != fi for - cross > 0) is repointed; a real file/dir at the link path is left alone. - """ - os.makedirs(eval_dir, exist_ok=True) - for name, target in ( - ("test", f"fold{test_fold}"), - ("test.out", f"fold{test_fold}.out"), - ): - link = f"{eval_dir}/{name}" - if os.path.islink(link): - if os.readlink(link) == target: - continue # already correct - os.unlink(link) # stale link → repoint below - elif os.path.lexists(link): - continue # real file/dir squatting the name → don't clobber - os.symlink(target, link) - - ################################################################################ # __main__ ################################################################################ diff --git a/src/baskerville/scripts/hound_train_folds.py b/src/baskerville/scripts/hound_train_folds.py index d4f6aec..b2359f8 100755 --- a/src/baskerville/scripts/hound_train_folds.py +++ b/src/baskerville/scripts/hound_train_folds.py @@ -25,7 +25,7 @@ import threading import time -from baskerville import utils +from baskerville import dataset, utils try: import slurmrunner @@ -44,7 +44,7 @@ from gcprunner import resolve_image from gcprunner import run_identity from gcprunner.batch_spec import DataMount -from baskerville.helpers import stage_cache +from baskerville.helpers import fold_utils, stage_cache from baskerville.helpers.gcs_utils import ( download_folder_from_gcs, gcs_file_exist, @@ -83,6 +83,13 @@ def main(): help="Conda environment to activate in each job (sources " "$BASKERVILLE_CONDA first if set). Default: the current environment.", ) + parser.add_argument( + "--eval", + default=False, + action="store_true", + help="Evaluate the test fold with hound_eval at the end of each training " + "job, into /eval*/fold as hound_eval_folds does [Default: %(default)s]", + ) parser.add_argument( "-f", "--fold_subset", @@ -114,6 +121,13 @@ def main(): default="titan_rtx", help="SLURM queue on which to run the jobs [Default: %(default)s]", ) + parser.add_argument( + "--rc", + default=False, + action="store_true", + help="--eval: average forward and reverse complement predictions " + "[Default: %(default)s]", + ) parser.add_argument( "-r", "--restart", @@ -186,6 +200,8 @@ def main(): container_transfer_dir = None if (args.conclude or args.extend) and not gcp_backend: parser.error("--conclude/--extend are only supported with --backend gcp") + if args.eval and args.whole: + parser.error("--eval needs a held-out test fold; drop --whole") if args.extend: # Edited params hash to a new run id, so take the run dir and image # from the marker instead. Before resolve_image_arg, so the pinned @@ -296,8 +312,9 @@ def main(): with open(data_stats_file) as data_stats_open: data_stats = json.load(data_stats_open) - # count folds + # count folds (all of them set the test fold, even under --fold_subset) num_folds = len([dkey for dkey in data_stats if dkey.startswith("fold")]) + args.num_folds_total = num_folds # subset folds if args.fold_subset is not None: @@ -370,6 +387,17 @@ def main(): cmd += " -w" cmd += f" {rep_dir}/params.json" cmd += f" {' '.join(args.data_dirs)}" + evals = _test_evals(args, fold_cross) + cmd += _test_eval_cmd( + args, + evals, + f"{rep_dir}/params.json", + f"{train_dir}/model_best.pth", + args.data_dirs, + rep_dir, + ) + for eval_sub, test_fold in evals: + fold_utils.link_test_metrics(f"{rep_dir}/{eval_sub}", test_fold) if slurmrunner is None: # Run locally @@ -612,7 +640,7 @@ def _read_progress_epoch(gcs_progress_uri): return None -def _fold_state(run_dir, fold_cross, active, epochs_max): +def _fold_state(run_dir, fold_cross, active, epochs_max, evals=()): """Infer a fold's state from GCS markers + the set of active Batch jobs. Returns (state, detail) where state is one of: @@ -623,12 +651,19 @@ def _fold_state(run_dir, fold_cross, active, epochs_max): new — nothing yet (fresh start) A COMPLETE fold that stopped at max_epochs below ``epochs_max`` counts as - incomplete, so raising train_epochs_max extends it. + incomplete, so raising train_epochs_max extends it. So does one missing any + ``evals`` (eval_sub, test_fold) acc.txt; the relaunch resumes a finished + trainer, which exits at once, and reruns the eval. """ base = f"{run_dir}/{fold_cross}/train" if gcs_file_exist(f"{base}/COMPLETE"): done = read_json_gcs(f"{base}/COMPLETE") - if done.get("reason") != "max_epochs" or done["epoch"] + 1 >= epochs_max: + trained = done.get("reason") != "max_epochs" or done["epoch"] + 1 >= epochs_max + evaluated = all( + gcs_file_exist(f"{run_dir}/{fold_cross}/{sub}/fold{t}/acc.txt") + for sub, t in evals + ) + if trained and evaluated: return ("complete", None) if gcs_file_exist(f"{base}/FAILED"): return ("failed", None) @@ -681,6 +716,14 @@ def _build_gcp_train_job( if args.whole: cmd += " -w" cmd += f" {container_params} {' '.join(container_data_dirs)}" + cmd += _test_eval_cmd( + args, + _test_evals(args, fold_cross), + container_params, + "/workspace/out/train/model_best.pth", + container_data_dirs, + "/workspace/out", + ) fold_gcs = f"{args.gcp_output_dir}/{fold_cross}" run_id = run_identity.run_id_from_gcs_dir(args.gcp_output_dir) @@ -805,6 +848,11 @@ def _run_gcp_training( try: if completed: _pull_full_run(args.gcp_output_dir, args.out_dir) + for fc in fold_crosses: + for eval_sub, test_fold in _test_evals(args, fc): + fold_utils.link_test_metrics( + f"{args.out_dir}/{fc}/{eval_sub}", test_fold + ) else: # light catch-up only; --conclude pulls an interrupted run _mirror_once(args.gcp_output_dir, args.out_dir) @@ -836,13 +884,15 @@ def _resubmit_loop( project = resolve_project(args.gcp_project) region = resolve_region(args.gcp_region) epochs_max = params_train.get("train_epochs_max", 10000) # Trainer default + evals = {fc: _test_evals(args, fc) for fc in fold_crosses} previous_states = {} for round_i in range(args.gcp_max_rounds): active = _active_fold_crosses(project, region, fold_crosses, run_id) states = { - fc: _fold_state(run_dir, fc, active, epochs_max) for fc in fold_crosses + fc: _fold_state(run_dir, fc, active, epochs_max, evals[fc]) + for fc in fold_crosses } _print_status_table(states, round_i, args.gcp_max_rounds) @@ -927,7 +977,10 @@ def _resubmit_loop( # exhausted rounds active = _active_fold_crosses(project, region, fold_crosses, run_id) - states = {fc: _fold_state(run_dir, fc, active, epochs_max) for fc in fold_crosses} + states = { + fc: _fold_state(run_dir, fc, active, epochs_max, evals[fc]) + for fc in fold_crosses + } incomplete = [fc for fc, (s, _) in states.items() if s != "complete"] if incomplete: print( @@ -940,6 +993,30 @@ def _resubmit_loop( return True +def _test_evals(args, fold_cross): + """[(eval_sub, test_fold)] per dataset under --eval, else [].""" + if not args.eval: + return [] + fi, ci = map(int, fold_cross[1:].split("c")) + test_fold = dataset.compute_fold_splits(args.num_folds_total, fi, ci)["test"] + num_data = len(args.data_dirs) + return [ + ("eval" if num_data == 1 else f"eval{di}", test_fold) for di in range(num_data) + ] + + +def _test_eval_cmd(args, evals, params_file, model_file, data_dirs, out_root): + """hound_eval commands to chain after hound_train, one per dataset.""" + cmd = "" + for di, (eval_sub, test_fold) in enumerate(evals): + cmd += f" && hound_eval --dataset {di} -o {out_root}/{eval_sub}/fold{test_fold}" + cmd += f" --split fold{test_fold}" + if args.rc: + cmd += " --rc" + cmd += f" {params_file} {model_file} {data_dirs[di]}" + return cmd + + def make_rep_params(params_file, rep_dir, pretrained_model): """Copy params file, including pretained model path. diff --git a/src/tests/test_train_folds_gcp.py b/src/tests/test_train_folds_gcp.py index f4907e6..3827cbf 100644 --- a/src/tests/test_train_folds_gcp.py +++ b/src/tests/test_train_folds_gcp.py @@ -19,6 +19,7 @@ class MockArgs: def __init__(self, **kwargs): self.params_file = "params.json" self.transfer = None + self.eval = False for key, value in kwargs.items(): setattr(self, key, value) @@ -633,3 +634,84 @@ def stop(args): htf.main() assert seen["gcp_output_dir"] == "gs://b/train/old-run" assert seen["gcp_image"] == "reg/img@sha256:abc" + + +def _eval_args(data_dirs, rc=False, eval=True): + return MockArgs( + eval=eval, + rc=rc, + num_folds_total=8, + data_dirs=data_dirs, + gcp_output_dir="gs://bucket/train/run123", + name="b32", + queue="l4", + whole=False, + gcp_retry_count=0, + gcp_data_dir="gs://bucket/data", + gcp_data_local="/workspace/data", + ) + + +def test_eval_chains_hound_eval_per_dataset(monkeypatch): + """--eval appends one hound_eval per dataset on the held-out test fold, + written where hound_eval_folds writes (eval/fold).""" + import gcprunner + + monkeypatch.setattr(htf, "slurmrunner", gcprunner) + args = _eval_args(["hg38", "mm10"], rc=True) + data = ["/workspace/data/hg38", "/workspace/data/mm10"] + job = htf._build_gcp_train_job(args, {}, "/c/params.json", data, "f2c1") + train, eval0, eval1 = job.cmd.split(" && ") + assert train.startswith("hound_train -o /workspace/out/train --fold 2 --cross 1") + model = "/workspace/out/train/model_best.pth" + # with 8 folds, fold 2 cross 1 tests on fold 2 (test = fold, valid = fold+1+cross) + assert eval0 == ( + "hound_eval --dataset 0 -o /workspace/out/eval0/fold2 --split fold2 --rc " + f"/c/params.json {model} /workspace/data/hg38" + ) + assert eval1.startswith("hound_eval --dataset 1 -o /workspace/out/eval1/fold2") + assert eval1.endswith("/workspace/data/mm10") + + +def test_eval_single_dataset_and_off(): + assert htf._test_evals(_eval_args(["hg38"]), "f3c0") == [("eval", 3)] + assert htf._test_evals(_eval_args(["hg38"], eval=False), "f3c0") == [] + assert htf._test_eval_cmd(_eval_args(["hg38"]), [], "p", "m", ["d"], "o") == "" + + +def test_eval_fold_uses_all_folds(): + """The test fold comes from the full fold count, not the --fold_subset cut.""" + args = _eval_args(["hg38"]) + expected = htf.dataset.compute_fold_splits(8, 1, 3)["test"] + assert htf._test_evals(args, "f1c3") == [("eval", expected)] + + +@pytest.mark.parametrize( + "has_acc, expected", [(True, "complete"), (False, "incomplete")] +) +def test_fold_state_requires_test_eval(monkeypatch, has_acc, expected): + """COMPLETE without the test acc.txt relaunches (resume is a no-op, eval reruns).""" + base = "gs://b/run/f0c0" + files = {f"{base}/train/COMPLETE", f"{base}/train/checkpoint.pth"} + if has_acc: + files.add(f"{base}/eval0/fold0/acc.txt") + monkeypatch.setattr(htf, "gcs_file_exist", lambda uri: uri in files) + monkeypatch.setattr( + htf, "read_json_gcs", lambda uri: {"epoch": 49, "reason": "early_stop"} + ) + state, _ = htf._fold_state("gs://b/run", "f0c0", set(), 50, [("eval0", 0)]) + assert state == expected + + +def test_success_links_test_eval(tmp_path, monkeypatch): + """A completed --eval run links eval*/test -> fold after the pull.""" + _patch_gcp_training_env(monkeypatch, []) + monkeypatch.setattr(htf, "_resubmit_loop", lambda *a, **k: True) + monkeypatch.setattr(htf, "_pull_full_run", lambda *a: None) + args = _gcp_training_args(tmp_path) + for k, v in vars(_eval_args(["hg38", "mm10"])).items(): + if not hasattr(args, k) or k in ("eval", "data_dirs"): + setattr(args, k, v) + htf._run_gcp_training(args, {}, {}, [], ["f1c0"]) + for di in (0, 1): + assert (tmp_path / f"f1c0/eval{di}/test").readlink().name == "fold1"