14. Predict with the score model¶

This page runs the trained AP-1 score model (previous page) forward. It is the same prediction_task as the accessibility page — DNA window + cell state → base-resolution accessibility — with one addition: each cell state also carries a chromVAR AP-1 activity score, and the model predicts the accessibility that matches it. The goal here is the I/O contract: how the score is supplied and what the task writes out — not benchmarking.

This uses the bolerodata model zoo (the paper route, like the variant-effect page). The AP-1 model is the published Bolero-AP1-Score model, registered as an ATAC score model (Score Input = AP1, no gene-count head). The notebook downloads its checkpoint, the HumanBrainDev cell table and pseudobulk records, and the hg38 genome.

Setup¶

In [1]:

Copied!





import logging

import joblib
import numpy as np
import pandas as pd

from bolero import init
from bolerodata import DATASETS, MODELS, metadata

# quiet the per-file download checks
logging.getLogger("httpx").setLevel(logging.WARNING)
init(num_cpus=16)
import logging

import joblib
import numpy as np
import pandas as pd

from bolero import init
from bolerodata import DATASETS, MODELS, metadata

# quiet the per-file download checks
logging.getLogger("httpx").setLevel(logging.WARNING)
init(num_cpus=16)

2026-10-27 16:06:26,479	INFO worker.py:1781 -- Started a local Ray instance.

Pick the model and the dataset to predict for¶

MODELS["Bolero10M+AP1+260104"] is the trained AP-1 ATAC score model. Because it is a multi-dataset model, you must tell it which dataset's cell states you are predicting (model.dataset) — that selects the per-dataset conditioning encoder and the peak set.

In [2]:

Copied!





MODEL_NAME = "Bolero10M+AP1+260104"
DATASET_NAME = "HumanBrainDev"

model = MODELS[MODEL_NAME]
model.dataset = DATASET_NAME  # multi-dataset model: set the dataset before predicting
print("model group:", model.model_group, "| score input:", model._tf_score_input)
MODEL_NAME = "Bolero10M+AP1+260104"
DATASET_NAME = "HumanBrainDev"

model = MODELS[MODEL_NAME]
model.dataset = DATASET_NAME  # multi-dataset model: set the dataset before predicting
print("model group:", model.model_group, "| score input:", model._tf_score_input)

model group: BorzoiSignalMultiDS | score input: AP1

The score rides in with the cell states¶

A prediction is made for a (region, cell state) pair. For the score model each cell state is a pseudobulk record that carries two things: its cell-state embedding and its __shared_data__ — the AP-1 score for that state. The model reads __shared_data__ and folds it into the same conditioning vector it was trained with.

We build these records here, from published files:

Cell states. Thirteen prenatal cortical neuron subclasses of the human brain-development dataset (HumanBrainDev). Each state's embedding is the embedding of its most central cell: the cell closest to the mean embedding of the subclass's prenatal cells.
AP-1 scores. Each state is predicted twice, once at a low AP-1 score (-1.5) and once at a high one (6.0). The published score file shows where these values sit.
The record file. The dataset's published pseudobulk records file supplies the rest of the format (the condition embeddings); we swap in our states and save it.

1. Cell states: the central cell of each prenatal neuron subclass¶

In [3]:

Copied!





dataset = DATASETS[DATASET_NAME]
# cell metadata + embedding (obsm["X_emb"])
cell_adata = dataset.get_meta_cell_adata("cell")

NEURON_SUBCLASSES = [
    "EN-L2_3-IT", "EN-L4-IT", "EN-L5-IT", "EN-L6-IT", "EN-L5-ET", "EN-L5_6-NP",
    "EN-L6-CT", "EN-L6b", "IN-MGE-PV", "IN-MGE-SST", "IN-CGE-VIP", "IN-CGE-SNCG",
    "IN-Mix-LAMP5",
]  # fmt: skip

# prenatal = the three trimesters; the other ages are postnatal
prenatal = cell_adata.obs["age"].astype(str).str.endswith("_trimester")
subclass = cell_adata.obs["subclass"].astype(str)


