pipeline step function creation: a developer’s guide#

we first import all script dependencies and load the provided data in GSE96583 into an AnnData object

in this vignette, we will explore how to write pipeline components. for our example demonstration, we will demonstrate how one could use the LEMUR framework within a HiDDEN pipeline.

as this vignette is entirely for demonstration purposes, and the LEMUR framework can be computationally intensive, we will also randomly subset our anndata object before proceeding

# subset to 5% of data
adata = adata[np.random.default_rng(RANDOM_STATE).choice(adata.n_obs, round(adata.n_obs * 0.005), replace=False)].copy()
print(adata)

# below is a naive implementation of inserting the LEMUR workflow into HiDDEN:
# we run LEMUR, optionally apply the harmony-driven alignment, then return the embedding matrix
def run_lemur(adata: ad.AnnData, k: int, lemur_design: str, lemur_grouping: pd.Series) -> FloatMtx:
    mod = LEMUR(adata, lemur_design, n_embedding=k)
    mod.fit(verbose=False)
    mod.align_with_grouping(lemur_grouping, verbose=False)

    return mod.embedding  # ty:ignore[invalid-return-type]


p_hat, labs = found.HiDDEN(
    adata,
    "stim",
    "ctrl",
    Pipeline(run_lemur, m.reg_logit, m.bin_kmeans),
    adata=adata,
    k=10,
    lemur_design="~ stim",
    lemur_grouping=adata.obs["cell"],
)

however, with this implementation of run_lemur, further pipeline steps become unable to utilize the LEMUR model.

say that we want to conduct hyperparameter tuning, and during the scoring process, we wish to access the LEMUR counterfactual predictions. this would require having the score function access the LEMUR fitted object to generate such counterfactuals. however, this object is generated during the dimensionality reduction step, and not returned by run_lemur, which returns only the embeddings.

luckily, it is actually possible to have step functions return multiple values! this requires the use of the step_fn() decorator provided in adapters.

# ⚠️ notice the use of the `step_fn` decorator!
#           |
#           V
@step_fn("Z", "lemur_mod")
def run_lemur_w_model_out(adata: ad.AnnData, k: int, lemur_design: str, lemur_grouping: pd.Series) -> tuple[FloatMtx, LEMUR]:
    mod = LEMUR(adata, lemur_design, n_embedding=k)
    mod.fit(verbose=False)
    mod.align_with_grouping(lemur_grouping, verbose=False)

    # ⚠️ we can now return both the embeddings and the model itself
    # ⚠️ important: the order in the tuple has to match the argument provided to the `step_fn` decorator!
    return (mod.embedding, mod)  # ty:ignore[invalid-return-type]


# ⚠️ we use the "named" lemur_mod output of `run_lemur_w_model_out` in our score function!
#                                 |
#                                 V
def score_lemur_counters(lemur_mod: LEMUR, V: BoolArr, W: BoolArr) -> np.floating:
    only_relab = V & ~W
    only_kept = V & W

    # we compute a score which attempts to:
    # maximize the mean squared difference between LEMUR counterfactuals and actual expression for the HiDDEN-kept cells
    # while minimizing the mean squared difference between LEMUR counterfactuals and actual expression for the HiDDEN-relabeled cells

    relab_dist = lemur_mod.predict(new_condition=lemur_mod.cond(stim="ctrl"))[only_relab] - lemur_mod.data_matrix[only_relab]
    kept_dist = lemur_mod.predict(new_condition=lemur_mod.cond(stim="ctrl"))[only_kept] - lemur_mod.data_matrix[only_kept]

    return np.mean(np.square(kept_dist)) - np.mean(np.square(relab_dist))


sel, out = found.HiDDENt(
    adata,
    "stim",
    "ctrl",
    NaiveMaxScoreTuner(score_lemur_counters, range(4, 17, 2)),
    Pipeline(run_lemur_w_model_out, m.reg_logit, m.bin_kmeans),
    adata=adata,
    lemur_design="~ stim",
    lemur_grouping=adata.obs["cell"],
)

exploring the HiDDENt outputs, we see that using this metric, we select for k=10, with a drop off in score for k>10. however, given that our data was subset to a very low number of cells for demonstration purposes, these results cannot be interpreted further.

pl.PlotTunerOutput(adata, sel, out).score_line().show()

this step_fn() wrapping mechanism is actually used internally during Pipeline construction! if a provided step function hasn’t already been decorated with step_fn(), Pipeline will wrap it in step_fn(). the return value will be named “Z” for Pipeline.dimr_fn, “Y” for Pipeline.regr_fn, and “W” for Pipeline.binr_fn.

some final notes for developers:

if a function inserted into a pipeline needs dynamic access to all pipeline values, this can be done by adding a variable keyword argument (i.e. **kwargs form). for example, this is used by the score_null_dist() function to rerun the pipeline regression step on randomly permuted case labels to approximate a “null” p_hat distribution.
pipeline functions cannot have positional only arguments (either a named positional only argument via the arg, /, construct or a variable-length positional argument via the *args construct). calling a Pipeline where one of the steps is a function with positional-only arguments will raise a TypeError during the validation step executed prior to computation (see found.types.check_step()).
pipeline functions should always use the value of get_seed() whenever a random seed is necessary, to allow for downstream users to generate fully reproducible results via set_seed()