utils.utils

Dataset, encoding utilities, and learning-rate scheduling helpers for BRIDGE.

This module groups together several utilities used across BRIDGE training/inference:

1. sequence preprocessing / encoding,

2. multimodal Dataset wrappers,

3. stratified train/test splitting for aligned modalities, and

4. a warmup learning-rate scheduler compatible with both step-based schedulers and ReduceLROnPlateau.

Provided functionality

Sequence / feature utilities

• seq2kmer(): Convert a raw sequence into overlapping k-mers (stride 1). (Note: current implementation returns a list of k-mers, not a single space-delimited string.)

• convert_one_hot(): Standard A/C/G/U(T) one-hot encoding with optional centered zero-padding.

• convert_one_hot2(): Attention-weighted one-hot encoding (writes attention[i] instead of 1.0).

Dataset splitting

• split_dataset(): Stratified split by a binary threshold on targets (targets < 0.5 vs >= 0.5), applied consistently across multiple aligned modalities.

Model inspection

• param_num(): Print total/trainable parameter counts for a PyTorch model.

Dataset wrappers (multimodal, BRIDGE-style)

• BaseRBPDataset: A generic dataset that returns an ordered tuple of tensors defined by the subclass attribute modalities.

• RBPTrainDataset: Training/validation dataset returning 6 items including labels.

• RBPInferDataset: Inference dataset returning 5 items without labels.

• myDataset, myDataset2: Legacy dataset wrappers kept for backward compatibility.

Tabular dataset readers

• read_csv(): Read a TSV file and return (sequences, structs, targets).

• read_csv_with_name(): Read a TSV file and also return record identifiers.

Learning-rate scheduling

• GradualWarmupScheduler: Warm up learning rate from base_lr to base_lr * multiplier over a fixed number of epochs, then delegate to another scheduler (including ReduceLROnPlateau).

Dynamic model name helper

• resolve_dynamic_model_name(): Heuristic mapping between cell-line suffixes (e.g., swapping to HepG2 or K562) for “dynamic” prediction settings.

BRIDGE batch conventions (important)

BRIDGE training/inference code typically expects each sample to be represented by multiple aligned modalities, returned as a tuple by the Dataset. The standard order is:

Training / labeled evaluation (6-tuple)

(embedding, attn, struct, motif, plfold, label)

Inference only (5-tuple)

(embedding, attn, struct, motif, biochem)

This module provides both modern dataset classes that follow this contract:

• RBPTrainDataset: modalities = (“embedding”, “attn”, “struct”, “motif”, “plfold”, “label”)

• RBPInferDataset: modalities = (“embedding”, “attn”, “struct”, “motif”, “biochem”)

and legacy classes (myDataset, myDataset2) with fixed attribute names.

BaseRBPDataset design (singular key -> plural attribute)

BaseRBPDataset expects keyword tensors named exactly as in modalities. It stores each tensor as a pluralized attribute: key="embedding" becomes self.embeddings. Therefore, the keys you pass must be singular:

ds = RBPTrainDataset(
    embedding=...,  # becomes self.embeddings
    attn=...,       # becomes self.attns
    struct=...,     # becomes self.structs
    motif=...,      # becomes self.motifs
    plfold=...,     # becomes self.plfolds
    label=...,      # becomes self.labels
)

If you accidentally pass already-plural keys (e.g., embeddings=…), the dataset will create attributes like self.embeddingss and __getitem__ will fail.

File format expectations

read_csv / read_csv_with_name

These readers expect a TSV with at least 6 columns and a header-like row where column 0 equals the literal string “Type” (that row is dropped).

• Column 2: sequence string

• Column 3: structure string (raw string; downstream may parse it)

• Column 5: label/target (cast to float32 and reshaped to (N, 1))

read_csv_with_name() additionally returns column 1 as the record identifier.

Note: these readers do not validate sequence alphabet or structure length. If you need strict validation, validate upstream (e.g., in a FASTA/structure parser) or add checks here.

Encoding notes and caveats

convert_one_hot

• Encodes A/C/G/U(T) into 4 channels.

• Unknown characters (e.g., N) remain all-zeros unless handled elsewhere.

• Optional centered padding to max_length may be used to match fixed-length models.

convert_one_hot2

• Same channel mapping as convert_one_hot but writes attention weights.

• Assumes attention indexing is valid for the sequence length. If sequences vary in length, you likely need per-sequence attention vectors.

seq2kmer

• Returns a list of k-mers, not a joined string. If you need a tokenizer-friendly representation, join with spaces in a separate helper (or adjust this function).

