Prediction & Variant Effect

The inference engine and its high-level task methods, the variant-effect managers, and the directed-evolution / enhancer-design engine.

BorzoiPredictor owns a trained model plus a data manager and exposes the task methods that map to the paper (prediction_task, caqtl_task, eqtl_task, peak_task, attribution_task, …). BorzoiSignalPredictor is the signal/flow variant used across the prediction tutorials.

Predictors

BorzoiPredictor

Bases: GenericPredictor

Inference engine for a trained Bolero (BorzoiLoRA) model.

Owns a model plus a :class:~bolero.tl.predict.datamanager.GenericGenomeDataManager and exposes high-level task methods that map to the paper's analyses: :meth:prediction_task / :meth:inference_task (accessibility and gene-count tracks), :meth:caqtl_task / :meth:qtl_task / :meth:eqtl_task (variant effect), :meth:peak_task (per-peak aggregation) and :meth:attribution_task (per-base DNA attribution). The data manager assembles the cell-state conditioning and reference signal for each region; :class:BorzoiSignalPredictor extends this with the signal / velocity modes.

Parameters:

Name	Type	Description	Default
config		Predictor configuration dict (validated against default_config); see the prediction tutorials for the expected keys.	required

register_qtl_manager

register_qtl_manager(qtl_table, qtl_type, channel_weights_path=None)

Register the QTL manager with the given QTL table.

This is required in qtl task

register_peak_manager

register_peak_manager(peak_table)

Register the peak manager with the given peak table.

This is required in peak task

iter_debug_batch

iter_debug_batch(dna_key='dna', embedding_key='embedding', pseudobulk_ids=None, add_true_data=True, batch_size=2, mode='prediction', regions=None, trigger_model=False)

Iterable for debugging

get_prediction_dataloader

get_prediction_dataloader(regions, batch_dir, pseudobulk_ids=None, add_true_data=False, dna_key='dna', embedding_key='embedding', batch_size=32, verbose=True, mode='prediction', _dataloader_batch_size=None) -> Generator

Get the dataloader for prediction.

get_dna_model

get_dna_model(mode='prediction', pseudobulk_id=None, attr_kwargs=None, **emb_data)

Get the collapsed DNA only model for DNA experiments. Either provide **emb_data or a pseudobulk id to collapse the model.

Parameters:

Name	Type	Description	Default
mode	str	The mode to collapse the model. Default is 'prediction'. Supported modes: 'prediction', 'attribution', 'gene_count_attribution'.	'prediction'
pseudobulk_id	str	The pseudobulk id to collapse the model. If None, **emb_data should be provided.	None
attr_kwargs	dict	The keyword arguments to pass to the attribution model. Default is None.	None
**emb_data	dict	The embedding data to collapse the model. Default is None.	{}

Returns:

Name	Type	Description
model	Module	The collapsed DNA input only model.

get_attribution_dataloader

get_attribution_dataloader(regions_per_pseudobulk, pseudobulk_ids, dna_key='dna', embedding_key='embedding', batch_size=6, verbose=True, mode='attribution', qtl_mutations=False) -> Generator

Get the dataloader for attribution.

The main difference on data loader side is that prediction dataloader iterates all region and all pseudobulk together; attribution dataloader iterate all regions for one pseudobulk at a time. This is because 1. predition task fetch parquet only once by putting all pseudobulk together, attribution task will not use parquet data. 2. attribution task needs to collapse lora model into base, this can only be done one pseudobulk at a time.

prediction_task

prediction_task(output_dir: str, regions: str | DataFrame | list[str] = 'test_regions', downsample_regions: int | None = None, downsample_seed: int = 0, pseudobulk_ids: list[str] | None = None, batch_size: int = 16, save_keys: str | list[str] | None = 'default', stats_keys: list[str] | None = None, verbose: bool = True, save_first_batch: bool = False, mode: str = 'prediction', filter_valid_regions: bool = True, _dataloader_batch_size: int | None = None) -> None

Prediction task for Borzoi. Compute the prediction on a set of regions and pseudobulk records. Then compute the stats and save them to a file.

Parameters:

Name	Type	Description	Default
output_dir	str	The output directory to save the results.	required
regions	str | DataFrame | list[str]	The regions to predict. If "test_regions", use the borzoi test regions based on fold in config.	'test_regions'
downsample_regions	int | None	The number of regions to downsample. If None, use all regions.	None
downsample_seed	int	The seed for downsampling.	0
pseudobulk_ids	list[str] | None	The pseudobulk ids to use. If None, use all pseudobulk ids.	None
batch_size	int	The batch size for prediction.	16
save_keys	str | list[str] | None	The keys to save in the output batch file. If None, nothing will be saved.	'default'
verbose	bool	Whether to print the progress.	True
save_first_batch	bool	Whether to save the first full batch for debugging purposes.	False
mode	str	The mode of the task. One of "prediction", "gene_count_prediction". If "prediction", predict the genome tracks. If "gene_count_prediction", predict the gene counts along with genome tracks.	'prediction'
filter_valid_regions	bool	Whether to filter the regions to valid regions.	True
_dataloader_batch_size	int | None	The batch size for the dataloader. If None, use the batch size.	None

inference_task

inference_task(output_dir: str, regions: str | DataFrame | list[str], downsample_regions: int | None = None, downsample_seed: int = 0, pseudobulk_ids: list[str] | None = None, batch_size: int = 16, save_keys: str | list[str] | None = 'default', verbose: bool = True, save_first_batch: bool = False, mode: str = 'prediction', _dataloader_batch_size: int | None = None) -> None

Inference task for Borzoi. Compute the prediction on a set of regions and pseudobulk records. Then compute the stats and save them to a file.

Parameters:

Name	Type	Description	Default
output_dir	str	The output directory to save the results.	required
regions	str | DataFrame | list[str]	The regions to predict.	required
downsample_regions	int | None	The number of regions to downsample. If None, use all regions.	None
downsample_seed	int	The seed for downsampling.	0
pseudobulk_ids	list[str] | None	The pseudobulk ids to use. If None, use all pseudobulk ids.	None
batch_size	int	The batch size for prediction.	16
save_keys	str | list[str] | None	The keys to save in the output batch file. If None, nothing will be saved.	'default'
verbose	bool	Whether to print the progress.	True
save_first_batch	bool	Whether to save the first full batch for debugging purposes.	False
mode	str	The mode of the task. One of "prediction", "gene_count_prediction". If "prediction", predict the genome tracks. If "gene_count_prediction", predict the gene counts along with genome tracks.	'prediction'
_dataloader_batch_size	int | None	The batch size for the dataloader. If None, use the batch size.	None

select_top_std_genes

select_top_std_genes(gene_data_path, top_n=5000)

Select top variable genes, intersect with test fold.

caqtl_task

caqtl_task(output_dir: str, qtl_table: str | DataFrame, pseudobulk_ids: list[str] | None = None, batch_size: int = 16, save_keys: str | list[str] | None = 'default', add_true_data: bool = False, verbose: bool = True, save_first_batch: bool = False, qtl_type: str = 'caqtl', channel_weights_path: str | None = None, _dataloader_batch_size: int | None = None) -> None

QTL Prediction task for Borzoi. Compute the ref and alt prediction of a QTL dataset Then compute the stats and save them to a file.

Parameters:

Name	Type	Description	Default
output_dir	str	The output directory to save the results.	required
qtl_table	str | DataFrame	The QTL table path to use. QTL table should contain inference regions, qtl mutations and qtl peaks.	required
pseudobulk_ids	list[str] | None	The pseudobulk ids to use. If None, use all pseudobulk ids.	None
batch_size	int	The batch size for prediction.	16
save_keys	str | list[str] | None	The keys to save in the output file. If None, save all keys.	'default'
verbose	bool	Whether to print the progress.	True
save_first_batch	bool	Whether to save the first complete batch into output dir.	False

qtl_task

qtl_task(*args, **kwargs)

Alias for caqtl_task

eqtl_task

eqtl_task(output_dir: str, qtl_table: str | DataFrame, pseudobulk_ids: list[str] | None = None, batch_size: int = 16, save_keys: str | list[str] | None = 'default', add_true_data: bool = False, verbose: bool = True, save_first_batch: bool = False, _dataloader_batch_size: int | None = None) -> None

eQTL task for Borzoi - predict effect of variants on gene expression.

Parameters:

Name	Type	Description	Default
output_dir	str	The output directory to save the results.	required
qtl_table	str or DataFrame	eQTL table with variant info, gene info, and regions	required
pseudobulk_ids	list[str] | None	The pseudobulk ids to use. If None, use all pseudobulk ids.	None
batch_size	int	The batch size for prediction.	16
save_keys	str | list[str] | None	The keys to save in the output file. If None, save all keys.	'default'
verbose	bool	Whether to print the progress.	True
add_true_data	bool	Whether to add true data to the batch.	False
save_first_batch	bool	Whether to save the first complete batch into output dir.	False

