import os

import re

import copy

from typing import Dict, List, Optional, Union

from urllib.parse import urlparse

import scipy.io

from typing_extensions import Literal

from .config import get_bustools_binary_path, get_kallisto_binary_path

from .constants import (

ABUNDANCE_FILENAME,

ABUNDANCE_GENE_FILENAME,

ABUNDANCE_GENE_TPM_FILENAME,

ABUNDANCE_GENE_NAMES_FILENAME,

ABUNDANCE_TPM_FILENAME,

ADATA_PREFIX,

BUS_FILENAME,

CAPTURE_FILENAME,

CELLRANGER_BARCODES,

CELLRANGER_DIR,

CELLRANGER_GENES,

CELLRANGER_MATRIX,

CORRECT_CODE,

COUNTS_PREFIX,

ECMAP_FILENAME,

FEATURE_NAME,

FEATURE_PREFIX,

FILTER_WHITELIST_FILENAME,

FILTERED_CODE,

FILTERED_COUNTS_DIR,

FLD_FILENAME,

FLENS_FILENAME,

GENE_NAME,

GENES_FILENAME,

GENE_NAMES_FILENAME,

GENOMEBAM_FILENAME,

GENOMEBAM_INDEX_FILENAME,

INSPECT_FILENAME,

INSPECT_INTERNAL_FILENAME,

INSPECT_UMI_FILENAME,

INTERNAL_SUFFIX,

KALLISTO_INFO_FILENAME,

KB_INFO_FILENAME,

PROJECT_CODE,

REPORT_HTML_FILENAME,

REPORT_NOTEBOOK_FILENAME,

SAVED_INDEX_FILENAME,

SORT_CODE,

TCC_PREFIX,

TRANSCRIPT_NAME,

TXNAMES_FILENAME,

UMI_SUFFIX,

UNFILTERED_CODE,

UNFILTERED_COUNTS_DIR,

UNFILTERED_QUANT_DIR,

WHITELIST_FILENAME,

)

from .dry import dryable

from .dry import count as dry_count

from .logging import logger

from .report import render_report

from .utils import (

copy_whitelist,

create_10x_feature_barcode_map,

get_temporary_filename,

import_matrix_as_anndata,

import_tcc_matrix_as_anndata,

make_directory,

open_as_text,

overlay_anndatas,

read_t2g,

remove_directory,

run_executable,

stream_file,

sum_anndatas,

update_filename,

whitelist_provided,

obtain_gene_names,

write_list_to_file,

do_sum_matrices,

move_file,

)

from .stats import STATS

from .validate import validate_files

INSPECT_PARSER = re.compile(r'^.*?(?P<count>[0-9]+)')

def make_transcript_t2g(txnames_path: str, out_path: str) -> str:

"""Make a two-column t2g file from a transcripts file

Args:

txnames_path: Path to transcripts.txt

out_path: Path to output t2g file

Returns:

Path to output t2g file

"""

with open_as_text(txnames_path, 'r') as f, open_as_text(out_path,

'w') as out:

for line in f:

out.write(f'{line.strip()}\t{line.strip()}\n')

return out_path

def kallisto_bus(

fastqs: Union[List[str], str],

index_path: str,

technology: str,

out_dir: str,

threads: int = 8,

n: bool = False,

k: bool = False,

paired: bool = False,

genomebam: bool = False,

aa: bool = False,

strand: Optional[Literal['unstranded', 'forward', 'reverse']] = None,

gtf_path: Optional[str] = None,

chromosomes_path: Optional[str] = None,

inleaved: bool = False,

demultiplexed: bool = False,

batch_barcodes: bool = False,

numreads: int = None,

lr: bool = False,

lr_thresh: float = 0.8,

lr_error_rate: float = None,

union: bool = False,

no_jump: bool = False,

) -> Dict[str, str]:

"""Runs `kallisto bus`.

Args:

fastqs: List of FASTQ file paths, or a single path to a batch file

index_path: Path to kallisto index

technology: Single-cell technology used

out_dir: Path to output directory

threads: Number of threads to use, defaults to `8`

n: Include number of read in flag column (used when splitting indices),

defaults to `False`

k: Alignment is done per k-mer (used when splitting indices),

defaults to `False`

paired: Whether or not to supply the `--paired` flag, only used for

bulk and smartseq2 samples, defaults to `False`

genomebam: Project pseudoalignments to genome sorted BAM file, defaults to

`False`

aa: Align to index generated from a FASTA-file containing amino acid sequences,

defaults to `False`

strand: Strandedness, defaults to `None`

gtf_path: GTF file for transcriptome information (required for --genomebam),

defaults to `None`

chromosomes_path: Tab separated file with chromosome names and lengths

(optional for --genomebam, but recommended), defaults to `None`

inleaved: Whether input FASTQ is interleaved, defaults to `False`

demultiplexed: Whether FASTQs are demultiplexed, defaults to `False`

batch_barcodes: Whether sample ID should be in barcode, defaults to `False`

numreads: Maximum number of reads to process from supplied input

lr: Whether to use lr-kallisto in read mapping, defaults to `False`

lr_thresh: Sets the --threshold for lr-kallisto, defaults to `0.8`

lr_error_rate: Sets the --error-rate for lr-kallisto, defaults to `None`

union: Use set union for pseudoalignment, defaults to `False`

no_jump: Disable pseudoalignment "jumping", defaults to `False`

Returns:

Dictionary containing paths to generated files

"""

logger.info(

f'Using index {index_path} to generate BUS file to {out_dir} from'

)

results = {

'bus': os.path.join(out_dir, BUS_FILENAME),

'ecmap': os.path.join(out_dir, ECMAP_FILENAME),

'txnames': os.path.join(out_dir, TXNAMES_FILENAME),

'info': os.path.join(out_dir, KALLISTO_INFO_FILENAME)

}

is_batch = isinstance(fastqs, str)

for fastq in [fastqs] if is_batch else fastqs:

logger.info((' ' * 8) + fastq)

command = [get_kallisto_binary_path(), 'bus']

command += ['-i', index_path]

command += ['-o', out_dir]

if not demultiplexed:

if technology.upper() == "10XV4":

# TODO: REMOVE THIS WHEN KALLISTO IS UPDATED

command += ['-x', "10XV3"]

else:

command += ['-x', technology]

elif technology[0] == '-':

# User supplied a custom demuxed (no-barcode) technology

command += ['-x', technology]

else:

command += ['-x', 'BULK']

command += ['-t', threads]

if n:

command += ['--num']

if k:

command += ['--kmer']

if paired and not aa:

command += ['--paired']

results['flens'] = os.path.join(out_dir, FLENS_FILENAME)

if genomebam:

command += ['--genomebam']

if gtf_path is not None:

command += ['-g', gtf_path]

if chromosomes_path is not None:

command += ['-c', chromosomes_path]

results['genomebam'] = os.path.join(out_dir, GENOMEBAM_FILENAME)

results['genomebam_index'] = os.path.join(

out_dir, GENOMEBAM_INDEX_FILENAME

)

if numreads:

command += ['-N', numreads]

if aa:

command += ['--aa']

if paired:

logger.warning(

'`--paired` ignored since `--aa` only supports single-end reads'

)

if strand == 'unstranded':

command += ['--unstranded']

elif strand == 'forward':

command += ['--fr-stranded']

elif strand == 'reverse':

command += ['--rf-stranded']

if inleaved:

command += ['--inleaved']

if lr:

command += ['--long']

if lr and lr_thresh:

command += ['-r', str(lr_thresh)]

if lr and lr_error_rate:

command += ['-e', str(lr_error_rate)]

if union:

command += ['--union']

if no_jump:

command += ['--no-jump']

if batch_barcodes:

command += ['--batch-barcodes']

if is_batch:

command += ['--batch', fastqs]

else:

command += fastqs

run_executable(command)

if technology.upper() in ('BULK', 'SMARTSEQ3'):

results['saved_index'] = os.path.join(out_dir, SAVED_INDEX_FILENAME)

if os.path.exists(results['saved_index']):

os.remove(results['saved_index']) # TODO: Fix this in kallisto?

return results

@validate_files(pre=False)

