Pseudobulk data with fraction and median non-zero¶

import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)
warnings.simplefilter(action='ignore', category=RuntimeWarning)

from datetime import date
import hisepy
import numpy as np
import os
import pandas as pd
import pickle
import re
import scanpy as sc
import scipy.sparse as scs

if not os.path.exists('output'):
    os.mkdir('output')

out_files = []

def read_adata_uuid(h5ad_uuid):
    h5ad_path = '/home/jupyter/cache/{u}'.format(u = h5ad_uuid)
    if not os.path.isdir(h5ad_path):
        hise_res = hisepy.cache_files([h5ad_uuid])
    h5ad_filename = os.listdir(h5ad_path)[0]
    h5ad_file = '{p}/{f}'.format(p = h5ad_path, f = h5ad_filename)
    adata = sc.read_h5ad(h5ad_file)
    return adata

def sparse_nz_median(sparse_mat, feat_names):
    # transpose the matrix so values for each gene are together in memory
    sparse_mat = sparse_mat.transpose().tocsr()
    
    # get dimensions for calculations
    n_cells = sparse_mat.shape[1]
    
    # compute the fraction non-zero per gene (row = axis 1)
    pct_exp = sparse_mat.getnnz(axis = 1) / n_cells
    
    # split non-zero values from the matrix directly
    split_values = np.split(sparse_mat.data, sparse_mat.indptr)
    
    # remove first and last entry in split - indptr starts with 0 and ends with the last value
    # so the first and last split entries are not meaningful
    split_values = split_values[1:-1].copy()
    
    # compute medians for every split
    median_expression = [np.median(x) for x in split_values]

    # assemble a DataFrame with results
    res_df = pd.DataFrame({
        'gene': feat_names,
        'pct_exp': pct_exp,
        'median_expression': median_expression
    })
    
    # replace missing values from splits without any values
    # these were not detected so median expression is 0
    res_df['median_expression'] = res_df['median_expression'].fillna(0)
    
    return res_df

h5ad_uuid = '5b3a0cc8-1811-4126-90c7-e9cdd41459fd'

adata = read_adata_uuid(h5ad_uuid)

obs = adata.obs

levels = ['AIFI_L1', 'AIFI_L2', 'AIFI_L3']

level_combined = {}
for level in levels:
    all_level = pd.concat(level_results[level])
    all_level = all_level.reset_index(drop = False, names = [level, 'row_num'])
    all_level = all_level.drop('row_num', axis = 1)
    level_combined[level] = all_level

for level, df in level_combined.items():
    out_csv = 'output/pbmc-ref_{l}_nz-pct-median_pseudobulk_{d}.csv'.format(l = level, d = date.today())
    df.to_csv(out_csv)
    out_files.append(out_csv)
    out_parquet = 'output/pbmc-ref_{l}_nz-pct-median_pseudobulk_{d}.parquet'.format(l = level, d = date.today())
    df.to_parquet(out_parquet)
    out_files.append(out_parquet)

gene_results = {}
for level in levels:
    all_res = pd.concat(level_results[level])
    split_res = all_res.groupby('gene')
    split_dict = {}
    for gene, df in split_res:
        df = df.reset_index(drop = False)
        df = df.drop('level_1', axis = 1)
        df = df.rename({'level_0': 'obs_group'}, axis = 1)
        split_dict[gene] = df
    
    gene_results[level] = split_dict

type_pkl = 'output/pbmc-ref_type_nz-pct-median_pseudobulk_{d}.pkl'.format(d = date.today())
with open(type_pkl, 'wb') as out_file:
    pickle.dump(gene_results, out_file)
out_files.append(type_pkl)

level_combined = {}
for level in levels:
    all_level = pd.concat(level_results[level])
    all_level = all_level.reset_index(drop = False, names = ['group', 'row_num'])
    all_level[level] = [re.sub('.+_', '', x) for x in all_level['group']]
    all_level = all_level.drop(['group','row_num'], axis = 1)
    all_level = all_level[['cohort', level, 'gene', 'pct_exp', 'median_expression']]
    level_combined[level] = all_level

for level, df in level_combined.items():
    out_csv = 'output/pbmc-ref_cohort-{l}_nz-pct-median_pseudobulk_{d}.csv'.format(l = level, d = date.today())
    df.to_csv(out_csv)
    out_files.append(out_csv)
    out_parquet = 'output/pbmc-ref_cohort-{l}_nz-pct-median_pseudobulk_{d}.parquet'.format(l = level, d = date.today())
    df.to_parquet(out_parquet)
    out_files.append(out_parquet)

gene_results = {}
for level in levels:
    all_res = pd.concat(level_results[level])
    split_res = all_res.groupby('gene')
    
    split_dict = {}
    for gene, df in split_res:
        df = df.reset_index(drop = False)
        df = df.drop('level_1', axis = 1)
        df['obs_group'] = [re.sub('.+_', '', x) for x in df['level_0']]
        df = df.drop('level_0', axis = 1)
        df = df[['obs_group', 'gene', 'pct_exp', 'median_expression', 'cohort']]
        split_dict[gene] = df
    gene_results[level] = split_dict

cohort_type_pkl = 'output/pbmc-ref_cohort-type_nz-pct-median_pseudobulk_{d}.pkl'.format(d = date.today())
with open(cohort_type_pkl, 'wb') as out_file:
    pickle.dump(gene_results, out_file)
out_files.append(cohort_type_pkl)

study_space_uuid = '64097865-486d-43b3-8f94-74994e0a72e0'
title = 'PBMC Reference Pseudobulk Frac Medians {d}'.format(d = date.today())

in_files = [h5ad_uuid]