• The current implementation contains unused randomness (rand1/rand2) and a commented line that would add random flank bases; as written, randomness does not affect output.

Train/test splitting behavior

split_dataset() performs a stratified split by class, defined as:

negatives: targets < 0.5 positives: targets >= 0.5

It permutes indices within each class and then concatenates positives first, then negatives. All modalities are split using the same indices to preserve alignment.

If you need reproducibility, set NumPy RNG seed before calling:

np.random.seed(…)

Warmup scheduler usage

GradualWarmupScheduler increases LR linearly for total_epoch epochs until it reaches base_lr * multiplier. After warmup, it delegates to after_scheduler.

If after_scheduler is ReduceLROnPlateau, you must pass metrics:

warm = GradualWarmupScheduler(optimizer, multiplier=10, total_epoch=5,
                              after_scheduler=ReduceLROnPlateau(optimizer))
for epoch in range(num_epochs):
    train(...)
    val_loss = ...
    warm.step(metrics=val_loss)   # required for ReduceLROnPlateau

Otherwise, typical schedulers work as usual:

warm.step(epoch)

Dynamic model name resolution caveat

resolve_dynamic_model_name() currently returns a modified name only when a recognized suffix is found. If no suffix matches, it returns None implicitly. If you want identity behavior, add a final return name.

Functions

convert_one_hot(sequence[, max_length])	Convert DNA/RNA sequences to one-hot encoding with optional centered zero-padding.
convert_one_hot2(sequence, attention[, ...])	Convert sequences into an attention-weighted one-hot representation.
param_num(model)	Print the total/trainable/non-trainable parameter counts of a PyTorch model.
read_csv(path)	Read a tab-separated dataset file into sequences, structures, and labels.
read_csv_with_name(path)	Read a tab-separated dataset file and also return record identifiers/names.
resolve_dynamic_model_name(name)	Resolve an alternate cell-line model name for dynamic prediction.
seq2kmer(seq, k)	Convert a nucleotide sequence into overlapping k-mers (stride 1).
split_dataset(data1, data2, data3, ...[, ...])	Stratified train/test split for multiple aligned modalities.

Classes

BaseRBPDataset(**modal_tensors)	Base class for multimodal RBP datasets.
GradualWarmupScheduler(*args, **kwargs)	Gradually warm up (increase) the learning rate, then delegate to another scheduler.
RBPInferDataset(**modal_tensors)	Dataset for inference without labels.
RBPTrainDataset(**modal_tensors)	Dataset for training/validation that includes labels.
myDataset(*args, **kwargs)	Legacy training dataset wrapper for multiple modalities + label.
myDataset2(*args, **kwargs)	Legacy inference dataset wrapper for multiple modalities without labels.

utils.utils.seq2kmer(seq, k)

Convert a nucleotide sequence into overlapping k-mers (stride 1).

Parameters:

• seq (str) – Raw nucleotide sequence (DNA/RNA), e.g., “ACGT…” or “AUGC…”.

• k (int) – k-mer length. Typically 1 <= k <= len(seq).

Returns:

list[str] – List of k-mer substrings of length k. Example: seq=”ACGT”, k=2 -> [“AC”, “CG”, “GT”].

utils.utils.split_dataset(data1, data2, data3, data_motif, data_plfold, targets, valid_frac=0.2)

Stratified train/test split for multiple aligned modalities.

This function splits samples into train and test sets by thresholding targets at 0.5:

negatives: targets < 0.5 positives: targets >= 0.5

and sampling approximately valid_frac from each class into the test set.

Parameters:

• data1, data2, data3, data_motif, data_plfold (np.ndarray) – Input arrays for different modalities. Each must share the same first dimension N.

• targets (np.ndarray) – Target array aligned with the modalities along the first axis. Shape (N,) or (N,1).

• valid_frac (float, optional) – Fraction of each class assigned to the test split. Default: 0.2.

Returns:

tuple[list[np.ndarray], list[np.ndarray]] –

(train, test) where each is a list:

train = [X_train1, X_train2, X_train3, X_train4, X_train5, Y_train] test = [X_test1, X_test2, X_test3, X_test4, X_test5, Y_test]

Notes

• Indices are permuted independently within each class using np.random.permutation.

• The returned order concatenates positives first, then negatives (as implemented).

utils.utils.param_num(model)

Print the total/trainable/non-trainable parameter counts of a PyTorch model.

class utils.utils.BaseRBPDataset(**modal_tensors)

Bases: Dataset

Base class for multimodal RBP datasets.