def kallisto_quant_tcc(

mtx_path: str,

saved_index_path: str,

ecmap_path: str,

t2g_path: str,

out_dir: str,

flens_path: Optional[str] = None,

l: Optional[int] = None,

s: Optional[int] = None,

threads: int = 8,

bootstraps: int = 0,

matrix_to_files: bool = False,

matrix_to_directories: bool = False,

no_fragment: bool = False,

lr: bool = False,

lr_platform: str = 'ONT',

) -> Dict[str, str]:

"""Runs `kallisto quant-tcc`.

Args:

mtx_path: Path to counts matrix

saved_index_path: Path to index

ecmap_path: Path to ecmap

t2g_path: Path to T2G

out_dir: Output directory path

flens_path: Path to flens.txt, defaults to `None`

l: Mean fragment length, defaults to `None`

s: Standard deviation of fragment length, defaults to `None`

threads: Number of threads to use, defaults to `8`

bootstraps: Number of bootstraps to perform for quant-tcc, defaults to 0

matrix_to_files: Whether to write quant-tcc output to files, defaults to `False`

matrix_to_directories: Whether to write quant-tcc output to directories, defaults to `False`

no_fragment: Whether to disable quant-tcc effective length normalization, defaults to `False`

lr: Whether to use lr-kallisto in quantification, defaults to `False`

lr_platform: Sets the --platform for lr-kallisto, defaults to `ONT`

Returns:

Dictionary containing path to output files

"""

logger.info(

f'Quantifying transcript abundances to {out_dir} from mtx file {mtx_path}'

)

command = [get_kallisto_binary_path(), 'quant-tcc']

command += ['-o', out_dir]

command += ['-i', saved_index_path]

command += ['-e', ecmap_path]

command += ['-g', t2g_path]

command += ['-t', threads]

if lr:

command += ['--long']

if lr and lr_platform:

command += ['-P', lr_platform]

if flens_path and not no_fragment:

command += ['-f', flens_path]

if l and not no_fragment:

command += ['-l', l]

if s and not no_fragment:

command += ['-s', s]

if bootstraps and bootstraps != 0:

command += ['-b', bootstraps]

if matrix_to_files:

command += ['--matrix-to-files']

if matrix_to_directories:

command += ['--matrix-to-directories']

command += [mtx_path]

run_executable(command)

ret_dict = {

'genes': os.path.join(out_dir, GENES_FILENAME),

'gene_mtx': os.path.join(out_dir, ABUNDANCE_GENE_FILENAME),

'gene_tpm_mtx': os.path.join(out_dir, ABUNDANCE_GENE_TPM_FILENAME),

'mtx': os.path.join(out_dir, ABUNDANCE_FILENAME),

'tpm_mtx': os.path.join(out_dir, ABUNDANCE_TPM_FILENAME),

'txnames': os.path.join(out_dir, TXNAMES_FILENAME),

}

if flens_path or l or s:

ret_dict['fld'] = os.path.join(out_dir, FLD_FILENAME)

return ret_dict

@validate_files(pre=False)

def bustools_project(

bus_path: str, out_path: str, map_path: str, ecmap_path: str,

txnames_path: str

) -> Dict[str, str]:

"""Runs `bustools project`.

bus_path: Path to BUS file to sort

out_dir: Path to output directory

map_path: Path to file containing source-to-destination mapping

ecmap_path: Path to ecmap file, as generated by `kallisto bus`

txnames_path: Path to transcript names file, as generated by `kallisto bus`

Returns:

Dictionary containing path to generated BUS file

"""

logger.info('Projecting BUS file {} with map {}'.format(bus_path, map_path))

command = [get_bustools_binary_path(), 'project']

command += ['-o', out_path]

command += ['-m', map_path]

command += ['-e', ecmap_path]

command += ['-t', txnames_path]

command += ['--barcode']

command += [bus_path]

run_executable(command)

return {'bus': out_path}

def bustools_sort(

bus_path: str,

out_path: str,

temp_dir: str = 'tmp',

threads: int = 8,

memory: str = '2G',

flags: bool = False,

store_num: bool = False,

) -> Dict[str, str]:

"""Runs `bustools sort`.

Args:

bus_path: Path to BUS file to sort

out_dir: Path to output BUS path

temp_dir: Path to temporary directory, defaults to `tmp`

threads: Number of threads to use, defaults to `8`

memory: Amount of memory to use, defaults to `2G`

flags: Whether to supply the `--flags` argument to sort, defaults to

`False`

store_num: Whether to process BUS files with read numbers in flag,

defaults to `False`

Returns:

Dictionary containing path to generated index

"""

logger.info('Sorting BUS file {} to {}'.format(bus_path, out_path))

command = [get_bustools_binary_path(), 'sort']

command += ['-o', out_path]

command += ['-T', temp_dir]

command += ['-t', threads]

command += ['-m', memory]

if flags:

command += ['--flags']

if store_num:

command += ['--no-flags']

command += [bus_path]

run_executable(command)

return {'bus': out_path}

@validate_files(pre=False)

def bustools_inspect(

bus_path: str,

out_path: str,

whitelist_path: Optional[str] = None,

ecmap_path: Optional[str] = None,

) -> Dict[str, str]:

"""Runs `bustools inspect`.

Args:

bus_path: Path to BUS file to sort

out_path: Path to output inspect JSON file

whitelist_path: Path to whitelist

ecmap_path: Path to ecmap file, as generated by `kallisto bus`

Returns:

Dictionary containing path to generated index

"""

logger.info('Inspecting BUS file {}'.format(bus_path))

command = [get_bustools_binary_path(), 'inspect']

command += ['-o', out_path]

if whitelist_path:

command += ['-w', whitelist_path]

if ecmap_path:

command += ['-e', ecmap_path]

command += [bus_path]

run_executable(command)

return {'inspect': out_path}

def bustools_correct(

bus_path: str,

out_path: str,

whitelist_path: str,

replace: bool = False,

exact_barcodes: bool = False

) -> Dict[str, str]:

"""Runs `bustools correct`.

Args:

bus_path: Path to BUS file to correct

out_path: Path to output corrected BUS file

whitelist_path: Path to whitelist

replace: If whitelist is a replacement file, defaults to `False`

exact_barcodes: Use exact matching for 'correction', defaults to `False`

Returns:

Dictionary containing path to generated index

"""

logger.info(

'Correcting BUS records in {} to {} with on-list {}'.format(

bus_path, out_path, whitelist_path

)

)

command = [get_bustools_binary_path(), 'correct']

command += ['-o', out_path]

command += ['-w', whitelist_path]

command += [bus_path]

if replace:

command += ['--replace']

if exact_barcodes:

command += ['--nocorrect']

run_executable(command)

return {'bus': out_path}

@validate_files(pre=False)

def bustools_count(

bus_path: str,

out_prefix: str,

t2g_path: str,

ecmap_path: str,

txnames_path: str,

tcc: bool = False,

mm: bool = False,

cm: bool = False,

umi_gene: bool = True,

em: bool = False,

nascent_path: str = None,

batch_barcodes: bool = False,

) -> Dict[str, str]:

"""Runs `bustools count`.

Args:

bus_path: Path to BUS file to correct

out_prefix: Prefix of the output files to generate

t2g_path: Path to output transcript-to-gene mapping

ecmap_path: Path to ecmap file, as generated by `kallisto bus`

txnames_path: Path to transcript names file, as generated by `kallisto bus`

tcc: Whether to generate a TCC matrix instead of a gene count matrix,

defaults to `False`

mm: Whether to include BUS records that pseudoalign to multiple genes,

defaults to `False`

cm: Count multiplicities instead of UMIs. Used for chemitries

without UMIs, such as bulk and Smartseq2, defaults to `False`

umi_gene: Whether to use genes to deduplicate umis, defaults to `True`

em: Whether to estimate gene abundances using EM algorithm, defaults

to `False`

nascent_path: Path to list of nascent targets for obtaining

nascent/mature/ambiguous matrices, defaults to `None`

batch_barcodes: If sample ID is barcoded, defaults to `False`

Returns:

Dictionary containing path to generated index

"""

logger.info(

f'Generating count matrix {out_prefix} from BUS file {bus_path}'

)

command = [get_bustools_binary_path(), 'count']

command += ['-o', out_prefix]

command += ['-g', t2g_path]