def center_embedding(mask):
    # the embedding of the cell closest to the group's mean embedding
    emb = cell_adata.obsm["X_emb"][np.asarray(mask)]
    return emb[((emb - emb.mean(axis=0, keepdims=True)) ** 2).sum(axis=1).argmin()]


state_emb = {
    f"{name}+PreBirth": center_embedding((subclass == name) & prenatal)
    for name in sorted(NEURON_SUBCLASSES)
}
emb_width = next(iter(state_emb.values())).size
print(len(state_emb), "cell states, embedding width", emb_width)
dataset = DATASETS[DATASET_NAME]
# cell metadata + embedding (obsm["X_emb"])
cell_adata = dataset.get_meta_cell_adata("cell")

NEURON_SUBCLASSES = [
    "EN-L2_3-IT", "EN-L4-IT", "EN-L5-IT", "EN-L6-IT", "EN-L5-ET", "EN-L5_6-NP",
    "EN-L6-CT", "EN-L6b", "IN-MGE-PV", "IN-MGE-SST", "IN-CGE-VIP", "IN-CGE-SNCG",
    "IN-Mix-LAMP5",
]  # fmt: skip

# prenatal = the three trimesters; the other ages are postnatal
prenatal = cell_adata.obs["age"].astype(str).str.endswith("_trimester")
subclass = cell_adata.obs["subclass"].astype(str)


def center_embedding(mask):
    # the embedding of the cell closest to the group's mean embedding
    emb = cell_adata.obsm["X_emb"][np.asarray(mask)]
    return emb[((emb - emb.mean(axis=0, keepdims=True)) ** 2).sum(axis=1).argmin()]


state_emb = {
    f"{name}+PreBirth": center_embedding((subclass == name) & prenatal)
    for name in sorted(NEURON_SUBCLASSES)
}
emb_width = next(iter(state_emb.values())).size
print(len(state_emb), "cell states, embedding width", emb_width)

Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.

13 cell states, embedding width 30

2. AP-1 scores: a low and a high value¶

In [4]:

Copied!





score_data = joblib.load(metadata.get_chromvar_score_path("AP1"))
corpus_scores = pd.concat([df.iloc[:, 0] for df in score_data.values()])
pd.DataFrame(
    {
        DATASET_NAME: score_data[DATASET_NAME].iloc[:, 0].describe(),
        "all datasets": corpus_scores.describe(),
    }
).round(2)
score_data = joblib.load(metadata.get_chromvar_score_path("AP1"))
corpus_scores = pd.concat([df.iloc[:, 0] for df in score_data.values()])
pd.DataFrame(
    {
        DATASET_NAME: score_data[DATASET_NAME].iloc[:, 0].describe(),
        "all datasets": corpus_scores.describe(),
    }
).round(2)

Out[4]:

	HumanBrainDev	all datasets
count	535.00	71634.00
mean	-0.51	0.03
std	0.72	1.01
min	-1.43	-2.13
25%	-0.97	-0.76
50%	-0.81	-0.32
75%	-0.32	0.71
max	2.38	6.20

The HumanBrainDev pseudobulks span a narrow range, and -1.5 sits just below its bottom. Across all datasets the score reaches about 6, so 6.0 is an extreme AP-1 activity: near the top of the whole corpus and far above any HumanBrainDev cell state.

In [5]:

Copied!

AP1_SCORES = [-1.5, 6.0]  # low and high AP-1 activity
AP1_SCORES = [-1.5, 6.0]  # low and high AP-1 activity

3. The record file¶

Each record key is <cell state>:<score index>, so the 13 states × 2 scores give 26 records. The model also conditions on the dataset's own variables (age and brain region); condition_to_related_pseudobulk assigns states to a condition, and we put all 26 under one condition, second-trimester V1. __embedding_only__ tells the predictor that these records are embeddings, with no cells to read coverage from.

In [6]:

Copied!





records = joblib.load(dataset.get_meta_cell_cond_pseudobulk_records_path(5_000_000))