Subclasses define modalities as an ordered tuple of modality names (singular). The constructor expects keyword tensors whose keys match these modality names, and stores them as pluralized attributes (e.g., key=”embedding” -> self.embeddings).

modalities

Names and order of fields returned by __getitem__.

Type:

Tuple[str, …]

_length

Number of samples (N).

Type:

int

Example

>>> class RBPTrainDataset(BaseRBPDataset):
...     modalities = ("embedding", "attn", "struct", "motif", "plfold", "label")
...
>>> train_ds = RBPTrainDataset(
...     embedding=torch.randn(N, C, L),
...     attn=torch.randn(N, L, L),
...     struct=torch.randn(N, 1, L),
...     motif=torch.randn(N, M),
...     plfold=torch.randn(N, P),
...     label=torch.randint(0, 2, (N, 1)).float(),
... )
>>> batch = train_ds[0]
>>> len(batch)
6

>>> class RBPInferDataset(BaseRBPDataset):
...     modalities = ("embedding", "attn", "struct", "motif", "biochem")
...
>>> infer_ds = RBPInferDataset(
...     embedding=torch.randn(N, C, L),
...     attn=torch.randn(N, L, L),
...     struct=torch.randn(N, 1, L),
...     motif=torch.randn(N, M),
...     biochem=torch.randn(N, B),
... )
>>> len(infer_ds[0])
5

modalities: Tuple[str, ...] = ()

class utils.utils.RBPTrainDataset(**modal_tensors)

Bases: BaseRBPDataset

Dataset for training/validation that includes labels.

modalities: Tuple[str, ...] = ('embedding', 'attn', 'struct', 'motif', 'plfold', 'label')

class utils.utils.RBPInferDataset(**modal_tensors)

Bases: BaseRBPDataset

Dataset for inference without labels.

modalities: Tuple[str, ...] = ('embedding', 'attn', 'struct', 'motif', 'biochem')

class utils.utils.myDataset(*args: Any, **kwargs: Any)

Bases: Dataset

Legacy training dataset wrapper for multiple modalities + label.

Parameters:

• bert_embedding (np.ndarray or torch.Tensor) – Token/channel embeddings per sample.

• attn (np.ndarray or torch.Tensor) – Attention matrices per sample.

• structure (np.ndarray or torch.Tensor) – Structure features per sample.

• motif (np.ndarray or torch.Tensor) – Motif features per sample.

• plfold (np.ndarray or torch.Tensor) – RNAplfold/biochemical-descriptors per sample.

• label (np.ndarray or torch.Tensor) – Labels per sample.

Returns:

tuple – (embedding, attn, struct, motif, plfold, label) for the given index.

Notes

• This class is functionally similar to RBPTrainDataset but uses fixed attribute names.

• Prefer RBPTrainDataset for clearer modality control and validation.

class utils.utils.myDataset2(*args: Any, **kwargs: Any)

Bases: Dataset

Legacy inference dataset wrapper for multiple modalities without labels.

Parameters:

• bert_embedding, attn, structure, motif – Same semantics as myDataset.

• phys_chem (np.ndarray or torch.Tensor) – Biochemical features per sample.

Returns:

tuple – (embedding, attn, struct, motif, phys_chem) for the given index.

utils.utils.read_csv(path)

Read a tab-separated dataset file into sequences, structures, and labels.

The input file is expected to be TSV with at least 6 columns, where:

col 0: Type (string; header row uses “Type”) col 2: Seq (sequence string) col 3: Str (structure string) col 5: label (0/1 or numeric)

Parameters:

path (str) – Path to the TSV file.

Returns:

tuple[np.ndarray, np.ndarray, np.ndarray] –

sequences:

Array of sequence strings, shape (N,), dtype object.

structs:

Array of structure strings, shape (N,), dtype object.

targets:

Float32 label array, shape (N, 1).

Notes

• Rows with df[0] == “Type” are dropped (header-like row).

• No validation of sequence alphabet or structure length is performed here.

utils.utils.read_csv_with_name(path)

Read a tab-separated dataset file and also return record identifiers/names.

Expected columns (TSV):

col 1: loc/name/identifier col 2: Seq col 3: Str col 5: label

Parameters:

path (str) – Path to the TSV file.

Returns:

tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray] –

name:

Array of identifiers (col 1), shape (N,).

sequences:

Array of sequence strings, shape (N,).

structs:

Array of structure strings, shape (N,).

targets:

Float32 label array, shape (N, 1).

utils.utils.convert_one_hot(sequence, max_length=None)

