cherimoya.wrappers

Thin torch.nn.Module wrappers around a Cherimoya that expose a single tensor from its (profile, log-count) output, so the model can be dropped straight into attribution and design tools that expect a one-tensor forward pass, plus ControlWrapper, which supplies a zero control track to models that expect one. All four are re-exported from the top-level cherimoya namespace.

ControlWrapper

class cherimoya.wrappers.ControlWrapper(*args, **kwargs)

Bases: Module

A wrapper that supplies an all-zero control track when none is given.

Cherimoya models trained with control tracks expect a control tensor at every forward pass, but attribution and marginalization tools call the model with the sequence alone. This wrapper bridges that gap: when no control is passed it checks whether the model expects one and, if so, synthesizes a control track of all zeroes with the matching shape, dtype, and device. Models without control tracks are forwarded through unchanged. The wrapper returns the model’s full (profile, log-count) output, so it is meant to be the inner wrapper that ProfileWrapper, LogCountWrapper, or ExpectedCountsWrapper are layered on top of. This is a port of bpnetlite.bpnet.ControlWrapper so that Cherimoya does not depend on bpnet-lite for this behavior.

Parameters:

model (cherimoya.Cherimoya) – A Cherimoya model, which makes predictions for basepair resolution profiles and also for log counts.

Constructor

__init__(model)

forward(X, X_ctl=None)

Returns the model’s full (profile, log-count) output, synthesizing an all-zero control track when the model expects one but none is passed. It is the inner wrapper that the output wrappers below are layered on top of — for example LogCountWrapper(ControlWrapper(model)) — and is what cherimoya attribute and cherimoya marginalize use so a model can be called with the sequence alone. A drop-in port of bpnetlite.bpnet.ControlWrapper.

ProfileWrapper

class cherimoya.wrappers.ProfileWrapper(*args, **kwargs)

Bases: Module

A wrapper that returns the weighted-softmax of the profile logits.

This wrapper takes the predicted profile logits and returns the dot product between them and their softmaxed values, summed across positions. The mean-centering and softmax weighting collapse the per-position profile into a single number per example whose attribution reflects the predicted profile shape. This is a port of bpnetlite.bpnet.ProfileWrapper so that Cherimoya does not depend on bpnet-lite for attribution.

Parameters:

model (cherimoya.Cherimoya) – A Cherimoya model, which makes predictions for basepair resolution profiles and also for log counts.

Constructor

__init__(model)

forward(X, X_ctl=None)

Returns the mean-centered profile logits weighted by their own softmax and summed across positions, collapsing the predicted profile into a single shape-sensitive number per example. This is the wrapper used by cherimoya attribute when output is "profile" (see Attribution and Motif Analysis). It is a drop-in port of bpnetlite.bpnet.ProfileWrapper so that attribution does not require bpnet-lite. As with LogCountWrapper, pair it with ControlWrapper for models trained with control tracks.

LogCountWrapper

class cherimoya.wrappers.LogCountWrapper(*args, **kwargs)

Bases: Module

A wrapper that extracts the log count predictions.

This wraps a Cherimoya and slices out the second prediction, which is for the log counts. This is useful when you only care about the log count predictions, such as for feature attribution or design methods.

Parameters:

model (cherimoya.Cherimoya) – A Cherimoya model, which makes predictions for basepair resolution profiles and also for log counts.

Constructor

__init__(model)

forward(X, X_ctl=None)

Returns the model’s per-group log-count predictions. This is the wrapper used by cherimoya attribute when output is "counts" (see Attribution and Motif Analysis). Pair it with ControlWrapper when attributing a model that was trained with control tracks, so that zero controls are supplied automatically.

ExpectedCountsWrapper

class cherimoya.wrappers.ExpectedCountsWrapper(*args, **kwargs)

Bases: Module

A wrapper that provides the expected counts per basepair.

This wrapper combines the profile predictions and the log count predictions to give the expected number of reads mapping to each position. This is done by exponentiating the log count predictions and multiplying them by the softmaxed logit profiles. Essentially, we are distributing counts (not log counts) by the predicted probability distribution across positions.

The distribution is performed jointly within each signal group. A group’s profile channels and positions are softmaxed together so that the probabilities sum to one across the entire group, and the group’s counts are then spread across all of its channels and positions. For a stranded (+, -) pair this means the expected counts summed over both strands and all positions equals the predicted count for that group. The count head is trained against log(count + 1), so torch.expm1() is used to recover the counts.

Parameters:

model (cherimoya.Cherimoya) – A Cherimoya model, which makes predictions for basepair resolution profiles and also for log counts.

Constructor

__init__(model)

forward(X, X_ctl=None)

Combines the profile and log-count heads into the expected number of reads at each position. Within each signal group the profile channels and positions are softmaxed jointly and scaled by expm1 of that group’s log-count, so the expected counts summed over the whole group (e.g. both strands of a stranded pair) equal the predicted count for the group.