Project: s41467-024-46506-2 | Function: _grad (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 the gradient computation for tomographic reconstruction using automatic differentiation with forward and backward passes through the OCT signal propagation model

**KEY TERMS:** {
  "IM": "recorded OCM volumes",
  "vjp": "vector-Jacobian product for automatic differentiation",
  "step_loc": "forward propagation step with parameter partitioning",
  "bf_fwd": "forward pass saving intermediate variables",
  "bf_bwd": "backward pass computing gradients",
  "loss_fn": "data fidelity error metric"
}

**CONTEXT CODE:**         exp(- p.ratio / (2 * log(2)) * (pi * (x - p.x_C) / (2. * p.x_D))**2 )))

# adjoint motion for the mask 
def br(a, p):

    xr, yr, zr = RotMat_T_Vec(p.ψ, p.t, p.x[:,None,None], p.y[None,:,None], p.z[None,None], p.nx, p.ny, p.nz, p.Δx, p.Δy, p.Δz)

    return mc(a, (xr, yr, zr), order=1)

# interpolate (R, α, Δn) at the rotated coordinates
def ip(p, xim):

    xr, yr, zr = RotMat_Vec(p.ψ, p.t, xim, p.y[:,None], p.z[None], p.nx, p.ny, p.nz, p.Δx, p.Δy, p.Δz)
    rot = lambda a: mc(a, (xr, yr, zr), order=1)
    Rim, αim, Δnim = rot(p.R), rot(p.α), rot(p.Δn)

    return Rim, αim, Δnim, (xr, yr, zr)

# attenuate the signal locally by α
def att(p, Rim, αim, Aim, xi, xim):

    Ai = Aim * exp(- αim * (xi - xim))
    Iim = Rim * Ai * T(p, xi) * H(p, xi)

    return Iim, Ai

# full propagation step
def step(p, Aim, xim):
    # MASKED: Implements the layer-by-layer OCT signal propagation model with refractive index correction and attenuation according to Equations 1-2
    raise NotImplementedError("This function is masked for testing purposes.")

def _fwd(p):

    def bf(i, val):

        (Aim, xim, I) = val
        Iim, Ai, xi, _ = step(p, Aim, xim)
        I = I.at[i].set(Iim)
        return (Ai, xi, I)

    I = zeros((p.nx, p.ny, p.nz))
    A0 = ones((p.ny, p.nz))
    x0 = zeros((p.ny, p.nz))

    _start = (A0 ,x0, I)
    _result = fori_loop(0, p.nx_d, bf, _start)

    return _result[-1]   

# gradient function
def _grad(p, IM, loss_fn, filter_spec):
    # MASKED: Implements the gradient computation for tomographic reconstruction using automatic differentiation with forward and backward passes through the OCT signal propagation model
    raise NotImplementedError("This function is masked for testing purposes.")

# generate update function
def get_update(bounds, opt_par):

    def update(p, p_bar, lrs):

        return eqx.tree_at(lambda tree: tuple(getattr(tree, a) for a in opt_par), 
                            p,
                            replace = tuple(clip(getattr(p, a) - getattr(p_bar, a) * getattr(lrs, a), 
                                                 getattr(bounds, a)[0], 
                                                 getattr(bounds, a)[1])
                                            for a in opt_par))
    return update

# generate stochastic gradient update step (single view)
def get_sgd_step(loss_fn, bounds, reg_fn, filter_opt, opt_par):

    update = get_update(bounds, opt_par)
    grad_reg = value_and_grad(reg_fn, allow_int=True)

    def sgd_step(p, IMi, Mi, lrs):

        err, p_bar = _grad(p, IMi, loss_fn, filter_opt)
        Mi_r = br(Mi, p)
        err_rg, pr_bar = grad_reg(p, Mi_r)
        p_bar = tree_map(lambda a, b: a + b, p_bar, pr_bar)
        
        return err + err_rg, update(p, p_bar, lrs)

    return sgd_step

# one epoch of the stochastic gradient descent: iterate over whole dataset in each epoch
def get_sgd_epoch(sgd_step):

    def sgd_epoch(p_arr, lrs, IM, M, key):
        
        ns = IM.shape[0]
        ind = permutation(key, ns)
        
        def bf(i, val):

            err, p_arr = val
            j = ind[i]
            p_i = tree_at(lambda tree: tuple(getattr(tree, a) for a in ['ψ', 't', 'x_F']), 
                          p_arr, 
                          replace_fn = lambda tree: tree[j])
            e, p_i = sgd_step(p_i, IM[j], M[j], lrs)
            p_arr = tree_at(lambda tree: tuple(getattr(tree, a) for a in ['ψ', 't', 'x_F']), 
                            p_i, 
                            replace = tuple(getattr(p_arr, a) * (1. - one_hot(j, ns)[:,None,None,None]) + getattr(p_i, a)[None] * one_hot(j, ns)[:,None,None,None]  for a in ['ψ', 't', 'x_F']))

            return err + e, p_arr
        
        return fori_loop(0, ns, bf, (0, p_arr))

    return sgd_epoch












Complete the function implementation replacing "raise NotImplementedError".