Project: s41467-025-57717-6 | Function: CalcLossCD31 (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:** Calculates loss for CD31 vascular mask prediction using Dice loss. Takes model predictions and ground truth CD31 annotations, applies specific loss weighting for positive/negative pixels in medical imaging segmentation.

**KEY TERMS:** {
  "mask_pred": "CD31 vascular mask predictions",
  "cd31_mask": "ground truth CD31 annotations",
  "dice_loss": "segmentation loss function"
}

**CONTEXT CODE:** makePlots=False

if makePlots:
    tempBatch=next(iter(trainLoader))
    predMask,predScore,preMAngio=angioMixedModel(tempBatch['image'])
    for n in range(batchSizeMask):
        mask=tempBatch['mask'][n].detach().cpu().numpy()
        img=tempBatch['image'][n].detach().cpu().numpy().transpose(1,2,0)
        pMask=np.argmax(predMask[n].detach().cpu().numpy(),axis=0)
        plt.figure()
        plt.subplot(1,3,1)
        plt.imshow(img)
        plt.xticks([],[])
        plt.yticks([],[])

        plt.subplot(1,3,2)    
        if tempBatch['isScore'][0]==0:

            plt.imshow(img)
            plt.imshow(mask,vmin=0,vmax=1,cmap=colors.ListedColormap(['w','g']),alpha=0.5)
        else:
            plt.axis('square')
            plt.title('NO CD31 GT')
        plt.xticks([],[])
        plt.yticks([],[])
        
        plt.subplot(1,3,3)
        plt.imshow(img)
        plt.imshow(pMask,vmin=0,vmax=1,cmap=colors.ListedColormap(['w','g']),alpha=0.5)
        plt.xticks([],[])
        plt.yticks([],[])
# %%
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
angioMixedModel= angioMixedModel.to(device)
    
    
# %% Calculate Weights
numberOfClasses=2
classWeights=[]
for c in range(numberOfClasses):
    classWeights.append(np.sum(trainMasks==c))
classWeights=np.array(classWeights)
classWeights=torch.from_numpy(np.max(classWeights)/classWeights).type(torch.float32).to(device)

# %%
numEpochs=10
beta=0.9
cce=nn.CrossEntropyLoss(weight=classWeights)
dice=Dice.DiceLoss()

def CalcLossCD31(predMask, predMaskAngio,predAngio,trueMask, metrics):
    """
    Calculates loss for CD31 vascular mask prediction using mixed DL model approach.

    This function implements the CD31 loss calculation component of the mixed model
    described in the paper. It combines segmentation losses (cross-entropy and Dice)
    with consistency losses between mask-derived and direct RNA predictions to enforce
    coherence between vascular mask segmentation and angiogenesis score prediction.

    The loss combines:
    - Cross-entropy loss for pixel-wise classification
    - Dice loss for segmentation overlap optimization
    - Consistency loss between prediction pathways

    Args:
        predMask (torch.Tensor): Predicted CD31 vascular masks
        predMaskAngio (torch.Tensor): Angiogenesis scores derived from predicted masks
        predAngio (torch.Tensor): Direct angiogenesis score predictions
        trueMask (torch.Tensor): Ground truth CD31 annotations
        metrics (dict): Dictionary to accumulate loss metrics

    Returns:
        torch.Tensor: Combined CD31 loss for mixed model training

    Raises:
        NotImplementedError: Function has been masked for testing purposes
    """
    raise NotImplementedError("This function has been masked for testing purposes")

def CalcLossRNA(predMaskAngio,predAngio,trueAngio, metrics):
    angioLoss=torch.square(torch.mean(trueAngio)-torch.mean(predAngio))
    maskAngioLoss=torch.square(torch.mean(predMaskAngio)-torch.mean(trueAngio))
    metrics['angioLoss'] += angioLoss.data.cpu().numpy() * trueAngio.size(0)
    metrics['mAngioLoss'] += maskAngioLoss.data.cpu().numpy() * trueAngio.size(0)
    loss=angioLoss+maskAngioLoss
    return loss

def CalcLossRNAMulti(predMaskAngio,predAngio,trueAngio,samplesOh,metrics):
    sampleWeights=torch.sum(samplesOh,axis=0,keepdims=True)
    nSamples=torch.sum(sampleWeights)
    sampleWeights[sampleWeights==0]=1
    samplesOh=samplesOh/sampleWeights
    
    trueAngioSample=torch.einsum('b,bs->s', trueAngio, samplesOh)
    predAngioSample=torch.einsum('b,bs->s', predAngio, samplesOh)
    predMaskAngioSample=torch.einsum('b,bs->s', predMaskAngio, samplesOh)
    
    angioLoss=torch.sum(torch.square(trueAngioSample-predAngioSample))/nSamples
    maskAngioLoss=torch.sum(torch.square(trueAngioSample-predMaskAngioSample))/nSamples
    
    metrics['angioLoss'] += angioLoss.data.cpu().numpy() * trueAngio.size(0)
    metrics['mAngioLoss'] += maskAngioLoss.data.cpu().numpy() * trueAngio.size(0)
    loss=angioLoss+maskAngioLoss
    return loss



optimizer = optim.SGD(filter(lambda p: p.requires_grad, angioMixedModel.parameters()), lr=1e-4,momentum=0.9)
    

for epoch in range(numEpochs):
        print('Epoch {}/{}'.format(epoch, numEpochs - 1))
        print('-' * 10)
        
        since = time.time()
            
        angioMixedModel.train()  # Set model to training mode

        metrics = defaultdict(float)
        rnaSamples = 1
        cd31Samples=1


        with tqdm(trainLoader,unit='batch') as tepoch:
            for batch in tepoch:
                inputImg = batch['image'].to(device)
                trueMask = batch['mask'].type(torch.LongTensor).to(device)    
                trueAngio = batch['score'].to(device).float()
                samplesOh=batch['sample'].to(device).float()
                isRNA=batch['isScore'][0]
                # zero the parameter gradients
                optimizer.zero_grad()

                        # forward
                        # track history if only in train
                with torch.set_grad_enabled(True):
                    predMask,predAngio,predMAngio= angioMixedModel(inputImg)
                    if isRNA:
                        
                        #loss = CalcLossRNA(predMAngio,predAngio,trueAngio, metrics)
                        loss = CalcLossRNAMulti(predMAngio,predAngio,trueAngio,samplesOh, metrics)
                        rnaSamples+=inputImg.shape[0]
                    else:
                        
                        loss=CalcLossCD31(predMask, predMAngio,predAngio,trueMask, metrics)
                        cd31Samples+=inputImg.shape[0]

                    loss.backward()
                    optimizer.step()



Complete the function implementation replacing "raise NotImplementedError".