peak_task

peak_task(output_dir: str, peak_table: str | DataFrame, pseudobulk_ids: list[str] | None = None, batch_size: int = 16, save_keys: str | list[str] | None = 'default', verbose: bool = True, save_first_batch: bool = False, _dataloader_batch_size: int | None = None) -> None

Prediction task for Borzoi. Compute the prediction on a peak of interests.

Parameters:

Name	Type	Description	Default
output_dir	str	The output directory to save the results.	required
peak_table	str | DataFrame	The peak table path to use. Peak table should contain borzoi input regions and peak regions.	required
pseudobulk_ids	list[str] | None	The pseudobulk ids to use. If None, use all pseudobulk ids.	None
batch_size	int	The batch size for prediction.	16
save_keys	str | list[str] | None	The keys to save in the output file. If None, save all keys.	'default'
verbose	bool	Whether to print the progress.	True
save_first_batch	bool	Whether to save the first complete batch into output dir.	False
_dataloader_batch_size	int | None	The batch size for the dataloader. If None, use the batch size.	None

attribution_task

attribution_task(output_dir: str, regions_per_pseudobulk: dict | DataFrame | str, pseudobulk_ids: list[str] | None = None, batch_size: int = 6, save_keys: str | list[str] | None = 'default', verbose: bool = True, save_first_batch: bool = False, mode: str = 'attribution', qtl_table: str | DataFrame | None = None, qtl_type: str | None = None) -> None

Attribution task for Borzoi with optional QTL mutation analysis.

Parameters:

Name	Type	Description	Default
output_dir	str	The output directory to save the results.	required
regions_per_pseudobulk	str or DataFrame or dict	The regions to run attribution on. For gene_count_attribution mode with qtl_table, this should match the gene regions from the QTL table (will be auto-generated if None). For regular attribution, provide standard regions.	required
pseudobulk_ids	list[str]	The pseudobulk ids to use. If None, use all pseudobulk ids.	None
batch_size	int	The batch size for attribution computation.	6
save_keys	iterable of str	The keys to save in the output file. Can be a list, tuple, or any iterable of strings.	'default'
verbose	bool	Whether to print the progress.	True
save_first_batch	bool	Whether to save the first full batch for debugging purposes.	False
mode	str	The mode of the task. One of "attribution", "gene_count_attribution".	'attribution'
qtl_table	str or DataFrame	QTL table for mutation analysis. If provided, will compute attributions for both ref and alt alleles. The table should have columns: Chromosome, Start, End, Name (gene+variant), Strand, TSS, GeneStart, GeneEnd, MaskStart, MaskEnd, variant_id, Ref, Alt, gene_id, etc.	None
qtl_type	str	Type of QTL analysis ('eqtl' or 'caqtl'). Default None.	None

BorzoiSignalPredictor

Bases: BorzoiPairPredictor

BorzoiSignalPredictor is a predictor for Borzoi models that uses ODEs to predict the trajectory of the model given an initial condition. It is used for signal-based models.

get_evolution_dataloader

get_evolution_dataloader(batch_size, max_batches, dna_key, embedding_key, rank_pred_fn, top_k, forward_batch_size=16) -> Generator

Get the dataloader for directed DNA evolution.

Parameters:

Name	Type	Description	Default
batch_size		Number of regions to process in each batch.	required
max_batches		Number of batches to process.	required
dna_key		Key for the DNA one-hot encoding in the batch.	required
embedding_key		Key for the embedding in the batch.	required
rank_pred_fn		Function to rank the predicted values. This function should take a batch as input, calculate the rank based on the predicted values and return the rank. The rank should be a np.array of shape (n_region, ) with smaller values indicating better performance.	required
top_k		Number of top regions to select. selected based on rank < top_k.	required

Returns:

Name	Type	Description
dataloader	Generator	Generator that yields batches of DNA one-hot encoding.

Notes

This dataloader is used to perform directed DNA evolution. It will iterate over the batches and select the top k regions based on the predicted values. The selected regions' DNA one-hot encoding will be saved to the self.datamanager._onehot_encoder. The self.datamanager._onehot_encoder will be used to generate the next batch of DNA one-hot encoding.

evolution_task

evolution_task(output_dir: str, rank_pred_fn: callable, save_keys: list[str], top_k: int = 5, pseudobulk_ids: list[str] | None = None, batch_size: int = 32, max_batches: int = 50, dna_key: str = 'dna', embedding_key: str = 'embedding', n_experiments: int = 10) -> None

