refactor generate_ion_image CLI and use it in the snakemake workflow

src/depiction_targeted_preproc/workflow/vis/images.py → src/depiction/tools/cli/cli_generate_ion_images.py

@@ -1,25 +1,27 @@
 from pathlib import Path
+from typing import Annotated
+
 import polars as pl
 import typer
-from typing import Annotated
 
 from depiction.parallel_ops import ParallelConfig
 from depiction.persistence import ImzmlReadFile
 from depiction.tools.generate_ion_image import GenerateIonImage
-from depiction_targeted_preproc.pipeline_config.model import PipelineParameters
 
 
-def vis_images(
+# TODO for the cli module we should figure out a sane way to configure the default n_jobs value
+
+def generate_ion_images(
     imzml_path: Annotated[Path, typer.Option()],
-    config_path: Annotated[Path, typer.Option()],
     mass_list_path: Annotated[Path, typer.Option()],
     output_hdf5_path: Annotated[Path, typer.Option()],
+    n_jobs: Annotated[int, typer.Option()] = 16,
 ) -> None:
-    config = PipelineParameters.parse_yaml(config_path)
+    parallel_config = ParallelConfig(n_jobs=n_jobs)
+    gen_image = GenerateIonImage(parallel_config=parallel_config)
+
     mass_list_df = pl.read_csv(mass_list_path)
-    parallel_config = ParallelConfig(n_jobs=config.n_jobs, task_size=None)
 
-    gen_image = GenerateIonImage(parallel_config=parallel_config)
     image = gen_image.generate_ion_images_for_file(
         input_file=ImzmlReadFile(imzml_path),
         mz_values=mass_list_df["mass"],
@@ -28,14 +30,6 @@ def vis_images(
     )
     image.write_hdf5(output_hdf5_path)
 
-    # if output_stats_path:
-    #    stats = ImageChannelStatistics.compute_xarray(xarray)
-    #    stats.write_csv(output_stats_path)
-
-
-def main() -> None:
-    typer.run(vis_images)
-
 
 if __name__ == "__main__":
-    main()
+    typer.run(generate_ion_images)

src/depiction/tools/generate_ion_image.py

@@ -1,18 +1,14 @@
-import argparse
 from collections.abc import Sequence
-from pprint import pprint
 from typing import Optional
 
 import numpy as np
-import polars as pl
 from numpy.typing import NDArray
 from xarray import DataArray
 
 from depiction.image.multi_channel_image import MultiChannelImage
 from depiction.parallel_ops.parallel_config import ParallelConfig
 from depiction.parallel_ops.read_spectra_parallel import ReadSpectraParallel
 from depiction.persistence import ImzmlReadFile, ImzmlReader
-from depiction.tools.image_channel_statistics import ImageChannelStatistics
 
 
 class GenerateIonImage:
@@ -132,40 +128,3 @@ def _compute_for_mz_ranges(
                 mz_end = np.searchsorted(mz_arr, mz_max, side="right")
                 result[i_spectrum, i_range] = np.sum(int_arr[mz_begin:mz_end])
         return result
-
-
-def main_generate_ion_image(
-    input_imzml_path: str, mass_list_path: str, output_hdf5: str, n_jobs: int, output_stats_path: Optional[str]
-) -> None:
-    parallel_config = ParallelConfig(n_jobs=n_jobs, task_size=None)
-    gen_image = GenerateIonImage(parallel_config=parallel_config)
-    mass_list_df = pl.read_csv(mass_list_path)
-
-    image = gen_image.generate_ion_images_for_file(
-        input_file=ImzmlReadFile(input_imzml_path),
-        mz_values=mass_list_df["mass"],
-        tol=mass_list_df["tol"],
-        channel_names=mass_list_df["label"],
-    )
-    xarray = image.to_dense_xarray(bg_value=np.nan)
-    xarray.to_netcdf(output_hdf5)
-
-    if output_stats_path:
-        stats = ImageChannelStatistics.compute_xarray(xarray)
-        stats.write_csv(output_stats_path)
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--input-imzml", dest="input_imzml_path")
-    parser.add_argument("--mass-list", dest="mass_list_path")
-    parser.add_argument("--output-hdf5")
-    parser.add_argument("--n-jobs", default=20, type=int)
-    parser.add_argument("--output-stats", required=False, default=None, dest="output_stats_path")
-    args = vars(parser.parse_args())
-    pprint(args)
-    main_generate_ion_image(**args)
-
-
-if __name__ == "__main__":
-    main()

src/depiction_targeted_preproc/workflow/rules/rules_vis.smk

@@ -5,11 +5,11 @@ rule vis_images:
         mass_list="{sample}/mass_list.visualization.csv"
     output:
         hdf5="{sample}/images_default.hdf5"
+    # TODO how can i pass n-jobs nicely here
     shell:
-        "python -m depiction_targeted_preproc.workflow.vis.images "
-        " --imzml-path {input.imzml[0]} --mass-list-path {input.mass_list} "
-        " --output-hdf5-path {output.hdf5}"
-        " --config-path {input.config}"
+        "python -m depiction.tools.cli.cli_generate_ion_images"
+        " --imzml-path {input.imzml[0]} --mass-list-path {input.mass_list}"
+        " --output-hdf5-path {output.hdf5} --n-jobs 10"
 
 rule vis_images_norm:
     input: hdf5="{sample}/images_{label}.hdf5"