records["pseudobulk_records"] = {
    f"{state}:{idx}": {"embedding": emb, "__shared_data__": np.array([score])}
    for idx, score in enumerate(AP1_SCORES)
    for state, emb in state_emb.items()
}
records["__embedding_only__"] = True
records["condition_to_related_pseudobulk"]["Second_trimester+V1"] = list(
    records["pseudobulk_records"]
)

PSEUDOBULK_RECORDS = f"{DATASET_NAME}.AP1.pseudobulk_records.joblib"
joblib.dump(records, PSEUDOBULK_RECORDS)

_pid, _rec = next(iter(records["pseudobulk_records"].items()))
print(len(records["pseudobulk_records"]), "records; one of them:", _pid)
print("a cell-state record holds:", list(_rec))
print("  embedding shape :", _rec["embedding"].shape)
print("  __shared_data__ :", _rec["__shared_data__"], "(the AP-1 activity score)")
records = joblib.load(dataset.get_meta_cell_cond_pseudobulk_records_path(5_000_000))

records["pseudobulk_records"] = {
    f"{state}:{idx}": {"embedding": emb, "__shared_data__": np.array([score])}
    for idx, score in enumerate(AP1_SCORES)
    for state, emb in state_emb.items()
}
records["__embedding_only__"] = True
records["condition_to_related_pseudobulk"]["Second_trimester+V1"] = list(
    records["pseudobulk_records"]
)

PSEUDOBULK_RECORDS = f"{DATASET_NAME}.AP1.pseudobulk_records.joblib"
joblib.dump(records, PSEUDOBULK_RECORDS)

_pid, _rec = next(iter(records["pseudobulk_records"].items()))
print(len(records["pseudobulk_records"]), "records; one of them:", _pid)
print("a cell-state record holds:", list(_rec))
print("  embedding shape :", _rec["embedding"].shape)
print("  __shared_data__ :", _rec["__shared_data__"], "(the AP-1 activity score)")

26 records; one of them: EN-L2_3-IT+PreBirth:0
a cell-state record holds: ['embedding', '__shared_data__']
  embedding shape : (30,)
  __shared_data__ : [-1.5] (the AP-1 activity score)

Build the predictor and run the task¶

create_predictor assembles the model + data manager from the config; prediction_task does the GPU forward pass and writes the outputs. The toggles below match a signal-free, embedding-driven run:

nosignal=True — predict from DNA + conditioning only (no reference-coverage input). The AP-1 score is still applied; only the reference ATAC track branch is off.
use_ref_bw=False — no reference bigWig needed (paired with nosignal).
embedding_only_mode=True — drive purely from the pseudobulk embeddings/scores; skip loading ground-truth signal.
peak_bed_path — the dataset's peaks, so the base track is also summarized per peak.

regions="test_regions" predicts on the model's held-out chromosomes; we keep just a handful (downsample_regions=8) so the demo runs quickly. Drop that argument to predict the full set.

In [7]:

Copied!





predictor = model.create_predictor(
    subset_name=None,
    pseudobulk_records_path=PSEUDOBULK_RECORDS,
    sample_n_pseudobulks=None,
    nosignal=True,
    use_ref_bw=False,
    embedding_only_mode=True,
    peak_bed_path=model.dataset.peak_bed_path,
)
predictor = model.create_predictor(
    subset_name=None,
    pseudobulk_records_path=PSEUDOBULK_RECORDS,
    sample_n_pseudobulks=None,
    nosignal=True,
    use_ref_bw=False,
    embedding_only_mode=True,
    peak_bed_path=model.dataset.peak_bed_path,
)

In [8]:

Copied!





output_dir = f"predictions/{MODEL_NAME}:{DATASET_NAME}"
predictor.prediction_task(
    output_dir,
    regions="test_regions",
    downsample_regions=8,  # demo: a handful of held-out regions (drop to predict all)
    mode="prediction",
)
print("wrote predictions to", output_dir)
output_dir = f"predictions/{MODEL_NAME}:{DATASET_NAME}"
predictor.prediction_task(
    output_dir,
    regions="test_regions",
    downsample_regions=8,  # demo: a handful of held-out regions (drop to predict all)
    mode="prediction",
)
print("wrote predictions to", output_dir)