Evolution task for Borzoi. Perform directed DNA evolution.

Parameters:

Name	Type	Description	Default
output_dir	str	The output directory to save the results.	required
rank_pred_fn	callable	Function to rank the predicted values. This function should take a batch as input, calculate the rank based on the predicted values and return the rank. The rank should be a np.array of shape (n_region, ) with smaller values indicating better performance.	required
top_k	int	Number of top regions to select. selected based on rank < top_k.	5
pseudobulk_ids	list[str] | None	The pseudobulk ids to use. If None, use all pseudobulk ids.	None
batch_size	int	The batch size for evolution.	32
max_batches	int	Number of batches to process.	50
dna_key	str	Key for the DNA one-hot encoding in the batch.	'dna'
embedding_key	str	Key for the embedding in the batch.	'embedding'

attribution_task

attribution_task(output_dir: str, regions_per_pseudobulk: dict | DataFrame | str, pseudobulk_ids: list[str] | None = None, batch_size: int = 6, save_keys: str | list[str] | None = 'default', verbose: bool = True, mode: str = 'attribution', save_first_batch: bool = False, qtl_table: str | DataFrame | None = None, qtl_type: str | None = None) -> None

Attribution task for BorzoiSignalPredictor. Compute the attribution on a set of regions and pseudobulk records. Then save the results to a file.

qtl_task

qtl_task(output_dir: str, qtl_table: str | DataFrame, pseudobulk_ids: list[str] | None = None, batch_size: int = 16, save_keys: str | list[str] | None = 'default', verbose: bool = True) -> None

QTL task for BorzoiSignalPredictor.

TODO: since we do not use signal input for QTL task, we can disable add_true_data

this need later code assume x0 key can be missing

BorzoiInputXGradient

Per-base attribution for a collapsed, cell-state-specific Borzoi model.

Wraps model (typically a BorzoiLoRA collapsed onto one cell-state embedding) with Captum's :class:~captum.attr.InputXGradient. The forward hook sums the predicted signal over the centre peak_length bins, so the returned attribution is the input-times-gradient of that scalar peak score with respect to each input base, cropped to the centre attr_length bp.

Parameters:

Name	Type	Description	Default
model		A callable DNA(+signal) -> track model (e.g. a collapsed BorzoiLoRA).	required
peak_length		Width in bp of the centred region whose summed signal is attributed.	512
attr_length		Width in bp of the centred attribution window returned by __call__.	1024
model_dna_length		Input sequence length the model expects (524,288 bp).	524288
model_resolution		Output bin size in bp (32).	32

Variant-effect managers

caQTLManager

Bases: QTLMixIn

Prepare and score chromatin-accessibility QTLs (caQTLs).

Parses a caQTL table into Borzoi input regions, ref/alt substitutions, and the peak windows to aggregate over, then (via :class:QTLMixIn) builds ref/alt-mutated DNA batches and sums predicted accessibility over each variant's peak. Consumed by :meth:BorzoiPredictor.caqtl_task.

Parameters:

Name	Type	Description	Default
qtl_table	str	Path to the caQTL table (or a DataFrame accepted by :func:prepare_qtl_table).	required
resolution	int	Output bin size in bp. Default 32.	32
qtl_stats_cols	list[str]	Per-variant statistic columns to carry through. Default ["beta", "PIP"].	None
ypred_seq_len	int	Expected length of the prediction axis. Default 16384.	16384
channel_weights	optional	Optional per-channel weights applied when aggregating tracks.	None
qtl_type	str	Label for the QTL type. Default "caqtl".	'caqtl'

get_peak_sum

get_peak_sum(batch: dict, ypred_key)

Calculate the sum of predictions over the QTL peak regions for each region/mutation in the batch.

add_qtl_info

add_qtl_info(batch: dict)

Add peak information to the batch.

eQTLManager

Bases: QTLMixIn

Prepare and score expression QTLs (eQTLs).

Parses an eQTL table into gene/promoter regions (with strand), ref/alt substitutions, and target genes, then builds ref/alt-mutated DNA batches for the gene-count head. Consumed by :meth:BorzoiPredictor.eqtl_task.

mutate_dna

mutate_dna(batch: dict[str, Tensor], dna_key='__dna__', region_key='region_name')

Mutate dna for eQTL

add_qtl_info

add_qtl_info(batch: dict)

Add gene information to the batch.

PeakManager

Prepare and aggregate predictions over a peak table.