command += ['-e', ecmap_path]

command += ['-t', txnames_path]

if nascent_path:

command += ['-s', nascent_path]

if not tcc:

command += ['--genecounts']

if mm:

command += ['--multimapping']

if cm:

command += ['--cm']

if umi_gene and not cm:

command += ['--umi-gene']

if em:

command += ['--em']

command += [bus_path]

# There is currently a bug when a directory with the same path as `out_prefix`

# exists, the matrix is named incorrectly. So, to get around this, manually

# detect and remove such a directory should it exist.

if os.path.isdir(out_prefix):

remove_directory(out_prefix)

run_executable(command)

if nascent_path:

ret = {

'mtx0':

move_file(f'{out_prefix}.mtx', f'{out_prefix}.mature.mtx'),

'ec0' if tcc else 'genes0':

f'{out_prefix}.ec.txt' if tcc else f'{out_prefix}.genes.txt',

'barcodes0':

f'{out_prefix}.barcodes.txt',

'batch_barcodes0':

f'{out_prefix}.barcodes.prefix.txt' if batch_barcodes else None,

'mtx1':

move_file(f'{out_prefix}.2.mtx', f'{out_prefix}.nascent.mtx'),

'ec1' if tcc else 'genes1':

f'{out_prefix}.ec.txt' if tcc else f'{out_prefix}.genes.txt',

'barcodes1':

f'{out_prefix}.barcodes.txt',

'batch_barcodes1':

f'{out_prefix}.barcodes.prefix.txt' if batch_barcodes else None,

'mtx2':

f'{out_prefix}.ambiguous.mtx',

'ec2' if tcc else 'genes2':

f'{out_prefix}.ec.txt' if tcc else f'{out_prefix}.genes.txt',

'barcodes2':

f'{out_prefix}.barcodes.txt',

'batch_barcodes2':

f'{out_prefix}.barcodes.prefix.txt' if batch_barcodes else None,

}

if not batch_barcodes:

del ret['batch_barcodes0']

del ret['batch_barcodes1']

del ret['batch_barcodes2']

elif not os.path.exists(ret['batch_barcodes0']):

del ret['batch_barcodes0']

del ret['batch_barcodes1']

del ret['batch_barcodes2']

return ret

ret = {

'mtx':

f'{out_prefix}.mtx',

'ec' if tcc else 'genes':

f'{out_prefix}.ec.txt' if tcc else f'{out_prefix}.genes.txt',

'barcodes':

f'{out_prefix}.barcodes.txt',

'batch_barcodes':

f'{out_prefix}.barcodes.prefix.txt' if batch_barcodes else None,

}

if not batch_barcodes:

del ret['batch_barcodes']

elif not os.path.exists(ret['batch_barcodes']):

del ret['batch_barcodes']

return ret

@validate_files(pre=False)

def bustools_capture(

bus_path: str,

out_path: str,

capture_path: str,

ecmap_path: Optional[str] = None,

txnames_path: Optional[str] = None,

capture_type: Literal['transcripts', 'umis', 'barcode'] = 'transcripts',

complement: bool = True,

) -> Dict[str, str]:

"""Runs `bustools capture`.

Args:

bus_path: Path to BUS file to capture

out_path: Path to BUS file to generate

capture_path: Path transcripts-to-capture list

ecmap_path: Path to ecmap file, as generated by `kallisto bus`

txnames_path: Path to transcript names file, as generated by `kallisto bus`

capture_type: The type of information in the capture list. Can be one of

`transcripts`, `umis`, `barcode`.

complement: Whether or not to complement, defaults to `True`

Returns:

Dictionary containing path to generated index

"""

logger.info(

f'Capturing records from BUS file {bus_path} to {out_path} with capture list {capture_path}'

)

command = [get_bustools_binary_path(), 'capture']

command += ['-o', out_path]

command += ['-c', capture_path]

if ecmap_path:

command += ['-e', ecmap_path]

if txnames_path:

command += ['-t', txnames_path]

if complement:

command += ['--complement']

command += ['--{}'.format(capture_type)]

command += [bus_path]

run_executable(command)

return {'bus': out_path}

@validate_files(pre=False)

