Project: s41467-025-61628-x | Function: orient_score (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 the orientation score for a gene based on unspliced and spliced RNA time series using linear trees and intrinsic u-s relationship properties. This implements the core algorithm for inferring developmental direction at cluster level without ODE assumptions.

**KEY TERMS:** {
  "u": "unspliced RNA counts",
  "s": "spliced RNA counts",
  "rise_sec": "rising sections where both u and s increase",
  "fall_sec": "falling sections where both u and s decrease",
  "pos_score": "positive contribution from rising sections",
  "neg_score": "negative contribution from falling sections"
}

**CONTEXT CODE:** import numpy as np
import scanpy as sc
import matplotlib.pyplot as plt
from .linear_tree import create_tree, create_forecast



def get_slope(tree, data):
    if not isinstance(tree, dict):
        return tree[1]
    if data[tree['spInd']] > tree['spVal']:
        if isinstance(tree, dict):
            return get_slope(tree['left'], data)
        else:
            return tree[1]
    else:
        if isinstance(tree, dict):
            return get_slope(tree['right'], data)
        else:
            return tree[1]


def rle(inarray):
    # for a sequence, return the corresponding run-length, position and element
    ia = np.asarray(inarray)
    n = len(ia)
    if n == 0:
        return None, None, None
    else:
        y = np.array(ia[1:] != ia[:-1])
        i = np.append(np.where(y), n - 1)
        z = np.diff(np.append(-1, i))
        p = np.cumsum(np.append(0, z))[:-1]
        return z, p, ia[i]


def orient_score(u, s, rate=1e-10, dur=100, kernel_size=100, step=1):
    import os
    if os.getenv('TEST_GOLDEN', 'true') == 'false':
        raise NotImplementedError('This function is masked for testing purposes.')
    n_cells = len(u)

    # preprocess
    u = np.convolve(u, np.ones(kernel_size) / kernel_size, mode="same")
    s = np.convolve(s, np.ones(kernel_size) / kernel_size, mode="same")
    if np.sum(u) > 0:
        u = u / np.sum(u)
    if np.sum(s) > 0:
        s = s / np.sum(s)

    x = np.arange(n_cells)
    u_arr = np.concatenate((x[:, None], u[:, None]), axis=1)
    s_arr = np.concatenate((x[:, None], s[:, None]), axis=1)

    # create tree
    u_tree = create_tree(u_arr, rate, dur, step)
    s_tree = create_tree(s_arr, rate, dur, step)

    # calculate slope
    u_slope, s_slope = np.zeros(n_cells), np.zeros(n_cells)
    for i in range(n_cells):
        u_slope[i] = get_slope(u_tree, u_arr[i, :])
        s_slope[i] = get_slope(s_tree, s_arr[i, :])

    # select sections
    u_rise, s_rise = u_slope > 0, s_slope > 0
    u_fall, s_fall = u_slope < 0, s_slope < 0
    rise_sec = u_rise & s_rise
    fall_sec = u_fall & s_fall

    # calculate run length, normalized reads, weight of location for rising section
    zr, pr, er = rle(rise_sec)
    pr_ = np.append(pr, n_cells)
    rise_l, rise_w, u_norm, s_norm = np.zeros(n_cells), np.zeros(n_cells), np.zeros(n_cells), np.zeros(n_cells)
    for i in range(len(pr)):
        rise_l[pr_[i]: pr_[i+1]] = zr[i]
        rise_w[pr_[i]: pr_[i+1]] = pr[i]
        u_norm[pr_[i]: pr_[i+1]] = u[pr_[i]: pr_[i+1]] - u[pr_[i]]
        s_norm[pr_[i]: pr_[i + 1]] = s[pr_[i]: pr_[i + 1]] - s[pr_[i]]

    rise_diff = u - s

    # calculate run length, normalized reads, weight of location for falling section
    zf, pf, ef = rle(fall_sec)
    pf_ = np.append(pf, n_cells)
    fall_l, fall_w, u_norm, s_norm = np.zeros(n_cells), np.zeros(n_cells), np.zeros(n_cells), np.zeros(n_cells)
    for i in range(len(pf)):
        fall_l[pf_[i]: pf_[i+1]] = zf[i]
        fall_w[pf_[i]: pf_[i+1]] = pf[i]
        u_norm[pf_[i]: pf_[i+1]] = u[pf_[i]: pf_[i+1]] - u[pf_[i]]
        s_norm[pf_[i]: pf_[i + 1]] = s[pf_[i]: pf_[i + 1]] - s[pf_[i]]

    fall_diff = u - s

    # calculate orientation scores
    pos_score = np.sum(rise_sec * rise_l * rise_diff)
    neg_score = np.sum(fall_sec * fall_l * fall_diff)
    score = pos_score - neg_score

    # calculate section score
    rise_p_ = [pr[i] for i in np.arange(len(pr)) if er[i] == 1]
    rise_l_ = [zr[i] for i in np.arange(len(zr)) if er[i] == 1]
    pos_score_ = [np.sum(((u_rise & s_rise) * rise_l * rise_diff)[pr_[i]: pr_[i+1]]) for i in
                  np.arange(len(pr)) if er[i] == 1]
    fall_p_ = [pf[i] for i in np.arange(len(pf)) if ef[i] == 1]
    fall_l_ = [zf[i] for i in np.arange(len(zf)) if ef[i] == 1]
    neg_score_ = [np.sum(((u_fall & s_fall) * fall_l * fall_diff)[pf_[i]: pf_[i + 1]]) for i in
                   np.arange(len(pf)) if ef[i] == 1]

    sec_list = [rise_sec, fall_sec, rise_p_, fall_p_, rise_l_, fall_l_, pos_score_, neg_score_]

    return score, u_arr, s_arr, u_tree, s_tree, sec_list


def visualize(u_arr, s_arr, u_tree, s_tree, sec_list, fig_size=(12, 9)):
    u_hat = create_forecast(u_tree, u_arr[:, 0])
    s_hat = create_forecast(s_tree, s_arr[:, 0])
    rise_sec, fall_sec, rise_p_, fall_p_, rise_l_, fall_l_, pos_score_, neg_score_ = sec_list

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=fig_size, sharex=True)

    ax1.plot(u_arr[:, 0], u_arr[:, 1], c="blue", label="Mu")
    ax1.plot(s_arr[:, 0], s_arr[:, 1], c="red", label="Ms")

    # search for the beginning points
    n_cells = len(u_arr)
    u_slope, s_slope = np.zeros(n_cells), np.zeros(n_cells)
    for i in range(n_cells):
        u_slope[i] = get_slope(u_tree, u_arr[i, :])
        s_slope[i] = get_slope(s_tree, s_arr[i, :])
    zu, pu, eu = rle(u_slope)
    zs, ps, es = rle(s_slope)
    pu_ = np.append(pu, n_cells)
    ps_ = np.append(ps, n_cells)
    # plot
    for i in np.arange(len(pu)):


Complete the function implementation replacing "raise NotImplementedError".