Parses a peak table into Borzoi input regions plus the peak windows to score, and sums predicted signal over each peak. Used by :meth:BorzoiPredictor.peak_task.

Parameters:

Name	Type	Description	Default
peak_table	str	Path to the peak table (or a DataFrame accepted by :func:prepare_peak_table), containing Borzoi input regions and the peak regions to aggregate over.	required
resolution	int	Output bin size in bp. Default 32.	32
ypred_seq_len	int	Expected length of the prediction axis. Default 16384.	16384

get_peak_sum

get_peak_sum(batch: dict, ypred_key)

Calculate the sum of predictions over the QTL peak regions for each region/mutation in the batch.

add_peak_info

add_peak_info(batch: dict)

Add peak information to the batch.

Directed evolution & enhancer design

DNASynthesisFactory

get_fasta_region_sequence

get_fasta_region_sequence(chrom: str, start: int, end: int, genome: str = None) -> str

Get the DNA sequence from the fasta or 2bit file for a given region.

get_random_sequence

get_random_sequence(n_regions: int, seq_len: int, probs: tuple[float, float, float, float] = (0.29, 0.21, 0.21, 0.29), random_state=None) -> torch.Tensor

Get random DNA sequence from the factory.

Parameters:

Name	Type	Description	Default
n_regions	int	Number of regions to generate.	required
seq_len	int	Length of the sequence.	required
probs	tuple[float, float, float, float]	Probabilities of the four bases.	(0.29, 0.21, 0.21, 0.29)

substitute_sequence

substitute_sequence(input_one_hot: Tensor, pos_0base: int, seq: str) -> torch.Tensor

Substitute the DNA sequence at the given position.

mutate_single_sequence

mutate_single_sequence(input_one_hot: Tensor, pos_0base: int, ref: str, alt: str, validate_ref: bool = True) -> torch.Tensor

Mutate the DNA sequence at the given position.

Parameters:

Name	Type	Description	Default
input_one_hot	Tensor	Input DNA sequence one-hot encoding.	required
pos_0base	int	Position to mutate. Position is 0-based.	required
ref	str	Reference sequence.	required
alt	str	Alternative sequence.	required
validate_ref	bool	Whether to validate the reference sequence at the given position.	True

shuffle_sequence

shuffle_sequence(input_one_hot: Tensor, start: int, end: int, n_shuffle: int = 1, random_state: int = None, use_dinucleotide_shuffle: bool = False) -> torch.Tensor

Shuffle the DNA sequence at the given position.

input_one_hot shape should be (n_regions, 4, seq_len). output shape should be (n_regions * n_shuffle, 4, seq_len).

Parameters:

Name	Type	Description	Default
input_one_hot	Tensor	Input DNA sequence one-hot encoding.	required
start	int	Start position of the sequence.	required
end	int	End position of the sequence.	required
n_shuffle	int	Number of shuffle to perform.	1
random_state	int	Random state.	None
use_dinucleotide_shuffle	bool	Whether to perform dinucleotide shuffle (True) or simple shuffle (False).	False

Returns:

Type	Description
Tensor	Shuffled sequence batch. Shape: (n_regions * n_shuffle, 4, seq_len).

decode

decode(one_hot: Tensor) -> list[str]

Decode the one-hot encoding to a DNA sequence.

genome_mutation_task

genome_mutation_task(region_and_mutation: DataFrame, genome: str = None, validate_ref: bool = False, decode: bool = False) -> torch.Tensor | list[str]

Generate sequence batch using a table of region and mutation information.

Parameters:

Name	Type	Description	Default
region_and_mutation	DataFrame	Table of region and mutation information. Columns: chromosome, start, end, mut_pos_0base, ref, alt, genome (optional if each region has a different genome)	required
genome	str	Genome name. If not provided, the default genome will be used.	None
validate_ref	bool	Whether to validate the reference sequence at the given position.	False

Returns:

Type	Description
Tensor	Sequence batch. Shape: (n_regions, 4, seq_len).

genome_task

genome_task(region: DataFrame, genome: str = None, decode: bool = False) -> torch.Tensor | list[str]

Generate sequence batch using a table of region information.

Parameters:

Name	Type	Description	Default
region	DataFrame	Table of region information.	required
genome	str	Genome name. If not provided, the default genome will be used.	None
decode	bool	Whether to decode the sequence to a DNA sequence.	False

Returns:

Type	Description
Tensor or list[str]	Sequence batch. Shape: (n_regions, 4, seq_len). If decode is True, return list[str] of DNA sequences.

region_string_to_dataframe

