Project: s41467-025-59797-w | Function: compute_ctf (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:** Computes the 2D Contrast Transfer Function (CTF) for electron microscopy imaging using physical formulas. Takes defocus parameters, voltage, spherical aberration and other imaging parameters to calculate CTF values.

**KEY TERMS:** {
  "dfu": "DefocusU parameter",
  "dfv": "DefocusV parameter",
  "volt": "accelerating voltage",
  "cs": "spherical aberration",
  "gamma": "phase function",
  "ctf": "contrast transfer function"
}

**CONTEXT CODE:** import torch
import torch.nn as nn
import torch.nn.init as init
from torch.nn import functional as F
import numpy as np
import mrcfile


def compute_ctf(freqs, dfu, dfv, dfang, volt, cs, w, phase_shift=0, bfactor=None):
    """
    Compute the 2D CTF
    Input:
        freqs (np.ndarray) Nx2 or BxNx2 tensor of 2D spatial frequencies
        dfu (float or Bx1 tensor): DefocusU (Angstrom)
        dfv (float or Bx1 tensor): DefocusV (Angstrom)
        dfang (float or Bx1 tensor): DefocusAngle (degrees)
        volt (float or Bx1 tensor): accelerating voltage (kV)
        cs (float or Bx1 tensor): spherical aberration (mm)
        w (float or Bx1 tensor): amplitude contrast ratio
        phase_shift (float or Bx1 tensor): degrees
        bfactor (float or Bx1 tensor): envelope fcn B-factor (Angstrom^2)
    """
    # MASKED: Computes the 2D Contrast Transfer Function (CTF) for electron microscopy imaging using physical formulas implementation
    raise NotImplementedError("This function is masked for testing purposes.")

def discrete_radon_transform_3d(volume, rotation):
    # MASKED: Implements 3D discrete Radon transform for generating 2D projections from 3D volumes by applying rotation and summing along the Z-axis implementation
    raise NotImplementedError("This function is masked for testing purposes.")

def translation_2d(proj, trans):
    """
    Input:
        proj: Bx1xbsxbs tensor 
        trans: Bx2 tensor
    """
    
    b = trans.shape[0]
    
    eye = torch.eye(2).unsqueeze(0).repeat(b, 1, 1).to(proj.device)
    trans = trans.unsqueeze(-1)
    theta = torch.cat([eye, trans], dim=2)

    grid = F.affine_grid(theta, size=proj.shape)
    proj_trans = F.grid_sample(proj, grid, mode='bicubic')
    
    return proj_trans


def convolve_ctf(proj, ctf):
    Fproj = torch.fft.fftshift(torch.fft.fft2(proj), dim=(-2, -1))
    Fproj = Fproj * ctf
    Fproj = torch.fft.ifftshift(Fproj, dim=(-2, -1))
    proj = torch.fft.ifft2(Fproj)
    proj = proj.real
    
    return proj


def init_volume(box_size):
    volume = torch.zeros(box_size, box_size, box_size)
    volume[box_size // 4 : 3 * box_size // 4, 
           box_size // 4 : 3 * box_size // 4,
           box_size // 4 : 3 * box_size // 4] = 1
    
    return volume

def gaussian_analytical_kl(mu1, mu2, logsigma1, logsigma2):
    return -0.5 + logsigma2 - logsigma1 + 0.5 * (logsigma1.exp() ** 2 + (mu1 - mu2) ** 2) / (logsigma2.exp() ** 2)


def draw_gaussian_diag_samples(mu, logsigma):
    eps = torch.empty_like(mu).normal_(0., 1.)
    return torch.exp(logsigma) * eps + mu


def initialize_weights(net_l, scale=1):
    if not isinstance(net_l, list):
        net_l = [net_l]
    for net in net_l:
        for m in net.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, a=0, mode='fan_in')
                m.weight.data *= scale  # for residual block
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                init.kaiming_normal_(m.weight, a=0, mode='fan_in')
                m.weight.data *= scale
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias.data, 0.0)


class ProjNet(nn.Module):
    def __init__(self, nf=64, num_down=0):
        super(ProjNet, self).__init__()
        
        conv = []
        
        conv.append(nn.Conv2d(1, nf, 3, 1, 1))
        for _ in range(num_down):
            conv.append(nn.LeakyReLU(negative_slope=0.2, inplace=True))
            conv.append(nn.Conv2d(nf, nf, 4, 2, 1))
        
        conv.append(nn.Conv2d(nf, nf, 3, 1, 1))
        self.conv = nn.Sequential(*conv)


Complete the function implementation replacing "raise NotImplementedError".