Project: s41467-024-51026-0 | Function: fwd_bkw_prob2 (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:** Forward-backward algorithm for Hidden Markov Model implementing logarithmic probability calculations to avoid numerical underflow, enabling robust smoothing of copy number states along genomic coordinates in single-cell RNA-seq data

**KEY TERMS:** {
  "start_prob_log": "Initial state probability distribution in log space",
  "trans_prob_log": "State transition probability matrix in log space",
  "emm_prob_log": "Emission probability matrix in log space",
  "prev_f_sum_log": "Forward probability accumulator",
  "b_curr_log": "Current backward probability vector",
  "posterior_log": "Posterior state probability after forward-backward"
}

**CONTEXT CODE:**         if g == 0:
            prev_f_sum = start_prob
        else:
            prev_f_sum = np.dot(f_prev, trans_prob)        
        
        f_curr = emm_prob[g, :] * prev_f_sum
        f_curr = normalize(f_curr)
        # print("f_curr:", f_curr)

        if g == 0:
            fwd_mtx = f_curr
        else:
            fwd_mtx = np.vstack((fwd_mtx, f_curr))
        
        f_prev = f_curr.copy()

    # Backward part of the algorithm
    bkw = []
    for g in list(reversed(range(obs_num))):
        # b_curr = np.zeros((states_num))
        if g == obs_num-1:
            # base case for backward part
            b_curr = np.ones((states_num))
        else:
            b_curr = np.dot(trans_prob,(emm_prob[g+1,:] * b_prev).T)
        # print("normalize bkw") # actually, no need to normalize backward part
        b_curr = normalize(b_curr) # but add normalize can fix value problem, and not affect results
        bkw.insert(0,b_curr)
        b_prev = b_curr.copy()
    
    for cnt, b_curr in enumerate(bkw):
        if cnt == 0:
            bkw_mtx = b_curr
        else:
            bkw_mtx = np.vstack((bkw_mtx, b_curr))
    
    ## get two matrix, fwd_mtx and bkw_mtx: genes*states matrix
    ## Merging the two parts
    # print("fwd_mtx:", fwd_mtx)
    # print("bkw_mtx:", bkw_mtx)
    posterior_pre_mtx = fwd_mtx * bkw_mtx
    posterior_sum = posterior_pre_mtx.sum(axis=1).reshape(-1,1)
    posterior_mtx = posterior_pre_mtx/posterior_sum
    # print("posterior_pre_mtx:", posterior_pre_mtx)
    # print('posterior_mtx:', posterior_mtx)
    ## return genes*states matrix
    return posterior_mtx, fwd_mtx, bkw_mtx


## Forward-backward algorithm
def fwd_bkw_prob2(start_prob=None, trans_prob=None, emm_prob_log=None):
    """Forward–backward algorithm.

    fwd_bkw_prob: version 0.0.2
    try to fix Numerical issue by storing the log value, stable version

    # Example1:
    # --------
    # reference = np.array([10, 30, 50])
    # observations = np.array([8, 80, 50])
    # states = np.array([0.5, 1.0, 1.5])
    # start_prob = np.array([0.15, 0.7, 0.15])
    # trans_prob = np.array([[0.7, 0.15, 0.15],
    #                       [0.15, 0.7, 0.15],
    #                       [0.15, 0.15, 0.7]])
    # emm_prob_log = generate_nb_emmission_logprob(reference, observations, states)
    # states_num = len(states)
    # obs_num = len(observations)
    # fwd_bkw_prob2(start_prob, trans_prob, emm_prob_log)

    Example2:
    reference = np.array([10, 10, 10,10,10, 10,10])
    # reference = np.array([5., 9, 3, 5, 12, 6, 9])
    observations = np.array([5., 9, 3, 5, 12, 6, 9])
    states = np.array([0.5, 1.0])
    start_prob = np.array([0.8, 0.2])
    trans_prob = np.array([[0.7,0.3],[0.2,0.8]])

    emm_prob_log = generate_nb_emmission_logprob(reference, observations,states)
    # emm_prob_log = generate_poisson_emmission_logprob(reference, observations,states)
    # states_num = len(states)
    # obs_num = len(observations)
    fwd_bkw_prob2(start_prob, trans_prob, emm_prob_log)
    """
    # MASKED: Forward-backward algorithm for Hidden Markov Model implementing logarithmic probability calculations to avoid numerical underflow, enabling robust smoothing of copy number states along genomic coordinates in single-cell RNA-seq data
    raise NotImplementedError("Function has been masked for testing purposes")
 
def fwd_bkw_prob3(start_prob=None, trans_prob=None, emm_prob_log=None):
    """Forward–backward algorithm.

    fwd_bkw_prob: version 0.0.3 [to be tested]
    1)try to fix Numerical issue by storing the log value, stable version
    2)improve efficiency by applying to matrix.
      suitable for matrix calculation, broadcast to be fitted across cells.

    Parameters:
    ----------
    start_prob: np.array, states shape.
    trans_prob: np.array, states*states matrix.
    emm_prob_log: np.array, cells*genes*states shape.

    Return:
    -------
    # return cells*genes*states array
    posterior_mtx: 
    posterior_mtx_log:
    fwd_mtx_log:
    bkw_mtx_log:

    Example:
    -------
    reference = np.array([10, 10, 10,10,10, 10,10])
    # reference = np.array([5., 9, 3, 5, 12, 6, 9])
    observations = np.array([5., 9, 3, 5, 12, 6, 9])
    states = np.array([0.5, 1.0])
    start_prob = np.array([0.8, 0.2])
    trans_prob = np.array([[0.7,0.3],[0.2,0.8]])

    emm_prob_log = generate_nb_emmission_logprob(reference, observations, states)
    fwd_bkw_prob3(start_prob, trans_prob, emm_prob_log)
    """
    sample_num = emm_prob_log.shape[0] ## cells in XClone   
    obs_num = emm_prob_log.shape[-2] ## genes in XClone
    states_num = emm_prob_log.shape[-1]

    # Forward part of the algorithm
    trans_prob_log = np.log(trans_prob)
    start_prob_log = np.log(start_prob)

    for g in range(obs_num):
        if g == 0:
            prev_f_sum_log = start_prob_log
        else:
            prev_f_sum_log = logsumexp(f_prev_log.swapaxes(1,2) + trans_prob_log, axis=1)
            ## swap genes axis and states axis for calculation
        
        f_curr_log = emm_prob_log[:, g, :] + prev_f_sum_log
        # normalize
        f_curr_log = f_curr_log - logsumexp(f_curr_log)
        # keepdims
        f_curr_log = f_curr_log[:, np.newaxis, :]

        if g == 0:
            fwd_mtx_log = f_curr_log
        else:
            fwd_mtx_log = np.concatenate((fwd_mtx_log, f_curr_log), axis=1)
        
        f_prev_log = f_curr_log.copy()

    # Backward part of the algorithm


Complete the function implementation replacing "raise NotImplementedError".