region_string_to_dataframe(regions: list[str], region_names: list[str], mode: str = None)

Convert a list of regions and region names to a dataframe.

Parameters:

Name	Type	Description	Default
regions	list[str]	List of regions.	required
region_names	list[str]	List of region names. Region names should be in the format "region_name__additional_info". "additional_info" can be empty, or a string of additional information with fields separated by ":".	required
mode	str	Mode of the task. If "genome", the dataframe will have columns: chromosome, start, end, and optionally genome. If "mutation", the dataframe will have columns: chromosome, start, end, mut_pos_0base, ref, alt, and optionally genome.	None

get_regions_onehot

get_regions_onehot(regions: list[str], region_names: list[str]) -> torch.Tensor

Get the one-hot encoding for the regions.

DNAEvolutionFactory

Beam-search DNA evolution: maintain a pool of parent sequences and produce mutated batches for scoring, then update the pool with the top-scoring sequences.

Workflow

1. Initialize with initial one-hot sequence(s) and evolution parameters.
2. Call :meth:get_regions_onehot to get a batch of mutated sequences (batch_size, 4, seq_len) for the predictor to score.
3. Select the top sequences by your criterion and call :meth:update_current_one_hot with that tensor.
4. Repeat from step 2 for the next evolution round.

The number of mutations per sequence is derived from the evolution window length and mutation_rate. Mutations are applied only within the evolution window [evolution_window_start, evolution_window_end).

Parameters:

Name	Type	Description	Default
input_sequence	str	Initial DNA sequence used as the evolution parent.	required
evolution_window_start	int	Start (inclusive) of the evolution window. Used only to compute the number of mutations per sequence.	required
evolution_window_end	int	End (exclusive) of the evolution window. Must be > evolution_window_start.	required
mutation_rate	float	Fraction of the window length to mutate per sequence, in (0, 1). Number of mutations = max(1, mutation_rate * (eend - estart)). Default is 0.01.	0.01
device	str	Torch device for the one-hot tensors. Default is "cuda".	'cuda'
back_mutation_rate	float	Probability of reverting a differing position to the reference base. Default is 0.2.	0.2

Attributes:

Name	Type	Description
seq_len	int	Sequence length.
n_mutations	int	Number of random mutation positions applied per sequence per round.

get_regions_onehot

get_regions_onehot(batch_size: int = 32, add_ref: bool = True) -> torch.Tensor

Generate a batch of mutated sequences from the current parent pool.

Each of the batch_size sequences is produced by taking a parent from the pool (cycled by index) and applying n_mutations random mutations. Parents are selected in round-robin when batch_size > pool size.

Parameters:

Name	Type	Description	Default
batch_size	int	Number of mutated sequences to return per call.	32
add_ref	bool	Whether to add the current sequence to the batch.	True

Returns:

Name	Type	Description
new_batch	Tensor	One-hot batch of shape (batch_size, 4, seq_len), same device/dtype as the current pool.

get_evolution_sequence

get_evolution_sequence(new_batch: Tensor) -> list[str]

Get the DNA sequence at the evolution window positions for the given batch.

Parameters:

Name	Type	Description	Default
new_batch	Tensor	One-hot batch of shape (batch_size, 4, seq_len), same device/dtype as the current pool.	required

Returns:

Name	Type	Description
current_evolution_sequence	list[str]	List of DNA sequences at the evolution window positions for the given batch. Shape: (batch_size,).

update_current_one_hot

update_current_one_hot(one_hot: Tensor) -> None

Set the current parent pool to the given sequences (e.g. top-k from the last batch after scoring).

Parameters:

Name	Type	Description	Default
one_hot	Tensor	One-hot sequences to use as parents for the next round. Shape (n, 4, seq_len) with n typically equal to n_top_regions. seq_len must match self.seq_len.	required

reset_current_one_hot

reset_current_one_hot() -> None

Reset the current one-hot to the input one-hot.

Evaluation utilities

multi_level_peak_stats

multi_level_peak_stats(output_dir: str, precomputed_region_group_path: str | None = None) -> None

Generate five quantile cutoffs based on true value across sample STD. For each cutoff, select corresponding peaks and calculate profile/sample pearson corr and R2 metrics.

Parameters:

Name	Type	Description	Default
output_dir	str	The output directory containing the true and predicted peak data.	required
precomputed_region_group_path	str	The path to a precomputed region group file, if provided, peak group in this file will be used.	None

Returns:

Type	Description
None	Results are written to {output_dir}/region_group_and_stats.joblib.