Crop, orient and average genomic predictions

Your task

Function signature

aligned_ensemble_predictions(
    predictions: torch.Tensor,
    valid_mask: torch.Tensor,
    start: torch.Tensor,
    reverse: torch.Tensor,
    track_permutation: torch.Tensor,
    ensemble_weights: torch.Tensor,
    options: dict,
) -> dict[str, torch.Tensor]

Worked example

Reverse, shift, then average

import torch

predictions = torch.tensor([[[[1.0], [2.0], [3.0]]], [[[10.0], [20.0], [30.0]]]], dtype=torch.float64, requires_grad=True)
valid_mask = torch.tensor([[[True, True, True]], [[True, True, True]]], dtype=torch.bool)
start = torch.tensor([0, 1], dtype=torch.int64)
reverse = torch.tensor([False, True], dtype=torch.bool)
track_permutation = torch.tensor([0], dtype=torch.int64)
ensemble_weights = torch.tensor([1.0, 3.0], dtype=torch.float64, requires_grad=True)
options = {
    'target_length': 4,
}

result = aligned_ensemble_predictions(predictions=predictions, valid_mask=valid_mask, start=start, reverse=reverse, track_permutation=track_permutation, ensemble_weights=ensemble_weights, options=options)

Expected return

{
    'y': torch.tensor([[[1.0], [23.0], [15.75], [10.0]]], dtype=torch.float64),
    'aligned': torch.tensor([[[[1.0], [2.0], [3.0], [0.0]]], [[[0.0], [30.0], [20.0], [10.0]]]], dtype=torch.float64),
    'aligned_valid_mask': torch.tensor([[[True, True, True, False]], [[False, True, True, True]]], dtype=torch.bool),
    'weight_sum': torch.tensor([[1.0, 4.0, 4.0, 3.0]], dtype=torch.float64),
    'coordinates': torch.tensor([0, 1, 2, 3], dtype=torch.int64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• predictions [E,B,L,C] contains E member windows, each with L\ge 1. start [E] is the leftmost coordinate covered by that window on the forward-oriented output-bin axis, even when the member is reversed.
• Requested coordinates are \mathrm{target\_start} + \operatorname{arange}(\mathrm{target\_length}), default \mathrm{start}=0 and \mathrm{length}=L.
• For each member and target coordinate g, p=g-\mathrm{start}[e]; read original index p for forward members or L-1-p for reverse[e]=True. Exclude p outside [0,L) or an invalid source position.
• Reversed members additionally gather channels by track_permutation [C]; use the identity for SpliceAI because donor/acceptor classes do not exchange.
• Return aligned [E,B,T,C], aligned_valid_mask, and a position-wise weighted mean y [B,T,C] using nonnegative ensemble_weights [E] and a denominator containing only valid members. A zero denominator returns exactly zero; also return weight_sum [B,T] and integer coordinates [T].
• For a center crop of T\le L bins use every \mathrm{start}=-\operatorname{floor}((L-T)/2) and \mathrm{target\_start}=0; an odd excess discards \operatorname{floor}(\mathrm{excess}/2) on the left and the remainder on the right.
• Coordinates here are output bins, not implicit base pairs; genomic bin size and external genome origin are supplied by the caller's coordinate conversion.

Why this matters

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def aligned_ensemble_predictions(
    predictions: torch.Tensor,  # (E, B, L, C)
    valid_mask: torch.Tensor,  # (E, B, L)
    start: torch.Tensor,  # (E,)
    reverse: torch.Tensor,  # (E,)
    track_permutation: torch.Tensor,  # (C,)
    ensemble_weights: torch.Tensor,  # (E,)
    options: dict,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    y: (B, T, C).
    aligned: (E, B, T, C).
    aligned_valid_mask: (E, B, T).
    weight_sum: (B, T).
    coordinates: (T,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement aligned_ensemble_predictions")

Sources and further reading

• doi.org
• Illumina/SpliceAI · utils.py
• google-deepmind/deepmind-research · enformer.py
• www.nature.com