def bustools_whitelist(

bus_path: str,

out_path: str,

threshold: Optional[int] = None

) -> Dict[str, str]:

"""Runs `bustools allowlist`.

Args:

bus_path: Path to BUS file generate the on-list from

out_path: Path to output on-list

threshold: Barcode threshold to be included in on-list

Returns:

Dictionary containing path to generated index

"""

logger.info(

'Generating on-list {} from BUS file {}'.format(out_path, bus_path)

)

command = [get_bustools_binary_path(), 'allowlist']

command += ['-o', out_path]

if threshold:

command += ['--threshold', threshold]

command += [bus_path]

run_executable(command)

return {'whitelist': out_path}

def matrix_to_cellranger(

matrix_path: str,

barcodes_path: str,

genes_path: str,

t2g_path: str,

out_dir: str,

gzip: bool = False

) -> Dict[str, str]:

"""Convert bustools count matrix to cellranger-format matrix.

Args:

matrix_path: Path to matrix

barcodes_path: List of paths to barcodes.txt

genes_path: Path to genes.txt

t2g_path: Path to transcript-to-gene mapping

out_dir: Path to output matrix

gzip: Whether to gzip compress the output files, defaults to `False`

Returns:

Dictionary of matrix files

"""

make_directory(out_dir)

logger.info(f'Writing matrix in cellranger format to {out_dir}')

# Add .gz extension if gzip is enabled

gz_ext = '.gz' if gzip else ''

cr_matrix_path = os.path.join(out_dir, CELLRANGER_MATRIX + gz_ext)

cr_barcodes_path = os.path.join(out_dir, CELLRANGER_BARCODES + gz_ext)

cr_genes_path = os.path.join(out_dir, CELLRANGER_GENES + gz_ext)

# Cellranger outputs genes x cells matrix

mtx = scipy.io.mmread(matrix_path)

# Write to temporary file first, then optionally gzip

temp_mtx = os.path.join(out_dir, CELLRANGER_MATRIX)

scipy.io.mmwrite(temp_mtx, mtx.T, field='integer')

if gzip:

from .utils import compress_gzip

compress_gzip(temp_mtx, cr_matrix_path)

os.remove(temp_mtx)

opener = open_as_text if gzip else open

with open(barcodes_path, 'r') as f, opener(cr_barcodes_path, 'w') as out:

for line in f:

if line.isspace():

continue

out.write(f'{line.strip()}-1\n')

# Get all (available) gene names

gene_to_name = {}

with open(t2g_path, 'r') as f:

for line in f:

if line.isspace():

continue

split = line.strip().split('\t')

if len(split) > 2:

gene_to_name[split[1]] = split[2]

with opener(genes_path, 'r') as f, opener(cr_genes_path, 'w') as out:

for line in f:

if line.isspace():

continue

gene_id = line.strip()

gene_symbol = gene_to_name.get(gene_id, gene_id)

# CellRanger format: GENE_ID\tGENE_SYMBOL (or FEATURE_ID\tFEATURE_NAME for features)

out.write(f'{gene_id}\t{gene_symbol}\n')

return {

'mtx': cr_matrix_path,

'barcodes': cr_barcodes_path,

'genes': cr_genes_path

}

def convert_matrix(

counts_dir: str,

matrix_path: str,

barcodes_path: str,

batch_barcodes_path: Optional[str] = None,

genes_path: Optional[str] = None,

ec_path: Optional[str] = None,

t2g_path: Optional[str] = None,

txnames_path: Optional[str] = None,

name: str = 'gene',

loom: bool = False,

loom_names: List[str] = ['barcode', 'target_name'],

h5ad: bool = False,

by_name: bool = False,

tcc: bool = False,

threads: int = 8,

) -> Dict[str, str]:

"""Convert a gene count or TCC matrix to loom or h5ad.

Args:

counts_dir: Path to counts directory

matrix_path: Path to matrix

barcodes_path: List of paths to barcodes.txt

batch_barcodes_path: Path to barcodes prefixed with sample ID,

defaults to `None`

genes_path: Path to genes.txt, defaults to `None`

ec_path: Path to ec.txt, defaults to `None`

t2g_path: Path to transcript-to-gene mapping. If this is provided,

the third column of the mapping is appended to the anndata var,

defaults to `None`

txnames_path: Path to transcripts.txt, defaults to `None`

name: Name of the columns, defaults to "gene"

loom: Whether to generate loom file, defaults to `False`

loom_names: Names for col_attrs and row_attrs in loom file,

defaults to `['barcode','target_name']`

h5ad: Whether to generate h5ad file, defaults to `False`

by_name: Aggregate counts by name instead of ID.

tcc: Whether the matrix is a TCC matrix, defaults to `False`

threads: Number of threads to use, defaults to `8`

Returns:

Dictionary of generated files

"""

results = {}

logger.info(f'Reading matrix {matrix_path}')

adata = import_tcc_matrix_as_anndata(

matrix_path,

barcodes_path,

ec_path,

txnames_path,

threads=threads,

loom=loom,

loom_names=loom_names,

batch_barcodes_path=batch_barcodes_path

) if tcc else import_matrix_as_anndata(

matrix_path,

barcodes_path,

genes_path,

t2g_path=t2g_path,

name=name,

by_name=by_name,

loom=loom,

loom_names=loom_names,

batch_barcodes_path=batch_barcodes_path

)

if loom:

loom_path = os.path.join(counts_dir, f'{ADATA_PREFIX}.loom')

logger.info(f'Writing matrix to loom {loom_path}')

adata.write_loom(loom_path)

results.update({'loom': loom_path})

if h5ad:

h5ad_path = os.path.join(counts_dir, f'{ADATA_PREFIX}.h5ad')

logger.info(f'Writing matrix to h5ad {h5ad_path}')

adata.write(h5ad_path)

results.update({'h5ad': h5ad_path})

return results

def convert_matrices(

counts_dir: str,

matrix_paths: List[str],

barcodes_paths: List[str],

batch_barcodes_paths: Optional[List[str]] = None,

genes_paths: Optional[List[str]] = None,

ec_paths: Optional[List[str]] = None,

t2g_path: Optional[str] = None,

txnames_path: Optional[str] = None,

name: str = 'gene',

loom: bool = False,

loom_names: List[str] = ['barcode', 'target_name'],

h5ad: bool = False,

by_name: bool = False,

nucleus: bool = False,

tcc: bool = False,

threads: int = 8,

) -> Dict[str, str]:

"""Convert a gene count or TCC matrix to loom or h5ad.

Args:

counts_dir: Path to counts directory

matrix_paths: List of paths to matrices

barcodes_paths: List of paths to barcodes.txt

batch_barcodes_path: Paths to barcodes prefixed with sample ID,

defaults to `None`

genes_paths: List of paths to genes.txt, defaults to `None`

ec_paths: List of path to ec.txt, defaults to `None`

t2g_path: Path to transcript-to-gene mapping. If this is provided,

the third column of the mapping is appended to the anndata var,

defaults to `None`

txnames_path: List of paths to transcripts.txt, defaults to `None`

name: Name of the columns, defaults to "gene"

loom: Whether to generate loom file, defaults to `False`

loom_names: Names for col_attrs and row_attrs in loom file,

defaults to `['barcode','target_name']`

h5ad: Whether to generate h5ad file, defaults to `False`

by_name: Aggregate counts by name instead of ID.

nucleus: Whether the matrices contain single nucleus counts, defaults to `False`

tcc: Whether the matrix is a TCC matrix, defaults to `False`

threads: Number of threads to use, defaults to `8`

Returns:

Dictionary of generated files

"""

results = {}

adatas = []

matrix_paths = matrix_paths or []

barcodes_paths = barcodes_paths or []

batch_barcodes_paths = batch_barcodes_paths or []

if not batch_barcodes_paths:

batch_barcodes_paths = [None for x in matrix_paths]

genes_paths = genes_paths or []

ec_paths = ec_paths or []

for matrix_path, barcodes_path, batch_barcodes_path, genes_ec_path in zip(

matrix_paths, barcodes_paths, batch_barcodes_paths, ec_paths

if not genes_paths or None in genes_paths else genes_paths):

logger.info(f'Reading matrix {matrix_path}')