Saving batches to /tmp/bolero-docs/bolero/docs/tutorials/score_model/predictions/Bolero10M+AP1+260104:HumanBrainDev/batch
Saving stats to /tmp/bolero-docs/bolero/docs/tutorials/score_model/predictions/Bolero10M+AP1+260104:HumanBrainDev/summary_stats.joblib.gz
Using temporary directory /tmp/tmp38istnsn
0 regions has finished in /tmp/bolero-docs/bolero/docs/tutorials/score_model/predictions/Bolero10M+AP1+260104:HumanBrainDev/batch
8 regions to compute
Data loader batch size 16

Loading base model weights from: /tmp/bolero-docs/bdhome/flashzoi/fold_0.renamed.pt

Model does not have gene_count_output_head, skip
Loading checkpoint from /tmp/bolero-docs/bdhome/hf/models--arcinstitute--bolero-models/snapshots/eabdbb2e61072357ad0441d49eb66708fb858e33/BorzoiSignalMultiDS/Bolero10M+AP1+260104/checkpoint.pt

Skipping checkpoint encoders of datasets not in the config: ['CZIBodyAtlas', 'HYSocialRebound', 'Torpor+Control']

==========
Dataloader Batch Schema:
- __conditionemb__: <class 'torch.Tensor'> torch.Size([52, 1, 3]) torch.float32 cuda:0
- __conditionemb__:cond0: <class 'torch.Tensor'> torch.Size([26, 1, 3]) torch.float32 cuda:0
- __conditionemb__:cond1: <class 'torch.Tensor'> torch.Size([26, 1, 3]) torch.float32 cuda:0
- __conditionemb__:delta: <class 'torch.Tensor'> torch.Size([26, 1, 3]) torch.float32 cuda:0
- __dna__: <class 'torch.Tensor'> torch.Size([8, 4, 524288]) torch.uint8 cuda:0
- __embedding__: <class 'torch.Tensor'> torch.Size([52, 60]) torch.float32 cuda:0
- __embedding__:cond0: <class 'torch.Tensor'> torch.Size([26, 60]) torch.float32 cuda:0
- __embedding__:cond1: <class 'torch.Tensor'> torch.Size([26, 60]) torch.float32 cuda:0
- __embedding__:delta: <class 'torch.Tensor'> torch.Size([26, 60]) torch.float32 cuda:0
- __embedding_only__:cov_scale: <class 'torch.Tensor'> torch.Size([52]) torch.float64 cuda:0
- __shared_data__: <class 'torch.Tensor'> torch.Size([52, 1]) torch.float64 cuda:0
- __shared_data__:cond0: <class 'torch.Tensor'> torch.Size([26, 1]) torch.float64 cuda:0
- __shared_data__:cond1: <class 'torch.Tensor'> torch.Size([26, 1]) torch.float64 cuda:0
- __shared_data__:delta: <class 'torch.Tensor'> torch.Size([26, 1]) torch.float64 cuda:0
- __ypred__:cond1: <class 'torch.Tensor'> torch.Size([8, 26, 16352]) torch.float32 cuda:0
- __ypred__:peak:cond1: <class 'torch.Tensor'> torch.Size([26, 1793]) torch.float32 cuda:0
- condition_pairs: <class 'numpy.ndarray'> (26, 2) <U8
- peak: <class 'pandas.core.frame.DataFrame'> (1793, 7)
- pseudobulk_ids: <class 'numpy.ndarray'> (52,) <U25
- region: <class 'numpy.ndarray'> (8,) <U25
- region_name: <class 'numpy.ndarray'> (8,) <U25
==========

Total time: 8.36s
Inference time: 7.97s or 7.966s per batch
Callback time: 0.21s or 0.209s per batch
(total 1 batches)

