节点 n55

<<<<<<< SEARCH
N_HVG = 2000
N_PCS = 30
EPSILON = 1e-3
TAU_A = 0.95
TAU_B = 1.0
K_SMOOTH = 30
LAMBDA = 1.0     # 1 = continue the observed displacement at full rate
GROWTH = True    # resample output cells by g^dt_out (WOT birth-death model)
N_THREADS = 8


def weighted_rows(w: np.ndarray, n: int, rng: np.random.Generator) -> np.ndarray:
    p = w / w.sum()
    return np.sort(rng.choice(len(w), size=n, replace=n > len(w), p=p))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", required=True)
    parser.add_argument("--out", required=True)
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    manifest = load_manifest(args.data)
    genes = panel_genes(args.data, manifest)
    prev, last, mask, dt_in, last_e = stage_pair(args.data, manifest, genes)
    dt_out = float(manifest["target"]["time"] - last_e["time"])
    rng = np.random.default_rng(args.seed)
    n = target_n_cells(manifest, last.n_obs)

    if prev is None:
        (g_last,) = growth_rates([last], genes, mask, args.seed)
        rows = weighted_rows(g_last ** dt_out if GROWTH else np.ones(last.n_obs), n, rng)
        write_prediction(last.X[rows], genes, args.out, seed=args.seed)
        return

    import anndata as ad
    import pandas as pd
    import torch

（只摘前 40 行）