Convert DNA/RNA sequences to one-hot encoding with optional centered zero-padding.

Encoding:

Channel order: A, C, G, U/T - ‘A’ -> channel 0 - ‘C’ -> channel 1 - ‘G’ -> channel 2 - ‘U’ or ‘T’ -> channel 3

Parameters:

• sequence (Sequence[str]) – List/array of sequence strings. Characters are uppercased internally.

• max_length (int, optional) – If provided, sequences are padded (centered) to max_length with zeros.

Returns:

np.ndarray – Array of shape (N, 4, L) if max_length is None, otherwise (N, 4, max_length). dtype is float64 by default.

Notes

• Non-ACGU/T characters are left as all-zeros at that position.

• If you need explicit handling of ‘N’, add it upstream.

utils.utils.convert_one_hot2(sequence, attention, max_length=None)

Convert sequences into an attention-weighted one-hot representation.

Instead of writing 1.0 at nucleotide positions, this function writes attention[i] into the nucleotide channel at position i.

Parameters:

• sequence (Sequence[str]) – Sequence strings.

• attention (Sequence[float] or np.ndarray) – Per-position weights. This implementation uses attention[i] where i indexes the position in the sequence.

• max_length (int, optional) – If provided, output is padded (centered) to this length with zeros.

Returns:

np.ndarray – Array of shape (N, 4, L) or (N, 4, max_length) depending on padding.

Notes

• This function assumes attention is compatible with each sequence length. If sequences have different lengths, a single shared attention vector may not work.

class utils.utils.GradualWarmupScheduler(*args: Any, **kwargs: Any)

Bases: _LRScheduler

Gradually warm up (increase) the learning rate, then delegate to another scheduler.

During warmup (epochs 1..total_epoch):

lr = base_lr * ( (multiplier - 1) * epoch/total_epoch + 1 )

After warmup:

• If after_scheduler is provided, it is used for subsequent scheduling.

• If after_scheduler is ReduceLROnPlateau, use step(metrics=…).

Parameters:

• optimizer (torch.optim.Optimizer) – Wrapped optimizer.

• multiplier (float) – Target learning rate multiplier. Final warmup LR is base_lr * multiplier. Must be > 1.0.

• total_epoch (int) – Number of warmup epochs. Target LR is reached at total_epoch.

• after_scheduler (Optional[_LRScheduler or ReduceLROnPlateau]) – Scheduler used after warmup.

Raises:

ValueError – If multiplier <= 1.0.

get_lr()

Compute the learning rate(s) for the current epoch.

Returns:

List[float] – A list of learning rates, one per parameter group in the wrapped optimizer.

step_ReduceLROnPlateau(metrics, epoch=None)

Step function for the special case where after_scheduler is ReduceLROnPlateau.

Parameters:

• metrics (float) – Monitored metric value (e.g., validation loss) used by ReduceLROnPlateau.

• epoch (int, optional) – Epoch index. If None, uses self.last_epoch + 1.

Returns:

None.

Notes

• ReduceLROnPlateau is typically called at the end of an epoch, whereas most schedulers are called at the beginning. This method mirrors common warmup implementations:

  • During warmup, it manually sets optimizer.param_groups[*][‘lr’].

  • After warmup, it calls after_scheduler.step(metrics, epoch - total_epoch).

step(epoch=None, metrics=None)

Advance the scheduler by one step.

Parameters:

• epoch (int, optional) – Epoch index to step to. If None, scheduler increments internally.

• metrics (float, optional) – Metric value required when after_scheduler is ReduceLROnPlateau. Ignored for most other schedulers.

Returns:

None.

Notes

• If after_scheduler is not ReduceLROnPlateau:

  • During warmup, this delegates to _LRScheduler.step.

  • After warmup, this delegates to after_scheduler.step, shifting epochs by total_epoch.

• If after_scheduler is ReduceLROnPlateau:

  • This calls step_ReduceLROnPlateau(metrics, epoch).

utils.utils.resolve_dynamic_model_name(name)

Resolve an alternate cell-line model name for dynamic prediction.

Heuristic behavior:

• If name ends with one of: {“K562”,”HEK293”,”HEK293T”,”Hela”,”H9”}, replace that suffix with “HepG2”.

• If name ends with “HepG2”, replace it with “K562”.

Parameters:

name (str) – Model or experiment identifier string, typically with a cell-line suffix.

Return type:

str

Returns:

str – The modified name if a known suffix is found.

Notes

• If name does not match any recognized suffix, the current implementation returns None implicitly (because there is no final return name). If you want identity behavior, add return name at the end.