Project: s41467-025-60048-1 | Function: constraint_matrix_batch (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implements canonical base pairing constraints for RNA secondary structure prediction by computing valid pairing probabilities between nucleotides. Calculates Watson-Crick and wobble base pair matrices (AU, CG, GU) using tensor operations for batch processing.

**KEY TERMS:** {
  "x": "nucleotide sequence features",
  "base_a": "adenine probabilities",
  "base_u": "uracil probabilities",
  "base_c": "cytosine probabilities",
  "base_g": "guanine probabilities",
  "au_ua": "A-U/U-A base pair matrix"
}

**CONTEXT CODE:** import torch
import torch.nn.functional as F


def get_base_index():
    return {'A': 0, 'U': 1, 'C': 2, 'G': 3}


def constraint_matrix_batch(x, loop_min_len=2, is_nc=False):
    # MASKED: Implements canonical base pairing constraints for RNA secondary structure prediction by computing valid pairing probabilities between nucleotides implementation
    raise NotImplementedError("This function is masked for testing purposes.")


def contact_a(a_hat, m):
    a = a_hat * a_hat
    a = (a + torch.transpose(a, -1, -2)) / 2
    a = a * m
    return a


def sign(x):
    return (x > 0).type(x.dtype)


def soft_sign(x):
    k = 1
    return 1.0/(1.0+torch.exp(-2*k*x))


def apply_constraints(u, x, lr_min, lr_max, num_itr, rho=0.0, with_l1=False, s=2.2, is_nc=False):
    """
    :param u: utility matrix, u is assumed to be symmetric, in batch
    :param x: RNA sequence, in batch
    :param lr_min: learning rate for minimization step
    :param lr_max: learning rate for maximization step (for lagrangian multiplier)
    :param num_itr: number of iterations
    :param rho: sparsity coefficient
    :param with_l1:
    :return:
    """
    m = constraint_matrix_batch(x, is_nc=is_nc).float()
    # u with threshold
    # equivalent to sigmoid(u) > 1/(e^(-2.2)+1)
    # u = (u > 2.2).type(torch.FloatTensor) * u
    u = soft_sign(u - s) * u

    # initialization
    a_hat = (torch.sigmoid(u)) * soft_sign(u - s).detach()
    lmbd = F.relu(torch.sum(contact_a(a_hat, m), dim=-1) - 1).detach()

    # gradient descent
    for t in range(num_itr):
        grad_a = (lmbd * soft_sign(torch.sum(contact_a(a_hat, m), dim=-1) - 1)).unsqueeze_(-1).expand(u.shape) - u / 2
        grad = a_hat * m * (grad_a + torch.transpose(grad_a, -1, -2))
        a_hat -= lr_min * grad
        lr_min = lr_min * 0.99

        if with_l1:
            a_hat = F.relu(torch.abs(a_hat) - rho * lr_min)

        lmbd_grad = F.relu(torch.sum(contact_a(a_hat, m), dim=-1) - 1)
        lmbd += lr_max * lmbd_grad
        lr_max = lr_max * 0.99
    out_a = contact_a(a_hat, m)
    return out_a


def postprocess(pred_batch, seq_onehot, nc_map, return_nc=False, return_score=False):
    ''' nc: non-canonical '''
    def get_mat_score(pred_ori, pred_post):
        ret = []
        for row_ori, row_post in zip(pred_ori, pred_post):
            val = row_ori[row_post.argmax()]
            new_row = row_ori[:]
            new_row[new_row>val] = new_row.min()
            ret.append(new_row)
        return ret

    ret_pred = []
    ret_pred_nc = []
    ret_score = []
    ret_score_nc = []
    for i in range(pred_batch.shape[0]):
        pred = apply_constraints(pred_batch[i:i+1], seq_onehot[i:i+1], 0.01, 0.1, 100, 1.6, True, 1.5)
        pred = (pred > 0.5).float()
        ret_pred.append(pred[0])

        if return_score:
            ret_score.append(get_mat_score(pred_batch[i], pred[0]))
        if return_nc:
            pred_nc = apply_constraints(pred_batch[i:i+1], seq_onehot[i:i+1], 0.01, 0.1, 100, 0.5, True, 0.5, is_nc=True)
            pred_nc =  nc_map * pred_nc
            pred_nc = (pred_nc > 0.5).float()
            ret_pred_nc.append(pred_nc[0])
            if return_score:
                ret_score_nc.append(get_mat_score(pred_batch[i], pred_nc[0]))
    return ret_pred, ret_pred_nc, ret_score, ret_score_nc


Complete the function implementation replacing "raise NotImplementedError".