script.py

import anndata as ad
from fadvi import FADVI

## VIASH START
# Note: this section is auto-generated by viash at runtime. To edit it, make changes
# in config.vsh.yaml and then run `viash config inject config.vsh.yaml`.
par = {
    'input': 'resources_test/task_batch_integration/cxg_immune_cell_atlas/dataset.h5ad',
    'output': 'output.h5ad',
    'n_hvg': 2000,
    'n_latent_l': 30,
    'n_latent_b': 30,
    'n_layers': 2,
    'max_epochs': 30
}
meta = {
  'name': 'fadvi'
}
## VIASH END

print('Reading input files', flush=True)
adata = ad.read_h5ad(par['input'])

if par["n_hvg"]:
    print(f"Select top {par['n_hvg']} high variable genes", flush=True)
    idx = adata.var["hvg_score"].to_numpy().argsort()[::-1][:par["n_hvg"]]
    adata = adata[:, idx].copy()


print('Preprocess data', flush=True)
FADVI.setup_anndata(adata, batch_key="batch",labels_key="cell_type",
                    unlabeled_category='Unknown', layer='counts')
model = FADVI(adata, n_latent_l=par["n_latent_l"], 
              n_latent_b=par["n_latent_b"], 
              n_layers=par["n_layers"]) 

print('Train model', flush=True)
model.train(max_epochs=par["max_epochs"])

print('Generate predictions', flush=True)
# ... generate predictions ...

print("Write output AnnData to file", flush=True)
output = ad.AnnData(
    obs=adata.obs[[]],
    var=adata.var[[]],
    obsm={
        "X_emb": model.get_latent_representation(),
    },
    uns={
        "dataset_id": adata.uns["dataset_id"],
        "normalization_id": adata.uns["normalization_id"],
        "method_id": meta["name"],
    },
)
output.write_h5ad(par['output'], compression='gzip')