==========
Saved Batch Schema:
- __embedding__: <class 'numpy.ndarray'> (52, 60) float32
- __ypred__:peak:cond1: <class 'numpy.ndarray'> (26, 1793) float32
- condition_pairs: <class 'numpy.ndarray'> (26, 2) <U8
- peak: <class 'pandas.core.frame.DataFrame'> (1793, 7)
- pseudobulk_ids: <class 'numpy.ndarray'> (52,) <U25
- region: <class 'numpy.ndarray'> (8,) <U25
- region_name: <class 'numpy.ndarray'> (8,) <U25
==========

==========
Final Stats Batch Schema:
- condition_pairs: <class 'numpy.ndarray'> (26, 2) <U8
- peak: <class 'pandas.core.frame.DataFrame'> (1793, 7)
- pseudobulk_ids: <class 'numpy.ndarray'> (52,) <U25
- region: <class 'numpy.ndarray'> (8,) <U25
==========

Removed temporary files in /tmp/tmp38istnsn
1 batches in /tmp/bolero-docs/bolero/docs/tutorials/score_model/predictions/Bolero10M+AP1+260104:HumanBrainDev
True peak data not found in batch, skip true peak data saving.
wrote predictions to predictions/Bolero10M+AP1+260104:HumanBrainDev

What you get¶

prediction_task writes an output_dir/ with the per-batch predictions plus aggregated files: a per-base 32 bp accessibility track for each region × cell-state, a per-peak accessibility value, and summary_stats. Rows are the cell states you supplied (their pseudobulk ids); each was predicted with its own AP-1 score baked in.

In [9]:

Copied!





peak_pred = pd.read_feather(f"{output_dir}/pred_peak_data.feather")
print(peak_pred.shape, "(peaks x cell states)")

# the same cell state at low vs. high AP-1 score
low = peak_pred.filter(like=":0").rename(columns=lambda c: c[:-2])
high = peak_pred.filter(like=":1").rename(columns=lambda c: c[:-2])
print("mean predicted accessibility per peak, low vs. high AP-1:")
pd.DataFrame({"low AP-1": low.mean(), "high AP-1": high.mean()}).round(3)
peak_pred = pd.read_feather(f"{output_dir}/pred_peak_data.feather")
print(peak_pred.shape, "(peaks x cell states)")

# the same cell state at low vs. high AP-1 score
low = peak_pred.filter(like=":0").rename(columns=lambda c: c[:-2])
high = peak_pred.filter(like=":1").rename(columns=lambda c: c[:-2])
print("mean predicted accessibility per peak, low vs. high AP-1:")
pd.DataFrame({"low AP-1": low.mean(), "high AP-1": high.mean()}).round(3)

(1793, 26) (peaks x cell states)
mean predicted accessibility per peak, low vs. high AP-1:

Out[9]:

	low AP-1	high AP-1
EN-L2_3-IT+PreBirth	14.490000	12.332
EN-L4-IT+PreBirth	12.989000	11.901
EN-L5-ET+PreBirth	14.439000	13.177
EN-L5-IT+PreBirth	13.129000	12.477
EN-L5_6-NP+PreBirth	14.153000	13.277
EN-L6-CT+PreBirth	15.056000	14.292
EN-L6-IT+PreBirth	12.996000	12.347
EN-L6b+PreBirth	15.552000	15.663
IN-CGE-SNCG+PreBirth	16.236000	15.722
IN-CGE-VIP+PreBirth	16.759001	16.657
IN-MGE-PV+PreBirth	16.584999	16.740
IN-MGE-SST+PreBirth	15.722000	16.403
IN-Mix-LAMP5+PreBirth	16.475000	15.863

That is the whole contract: supply regions + score-carrying cell states, get score-aware accessibility back. The two score columns above already show the idea: predict the same regions and states while changing __shared_data__ and compare. Measuring a TF's sequence effect that way, and benchmarking the predictions against held-out activity data, is beyond this I/O walkthrough.