From 994c58d5d2dde310e344a8dfc0c1b7a5933cbacc Mon Sep 17 00:00:00 2001 From: Kun Qian Date: Sat, 11 Apr 2026 00:09:27 +0800 Subject: [PATCH] feat: port current v3 topic modeling changes --- pyproject.toml | 1 + .../cli/subcommand/topic_modeling.py | 662 ++++++++---------- .../topic_modeling/create_anndata.py | 74 +- .../topic_modeling/mallet_models.py | 521 ++++---------- src/pycisTopic/topic_modeling/plot_stats.py | 69 +- src/pycisTopic/topic_modeling/stats.py | 186 ++--- .../topic_modeling/tomotopy_models.py | 195 ++++++ src/pycisTopic/topic_modeling/topic_models.py | 151 ++++ 8 files changed, 929 insertions(+), 930 deletions(-) create mode 100644 src/pycisTopic/topic_modeling/tomotopy_models.py create mode 100644 src/pycisTopic/topic_modeling/topic_models.py diff --git a/pyproject.toml b/pyproject.toml index e84dec9..5cc0f78 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,7 @@ dependencies = [ "pybiomart>=0.2.0", "requests>=2.32.4", "scipy>=1.15.3", + "tomotopy>=0.14.0", "tmtoolkit>=0.12.0", ] diff --git a/src/pycisTopic/cli/subcommand/topic_modeling.py b/src/pycisTopic/cli/subcommand/topic_modeling.py index 4f4bfc1..b3083f0 100644 --- a/src/pycisTopic/cli/subcommand/topic_modeling.py +++ b/src/pycisTopic/cli/subcommand/topic_modeling.py @@ -1,10 +1,9 @@ from __future__ import annotations +import gzip import logging import os -import pickle import sys -import tempfile from argparse import ArgumentTypeError from typing import TYPE_CHECKING @@ -12,142 +11,155 @@ from argparse import ArgumentParser, _SubParsersAction -def run_topic_modeling_with_mallet(args): - from pycisTopic.topic_modeling.mallet_models import LDAMallet - - check_java_exists() - - mallet_corpus_filename = args.mallet_corpus_filename - output_prefix = args.output_prefix - n_topics_list = [args.topics] if isinstance(args.topics, int) else args.topics - alpha = args.alpha - alpha_by_topic = args.alpha_by_topic - eta = args.eta - eta_by_topic = args.eta_by_topic - n_iter = args.iterations - optimize_interval = args.optimize_interval - optimize_burn_in = args.optimize_burn_in - n_threads = args.parallel - random_seed = args.seed - memory_in_gb = f"{args.memory_in_gb}G" - mallet_path = args.mallet_path - - if args.verbose: - level = logging.INFO - log_format = "%(asctime)s %(name)-12s %(levelname)-8s %(message)s" - handlers = [logging.StreamHandler(stream=sys.stdout)] - logging.basicConfig(level=level, format=log_format, handlers=handlers) - - print("Run topic modeling with Mallet with the following settings:") - print(f" - Mallet corpus filename: {mallet_corpus_filename}") - print(f" - Output prefix: {output_prefix}") - print(f" - Number of topics to run topic modeling for: {n_topics_list}") - print(f" - Alpha: {alpha}") - print(f" - Divide alpha by the number of topics: {alpha_by_topic}") - print(f" - Eta: {eta}") - print(f" - Divide eta by the number of topics: {eta_by_topic}") - print(f" - Number of iterations of Gibbs sampling: {n_iter}") - print(f" - Optimize interval for hyperparameters: {optimize_interval}") - print(f" - Number of burn-in iterations: {optimize_burn_in}") - print(f" - Number threads Mallet is allowed to use: {n_threads}") - print(f" - Seed: {random_seed}") - print(f" - Amount of memory Mallet is allowed to use: {memory_in_gb}") - print(f" - Mallet binary: {mallet_path}") - - os.environ["MALLET_MEMORY"] = memory_in_gb - - for n_topics in n_topics_list: - # Run models - print(f"\nRunning Mallet topic modeling for {n_topics} topics.") - print(f"----------------------------------{'-' * len(str(n_topics))}--------") - - LDAMallet.run_mallet_topic_modeling( - mallet_corpus_filename=mallet_corpus_filename, - output_prefix=output_prefix, - n_topics=n_topics, - alpha=alpha, - alpha_by_topic=alpha_by_topic, - eta=eta, - eta_by_topic=eta_by_topic, - n_threads=n_threads, - iterations=n_iter, - optimize_interval=optimize_interval, - optimize_burn_in=optimize_burn_in, - topic_threshold=0.0, - random_seed=random_seed, - mallet_path=mallet_path, - ) - - print( - f'\nWriting Mallet topic modeling output files to "{output_prefix}.{n_topics}_topics.*"...' - ) - - -def run_convert_binary_matrix_to_mallet_corpus_file(args): +def run_create_corpus(args) -> None: import scipy from pycisTopic.topic_modeling.mallet_models import LDAMallet check_java_exists() - - binary_accessibility_matrix_filename = args.binary_accessibility_matrix_filename - mallet_corpus_filename = args.mallet_corpus_filename - mallet_path = args.mallet_path - memory_in_gb = f"{args.memory_in_gb}G" - - if args.verbose: - level = logging.INFO - log_format = "%(asctime)s %(name)-12s %(levelname)-8s %(message)s" - handlers = [logging.StreamHandler(stream=sys.stdout)] - logging.basicConfig(level=level, format=log_format, handlers=handlers) + _configure_logging(verbose=args.verbose) print( - f'Read binary accessibility matrix from "{binary_accessibility_matrix_filename}" Matrix Market file.' + f'Read binary accessibility matrix from "{args.binary_accessibility_matrix_filename}" Matrix Market file.' + ) + binary_accessibility_matrix = scipy.io.mmread( + args.binary_accessibility_matrix_filename ) - binary_accessibility_matrix = scipy.io.mmread(binary_accessibility_matrix_filename) - - os.environ["MALLET_MEMORY"] = memory_in_gb + os.environ["MALLET_MEMORY"] = f"{args.memory_in_gb}G" print( - f'Convert binary accessibility matrix to Mallet serialized corpus file "{mallet_corpus_filename}".' + f'Convert binary accessibility matrix to Mallet serialized corpus file "{args.mallet_corpus_filename}".' ) LDAMallet.convert_binary_matrix_to_mallet_corpus_file( binary_accessibility_matrix=binary_accessibility_matrix, - mallet_corpus_filename=mallet_corpus_filename, - mallet_path=mallet_path, - ) + mallet_corpus_filename=args.mallet_corpus_filename, + mallet_path=args.mallet_path, + ) + + +def run_topic_modeling(args) -> None: + n_topics_list = args.topics if isinstance(args.topics, list) else [args.topics] + _configure_logging(verbose=args.verbose) + + if args.backend == "mallet": + from pycisTopic.topic_modeling.mallet_models import LDAMallet + + check_java_exists() + os.environ["MALLET_MEMORY"] = f"{args.memory_in_gb}G" + + print("Run topic modeling with Mallet with the following settings:") + print(f" - Mallet corpus filename: {args.input_filename}") + print(f" - Output prefix: {args.output_prefix}") + print(f" - Topics: {n_topics_list}") + print(f" - Threads: {args.threads}") + print(f" - Alpha: {args.alpha}") + print(f" - Alpha by topic: {args.alpha_by_topic}") + print(f" - Eta: {args.eta}") + print(f" - Eta by topic: {args.eta_by_topic}") + print(f" - Iterations: {args.iterations}") + print(f" - Optimize interval: {args.optimize_interval}") + print(f" - Optimize burn-in: {args.optimize_burn_in}") + print(f" - Seed: {args.seed}") + print(f" - Mallet memory: {args.memory_in_gb}G") + print(f" - Mallet binary: {args.mallet_path}") + + for n_topics in n_topics_list: + print(f"\nRunning Mallet topic modeling for {n_topics} topics.") + LDAMallet.run_topic_modeling( + mallet_corpus_filename=args.input_filename, + output_prefix=args.output_prefix, + n_topics=n_topics, + alpha=args.alpha, + alpha_by_topic=args.alpha_by_topic, + eta=args.eta, + eta_by_topic=args.eta_by_topic, + n_threads=args.threads, + iterations=args.iterations, + optimize_interval=args.optimize_interval, + optimize_burn_in=args.optimize_burn_in, + topic_threshold=0.0, + random_seed=args.seed, + mallet_path=args.mallet_path, + ) + return + + import scipy + from pycisTopic.topic_modeling.tomotopy_models import LDATomotopy + + if args.cell_barcodes_filename is None or args.region_ids_filename is None: + raise ValueError( + "`--cb` and `--regions` should be provided when backend is `tomotopy`." + ) -def run_mallet_calculate_model_evaluation_stats(args): + print( + f'Read binary accessibility matrix from "{args.input_filename}" Matrix Market file.' + ) + binary_accessibility_matrix = scipy.io.mmread(args.input_filename) + cell_names = _read_names_file(args.cell_barcodes_filename) + region_names = _read_names_file(args.region_ids_filename) + + print("Run topic modeling with tomotopy with the following settings:") + print(f" - Binary accessibility matrix: {args.input_filename}") + print(f" - Output prefix: {args.output_prefix}") + print(f" - Topics: {n_topics_list}") + print(f" - Threads: {args.threads}") + print(f" - Alpha: {args.alpha}") + print(f" - Alpha by topic: {args.alpha_by_topic}") + print(f" - Eta: {args.eta}") + print(f" - Eta by topic: {args.eta_by_topic}") + print(f" - Iterations: {args.iterations}") + print(f" - Optimize interval: {args.optimize_interval}") + print(f" - Seed: {args.seed}") + + for n_topics in n_topics_list: + print(f"\nRunning tomotopy topic modeling for {n_topics} topics.") + LDATomotopy.run_topic_modeling( + binary_accessibility_matrix=binary_accessibility_matrix, + cell_names=cell_names, + region_names=region_names, + output_prefix=args.output_prefix, + n_topics=n_topics, + alpha=args.alpha, + alpha_by_topic=args.alpha_by_topic, + eta=args.eta, + eta_by_topic=args.eta_by_topic, + n_threads=args.threads, + iterations=args.iterations, + optimize_interval=args.optimize_interval, + random_seed=args.seed, + ) + + +def run_calculate_model_evaluation_stats(args) -> None: import scipy from pycisTopic.topic_modeling.stats import calculate_model_evaluation_stats - binary_accessibility_matrix_filename = args.binary_accessibility_matrix_filename - output_prefix = args.output_prefix - n_topics_list = [args.topics] if isinstance(args.topics, int) else args.topics - print( - f'Read binary accessibility matrix from "{binary_accessibility_matrix_filename}" Matrix Market file.' + f'Read binary accessibility matrix from "{args.binary_accessibility_matrix_filename}" Matrix Market file.' + ) + binary_accessibility_matrix = scipy.io.mmread( + args.binary_accessibility_matrix_filename ) - binary_accessibility_matrix = scipy.io.mmread(binary_accessibility_matrix_filename) + n_topics_list = args.topics if isinstance(args.topics, list) else [args.topics] for n_topics in n_topics_list: print( - f'Calculate model evaluation statistics for {n_topics} topics from "{output_prefix}.{n_topics}_topics.*"...' + f'Calculate model evaluation statistics for {n_topics} topics from "{args.output_prefix}.{n_topics}_topics.*"...' ) calculate_model_evaluation_stats( binary_accessibility_matrix=binary_accessibility_matrix, - output_prefix=output_prefix, + output_prefix=args.output_prefix, n_topics=n_topics, top_topics_coh=5, ) -def run_mallet_plot_model_evaluation_stats(args): +def run_plot_model_evaluation_stats(args) -> None: from pycisTopic.topic_modeling.plot_stats import plot_stats - n_topics_list = [args.n_topics] if isinstance(args.n_topics, int) else args.n_topics + n_topics_list = args.n_topics if isinstance(args.n_topics, list) else [args.n_topics] plot_stats( output_prefix=args.output_prefix, n_topics=n_topics_list, @@ -155,167 +167,138 @@ def run_mallet_plot_model_evaluation_stats(args): ) -def binarize_cell_or_region_topic(args): +def binarize_cell_or_region_topic(args) -> None: """Binarize cell-topics or region-topics.""" - target = args.target - method = args.method - ntop = args.ntop - smooth_topics = args.smooth_topics - nbins = args.nbins - cell_barcodes_filename = args.cell_barcodes_filename - region_ids_filename = args.region_ids_filename - output_prefix = args.output_prefix - n_topics = args.n_topics - out_dir = args.out_dir - - # input validation - if target == "cell" and cell_barcodes_filename is None: + from pycisTopic.topic_binarization import binarize_topics + from pycisTopic.topic_modeling.topic_models import TopicModelFilenames, load_topic_model_backend + + if args.target == "cell" and args.cell_barcodes_filename is None: raise ValueError( - "`cell_barcodes_filename` using `--cb` should be provided when target is `cell`" + "`cell_barcodes_filename` using `--cb` should be provided when target is `cell`." ) - if target == "region" and region_ids_filename is None: + if args.target == "region" and args.region_ids_filename is None: raise ValueError( - "`region_ids_filename` using `--regions` should be provided when target is `region`" + "`region_ids_filename` using `--regions` should be provided when target is `region`." ) - import os + if not os.path.exists(args.out_dir): + print(f"Making directory: {args.out_dir}") + os.makedirs(args.out_dir) - if not os.path.exists(out_dir): - print(f"Making directory: {out_dir}") - os.makedirs(out_dir) - - from pycisTopic.fragments import read_barcodes_file_to_polars_series - from pycisTopic.topic_binarization import binarize_topics - from pycisTopic.topic_modeling.mallet_models import LDAMallet, LDAMalletFilenames - - lda_mallet_filenames = LDAMalletFilenames( - output_prefix=output_prefix, n_topics=n_topics + filenames = TopicModelFilenames(output_prefix=args.output_prefix, n_topics=args.n_topics) + backend_cls = load_topic_model_backend( + output_prefix=args.output_prefix, + n_topics=args.n_topics, ) - if target == "cell": - print(f'Read cell barcodes filename "{cell_barcodes_filename}".') - cell_or_region_names = read_barcodes_file_to_polars_series( - barcodes_tsv_filename=cell_barcodes_filename, - sample_id=None, - cb_end_to_remove=None, - cb_sample_separator=None, - ).to_list() + if args.target == "cell": + print(f'Read cell barcodes filename "{args.cell_barcodes_filename}".') + cell_or_region_names = _read_names_file(args.cell_barcodes_filename) print( - f'Read cell-topic probabilities filename "{lda_mallet_filenames.cell_topic_probabilities_parquet_filename}".' + f'Read cell-topic probabilities filename "{filenames.cell_topic_probabilities_parquet_filename}".' ) - cell_or_region_topic_prob = LDAMallet.read_cell_topic_probabilities_parquet_file( - mallet_cell_topic_probabilities_parquet_filename=lda_mallet_filenames.cell_topic_probabilities_parquet_filename + cell_or_region_topic_prob = backend_cls.read_cell_topic_probabilities_parquet_file( + cell_topic_probabilities_parquet_filename=filenames.cell_topic_probabilities_parquet_filename ) - - if target == "region": - print(f'Read region IDs filename "{region_ids_filename}".') - cell_or_region_names = read_barcodes_file_to_polars_series( - barcodes_tsv_filename=region_ids_filename, - sample_id=None, - cb_end_to_remove=None, - cb_sample_separator=None, - ).to_list() + else: + print(f'Read region IDs filename "{args.region_ids_filename}".') + cell_or_region_names = _read_names_file(args.region_ids_filename) print( - f'Read region-topic probabilities filename "{lda_mallet_filenames.region_topic_counts_parquet_filename}".' + f'Read region-topic probabilities filename "{filenames.region_topic_counts_parquet_filename}".' + ) + cell_or_region_topic_prob = ( + backend_cls.read_region_topic_counts_parquet_file_to_region_topic_probabilities( + region_topic_counts_parquet_filename=filenames.region_topic_counts_parquet_filename + ).T ) - cell_or_region_topic_prob = LDAMallet.read_region_topic_counts_parquet_file_to_region_topic_probabilities( - mallet_region_topic_counts_parquet_filename=lda_mallet_filenames.region_topic_counts_parquet_filename - ).T print("Binarizing topics ...") cell_or_region_names_per_topic, scores_per_topic, thresholds = binarize_topics( cell_or_region_topic_prob=cell_or_region_topic_prob, cell_or_region_names=cell_or_region_names, - method=method, - smooth_topics=smooth_topics, - ntop=ntop, - nbins=nbins, + method=args.method, + smooth_topics=args.smooth_topics, + ntop=args.ntop, + nbins=args.nbins, ) - print(f'Saving results to "{out_dir}".') - - with open(os.path.join(out_dir, f"{target}_thresholds.tsv"), "w") as f: - for topic, thr in enumerate(thresholds): - f.write(f"{topic + 1}\t{thr}\n") + print(f'Saving results to "{args.out_dir}".') + with open(os.path.join(args.out_dir, f"{args.target}_thresholds.tsv"), "w") as fh: + for topic, threshold in enumerate(thresholds): + fh.write(f"{topic + 1}\t{threshold}\n") - if target == "cell": + if args.target == "cell": for topic, (cells, scores) in enumerate( zip(cell_or_region_names_per_topic, scores_per_topic) ): with open( - os.path.join(out_dir, f"{target}_Topic_{topic + 1}_binarized.txt"), "w" - ) as f: + os.path.join(args.out_dir, f"cell_Topic_{topic + 1}_binarized.txt"), + "w", + ) as fh: for cell, score in zip(cells, scores): - f.write(f"{cell}\t{score}\n") - - elif target == "region": + fh.write(f"{cell}\t{score}\n") + else: for topic, (regions, scores) in enumerate( zip(cell_or_region_names_per_topic, scores_per_topic) ): with open( - os.path.join(out_dir, f"{target}_Topic_{topic + 1}_binarized.bed"), "w" - ) as f: + os.path.join(args.out_dir, f"region_Topic_{topic + 1}_binarized.bed"), + "w", + ) as fh: for region, score in zip(regions, scores): chrom, start, end = region.replace(":", "-").split("-") - f.write(f"{chrom}\t{start}\t{end}\tTopic_{topic + 1}\t{score}\n") - - -def run_create_anndata_from_mallet(args): - from pycisTopic.topic_modeling.create_anndata import create_anndata_from_mallet + fh.write(f"{chrom}\t{start}\t{end}\tTopic_{topic + 1}\t{score}\n") - cell_barcodes: list[str] = [] - with open(args.cell_barcodes) as f: - for line in f: - cell_barcodes.append(line.strip()) - region_ids: list[str] = [] - with open(args.region_ids) as f: - for line in f: - region_ids.append(line.strip()) +def run_create_anndata(args) -> None: + from pycisTopic.topic_modeling.create_anndata import create_anndata_from_topic_model - create_anndata_from_mallet( + create_anndata_from_topic_model( output_prefix=args.output_prefix, n_topics=args.n_topics, - cell_barcodes=cell_barcodes, - region_ids=region_ids, + cell_barcodes=_read_names_file(args.cell_barcodes), + region_ids=_read_names_file(args.region_ids), ) -def str_to_bool(v: str) -> bool: - """ - Convert string representation of a boolean value to a boolean. +def _configure_logging(verbose: bool) -> None: + if not verbose: + return - Parameters - ---------- - v - String representation of a boolean value. - After conversion to lowercase, the following string values can be converted: - - "yes", "true", "t", "y", "1" -> True - - "no", "false", "f", "n", "0" -> False + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(name)-12s %(levelname)-8s %(message)s", + handlers=[logging.StreamHandler(stream=sys.stdout)], + force=True, + ) + + +def _read_names_file(filename: str) -> list[str]: + opener = gzip.open if filename.endswith(".gz") else open + with opener(filename, "rt", encoding="utf-8") as fh: + return [line.strip() for line in fh if line.strip()] - Returns - ------- - True or False - """ - if isinstance(v, str): - v = v.lower() - if v.lower() in ("yes", "true", "t", "y", "1"): +def str_to_bool(value: str) -> bool: + """Convert a string representation of a boolean value to a boolean.""" + if isinstance(value, str): + lowered = value.lower() + if lowered in ("yes", "true", "t", "y", "1"): return True - elif v.lower() in ("no", "false", "f", "n", "0"): + if lowered in ("no", "false", "f", "n", "0"): return False raise ArgumentTypeError("Boolean value expected.") -def check_java_exists(): +def check_java_exists() -> None: import subprocess print("Checking whether Java exists.") subprocess.run(["java", "--version"], shell=False, stdout=subprocess.DEVNULL) -def add_parser_topic_modeling(subparsers: _SubParsersAction[ArgumentParser]): - """Creates an ArgumentParser to read the options for this script.""" +def add_parser_topic_modeling(subparsers: _SubParsersAction[ArgumentParser]) -> None: + """Create the topic modeling CLI parser.""" parser_topic_modeling = subparsers.add_parser( "topic_modeling", help="Run LDA topic modeling.", @@ -328,62 +311,46 @@ def add_parser_topic_modeling(subparsers: _SubParsersAction[ArgumentParser]): ) subparser_topic_modeling.required = True - parser_topic_modeling_mallet = subparser_topic_modeling.add_parser( - "mallet", help="Run LDA topic modeling with Mallet." - ) - - subparser_topic_modeling_mallet = parser_topic_modeling_mallet.add_subparsers( - title="Topic modeling with Mallet", - dest="mallet", - help="List of Mallet topic modeling subcommands.", - ) - subparser_topic_modeling_mallet.required = True - - parser_topic_modeling_mallet_create_corpus = subparser_topic_modeling_mallet.add_parser( + parser_create_corpus = subparser_topic_modeling.add_parser( "create_corpus", - help="Convert binary accessibility matrix to Mallet serialized corpus file.", - ) - parser_topic_modeling_mallet_create_corpus.set_defaults( - func=run_convert_binary_matrix_to_mallet_corpus_file + help="Convert a binary accessibility matrix to a Mallet serialized corpus file.", ) - - parser_topic_modeling_mallet_create_corpus.add_argument( + parser_create_corpus.set_defaults(func=run_create_corpus) + parser_create_corpus.add_argument( "-i", "--input", dest="binary_accessibility_matrix_filename", - action="store", type=str, required=True, - help="Binary accessibility matrix (region IDs vs cell barcodes) in Matrix Market format.", + help="Binary accessibility matrix in Matrix Market format.", ) - parser_topic_modeling_mallet_create_corpus.add_argument( + parser_create_corpus.add_argument( "-o", "--output", dest="mallet_corpus_filename", - action="store", type=str, required=True, help="Mallet serialized corpus filename.", ) - parser_topic_modeling_mallet_create_corpus.add_argument( + parser_create_corpus.add_argument( "-m", "--memory", dest="memory_in_gb", type=int, required=False, default=10, - help='Amount of memory (in GB) Mallet is allowed to use. Default: "10".', + help='Amount of memory in GB that Mallet is allowed to use. Default: "10".', ) - parser_topic_modeling_mallet_create_corpus.add_argument( + parser_create_corpus.add_argument( "-b", "--mallet_path", dest="mallet_path", type=str, required=False, default="mallet", - help='Path to Mallet binary (e.g. "/xxx/Mallet/bin/mallet"). Default: "mallet".', + help='Path to the Mallet binary. Default: "mallet".', ) - parser_topic_modeling_mallet_create_corpus.add_argument( + parser_create_corpus.add_argument( "-v", "--verbose", dest="verbose", @@ -392,31 +359,36 @@ def add_parser_topic_modeling(subparsers: _SubParsersAction[ArgumentParser]): help="Enable verbose mode.", ) - parser_topic_modeling_mallet_run = subparser_topic_modeling_mallet.add_parser( + parser_run = subparser_topic_modeling.add_parser( "run", - help="Run LDA topic modeling with Mallet.", + help="Run topic modeling with the selected backend.", ) - parser_topic_modeling_mallet_run.set_defaults(func=run_topic_modeling_with_mallet) - - parser_topic_modeling_mallet_run.add_argument( + parser_run.set_defaults(func=run_topic_modeling) + parser_run.add_argument( + "--backend", + dest="backend", + type=str, + choices=("mallet", "tomotopy"), + required=True, + help="Topic modeling backend to use.", + ) + parser_run.add_argument( "-i", "--input", - dest="mallet_corpus_filename", - action="store", + dest="input_filename", type=str, required=True, - help="Mallet corpus filename.", + help="Input filename. Use a Mallet corpus for `mallet` or a Matrix Market matrix for `tomotopy`.", ) - parser_topic_modeling_mallet_run.add_argument( + parser_run.add_argument( "-o", "--output", dest="output_prefix", - action="store", type=str, required=True, help="Topic model output prefix.", ) - parser_topic_modeling_mallet_run.add_argument( + parser_run.add_argument( "-t", "--topics", dest="topics", @@ -425,15 +397,15 @@ def add_parser_topic_modeling(subparsers: _SubParsersAction[ArgumentParser]): nargs="+", help="Number(s) of topics to create during topic modeling.", ) - parser_topic_modeling_mallet_run.add_argument( + parser_run.add_argument( "-p", - "--parallel", - dest="parallel", + "--threads", + dest="threads", type=int, required=True, - help="Number of threads Mallet is allowed to use.", + help="Number of threads the backend is allowed to use.", ) - parser_topic_modeling_mallet_run.add_argument( + parser_run.add_argument( "-n", "--iterations", dest="iterations", @@ -442,36 +414,32 @@ def add_parser_topic_modeling(subparsers: _SubParsersAction[ArgumentParser]): default=150, help="Number of iterations of Gibbs sampling. Default: 150.", ) - parser_topic_modeling_mallet_run.add_argument( + parser_run.add_argument( "--optimize-interval", dest="optimize_interval", type=int, required=False, default=0, - help="Optimize hyperparameters every `optimize_interval` iterations. " - "Only takes effect after running `optimize_burn_in` iterations. " - "Disable optimizing hyperparameters by setting this option to 0. " - "Default: 0.", + help="Optimize hyperparameters every `optimize_interval` iterations. Default: 0.", ) - parser_topic_modeling_mallet_run.add_argument( + parser_run.add_argument( "--optimize-burn-in", dest="optimize_burn_in", type=int, required=False, default=50, - help="The number of iterations before starting hyperparameter optimization. " - "Default: 50.", + help="Number of iterations before starting hyperparameter optimization. Default: 50.", ) - parser_topic_modeling_mallet_run.add_argument( + parser_run.add_argument( "-a", "--alpha", dest="alpha", - type=int, + type=float, required=False, default=50, help="Alpha value. Default: 50.", ) - parser_topic_modeling_mallet_run.add_argument( + parser_run.add_argument( "-A", "--alpha_by_topic", dest="alpha_by_topic", @@ -479,9 +447,9 @@ def add_parser_topic_modeling(subparsers: _SubParsersAction[ArgumentParser]): choices=(True, False), required=False, default=True, - help="Whether the alpha value should by divided by the number of topics. Default: True.", + help="Whether alpha should be divided by the number of topics. Default: True.", ) - parser_topic_modeling_mallet_run.add_argument( + parser_run.add_argument( "-e", "--eta", dest="eta", @@ -490,7 +458,7 @@ def add_parser_topic_modeling(subparsers: _SubParsersAction[ArgumentParser]): default=0.1, help="Eta value. Default: 0.1.", ) - parser_topic_modeling_mallet_run.add_argument( + parser_run.add_argument( "-E", "--eta_by_topic", dest="eta_by_topic", @@ -498,38 +466,52 @@ def add_parser_topic_modeling(subparsers: _SubParsersAction[ArgumentParser]): choices=(True, False), required=False, default=False, - help="Whether the eta value should by divided by the number of topics. Default: False.", + help="Whether eta should be divided by the number of topics. Default: False.", ) - parser_topic_modeling_mallet_run.add_argument( + parser_run.add_argument( "-s", "--seed", dest="seed", type=int, required=False, default=555, - help="Seed for ensuring reproducibility. " - "To get reproducible output, Mallet also has to be run with the same number of threads. " - "Default: 555.", + help="Seed for ensuring reproducibility. Default: 555.", ) - parser_topic_modeling_mallet_run.add_argument( + parser_run.add_argument( "-m", "--memory", dest="memory_in_gb", type=int, required=False, default=100, - help='Amount of memory (in GB) Mallet is allowed to use. Default: "100".', + help='Amount of memory in GB Mallet is allowed to use. Default: "100".', ) - parser_topic_modeling_mallet_run.add_argument( + parser_run.add_argument( "-b", "--mallet_path", dest="mallet_path", type=str, required=False, default="mallet", - help='Path to Mallet binary (e.g. "/xxx/Mallet/bin/mallet"). Default: "mallet".', + help='Path to the Mallet binary. Default: "mallet".', + ) + parser_run.add_argument( + "-c", + "--cb", + dest="cell_barcodes_filename", + type=str, + required=False, + help="Filename with cell barcodes. Required for `tomotopy`.", ) - parser_topic_modeling_mallet_run.add_argument( + parser_run.add_argument( + "-r", + "--regions", + dest="region_ids_filename", + type=str, + required=False, + help="Filename with region IDs. Required for `tomotopy`.", + ) + parser_run.add_argument( "-v", "--verbose", dest="verbose", @@ -538,232 +520,194 @@ def add_parser_topic_modeling(subparsers: _SubParsersAction[ArgumentParser]): help="Enable verbose mode.", ) - parser_topic_modeling_mallet_calculate_stats = ( - subparser_topic_modeling_mallet.add_parser( - "stats", - help="Calculate model evaluation statistics.", - ) - ) - parser_topic_modeling_mallet_calculate_stats.set_defaults( - func=run_mallet_calculate_model_evaluation_stats + parser_stats = subparser_topic_modeling.add_parser( + "stats", + help="Calculate model evaluation statistics.", ) - - parser_topic_modeling_mallet_calculate_stats.add_argument( + parser_stats.set_defaults(func=run_calculate_model_evaluation_stats) + parser_stats.add_argument( "-i", "--input", dest="binary_accessibility_matrix_filename", - action="store", type=str, required=True, - help="Binary accessibility matrix (region IDs vs cell barcodes) in Matrix Market format.", + help="Binary accessibility matrix in Matrix Market format.", ) - parser_topic_modeling_mallet_calculate_stats.add_argument( + parser_stats.add_argument( "-o", "--output", dest="output_prefix", - action="store", type=str, required=True, help="Topic model output prefix.", ) - parser_topic_modeling_mallet_calculate_stats.add_argument( + parser_stats.add_argument( "-t", "--topics", dest="topics", type=int, required=True, nargs="+", - help="Topic number(s) to create the model evaluation statistics for.", - ) - parser_topic_modeling_mallet_calculate_stats.add_argument( - "-v", - "--verbose", - dest="verbose", - action="store_true", - required=False, - help="Enable verbose mode.", + help="Topic number(s) to calculate model evaluation statistics for.", ) - parser_topic_modeling_mallet_plot_stats = ( - subparser_topic_modeling_mallet.add_parser( - "plot_stats", - help="Plot evaluation statistics.", - ) + parser_plot_stats = subparser_topic_modeling.add_parser( + "plot_stats", + help="Plot evaluation statistics.", ) - parser_topic_modeling_mallet_plot_stats.set_defaults( - func=run_mallet_plot_model_evaluation_stats - ) - parser_topic_modeling_mallet_plot_stats.add_argument( + parser_plot_stats.set_defaults(func=run_plot_model_evaluation_stats) + parser_plot_stats.add_argument( "-o", "--output", dest="output_prefix", - action="store", type=str, required=True, help="Topic model output prefix.", ) - parser_topic_modeling_mallet_plot_stats.add_argument( + parser_plot_stats.add_argument( "-t", "--topics", dest="n_topics", type=int, required=True, nargs="+", - help="Topic number(s) to create the model evaluation statistics for.", + help="Topic number(s) to plot model evaluation statistics for.", ) - parser_topic_modeling_mallet_plot_stats.add_argument( + parser_plot_stats.add_argument( "-q", "--format", dest="plot_file_format", - action="store", type=str, required=False, default="png", - help="File format of the resulting plots. Default: png.", + help="File format of the resulting plot. Default: png.", ) - parser_topic_modeling_mallet_binarize = subparser_topic_modeling_mallet.add_parser( + parser_binarize = subparser_topic_modeling.add_parser( "binarize", help="Binarize cell- or region-topic probabilities.", ) - parser_topic_modeling_mallet_binarize.set_defaults( - func=binarize_cell_or_region_topic - ) - parser_topic_modeling_mallet_binarize.add_argument( + parser_binarize.set_defaults(func=binarize_cell_or_region_topic) + parser_binarize.add_argument( "-a", "--target", dest="target", - action="store", type=str, choices=["region", "cell"], required=True, help='Choose between "region" or "cell" topic binarization.', ) - parser_topic_modeling_mallet_binarize.add_argument( + parser_binarize.add_argument( "-m", "--method", dest="method", - action="store", type=str, choices=("ntop", "otsu", "aucell", "li", "yen"), required=True, - help='Binarization method. Choose between "ntop", "otsu", "aucell", "li" or "yen" for cell-or region-topic binarization.', + help="Binarization method.", ) - parser_topic_modeling_mallet_binarize.add_argument( + parser_binarize.add_argument( "-n", "--ntop", dest="ntop", - action="store", type=int, required=False, - help="Number of top regions to select. Can only be used when `--method` is set to `ntop`.", + help="Number of top regions to select when `--method` is `ntop`.", ) - parser_topic_modeling_mallet_binarize.add_argument( + parser_binarize.add_argument( "-s", "--smooth", dest="smooth_topics", - action="store", type=str_to_bool, choices=(True, False), required=False, default=True, help="Whether to smooth the cell- or region-topic probabilities.", ) - parser_topic_modeling_mallet_binarize.add_argument( + parser_binarize.add_argument( "-b", "--nbins", dest="nbins", - action="store", type=int, required=False, default=100, - help="Number of bins to use in the histogram used for `otsu`, `yen` and `li` thresholding.", + help="Number of bins to use in thresholding histograms.", ) - parser_topic_modeling_mallet_binarize.add_argument( + parser_binarize.add_argument( "-c", "--cb", dest="cell_barcodes_filename", - action="store", type=str, required=False, help="Filename with cell barcodes.", ) - parser_topic_modeling_mallet_binarize.add_argument( + parser_binarize.add_argument( "-r", "--regions", dest="region_ids_filename", - action="store", type=str, required=False, help="Filename with region IDs.", ) - parser_topic_modeling_mallet_binarize.add_argument( + parser_binarize.add_argument( "-o", "--output", dest="output_prefix", - action="store", type=str, required=True, help="Topic model output prefix.", ) - parser_topic_modeling_mallet_binarize.add_argument( + parser_binarize.add_argument( "-t", "--n_topics", dest="n_topics", type=int, required=True, - help="Model with `topic` number of topics to binarize.", + help="Model topic count to binarize.", ) - parser_topic_modeling_mallet_binarize.add_argument( + parser_binarize.add_argument( "-p", "--output_dir", dest="out_dir", - action="store", type=str, required=True, help="Directory to store results.", ) - parser_topic_modeling_mallet_anndata = subparser_topic_modeling_mallet.add_parser( + parser_anndata = subparser_topic_modeling.add_parser( "create_anndata", - help="Generate AnnData h5ad file from Mallet result.", - ) - parser_topic_modeling_mallet_anndata.set_defaults( - func=run_create_anndata_from_mallet + help="Generate AnnData h5ad files from topic modeling results.", ) - - parser_topic_modeling_mallet_anndata.add_argument( + parser_anndata.set_defaults(func=run_create_anndata) + parser_anndata.add_argument( "-c", "--cb", dest="cell_barcodes", - action="store", type=str, - required=False, + required=True, help="Filename with cell barcodes.", ) - parser_topic_modeling_mallet_anndata.add_argument( + parser_anndata.add_argument( "-r", "--regions", dest="region_ids", - action="store", type=str, - required=False, + required=True, help="Filename with region IDs.", ) - parser_topic_modeling_mallet_anndata.add_argument( + parser_anndata.add_argument( "-o", "--output", dest="output_prefix", - action="store", type=str, required=True, help="Topic model output prefix.", ) - parser_topic_modeling_mallet_anndata.add_argument( + parser_anndata.add_argument( "-t", "--n_topics", dest="n_topics", type=int, required=True, - help="Model with `topic` number of topics to generate AnnData from.", + help="Model topic count to convert to AnnData.", ) diff --git a/src/pycisTopic/topic_modeling/create_anndata.py b/src/pycisTopic/topic_modeling/create_anndata.py index b75d420..316f8b2 100644 --- a/src/pycisTopic/topic_modeling/create_anndata.py +++ b/src/pycisTopic/topic_modeling/create_anndata.py @@ -1,74 +1,52 @@ +from __future__ import annotations + import anndata import pandas as pd -from pycisTopic.topic_modeling.mallet_models import LDAMallet, LDAMalletFilenames +from pycisTopic.topic_modeling.topic_models import TopicModelFilenames, load_topic_model_backend -def create_anndata_from_mallet( +def create_anndata_from_topic_model( output_prefix: str, n_topics: int, cell_barcodes: list[str], - region_ids: list[str] -): - """ - Create an AnnData object from mallet topic modeling results. - - Parameters - ---------- - output_prefix - Output prefix used for running topic modeling with Mallet. - n_topics - Number of topics used in the topic model created by Mallet. - In combination with output_prefix, this allows to load the correct region - topic counts and cell topic probabilties parquet files. - cell_barcodes - List containing cell names as ordered in the binary matrix columns. - region_ids - List containing region names as ordered in the binary matrix rows. - - Return - ------ - None - - """ - # Get distributions - print("Reading Mallet results ...") - lda_mallet_filenames = LDAMalletFilenames( - output_prefix=output_prefix, n_topics=n_topics + region_ids: list[str], +) -> None: + """Create AnnData objects from backend-agnostic v3 topic modeling artifacts.""" + filenames = TopicModelFilenames(output_prefix=output_prefix, n_topics=n_topics) + backend_cls = load_topic_model_backend(output_prefix=output_prefix, n_topics=n_topics) + + print(f"Reading {backend_cls.backend} results ...") + topic_word_distrib = ( + backend_cls.read_region_topic_counts_parquet_file_to_region_topic_probabilities( + region_topic_counts_parquet_filename=filenames.region_topic_counts_parquet_filename + ) ) - topic_word_distrib = LDAMallet.read_region_topic_counts_parquet_file_to_region_topic_probabilities( - mallet_region_topic_counts_parquet_filename=lda_mallet_filenames.region_topic_counts_parquet_filename - ) - doc_topic_distrib = LDAMallet.read_cell_topic_probabilities_parquet_file( - mallet_cell_topic_probabilities_parquet_filename=lda_mallet_filenames.cell_topic_probabilities_parquet_filename + doc_topic_distrib = backend_cls.read_cell_topic_probabilities_parquet_file( + cell_topic_probabilities_parquet_filename=filenames.cell_topic_probabilities_parquet_filename ) cell_topic = pd.DataFrame.from_records( doc_topic_distrib, index=cell_barcodes, - columns=["Topic" + str(i) for i in range(1, n_topics + 1)], + columns=[f"Topic{i}" for i in range(1, n_topics + 1)], ) - region_topic = pd.DataFrame.from_records( topic_word_distrib, columns=region_ids, - index=["Topic" + str(i) for i in range(1, n_topics + 1)], + index=[f"Topic{i}" for i in range(1, n_topics + 1)], ).transpose() - print("Generating cell_topic AnnData object") - adata_cell_topic = anndata.AnnData( - X=cell_topic - ) + print("Generating cell topic AnnData object") + adata_cell_topic = anndata.AnnData(X=cell_topic) print(f"Done, shape is: {adata_cell_topic.shape}") print("Generating region topic AnnData object") - adata_region_topic = anndata.AnnData( - X=region_topic - ) + adata_region_topic = anndata.AnnData(X=region_topic) print(f"Done, shape is: {adata_region_topic.shape}") - print(f"Writing to: {lda_mallet_filenames.anndata_cell_topic_filename}") - adata_cell_topic.write(lda_mallet_filenames.anndata_cell_topic_filename) + print(f"Writing to: {filenames.anndata_cell_topic_filename}") + adata_cell_topic.write(filenames.anndata_cell_topic_filename) - print(f"Writing to: {lda_mallet_filenames.anndata_region_topic_filename}") - adata_region_topic.write(lda_mallet_filenames.anndata_region_topic_filename) + print(f"Writing to: {filenames.anndata_region_topic_filename}") + adata_region_topic.write(filenames.anndata_region_topic_filename) diff --git a/src/pycisTopic/topic_modeling/mallet_models.py b/src/pycisTopic/topic_modeling/mallet_models.py index fd97fb4..488e320 100644 --- a/src/pycisTopic/topic_modeling/mallet_models.py +++ b/src/pycisTopic/topic_modeling/mallet_models.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import json import logging import os @@ -8,80 +10,66 @@ import polars as pl import scipy +from pycisTopic.topic_modeling.topic_models import LDAModel, TopicModelFilenames + -class LDAMallet: - """Class for running LDA models with Mallet.""" +class LDAMallet(LDAModel): + """Run LDA models with Mallet and write v3 artifacts.""" + + backend = "mallet" @staticmethod def convert_binary_matrix_to_mallet_corpus_file( - binary_accessibility_matrix: scipy.sparse.csr, + binary_accessibility_matrix: scipy.sparse.csr_matrix, mallet_corpus_filename: str, mallet_path: str = "mallet", ) -> None: """ - Convert binary matrix to Mallet serialized corpus file. + Convert a binary accessibility matrix to a serialized Mallet corpus file. Parameters ---------- binary_accessibility_matrix - Binary accessibility matrix (region IDs vs cell barcodes) + Binary accessibility matrix (region IDs vs cell barcodes). mallet_corpus_filename - Mallet serialized corpus filename + Mallet serialized corpus filename. mallet_path - Path to Mallet binary. - - Returns - ------- - None. + Path to the Mallet binary. """ logger = logging.getLogger("LDAMallet") - # Convert binary accessibility matrix to compressed sparse column matrix format - # and eliminate zeros as we assume later that for each found index, the - # associated value is 1. - binary_accessibility_matrix_csc = binary_accessibility_matrix.tocsc() - binary_accessibility_matrix_csc.eliminate_zeros() - - mallet_corpus_txt_filename = f"{mallet_corpus_filename}.txt" - - logger.info( - f'Serializing binary accessibility matrix to Mallet text corpus to "{mallet_corpus_txt_filename}".' - ) + matrix = binary_accessibility_matrix.tocsc() + matrix.eliminate_zeros() - if binary_accessibility_matrix_csc.shape[0] == 0: + if matrix.shape[0] == 0: + raise ValueError("Binary accessibility matrix does not contain any regions.") + if matrix.shape[1] == 0: raise ValueError( "Binary accessibility matrix does not contain any cell barcodes." ) - if binary_accessibility_matrix_csc.shape[1] == 0: - raise ValueError( - "Binary accessibility matrix does not contain any regions." - ) + mallet_corpus_txt_filename = f"{mallet_corpus_filename}.txt" + logger.info( + 'Serializing binary accessibility matrix to Mallet text corpus "%s".', + mallet_corpus_txt_filename, + ) - with open(mallet_corpus_txt_filename, "w") as mallet_corpus_txt_fh: - # Iterate over each column (cell barcode index) of the sparse binary - # accessibility matrix in compressed sparse column matrix format and get - # all index positions (region IDs indices) for that cell barcode index. + with open(mallet_corpus_txt_filename, "w", encoding="utf-8") as fh: + buffered_lines: list[str] = [] for cell_barcode_idx, (indptr_start, indptr_end) in enumerate( - zip( - binary_accessibility_matrix_csc.indptr, - binary_accessibility_matrix_csc.indptr[1:], - ) + zip(matrix.indptr, matrix.indptr[1:]) ): - # Get all region ID indices (assume all have an associated value of 1) - # for the current cell barcode index. - region_ids_idx = binary_accessibility_matrix_csc.indices[ - indptr_start:indptr_end - ] - - # Write Mallet text corpus for the current cell barcode index: - # - column 1: cell barcode index. - # - column 2: document number (always 0). - # - column 3: region IDs indices accessible in the current cell barcode. - mallet_corpus_txt_fh.write( - f"{cell_barcode_idx}\t0\t{' '.join([str(x) for x in region_ids_idx])}\n" + region_ids_idx = matrix.indices[indptr_start:indptr_end] + buffered_lines.append( + f"{cell_barcode_idx}\t0\t{' '.join(map(str, region_ids_idx))}\n" ) + if len(buffered_lines) >= 1024: + fh.writelines(buffered_lines) + buffered_lines.clear() + + if buffered_lines: + fh.writelines(buffered_lines) mallet_import_file_cmd = [ mallet_path, @@ -95,21 +83,22 @@ def convert_binary_matrix_to_mallet_corpus_file( "--output", mallet_corpus_filename, ] - logger.info( - f"Converting Mallet text corpus to Mallet serialised corpus with: {' '.join(mallet_import_file_cmd)}" + "Converting Mallet text corpus to Mallet serialized corpus with: %s", + " ".join(mallet_import_file_cmd), ) try: subprocess.check_output( - args=mallet_import_file_cmd, shell=False, stderr=subprocess.STDOUT + args=mallet_import_file_cmd, + shell=False, + stderr=subprocess.STDOUT, ) - except subprocess.CalledProcessError as e: + except subprocess.CalledProcessError as exc: raise RuntimeError( - f"command '{e.cmd}' return with error (code {e.returncode}): {e.output}" - ) + f"command '{exc.cmd}' return with error (code {exc.returncode}): {exc.output}" + ) from exc - # Remove Mallet text corpus as only Mallet serialised corpus file is needed. if os.path.exists(mallet_corpus_txt_filename): os.remove(mallet_corpus_txt_filename) @@ -118,231 +107,73 @@ def convert_cell_topic_probabilities_txt_to_parquet( mallet_cell_topic_probabilities_txt_filename: str, mallet_cell_topic_probabilities_parquet_filename: str, ) -> None: - """ - Convert cell-topic probabilities from Mallet output to Parquet file. - - Parameters - ---------- - mallet_cell_topic_probabilities_txt_filename - Mallet cell-topic probabilities text file. - mallet_cell_topic_probabilities_parquet_filename - Parquet output file with cell-topic probabilities. - - Returns - ------- - None - - """ - # Read cell-topic probabilities Mallet output file and extract for each cell - # barcode the probability for the cell barcode to belong to a certain topic. - # - # Column 0: order in which cell barcode idx was seen in the input corpus file. - # Column 1: cell barcode idx - # Column 3-n: "topic probability" for each topic - # - # Mallet cell-topic probabilities file example: - # --------------------------------------------- - # - # 0 0 0.06355276993175679 0.1908026307651073 0.06691338680081645 0.007391295383790694 0.07775807681999052 0.08091252087499742 0.08262375523163516 9.793208667505102E-4 0.007721171886275076 0.01605055357400573 0.014071294559099437 0.025307712924973712 0.020524503638950167 0.061903387419334883 0.07344906500628827 0.02866832979403336 0.03520400799950518 0.07608807702616333 0.047656845968290625 0.022421293528235367 - # 1 1 0.10109016579604815 0.0016579604814898933 0.033499886441062915 0.003792868498750852 0.06665909607086078 0.19216443334090394 0.023143311378605497 0.0011128775834658188 0.08719055189643425 0.00401998637292755 0.0030206677265500795 0.03617987735634794 0.02473313649784238 0.255984555984556 0.004383374971610266 0.037179196002725415 0.023143311378605497 0.06202589143765614 0.009379968203497615 0.02963888258005905 - # 2 2 0.08937104175357427 0.03120615116234973 0.11623971329970799 0.03952083886381736 0.034562364898175886 0.08415658538435283 0.03002104744207213 0.040440479350752775 0.02172532140012894 0.025119458455003983 0.01332530623080132 0.06196196291099397 0.07174617922560582 0.03189825173499185 0.05144772270469111 0.00540881337934696 0.08696291099397019 0.07489381470666313 0.04997819409154689 0.04001384201145284 - # 3 3 0.05694870514375401 0.003620603552828708 0.07264393236783906 0.11541342655347078 0.005546835984875508 0.025451237782692444 0.010790468716558465 0.377309695369908 0.03540343868160091 0.007580081329813798 0.023453663408717986 0.02869729614040094 0.08166868802168795 0.01703288863522865 0.006153242491260612 0.0172112434900478 0.06311978312049654 0.02124206320896055 0.012895056003424414 0.017817649996432903 - # 4 4 0.08079825190344497 0.002168049438355697 0.06058588548601864 0.002919184676841135 0.07448188739799926 0.12989518249172044 0.15225852709208235 0.008962409095564889 0.02753593499265936 0.001519341732391 0.011386527365222438 0.012376660179589606 0.015108061046809382 0.1424596264809314 0.015449486155211854 0.027740790057700842 0.068370377957595 0.1540339376557752 0.002168049438355697 0.00978182935573082 - cell_topic_probabilities_ldf = pl.scan_csv( - mallet_cell_topic_probabilities_txt_filename, - separator="\t", - has_header=False, - with_column_names=lambda cols: [ - f"topic_{idx - 1}" if idx > 1 else f"cell_idx{idx}" - for idx, col in enumerate(cols) - ], - ) - # Get cell-topic probabilities as numpy matrix. + """Convert Mallet cell-topic probabilities text output to Parquet.""" cell_topic_probabilities = ( - cell_topic_probabilities_ldf.select( - pl.col("^topic_[0-9]+$").cast(pl.Float32) + pl.scan_csv( + mallet_cell_topic_probabilities_txt_filename, + separator="\t", + has_header=False, + with_column_names=lambda cols: [ + f"topic_{idx - 1}" if idx > 1 else f"cell_idx{idx}" + for idx, _ in enumerate(cols) + ], ) + .select(pl.col("^topic_[0-9]+$").cast(pl.Float32)) .collect() .to_numpy() ) - # Write cell-topic probabilities matrix to one column of a Parquet file. pl.Series( - "cell_topic_probabilities", cell_topic_probabilities - ).to_frame().write_parquet( - f"{mallet_cell_topic_probabilities_parquet_filename}" - ) - - @staticmethod - def read_cell_topic_probabilities_parquet_file( - mallet_cell_topic_probabilities_parquet_filename: str, - ) -> np.ndarray: - """ - Read cell-topic probabilities Parquet file to cell-topic probabilities matrix. - - Parameters - ---------- - mallet_cell_topic_probabilities_parquet_filename - Mallet cell-topic probabilities Parquet filename. - - Returns - ------- - Cell-topic probabilities matrix. - - """ - return ( - pl.read_parquet(mallet_cell_topic_probabilities_parquet_filename) - .get_column("cell_topic_probabilities") - .to_numpy() - ) + "cell_topic_probabilities", + cell_topic_probabilities, + ).to_frame().write_parquet(mallet_cell_topic_probabilities_parquet_filename) @staticmethod def convert_region_topic_counts_txt_to_parquet( mallet_region_topic_counts_txt_filename: str, mallet_region_topic_counts_parquet_filename: str, ) -> None: - """ - Convert region-topic counts from Mallet output to Parquet file. - - Parameters - ---------- - mallet_region_topic_counts_txt_filename - Mallet region-topic counts text file. - mallet_region_topic_counts_parquet_filename - Parquet output file with region-topic counts. - - Returns - ------- - None - - """ + """Convert Mallet region-topic counts text output to Parquet.""" n_region_ids = -1 n_topics = -1 - region_id_topic_counts = [] - - with open(mallet_region_topic_counts_txt_filename) as fh: - # Column 0: order in which region ID idx was seen in the input corpus file. - # Column 1: region ID idx - # Column 3-n: "topic:count" pairs - # - # Mallet region-topics count file example: - # ---------------------------------------- - # - # 0 12 3:94 11:84 1:84 18:75 17:36 0:31 13:25 4:23 6:22 12:16 9:10 10:6 15:3 7:2 8:1 - # 1 28 8:368 15:267 3:267 17:255 0:245 10:227 16:216 19:201 7:92 18:85 1:58 14:52 9:31 6:17 13:6 2:3 - # 2 33 8:431 16:418 10:354 3:257 17:211 12:146 7:145 9:115 4:108 13:106 18:66 1:60 15:45 6:45 19:33 5:19 14:12 0:1 - # 3 35 7:284 18:230 15:199 10:191 16:164 0:114 4:112 19:107 12:104 13:68 3:49 9:35 1:28 11:25 5:20 17:17 6:11 14:2 8:1 - # 4 57 8:192 3:90 19:88 1:69 18:67 2:63 10:62 17:38 15:37 13:10 4:9 12:2 9:1 - for line in fh: + region_id_topic_counts: list[tuple[int, np.ndarray, np.ndarray]] = [] + + with open( + mallet_region_topic_counts_txt_filename, encoding="utf-8" + ) as text_file: + for line in text_file: columns = line.rstrip().split() - # Get region ID index from second column. region_id_idx = int(columns[1]) - # Get topic index and counts from column 3 till the end by splitting - # "topic:count" pairs. topics_counts = [ (int(topic), int(count)) - for topic, count in [ - topic_counts.split(":", 1) for topic_counts in columns[2:] - ] + for topic, count in ( + topic_count.split(":", 1) for topic_count in columns[2:] + ) ] - # Get topic indices. - topics_idx = np.array([topic for topic, count in topics_counts]) - # Get counts. - counts = np.array([count for topic, count in topics_counts]) - # Store region ID index, topics indices and counts till we know how many - # regions and topics we have. + + topics_idx = np.asarray( + [topic for topic, _ in topics_counts], dtype=np.int32 + ) + counts = np.asarray([count for _, count in topics_counts], dtype=np.int32) region_id_topic_counts.append((region_id_idx, topics_idx, counts)) - # Keep track of the highest seen region ID index and topic index - # (0-based). n_region_ids = max(region_id_idx, n_region_ids) - n_topics = max(topics_idx.max(), n_topics) + if topics_idx.size > 0: + n_topics = max(int(topics_idx.max()), n_topics) - # Add 1 to region IDs and topics counts to account for start at 0. n_region_ids += 1 n_topics += 1 - # Create region-topic counts matrix and populate it. - regions_topic_counts = np.zeros((n_topics, n_region_ids), dtype=np.int32) + region_topic_counts = np.zeros((n_topics, n_region_ids), dtype=np.int32) for region_idx, topics_idx, counts in region_id_topic_counts: - regions_topic_counts[topics_idx, region_idx] = counts + region_topic_counts[topics_idx, region_idx] = counts - # Write region-topic counts matrix to one column of a Parquet file. - pl.Series("region_topic_counts", regions_topic_counts).to_frame().write_parquet( + pl.Series("region_topic_counts", region_topic_counts).to_frame().write_parquet( mallet_region_topic_counts_parquet_filename ) @staticmethod - def read_region_topic_counts_parquet_file( - mallet_region_topic_counts_parquet_filename: str, - ) -> np.ndarray: - """ - Read region-topic counts Parquet file to region-topic counts matrix. - - Parameters - ---------- - mallet_region_topic_counts_parquet_filename - Mallet region-topic counts Parquet filename. - - Returns - ------- - Region-topic counts matrix. - - """ - return ( - pl.read_parquet(mallet_region_topic_counts_parquet_filename) - .get_column("region_topic_counts") - .to_numpy() - ) - - @staticmethod - def read_region_topic_counts_parquet_file_to_region_topic_probabilities( - mallet_region_topic_counts_parquet_filename: str, - ) -> np.ndarray: - """ - Get the region-topic probabilities matrix learned during inference. - - Returns - ------- - The probability for each region in each topic, shape (n_regions, n_topics). - - """ - region_topic_counts = np.asarray( - LDAMallet.read_region_topic_counts_parquet_file( - mallet_region_topic_counts_parquet_filename=mallet_region_topic_counts_parquet_filename, - ), - np.float64, - ) - - # Create region-topic probabilities matrix by dividing all count values for a - # topic by total counts for that topic. - region_topic_probabilities = ( - region_topic_counts / region_topic_counts.sum(axis=1)[:, None] - ).astype(np.float32) - - return region_topic_probabilities - - @staticmethod - def read_parameters_json_filename(parameters_json_filename: str) -> dict: - """ - Read parameters from JSON file which gets written by `LDAMallet.run_mallet_topic_modeling`. - - Parameters - ---------- - parameters_json_filename - Parameters JSON filename created by `LDAMallet.run_mallet_topic_modeling`. - - Returns - ------- - Dictionary with Mallet LDA parameters and settings. - - """ - with open(parameters_json_filename) as fh: - mallet_train_topics_parameters = json.load(fh) - return mallet_train_topics_parameters - - @staticmethod - def run_mallet_topic_modeling( + def run_topic_modeling( mallet_corpus_filename: str, output_prefix: str, n_topics: int, @@ -357,62 +188,20 @@ def run_mallet_topic_modeling( topic_threshold: float = 0.0, random_seed: int = 555, mallet_path: str = "mallet", - ): - """ - Run Mallet LDA. - - Parameters - ---------- - mallet_corpus_filename - Mallet corpus file. - output_prefix - Output prefix. - n_topics - The number of topics to use in the model. - alpha - Scalar value indicating the symmetric Dirichlet hyperparameter for topic - proportions. Default: 50. - alpha_by_topic - Boolean indicating whether the scalar given in alpha has to be divided by - the number of topics. Default: True. - eta - Scalar value indicating the symmetric Dirichlet hyperparameter for topic - multinomials. Default: 0.1. - eta_by_topic - Boolean indicating whether the scalar given in beta has to be divided by - the number of topics. Default: False. - n_threads - Number of threads that will be used for training. Default: 1. - iterations - Number of training iterations of Gibbs sampling. Default: 150. - optimize_interval - Optimize hyperparameters every `optimize_interval` iterations (sometimes - leads to Java exception, 0 to switch off hyperparameter optimization). - Only takes effect after running `optimize_burn_in` iterations. - Default: 0. - optimize_burn_in - The number of iterations before hyperparameter optimization begins. - Default: 50. - topic_threshold - Threshold of the probability above which we consider a topic. Default: 0.0. - random_seed - Random seed to ensure consistent results, if 0 - use system clock. - Default: 555. - mallet_path - Path to the mallet binary (e.g. /xxx/Mallet/bin/mallet). Default: "mallet". - - """ + ) -> None: + """Run Mallet LDA and write the standard v3 artifacts.""" logger = logging.getLogger("LDAMallet") - # Mallet divides alpha value by default by the number of topics, so in case - # alpha_by_topic=False, input alpha needs to be multiplied by n_topics. - mallet_alpha = alpha if alpha_by_topic else alpha * n_topics + if topic_threshold != 0.0: + raise ValueError( + "topic_threshold must be 0.0 because pycisTopic stores dense " + "cell-topic probability tables and cannot parse Mallet sparse " + "doc-topic output." + ) + mallet_alpha = alpha if alpha_by_topic else alpha * n_topics mallet_beta = eta / n_topics if eta_by_topic else eta - - lda_mallet_filenames = LDAMalletFilenames( - output_prefix=output_prefix, n_topics=n_topics - ) + filenames = TopicModelFilenames(output_prefix=output_prefix, n_topics=n_topics) if not os.path.exists(mallet_corpus_filename): raise FileNotFoundError( @@ -439,9 +228,9 @@ def run_mallet_topic_modeling( "--num-iterations", str(iterations), "--word-topic-counts-file", - lda_mallet_filenames.region_topic_counts_txt_filename, + filenames.region_topic_counts_txt_filename, "--output-doc-topics", - lda_mallet_filenames.cell_topic_probabilities_txt_filename, + filenames.cell_topic_probabilities_txt_filename, "--doc-topics-threshold", str(topic_threshold), "--random-seed", @@ -449,113 +238,51 @@ def run_mallet_topic_modeling( ] start_time = time.time() - logger.info(f"Train topics with Mallet LDA: {' '.join(cmd)}") + logger.info("Train topics with Mallet LDA: %s", " ".join(cmd)) try: subprocess.check_output(args=cmd, shell=False, stderr=subprocess.STDOUT) - except subprocess.CalledProcessError as e: - raise RuntimeError( # noqa: B904 - f"command '{e.cmd}' return with error (code {e.returncode}): {e.output}" - ) + except subprocess.CalledProcessError as exc: + raise RuntimeError( + f"command '{exc.cmd}' return with error (code {exc.returncode}): {exc.output}" + ) from exc - # Convert cell-topic probabilities text version to parquet. logger.info( - f'Write cell-topic probabilities to "{lda_mallet_filenames.cell_topic_probabilities_parquet_filename}".' + 'Write cell-topic probabilities to "%s".', + filenames.cell_topic_probabilities_parquet_filename, ) LDAMallet.convert_cell_topic_probabilities_txt_to_parquet( - mallet_cell_topic_probabilities_txt_filename=lda_mallet_filenames.cell_topic_probabilities_txt_filename, - mallet_cell_topic_probabilities_parquet_filename=lda_mallet_filenames.cell_topic_probabilities_parquet_filename, + mallet_cell_topic_probabilities_txt_filename=filenames.cell_topic_probabilities_txt_filename, + mallet_cell_topic_probabilities_parquet_filename=filenames.cell_topic_probabilities_parquet_filename, ) - # Convert region-topic counts text version to parquet. logger.info( - f'Write region-topic counts to "{lda_mallet_filenames.region_topic_counts_parquet_filename}".' + 'Write region-topic counts to "%s".', + filenames.region_topic_counts_parquet_filename, ) LDAMallet.convert_region_topic_counts_txt_to_parquet( - mallet_region_topic_counts_txt_filename=lda_mallet_filenames.region_topic_counts_txt_filename, - mallet_region_topic_counts_parquet_filename=lda_mallet_filenames.region_topic_counts_parquet_filename, - ) - - total_time = time.time() - start_time - - # Write JSON file with all used parameters. - logger.info( - f'Write JSON parameters file to "{lda_mallet_filenames.parameters_json_filename}".' - ) - with open(lda_mallet_filenames.parameters_json_filename, "w") as fh: - mallet_train_topics_parameters = { - "mallet_corpus_filename": mallet_corpus_filename, - "output_prefix": output_prefix, - "n_topics": n_topics, - "alpha": alpha, - "alpha_by_topic": alpha_by_topic, - "eta": eta, - "eta_by_topic": eta_by_topic, - "n_threads": n_threads, - "iterations": iterations, - "optimize_interval": optimize_interval, - "optimize_burn_in": optimize_burn_in, - "random_seed": random_seed, - "mallet_path": mallet_path, - "time": total_time, - "mallet_cmd": cmd, - } - json.dump(mallet_train_topics_parameters, fh) - - -class LDAMalletFilenames: - """Class to generate output filenames when running functions of LDAMallet.""" - - def __init__(self, output_prefix: str, n_topics: int): - """ - Generate output filenames when running functions of LDAMallet. - - Parameters - ---------- - output_prefix - Output prefix. - n_topics - The number of topics used in the model. - - """ - self.output_prefix = output_prefix - self.n_topics = n_topics - - @property - def parameters_json_filename(self): - return os.path.join( - f"{self.output_prefix}.{self.n_topics}_topics.parameters.json" - ) - - @property - def cell_topic_probabilities_txt_filename(self): - return os.path.join( - f"{self.output_prefix}.{self.n_topics}_topics.cell_topic_probabilities.txt" + mallet_region_topic_counts_txt_filename=filenames.region_topic_counts_txt_filename, + mallet_region_topic_counts_parquet_filename=filenames.region_topic_counts_parquet_filename, ) - @property - def cell_topic_probabilities_parquet_filename(self): - return os.path.join( - f"{self.output_prefix}.{self.n_topics}_topics.cell_topic_probabilities.parquet" - ) - - @property - def region_topic_counts_txt_filename(self): - return f"{self.output_prefix}.{self.n_topics}_topics.region_topic_counts.txt" - - @property - def region_topic_counts_parquet_filename(self): - return ( - f"{self.output_prefix}.{self.n_topics}_topics.region_topic_counts.parquet" - ) - - @property - def model_stats_filename(self): - return f"{self.output_prefix}.{self.n_topics}_topics.stats.json" - - @property - def anndata_cell_topic_filename(self): - return f"{self.output_prefix}.{self.n_topics}_topics_cell_topic_adata.h5ad" - - @property - def anndata_region_topic_filename(self): - return f"{self.output_prefix}.{self.n_topics}_topics_region_topic_adata.h5ad" + parameters = { + "backend": LDAMallet.backend, + "mallet_corpus_filename": mallet_corpus_filename, + "output_prefix": output_prefix, + "n_topics": n_topics, + "alpha": alpha, + "alpha_by_topic": alpha_by_topic, + "eta": eta, + "eta_by_topic": eta_by_topic, + "n_threads": n_threads, + "iterations": iterations, + "optimize_interval": optimize_interval, + "optimize_burn_in": optimize_burn_in, + "topic_threshold": topic_threshold, + "random_seed": random_seed, + "mallet_path": mallet_path, + "time": time.time() - start_time, + "mallet_cmd": cmd, + } + logger.info('Write JSON parameters file to "%s".', filenames.parameters_json_filename) + with open(filenames.parameters_json_filename, "w", encoding="utf-8") as fh: + json.dump(parameters, fh, indent=2) diff --git a/src/pycisTopic/topic_modeling/plot_stats.py b/src/pycisTopic/topic_modeling/plot_stats.py index fef4f33..41312de 100644 --- a/src/pycisTopic/topic_modeling/plot_stats.py +++ b/src/pycisTopic/topic_modeling/plot_stats.py @@ -1,14 +1,21 @@ +from __future__ import annotations + import json from dataclasses import dataclass import matplotlib.pyplot as plt import numpy as np -from pycisTopic.topic_modeling.mallet_models import LDAMalletFilenames +from pycisTopic.topic_modeling.topic_models import TopicModelFilenames, load_topic_model_backend + +def scale(values: np.ndarray) -> np.ndarray: + min_value = values.min() + max_value = values.max() + if min_value == max_value: + return np.ones_like(values, dtype=np.float64) + return (values - min_value) / (max_value - min_value) -def scale(X: np.ndarray) -> np.ndarray: - return (X - X.min()) / (X.max() - X.min()) @dataclass class TopicModelMetrics: @@ -45,54 +52,40 @@ def plot_stats( """ metrics_per_topic: dict[int, TopicModelMetrics] = {} for n_topic in n_topics: - lda_mallet_filenames = LDAMalletFilenames( - output_prefix=output_prefix, n_topics=n_topic - ) - with open(lda_mallet_filenames.model_stats_filename) as infile: - metrics_per_topic[n_topic] = TopicModelMetrics( - **json.load(infile) - ) + load_topic_model_backend(output_prefix=output_prefix, n_topics=n_topic) + filenames = TopicModelFilenames(output_prefix=output_prefix, n_topics=n_topic) + with open(filenames.model_stats_filename, encoding="utf-8") as infile: + metrics_per_topic[n_topic] = TopicModelMetrics(**json.load(infile)) + sorted_topics = sorted(n_topics) metrics = { "Inv_Arun_2010": scale( - -np.array([ - metrics_per_topic[t].Arun_2010 - for t in sorted(n_topics) - ]) + -np.asarray([metrics_per_topic[t].Arun_2010 for t in sorted_topics]) ), "Inv_Cao_Juan_2009": scale( - -np.array([ - metrics_per_topic[t].Cao_Juan_2009 - for t in sorted(n_topics) - ]) + -np.asarray([metrics_per_topic[t].Cao_Juan_2009 for t in sorted_topics]) ), "Mimno_2011": scale( - np.array([ - metrics_per_topic[t].Mimno_2011 - for t in sorted(n_topics) - ]) + np.asarray([metrics_per_topic[t].Mimno_2011 for t in sorted_topics]) ), "Loglikelihood": scale( - np.array([ - metrics_per_topic[t].loglikelihood - for t in sorted(n_topics) - ]) - ) + np.asarray([metrics_per_topic[t].loglikelihood for t in sorted_topics]) + ), } - fig, ax = plt.subplots(figsize = (8, 8)) + figure, axis = plt.subplots(figsize=(8, 8)) for metric, values in metrics.items(): - _ = ax.plot( - sorted(n_topics), + axis.plot( + sorted_topics, values, linestyle="--", marker="o", - label=metric + label=metric, ) - ax.grid(True) - ax.set_axisbelow(True) - _ = ax.set_xlabel("Number of topics") - _ = ax.set_ylabel("Scaled metric") - _ = ax.legend() - fig.tight_layout() - fig.savefig(f"{output_prefix}.model_evaluation_stats.{plot_file_format}") + axis.grid(True) + axis.set_axisbelow(True) + axis.set_xlabel("Number of topics") + axis.set_ylabel("Scaled metric") + axis.legend() + figure.tight_layout() + figure.savefig(f"{output_prefix}.model_evaluation_stats.{plot_file_format}") diff --git a/src/pycisTopic/topic_modeling/stats.py b/src/pycisTopic/topic_modeling/stats.py index 0220168..770db63 100644 --- a/src/pycisTopic/topic_modeling/stats.py +++ b/src/pycisTopic/topic_modeling/stats.py @@ -1,12 +1,14 @@ +from __future__ import annotations + import json -import math from itertools import chain import numpy as np import scipy +from scipy.special import gammaln import pycisTopic.topic_modeling.tmtoolkit_lite as tmtoolkit_lite -from pycisTopic.topic_modeling.mallet_models import LDAMallet, LDAMalletFilenames +from pycisTopic.topic_modeling.topic_models import TopicModelFilenames, load_topic_model_backend class JsonNumpyEncode(json.JSONEncoder): @@ -20,38 +22,49 @@ def default(self, obj): return super().default(obj) -def loglikelihood(nzw, ndz, alpha, eta): - D = ndz.shape[0] - n_topics = ndz.shape[1] - vocab_size = nzw.shape[1] - - const_prior = (n_topics * math.lgamma(alpha) - math.lgamma(alpha * n_topics)) * D - const_ll = ( - vocab_size * math.lgamma(eta) - math.lgamma(eta * vocab_size) - ) * n_topics - - # calculate log p(w|z) - topic_ll = 0 - for k in range(n_topics): - sum = eta * vocab_size - for w in range(vocab_size): - if nzw[k, w] > 0: - topic_ll = math.lgamma(nzw[k, w] + eta) - sum += nzw[k, w] - topic_ll -= math.lgamma(sum) - - # calculate log p(z) - doc_ll = 0 - for d in range(D): - sum = alpha * n_topics - for k in range(n_topics): - if ndz[d, k] > 0: - doc_ll = math.lgamma(ndz[d, k] + alpha) - sum += ndz[d, k] - doc_ll -= math.lgamma(sum) - - ll = doc_ll - const_prior + topic_ll - const_ll - return ll +def loglikelihood( + nzw: np.ndarray, + ndz: np.ndarray, + alpha: float | list[float] | np.ndarray, + eta: float | list[float] | np.ndarray, +) -> float: + """Calculate the LDA log likelihood for scalar or vector hyperparameters.""" + nzw = np.asarray(nzw, dtype=np.float64) + ndz = np.asarray(ndz, dtype=np.float64) + + alpha_array = np.asarray(alpha, dtype=np.float64).reshape(-1) + if alpha_array.size == 1: + alpha_array = np.repeat(alpha_array.item(), ndz.shape[1]) + elif alpha_array.size != ndz.shape[1]: + raise ValueError( + f"alpha length {alpha_array.size} does not match topic count {ndz.shape[1]}." + ) + + eta_array = np.asarray(eta, dtype=np.float64).reshape(-1) + if eta_array.size == 1: + eta_array = np.repeat(eta_array.item(), nzw.shape[1]) + elif eta_array.size != nzw.shape[1]: + raise ValueError( + f"eta length {eta_array.size} does not match vocab size {nzw.shape[1]}." + ) + + alpha_sum = float(alpha_array.sum()) + eta_sum = float(eta_array.sum()) + doc_lengths = ndz.sum(axis=1) + topic_lengths = nzw.sum(axis=1) + + doc_ll = np.sum( + gammaln(alpha_sum) + - gammaln(doc_lengths + alpha_sum) + + np.sum(gammaln(ndz + alpha_array) - gammaln(alpha_array), axis=1) + ) + topic_ll = np.sum( + gammaln(eta_sum) + - gammaln(topic_lengths + eta_sum) + + np.sum(gammaln(nzw + eta_array) - gammaln(eta_array), axis=1) + ) + + return float(doc_ll + topic_ll) def calculate_model_evaluation_stats( @@ -61,67 +74,49 @@ def calculate_model_evaluation_stats( top_topics_coh: int = 5, ) -> None: """ - Calculate model evaluation statistics after running Mallet (McCallum, 2002) topic modeling. + Calculate topic model evaluation statistics from backend-agnostic v3 artifacts. Parameters ---------- binary_accessibility_matrix - Binary accessibility sparse matrix with cells as columns, regions as rows, - and 1 as value if a region is considered accessible in a cell (otherwise, 0). + Binary accessibility sparse matrix with cells as columns and regions as rows. output_prefix - Output prefix used for running topic modeling with Mallet. + Output prefix used for running topic modeling. n_topics - Number of topics used in the topic model created by Mallet. - In combination with output_prefix, this allows to load the correct region - topic counts and cell topic probabilties parquet files. + Number of topics used in the topic model. top_topics_coh - Number of topics to use to calculate the model coherence. For each model, - the coherence will be calculated as the average of the top coherence values. - Default: 5. - - Return - ------ - None + Number of topics to use to calculate the model coherence. - References - ---------- - McCallum, A. K. (2002). Mallet: A machine learning for language toolkit. http://mallet.cs.umass.edu. + """ + filenames = TopicModelFilenames(output_prefix=output_prefix, n_topics=n_topics) + backend_cls = load_topic_model_backend(output_prefix=output_prefix, n_topics=n_topics) - """ # noqa: W505 - # Get distributions - lda_mallet_filenames = LDAMalletFilenames( - output_prefix=output_prefix, n_topics=n_topics + topic_word_distrib = ( + backend_cls.read_region_topic_counts_parquet_file_to_region_topic_probabilities( + region_topic_counts_parquet_filename=filenames.region_topic_counts_parquet_filename + ) ) - topic_word_distrib = LDAMallet.read_region_topic_counts_parquet_file_to_region_topic_probabilities( - mallet_region_topic_counts_parquet_filename=lda_mallet_filenames.region_topic_counts_parquet_filename + doc_topic_distrib = backend_cls.read_cell_topic_probabilities_parquet_file( + cell_topic_probabilities_parquet_filename=filenames.cell_topic_probabilities_parquet_filename ) - doc_topic_distrib = LDAMallet.read_cell_topic_probabilities_parquet_file( - mallet_cell_topic_probabilities_parquet_filename=lda_mallet_filenames.cell_topic_probabilities_parquet_filename + topic_word_counts = backend_cls.read_region_topic_counts_parquet_file( + region_topic_counts_parquet_filename=filenames.region_topic_counts_parquet_filename ) - topic_word_counts = LDAMallet.read_region_topic_counts_parquet_file( - mallet_region_topic_counts_parquet_filename=lda_mallet_filenames.region_topic_counts_parquet_filename + parameters = backend_cls.read_parameters_json_filename( + filenames.parameters_json_filename ) - # Read used Mallet LDA parameters from JSON file. - mallet_train_topics_parameters = LDAMallet.read_parameters_json_filename( - lda_mallet_filenames.parameters_json_filename - ) - - if mallet_train_topics_parameters["n_topics"] != n_topics: + if parameters["n_topics"] != n_topics: raise ValueError( - f"Number of topics does not match: {n_topics} vs {mallet_train_topics_parameters['n_topics']}." + f"Number of topics does not match: {n_topics} vs {parameters['n_topics']}." ) - alpha = mallet_train_topics_parameters["alpha"] - alpha_by_topic = mallet_train_topics_parameters["alpha_by_topic"] - eta = mallet_train_topics_parameters["eta"] - eta_by_topic = mallet_train_topics_parameters["eta_by_topic"] - - ll_alpha = alpha / n_topics if alpha_by_topic else alpha - ll_eta = eta / n_topics if eta_by_topic else eta + ll_alpha, ll_eta = _resolve_loglikelihood_hyperparameters( + parameters=parameters, + n_topics=n_topics, + ) - # Model evaluation - cell_cov = np.asarray(binary_accessibility_matrix.sum(axis=0)).astype(float) + cell_cov = np.asarray(binary_accessibility_matrix.sum(axis=0)).astype(float).ravel() arun_2010 = tmtoolkit_lite.topicmod.evaluate.metric_arun_2010( topic_word_distrib=topic_word_distrib, doc_topic_distrib=doc_topic_distrib, @@ -130,24 +125,24 @@ def calculate_model_evaluation_stats( cao_juan_2009 = tmtoolkit_lite.topicmod.evaluate.metric_cao_juan_2009( topic_word_distrib=topic_word_distrib, ) + coherence_top_n = max(1, min(20, topic_word_distrib.shape[1])) mimno_2011 = tmtoolkit_lite.topicmod.evaluate.metric_coherence_mimno_2011( topic_word_distrib=topic_word_distrib, dtm=binary_accessibility_matrix.transpose(), - top_n=20, + top_n=coherence_top_n, eps=1e-12, normalize=True, return_mean=False, ) - doc_topic_counts = (doc_topic_distrib.T * (cell_cov)).T + doc_topic_counts = (doc_topic_distrib.T * cell_cov).T ll = loglikelihood(topic_word_counts, doc_topic_counts, ll_alpha, ll_eta) marg_topic = tmtoolkit_lite.topicmod.model_stats.marginal_topic_distrib( doc_topic_distrib=doc_topic_distrib, doc_lengths=cell_cov, ) - - topic_ass = list(chain.from_iterable(topic_word_counts.sum(axis=1)[:, None])) + topic_assignments = list(chain.from_iterable(topic_word_counts.sum(axis=1)[:, None])) metrics = { "Arun_2010": arun_2010, @@ -164,12 +159,27 @@ def calculate_model_evaluation_stats( "loglikelihood": ll, "coherence": mimno_2011, "marg_topic": marg_topic, - "assignments": topic_ass, + "assignments": topic_assignments, } - with open(lda_mallet_filenames.model_stats_filename, "w") as fh: - json.dump( - metrics, - fh, - cls=JsonNumpyEncode, - ) + with open(filenames.model_stats_filename, "w", encoding="utf-8") as fh: + json.dump(metrics, fh, cls=JsonNumpyEncode) + + +def _resolve_loglikelihood_hyperparameters( + parameters: dict, + n_topics: int, +) -> tuple[float | list[float], float | list[float]]: + backend_name = str(parameters.get("backend", "mallet")).lower() + + if backend_name == "tomotopy": + return parameters["alpha"], parameters["eta"] + + alpha = parameters["alpha"] + eta = parameters["eta"] + alpha_by_topic = parameters.get("alpha_by_topic", True) + eta_by_topic = parameters.get("eta_by_topic", False) + + ll_alpha = alpha / n_topics if alpha_by_topic else alpha + ll_eta = eta / n_topics if eta_by_topic else eta + return ll_alpha, ll_eta diff --git a/src/pycisTopic/topic_modeling/tomotopy_models.py b/src/pycisTopic/topic_modeling/tomotopy_models.py new file mode 100644 index 0000000..bc93ab6 --- /dev/null +++ b/src/pycisTopic/topic_modeling/tomotopy_models.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +import json +import logging +import time + +import numpy as np +import polars as pl +import tomotopy as tp +from scipy import sparse + +from pycisTopic.topic_modeling.topic_models import LDAModel, TopicModelFilenames + + +class LDATomotopy(LDAModel): + """Run LDA models with tomotopy and write v3 artifacts.""" + + backend = "tomotopy" + + @staticmethod + def run_topic_modeling( + binary_accessibility_matrix: sparse.csr_matrix, + cell_names: list[str], + region_names: list[str], + output_prefix: str, + n_topics: int, + alpha: float = 50.0, + alpha_by_topic: bool = True, + eta: float = 0.1, + eta_by_topic: bool = False, + n_threads: int = 1, + iterations: int = 150, + optimize_interval: int = 0, + random_seed: int = 555, + ) -> None: + """Run tomotopy LDA and write the standard v3 artifacts.""" + logger = logging.getLogger("LDATomotopy") + matrix = sparse.csr_matrix(binary_accessibility_matrix) + + if len(cell_names) != matrix.shape[1]: + raise ValueError( + "Number of cell names does not match the accessibility matrix columns." + ) + if len(region_names) != matrix.shape[0]: + raise ValueError( + "Number of region names does not match the accessibility matrix rows." + ) + + documents, valid_cell_indices = _build_documents( + binary_accessibility_matrix=matrix, + region_names=region_names, + ) + + effective_alpha = alpha / n_topics if alpha_by_topic else alpha + effective_eta = eta / n_topics if eta_by_topic else eta + + model = tp.LDAModel( + k=n_topics, + alpha=effective_alpha, + eta=effective_eta, + seed=random_seed, + min_cf=0, + ) + model.optim_interval = optimize_interval + + for document in documents: + model.add_doc(words=document) + + start_time = time.time() + logger.info( + "Training tomotopy model for %s topics, %s iterations and %s threads.", + n_topics, + iterations, + n_threads, + ) + model.train(iterations=iterations, workers=n_threads, show_progress=False) + elapsed = time.time() - start_time + + region_topic_counts, doc_topic_counts = _extract_exact_counts( + model=model, + region_names=region_names, + cell_count=matrix.shape[1], + valid_cell_indices=valid_cell_indices, + n_topics=n_topics, + ) + + cell_topic_probabilities = _counts_to_probabilities(doc_topic_counts) + filenames = TopicModelFilenames(output_prefix=output_prefix, n_topics=n_topics) + + pl.Series( + "cell_topic_probabilities", + cell_topic_probabilities.astype(np.float32), + ).to_frame().write_parquet(filenames.cell_topic_probabilities_parquet_filename) + pl.Series( + "region_topic_counts", + region_topic_counts.astype(np.int32), + ).to_frame().write_parquet(filenames.region_topic_counts_parquet_filename) + + learned_alpha = np.asarray(model.alpha, dtype=np.float64).reshape(-1) + parameters = { + "backend": LDATomotopy.backend, + "output_prefix": output_prefix, + "n_topics": n_topics, + "alpha": ( + learned_alpha.tolist() + if learned_alpha.size > 1 + else float(learned_alpha[0]) + ), + "alpha_input": alpha, + "alpha_by_topic": alpha_by_topic, + "eta": float(np.asarray(model.eta, dtype=np.float64).reshape(())), + "eta_input": eta, + "eta_by_topic": eta_by_topic, + "n_threads": n_threads, + "iterations": iterations, + "optimize_interval": optimize_interval, + "random_seed": random_seed, + "time": elapsed, + "tomotopy_version": tp.__version__, + } + with open(filenames.parameters_json_filename, "w", encoding="utf-8") as fh: + json.dump(parameters, fh, indent=2) + + +def _build_documents( + binary_accessibility_matrix: sparse.csr_matrix, + region_names: list[str], +) -> tuple[list[list[str]], list[int]]: + matrix = binary_accessibility_matrix.tocsc() + matrix.eliminate_zeros() + + if matrix.shape[0] == 0: + raise ValueError("Binary accessibility matrix does not contain any regions.") + if matrix.shape[1] == 0: + raise ValueError( + "Binary accessibility matrix does not contain any cell barcodes." + ) + + documents: list[list[str]] = [] + valid_cell_indices: list[int] = [] + for cell_idx, (indptr_start, indptr_end) in enumerate( + zip(matrix.indptr, matrix.indptr[1:]) + ): + region_indices = matrix.indices[indptr_start:indptr_end] + if region_indices.size == 0: + continue + documents.append([region_names[region_idx] for region_idx in region_indices]) + valid_cell_indices.append(cell_idx) + + if not documents: + raise ValueError( + "Binary accessibility matrix does not contain any accessible cells." + ) + + return documents, valid_cell_indices + + +def _extract_exact_counts( + model: tp.LDAModel, + region_names: list[str], + cell_count: int, + valid_cell_indices: list[int], + n_topics: int, +) -> tuple[np.ndarray, np.ndarray]: + region_name_to_index = { + region_name: idx for idx, region_name in enumerate(region_names) + } + word_id_to_region_index = np.asarray( + [region_name_to_index[token] for token in model.used_vocabs], + dtype=np.int64, + ) + + region_topic_counts = np.zeros((n_topics, len(region_names)), dtype=np.int64) + doc_topic_counts = np.zeros((cell_count, n_topics), dtype=np.int64) + + for doc_idx, doc in enumerate(model.docs): + cell_idx = valid_cell_indices[doc_idx] + word_ids = np.asarray(doc.words, dtype=np.int64) + topic_ids = np.asarray(doc.topics, dtype=np.int64) + region_indices = word_id_to_region_index[word_ids] + + np.add.at(region_topic_counts, (topic_ids, region_indices), 1) + doc_topic_counts[cell_idx] = np.bincount(topic_ids, minlength=n_topics) + + return region_topic_counts, doc_topic_counts + + +def _counts_to_probabilities(counts: np.ndarray) -> np.ndarray: + totals = counts.sum(axis=1, keepdims=True).astype(np.float64) + return np.divide( + counts, + totals, + out=np.zeros_like(counts, dtype=np.float64), + where=totals != 0, + ).astype(np.float32) diff --git a/src/pycisTopic/topic_modeling/topic_models.py b/src/pycisTopic/topic_modeling/topic_models.py new file mode 100644 index 0000000..abf1ad1 --- /dev/null +++ b/src/pycisTopic/topic_modeling/topic_models.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import json +import os +from abc import ABC, abstractmethod +from typing import ClassVar + +import numpy as np +import polars as pl + + +class TopicModelFilenames: + """Generate artifact filenames for v3 topic modeling outputs.""" + + def __init__(self, output_prefix: str, n_topics: int): + self.output_prefix = output_prefix + self.n_topics = n_topics + + @property + def parameters_json_filename(self) -> str: + return f"{self.output_prefix}.{self.n_topics}_topics.parameters.json" + + @property + def cell_topic_probabilities_txt_filename(self) -> str: + return ( + f"{self.output_prefix}.{self.n_topics}_topics.cell_topic_probabilities.txt" + ) + + @property + def cell_topic_probabilities_parquet_filename(self) -> str: + return ( + f"{self.output_prefix}.{self.n_topics}_topics.cell_topic_probabilities.parquet" + ) + + @property + def region_topic_counts_txt_filename(self) -> str: + return f"{self.output_prefix}.{self.n_topics}_topics.region_topic_counts.txt" + + @property + def region_topic_counts_parquet_filename(self) -> str: + return f"{self.output_prefix}.{self.n_topics}_topics.region_topic_counts.parquet" + + @property + def model_stats_filename(self) -> str: + return f"{self.output_prefix}.{self.n_topics}_topics.stats.json" + + @property + def anndata_cell_topic_filename(self) -> str: + return f"{self.output_prefix}.{self.n_topics}_topics_cell_topic_adata.h5ad" + + @property + def anndata_region_topic_filename(self) -> str: + return f"{self.output_prefix}.{self.n_topics}_topics_region_topic_adata.h5ad" + + +class LDAModel(ABC): + """Abstract base class for v3 LDA backends.""" + + backend: ClassVar[str] + + @staticmethod + def read_parameters_json_filename(parameters_json_filename: str) -> dict: + with open(parameters_json_filename, encoding="utf-8") as fh: + return json.load(fh) + + @staticmethod + def read_matrix_parquet_file( + parquet_filename: str, + column_name: str, + ) -> np.ndarray: + matrix_column = pl.read_parquet(parquet_filename).get_column(column_name) + values = matrix_column.to_numpy() + + if values.ndim == 2: + return values + if values.size == 0: + return np.empty((0, 0), dtype=np.float32) + + return np.stack([np.asarray(row) for row in values], axis=0) + + @classmethod + def read_cell_topic_probabilities_parquet_file( + cls, + cell_topic_probabilities_parquet_filename: str, + ) -> np.ndarray: + return cls.read_matrix_parquet_file( + parquet_filename=cell_topic_probabilities_parquet_filename, + column_name="cell_topic_probabilities", + ) + + @classmethod + def read_region_topic_counts_parquet_file( + cls, + region_topic_counts_parquet_filename: str, + ) -> np.ndarray: + return cls.read_matrix_parquet_file( + parquet_filename=region_topic_counts_parquet_filename, + column_name="region_topic_counts", + ) + + @classmethod + def read_region_topic_counts_parquet_file_to_region_topic_probabilities( + cls, + region_topic_counts_parquet_filename: str, + ) -> np.ndarray: + region_topic_counts = np.asarray( + cls.read_region_topic_counts_parquet_file( + region_topic_counts_parquet_filename=region_topic_counts_parquet_filename + ), + dtype=np.float64, + ) + topic_totals = region_topic_counts.sum(axis=1, keepdims=True) + return np.divide( + region_topic_counts, + topic_totals, + out=np.zeros_like(region_topic_counts, dtype=np.float64), + where=topic_totals != 0, + ).astype(np.float32) + + @staticmethod + @abstractmethod + def run_topic_modeling(*args, **kwargs) -> None: + """Run topic modeling and write the standard v3 artifacts.""" + + +def load_topic_model_backend(output_prefix: str, n_topics: int) -> type[LDAModel]: + """Resolve the backend used for a topic-model artifact bundle.""" + + filenames = TopicModelFilenames(output_prefix=output_prefix, n_topics=n_topics) + backend_name = "mallet" + + if os.path.exists(filenames.parameters_json_filename): + parameters = LDAModel.read_parameters_json_filename( + filenames.parameters_json_filename + ) + backend_name = str(parameters.get("backend", "mallet")).lower() + + if backend_name == "mallet": + from pycisTopic.topic_modeling.mallet_models import LDAMallet + + return LDAMallet + + if backend_name == "tomotopy": + from pycisTopic.topic_modeling.tomotopy_models import LDATomotopy + + return LDATomotopy + + raise ValueError( + f"Unsupported topic modeling backend {backend_name!r} in " + f'"{filenames.parameters_json_filename}".' + )