adatas.append(

import_tcc_matrix_as_anndata(

matrix_path,

barcodes_path,

genes_ec_path,

txnames_path,

threads=threads,

loom=loom,

loom_names=loom_names,

batch_barcodes_path=batch_barcodes_path

) if tcc else import_matrix_as_anndata(

matrix_path,

barcodes_path,

genes_ec_path,

t2g_path=t2g_path,

name=name,

by_name=by_name,

loom=loom,

loom_names=loom_names,

batch_barcodes_path=batch_barcodes_path

)

)

logger.info('Combining matrices')

adata = sum_anndatas(*adatas) if nucleus else overlay_anndatas(*adatas)

if loom:

loom_path = os.path.join(counts_dir, f'{ADATA_PREFIX}.loom')

logger.info(f'Writing matrices to loom {loom_path}')

adata.write_loom(loom_path)

results.update({'loom': loom_path})

if h5ad:

h5ad_path = os.path.join(counts_dir, f'{ADATA_PREFIX}.h5ad')

logger.info(f'Writing matrices to h5ad {h5ad_path}')

adata.write(h5ad_path)

results.update({'h5ad': h5ad_path})

return results

def count_result_to_dict(count_result: Dict[str, str]) -> List[Dict[str, str]]:

"""Converts count result dict to list.

Args:

count_result: Count result object returned by bustools_count

Returns:

List of count result dicts

"""

new_count_result = []

for i in range(len(count_result)):

if f'mtx{i}' not in count_result:

break

new_count_result.append({

'mtx':

count_result[f'mtx{i}'],

'ec' if f'ec{i}' in count_result else 'genes':

count_result[f'ec{i}' if f'ec{i}' in

count_result else f'genes{i}'],

'barcodes':

count_result[f'barcodes{i}'],

'batch_barcodes':

count_result[f'batch_barcodes{i}']

if f'batch_barcodes{i}' in count_result else None,

})

return new_count_result

def filter_with_bustools(

bus_path: str,

ecmap_path: str,

txnames_path: str,

t2g_path: str,

whitelist_path: str,

filtered_bus_path: str,

filter_threshold: Optional[int] = None,

counts_prefix: Optional[str] = None,

tcc: bool = False,

mm: bool = False,

kite: bool = False,

temp_dir: str = 'tmp',

threads: int = 8,

memory: str = '2G',

count: bool = True,

loom: bool = False,

loom_names: List[str] = ['barcode', 'target_name'],

h5ad: bool = False,

by_name: bool = False,

cellranger: bool = False,

gzip: bool = False,

umi_gene: bool = True,

em: bool = False,

) -> Dict[str, str]:

"""Generate filtered count matrices with bustools.

Args:

bus_path: Path to sorted, corrected, sorted BUS file

ecmap_path: Path to matrix ec file

txnames_path: Path to list of transcripts

t2g_path: Path to transcript-to-gene mapping

whitelist_path: Path to filter whitelist to generate

filtered_bus_path: Path to filtered BUS file to generate

filter_threshold: Barcode filter threshold for bustools, defaults

to `None`

counts_prefix: Prefix of count matrix, defaults to `None`

tcc: Whether to generate a TCC matrix instead of a gene count matrix,

defaults to `False`

mm: Whether to include BUS records that pseudoalign to multiple genes,

defaults to `False`

kite: Whether this is a KITE workflow

temp_dir: Path to temporary directory, defaults to `tmp`

threads: Number of threads to use, defaults to `8`

memory: Amount of memory to use, defaults to `2G`

count: Whether to run `bustools count`, defaults to `True`

loom: Whether to convert the final count matrix into a loom file,

defaults to `False`

loom_names: Names for col_attrs and row_attrs in loom file,

defaults to `['barcode','target_name']`

h5ad: Whether to convert the final count matrix into a h5ad file,

defaults to `False`

by_name: Aggregate counts by name instead of ID.

cellranger: Whether to convert the final count matrix into a

cellranger-compatible matrix, defaults to `False`

umi_gene: Whether to perform gene-level UMI collapsing, defaults to

`True`

em: Whether to estimate gene abundances using EM algorithm, defaults to

`False`

Returns:

Dictionary of generated files

"""

logger.info('Filtering with bustools')

results = {}