shift map visualization

src/depiction_targeted_preproc/example_compare/run_compare.py

@@ -17,7 +17,9 @@ def prepare_tasks(input_imzml_path: Path, work_dir: Path) -> list[Path]:
     requested_files = get_all_output_files(folders)
 
     combined_dir = work_dir / input_imzml_path.stem
-    exp_files = [combined_dir / "exp_compare_cluster_stats.pdf", combined_dir / "exp_plot_compare_peak_density.pdf"]
+    exp_files = [combined_dir / "exp_compare_cluster_stats.pdf",
+                 combined_dir / "exp_plot_compare_peak_density.pdf",
+                 combined_dir / "exp_plot_map_comparison.pdf"]
     return requested_files + exp_files
 
 
@@ -37,7 +39,7 @@ def main() -> None:
         requested_files += prepare_tasks(data_raw_dir / imzml, work_dir=work_dir)
 
     ## TODO quick hack
-    #requested_files = [f for f in requested_files if "mini" in str(f)]
+    # requested_files = [f for f in requested_files if "mini" in str(f)]
 
     SnakemakeInvoke(continue_on_error=True).invoke(work_dir=work_dir, result_files=requested_files, n_cores=4)
 

src/depiction_targeted_preproc/workflow/exp/plot_map_comparison.py

@@ -0,0 +1,26 @@
+from pathlib import Path
+from typing import Annotated
+import xarray
+import typer
+from matplotlib import pyplot as plt
+
+
+def exp_plot_map_comparison(
+    input_mass_shift_paths: Annotated[list[Path], typer.Argument()],
+    output_pdf_path: Annotated[Path, typer.Option()],
+) -> None:
+    # load all the inputs
+    mass_shifts = [xarray.open_dataarray(path) for path in input_mass_shift_paths]
+
+    fig, axs = plt.subplots(1, len(mass_shifts), figsize=(10 * len(mass_shifts), 10))
+    for i, shift_map in enumerate(mass_shifts):
+        shift_map.isel(c=0).plot(x="x", y="y", ax=axs[i], cmap="coolwarm", vmin=-1, vmax=+1)
+        axs[i].set_aspect("equal")
+        variant = input_mass_shift_paths[i].parent.name
+        test_mass = shift_map.coords["c"][0]
+        axs[i].set_title(f"{variant} computed shift for test mass {test_mass:.2f}")
+
+    plt.savefig(output_pdf_path, bbox_inches="tight")
+
+if __name__ == "__main__":
+    typer.run(exp_plot_map_comparison)

src/depiction_targeted_preproc/workflow/experimental.smk

@@ -82,3 +82,14 @@ rule qc_plot_marker_presence_mini:
         " --table-marker-distances-calib {input.table_marker_distances_calib}"
         " --output-pdf {output.pdf}"
 
+variants_with_map = ["mass_cluster", "reg_shift"]
+
+rule exp_plot_map_comparison:
+    input:
+        mass_shifts=expand("{{sample}}/{exp_variant}/test_mass_shifts.hdf5", exp_variant=variants_with_map)
+    output:
+        pdf="{sample}/exp_plot_map_comparison.pdf"
+    shell:
+        "python -m depiction_targeted_preproc.workflow.exp.plot_map_comparison"
+        " {input.mass_shifts}"
+        " --output-pdf-path {output.pdf}"

src/depiction_targeted_preproc/workflow/rules/rules_vis.smk

@@ -36,4 +36,17 @@ rule vis_clustering:
         png="{sample}/cluster_{label}.png"
     shell:
         "python -m depiction_targeted_preproc.workflow.vis.clustering "
-        " --input-netcdf-path {input.netcdf} --output-png-path {output.png}"
+        " --input-netcdf-path {input.netcdf} --output-png-path {output.png}"
+
+rule vis_test_mass_shifts:
+    input:
+        calib_hdf5="{sample}/calib_data.hdf5",
+        config="{sample}/pipeline_params.yml",
+        mass_list="{sample}/mass_list.calibration.csv"
+    output:
+        hdf5="{sample}/test_mass_shifts.hdf5"
+    shell:
+        "python -m depiction_targeted_preproc.workflow.vis.test_mass_shifts "
+        " --calib-hdf5-path {input.calib_hdf5} --mass-list-path {input.mass_list} "
+        " --config-path {input.config}"
+        " --output-hdf5-path {output.hdf5}"

src/depiction_targeted_preproc/workflow/vis/test_mass_shifts.py

@@ -0,0 +1,54 @@
+from pathlib import Path
+from typing import Annotated
+
+import numpy as np
+import polars as pl
+import typer
+import xarray as xr
+from typer import Option
+
+from depiction_targeted_preproc.pipeline_config.model import PipelineParameters
+from depiction_targeted_preproc.workflow.proc.calibrate import get_calibration_from_config
+
+
+def vis_test_mass_shifts(
+    calib_hdf5_path: Annotated[Path, Option()],
+    mass_list_path: Annotated[Path, Option()],
+    config_path: Annotated[Path, Option()],
+    output_hdf5_path: Annotated[Path, Option()],
+) -> None:
+    # load inputs
+    model_coefs = xr.open_dataarray(calib_hdf5_path, group="model_coefs")
+    config = PipelineParameters.parse_yaml(config_path)
+    mass_list = pl.read_csv(mass_list_path)
+    calibration = get_calibration_from_config(mass_list=mass_list, calib_config=config.calibration)
+
+    # define test masses
+    # to keep it simple for now only 1
+    test_masses = np.array([(mass_list["mass"].max() + mass_list["mass"].min()) / 2])
+    test_masses_int = np.ones_like(test_masses)
+
+    # compute the shifts
+    def compute_shifts(coef):
+        result = calibration.apply_spectrum_model(spectrum_mz_arr=test_masses, spectrum_int_arr=test_masses_int,
+                                                  model_coef=xr.DataArray(coef, dims=["c"]))
+        return xr.DataArray(result[0] - test_masses, dims=["m"])
+
+    shifts = xr.apply_ufunc(
+        compute_shifts,
+        model_coefs,
+        input_core_dims=[["c"]],
+        output_core_dims=[["m"]],
+        vectorize=True,
+    ).rename({"m": "c"})
+    shifts = shifts.assign_coords(c=test_masses)
+
+    shifts_2d = shifts.set_xindex(["x", "y"]).unstack("i")
+    shifts_2d.attrs["bg_value"] = np.nan
+
+    # save the result
+    shifts_2d.to_netcdf(output_hdf5_path)
+
+
+if __name__ == "__main__":
+    typer.run(vis_test_mass_shifts)