01_clone_counters.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f53ccb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp clone_counters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e892c4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3a39c6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "import re\n",
    "from functools import partial, reduce\n",
    "from glob import glob\n",
    "from typing import Callable\n",
    "\n",
    "import dask.array as da\n",
    "import dask.dataframe as dd\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import xarray as xr\n",
    "from matplotlib import pyplot as plt\n",
    "from skimage import measure\n",
    "\n",
    "from clonedetective import clone_analysis as ca\n",
    "from clonedetective.utils import (\n",
    "    RGB_image_from_CYX_img,\n",
    "    add_scale_regionprops_table_area_measurements,\n",
    "    calc_allfilt_from_thresholds,\n",
    "    calculate_corresponding_labels,\n",
    "    calculate_overlap,\n",
    "    check_channels_input_suitable_and_return_channels,\n",
    "    concat_list_of_thresholds_to_string,\n",
    "    extend_region_properties_list,\n",
    "    four_ch_CYX_img_to_three_ch_CYX_img,\n",
    "    get_all_labeled_clones_unmerged_and_merged,\n",
    "    img_path_to_xarr,\n",
    "    last2dims,\n",
    "    lazy_props,\n",
    "    plot_new_images,\n",
    "    reorder_df_to_put_ch_info_first,\n",
    "    update_1st_coord_and_dim_of_xarr,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "930256ed",
   "metadata": {},
   "source": [
    "# CloneCounter Classes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c98b4b1",
   "metadata": {},
   "source": [
    "## Parent Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d99d2b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class CloneCounter:\n",
    "    def __init__(\n",
    "        self,\n",
    "        exp_name: str,\n",
    "        img_name_regex: str,\n",
    "        pixel_size: float,\n",
    "        tot_seg_ch: str = \"C0\",\n",
    "    ):\n",
    "        self.exp_name = exp_name\n",
    "        self.img_name_regex = img_name_regex\n",
    "        self.pixel_size = pixel_size\n",
    "        self.tot_seg_ch = tot_seg_ch\n",
    "        self.results_clones_and_neighbour_counts = dict()\n",
    "        self.defined_thresholds = dict()\n",
    "\n",
    "    def add_images(self, **channel_path_globs):\n",
    "        return img_path_to_xarr(\n",
    "            self.img_name_regex,\n",
    "            self.pixel_size,\n",
    "            ch_name_for_first_dim=\"img_channels\",\n",
    "            **channel_path_globs,\n",
    "        )\n",
    "\n",
    "    def add_segmentations(\n",
    "        self,\n",
    "        additional_func_to_map: Callable = None,\n",
    "        ad_func_kwargs: dict = None,\n",
    "        **channel_path_globs,\n",
    "    ):\n",
    "        segmentations = img_path_to_xarr(\n",
    "            self.img_name_regex,\n",
    "            self.pixel_size,\n",
    "            ch_name_for_first_dim=\"seg_channels\",\n",
    "            **channel_path_globs,\n",
    "        )\n",
    "\n",
    "        if additional_func_to_map is not None:\n",
    "            segmentations.data = segmentations.data.map_blocks(\n",
    "                additional_func_to_map, **ad_func_kwargs, dtype=np.uint16\n",
    "            )\n",
    "\n",
    "        segmentations.data = segmentations.data.map_blocks(\n",
    "            last2dims(partial(measure.label)), dtype=np.uint16\n",
    "        )\n",
    "        return segmentations\n",
    "\n",
    "    def combine_C0_overlaps_and_measurements(self):\n",
    "        ov_df = (\n",
    "            self.results_overlaps.pivot(\n",
    "                index=[\"img_name\", \"C0_labels\"],\n",
    "                columns=[\"colocalisation_ch\"],\n",
    "                values=\"is_in_label\",\n",
    "            )\n",
    "            .query(\"C0_labels != 0\")\n",
    "            .copy()\n",
    "        )\n",
    "        sk_df = self.results_measurements.query(\"seg_ch== 'C0'\").set_index(\n",
    "            [\"seg_img\", \"label\"]\n",
    "        )\n",
    "        sk_df.index.rename([\"img_name\", \"C0_labels\"], inplace=True)\n",
    "        return pd.merge(ov_df, sk_df, left_index=True, right_index=True)\n",
    "\n",
    "    def _determine_seg_img_channel_pairs(\n",
    "        self, seg_channels: list = None, img_channels: list = None\n",
    "    ):\n",
    "        seg_channels = check_channels_input_suitable_and_return_channels(\n",
    "            channels=seg_channels,\n",
    "            available_channels=self.image_data.seg_channels.values.tolist(),\n",
    "        )\n",
    "\n",
    "        img_channels = check_channels_input_suitable_and_return_channels(\n",
    "            channels=img_channels,\n",
    "            available_channels=self.image_data.img_channels.values.tolist(),\n",
    "        )\n",
    "\n",
    "        seg_img_channel_pairs = pd.DataFrame()\n",
    "        seg_img_channel_pairs[\"image_channel\"] = pd.Series(img_channels)\n",
    "        seg_img_channel_pairs[\"segmentation_channel\"] = pd.Series(seg_channels)\n",
    "        self.seg_img_channel_pairs = seg_img_channel_pairs.fillna(method=\"ffill\")[\n",
    "            [\"segmentation_channel\", \"image_channel\"]\n",
    "        ]\n",
    "\n",
    "    def make_measurements(\n",
    "        self,\n",
    "        seg_channels: list = None,\n",
    "        img_channels: list = None,\n",
    "        extra_properties: list = None,\n",
    "        **kwargs,\n",
    "    ):\n",
    "\n",
    "        self._determine_seg_img_channel_pairs(seg_channels, img_channels)\n",
    "\n",
    "        properties = extend_region_properties_list(extra_properties)\n",
    "\n",
    "        results = list()\n",
    "        for _, seg_ch, img_ch in self.seg_img_channel_pairs.itertuples():\n",
    "            for seg, img in zip(\n",
    "                self.image_data[\"segmentations\"].loc[seg_ch],\n",
    "                self.image_data[\"images\"].loc[img_ch],\n",
    "            ):\n",
    "                results.append(\n",
    "                    lazy_props(\n",
    "                        seg.data,\n",
    "                        img.data,\n",
    "                        seg.seg_channels.item(),\n",
    "                        img.img_channels.item(),\n",
    "                        seg.img_name.item(),\n",
    "                        img.img_name.item(),\n",
    "                        properties,\n",
    "                        **kwargs,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "        df = dd.from_delayed(results).compute()\n",
    "        df = add_scale_regionprops_table_area_measurements(df, self.pixel_size)\n",
    "        self.results_measurements = reorder_df_to_put_ch_info_first(df)\n",
    "        self._determine_max_seg_label_levels()\n",
    "\n",
    "    def _determine_max_seg_label_levels(self):\n",
    "        self.tot_seg_ch_max_labels = (\n",
    "            self.image_data[\"segmentations\"]\n",
    "            .loc[self.tot_seg_ch]\n",
    "            .data.map_blocks(\n",
    "                lambda x: np.unique(x).shape[0], drop_axis=(1, 2), dtype=np.uint16,\n",
    "            )\n",
    "            .compute()\n",
    "            .max()\n",
    "        )\n",
    "\n",
    "    def _create_df_from_arr(self, arr):\n",
    "        return (\n",
    "            xr.DataArray(\n",
    "                np.moveaxis(arr, 1, 0),\n",
    "                coords=(\n",
    "                    self.image_data[\"segmentations\"].coords[\"seg_channels\"][1:],\n",
    "                    self.image_data[\"segmentations\"].coords[\"img_name\"],\n",
    "                    np.arange(self.tot_seg_ch_max_labels),\n",
    "                ),\n",
    "                dims=(\"colocalisation_ch\", \"img_name\", \"C0_labels\",),\n",
    "            )\n",
    "            .to_dataframe(\"is_in_label\")\n",
    "            .reset_index()\n",
    "            .dropna()\n",
    "        )\n",
    "\n",
    "    def measure_overlap(self):\n",
    "        self._determine_max_seg_label_levels()\n",
    "        arr = (\n",
    "            self.image_data[\"segmentations\"]\n",
    "            .data.map_blocks(\n",
    "                calculate_overlap,\n",
    "                drop_axis=[0],\n",
    "                dtype=np.float64,\n",
    "                num_of_segs=self.image_data[\"segmentations\"].shape[0],\n",
    "                preallocate_value=self.tot_seg_ch_max_labels,\n",
    "            )\n",
    "            .compute()\n",
    "        )\n",
    "\n",
    "        df = self._create_df_from_arr(arr)\n",
    "        df[\"is_in_label\"] = df[\"is_in_label\"].astype(np.uint16)\n",
    "        self.results_overlaps = df[\n",
    "            [\"img_name\", \"C0_labels\", \"colocalisation_ch\", \"is_in_label\"]\n",
    "        ]\n",
    "\n",
    "    def testing_possible_thresholds(\n",
    "        self,\n",
    "        int_img: str,\n",
    "        thresholds: list,\n",
    "        figure_shape=None,\n",
    "        figure_size=None,\n",
    "        **kwargs,\n",
    "    ):\n",
    "        seg = (\n",
    "            self.image_data[\"segmentations\"]\n",
    "            .sel(seg_channels=self.tot_seg_ch, img_name=int_img)\n",
    "            .compute()\n",
    "        ).data\n",
    "\n",
    "        df = self.results_measurements.query(\"int_img == @int_img\")\n",
    "\n",
    "        thresh_img_dict = dict()\n",
    "        for thresh in thresholds:\n",
    "            all_filt = calc_allfilt_from_thresholds(thresh, df)\n",
    "            to_keep = all_filt[all_filt].reset_index()[\"label\"].values\n",
    "            thresh = concat_list_of_thresholds_to_string(thresh)\n",
    "            thresh_img_dict[thresh] = np.isin(seg, to_keep) * seg\n",
    "\n",
    "        img = (self.image_data[\"images\"].sel(img_name=int_img).compute()).data\n",
    "\n",
    "        if img.shape[0] == 4:\n",
    "            img = four_ch_CYX_img_to_three_ch_CYX_img(img)\n",
    "\n",
    "        imgs_to_plot = [\n",
    "            RGB_image_from_CYX_img(red=img[2], green=img[1], blue=img[0]),\n",
    "            seg,\n",
    "        ] + list(thresh_img_dict.values())\n",
    "\n",
    "        labels_to_plot = [\n",
    "            f\"composite image\",\n",
    "            f\"{self.tot_seg_ch} segmentation\",\n",
    "        ] + list(thresh_img_dict.keys())\n",
    "\n",
    "        plot_new_images(\n",
    "            imgs_to_plot,\n",
    "            labels_to_plot,\n",
    "            figure_shape=figure_shape,\n",
    "            figure_size=figure_size,\n",
    "            **kwargs,\n",
    "        )\n",
    "\n",
    "    def _filter_labels(self, thresholds: list, thresh_name: str, calc_clone: bool):\n",
    "        all_filt = calc_allfilt_from_thresholds(\n",
    "            thresholds, self.results_measurements.copy()\n",
    "        )\n",
    "\n",
    "        if calc_clone:\n",
    "            all_filt.name = f\"{thresh_name}_clonepos\"\n",
    "\n",
    "        else:\n",
    "            all_filt.name = f\"{thresh_name}_pos\"\n",
    "\n",
    "        return (\n",
    "            all_filt[all_filt == True]\n",
    "            .reset_index()\n",
    "            .groupby(\"int_img\")\n",
    "            .agg({\"label\": lambda x: list(x)})[\"label\"]\n",
    "            .to_dict(),\n",
    "            all_filt,\n",
    "        )\n",
    "\n",
    "    def _update_measurements_df_with_all_filt(self, all_filt):\n",
    "        self.results_measurements = pd.merge(\n",
    "            self.results_measurements,\n",
    "            all_filt,\n",
    "            suffixes=(\"_unwanted\", None),\n",
    "            on=[\"int_img\", \"label\"],\n",
    "        )\n",
    "        self.results_measurements.drop(\n",
    "            columns=self.results_measurements.filter(regex=\"_unwanted\").columns,\n",
    "            inplace=True,\n",
    "        )\n",
    "\n",
    "    def add_clones_and_neighbouring_labels(\n",
    "        self,\n",
    "        thresholds: list,\n",
    "        thresh_name: str = \"C1\",\n",
    "        calc_clones: bool = True,\n",
    "    ):\n",
    "        new_coord = [\n",
    "            \"ext_tot_seg_labs\",\n",
    "            f\"{thresh_name}_neg_labs\",\n",
    "            f\"{thresh_name}_pos_labs\",\n",
    "            \"tot_nc\",\n",
    "            f\"{thresh_name}neg_nc\",\n",
    "            f\"{thresh_name}pos_nc\",\n",
    "        ]\n",
    "\n",
    "        if calc_clones:\n",
    "            new_coord.insert(3, f\"{thresh_name}_clone\")\n",
    "\n",
    "        clone_coords, clone_dims = update_1st_coord_and_dim_of_xarr(\n",
    "            self.image_data[\"images\"],\n",
    "            new_coord=new_coord,\n",
    "            new_dim=f\"{thresh_name}_neighbours\",\n",
    "        )\n",
    "\n",
    "        labels_to_keep, all_filt = self._filter_labels(\n",
    "            thresholds, thresh_name, calc_clones\n",
    "        )\n",
    "\n",
    "        self._update_measurements_df_with_all_filt(all_filt)\n",
    "\n",
    "        new_label_imgs = get_all_labeled_clones_unmerged_and_merged(\n",
    "            self.image_data[\"segmentations\"].loc[self.tot_seg_ch],\n",
    "            labels_to_keep,\n",
    "            calc_clones,\n",
    "        )\n",
    "\n",
    "        return xr.DataArray(\n",
    "            data=new_label_imgs,\n",
    "            coords=clone_coords,\n",
    "            dims=clone_dims,\n",
    "            attrs={f\"{self.tot_seg_ch}_labels_kept_query\": thresholds},\n",
    "        )\n",
    "\n",
    "    def _get_centroids_list(self):\n",
    "        df = self.results_measurements.query(\"int_img_ch == @self.tot_seg_ch\")\n",
    "        centroids_list = list()\n",
    "        for img_name in df[\"int_img\"].unique():\n",
    "            centroids_list.append(\n",
    "                (\n",
    "                    df.query(\"int_img == @img_name\")\n",
    "                    .loc[:, [\"centroid-0\", \"centroid-1\"]]\n",
    "                    .values.astype(int)\n",
    "                )\n",
    "            )\n",
    "        return centroids_list\n",
    "\n",
    "    def _colabels_to_df(self, colabels, thresh_name):\n",
    "        return (\n",
    "            xr.DataArray(\n",
    "                colabels,\n",
    "                coords=(\n",
    "                    self.image_data[thresh_name].coords[f\"{thresh_name}_neighbours\"],\n",
    "                    self.image_data[thresh_name].coords[\"img_name\"],\n",
    "                    range(1, colabels.shape[2] + 1),\n",
    "                ),\n",
    "                dims=(f\"{thresh_name}_neighbours\", \"img_name\", \"label\"),\n",
    "            )\n",
    "            .to_dataframe(\"colabel\")\n",
    "            .reset_index()\n",
    "            .dropna()\n",
    "            .pivot(\n",
    "                index=[\"img_name\", \"label\"],\n",
    "                columns=[f\"{thresh_name}_neighbours\"],\n",
    "                values=\"colabel\",\n",
    "            )\n",
    "            .astype(np.uint16)\n",
    "            .query(\"label == ext_tot_seg_labs\")\n",
    "        )\n",
    "\n",
    "    def mutually_exclusive_cell_types(self):\n",
    "        return (self.results_measurements.filter(regex=\"_pos\").sum(axis=1) <= 1).all()\n",
    "\n",
    "    def complete_set_of_cell_types(self):\n",
    "        return (self.results_measurements.filter(regex=\"_pos\").sum(axis=1) > 0).all()\n",
    "\n",
    "    def measure_clones_and_neighbouring_labels_for_ind_thresh(self, thresh_name):\n",
    "        colabels = calculate_corresponding_labels(\n",
    "            self.image_data[thresh_name].data,\n",
    "            self._get_centroids_list(),\n",
    "            self.image_data[thresh_name].shape[0],\n",
    "            self.tot_seg_ch_max_labels,\n",
    "        )\n",
    "\n",
    "        df = self._colabels_to_df(colabels, thresh_name)\n",
    "\n",
    "        df.index.rename([\"int_img\", \"label\"], inplace=True)\n",
    "\n",
    "        self.results_clones_and_neighbour_counts[thresh_name] = df.drop(\n",
    "            columns=df.filter(regex=r\"labs\").columns\n",
    "        )\n",
    "\n",
    "    def measure_all_clones_and_neighbouring_labels(self):\n",
    "        if not self.mutually_exclusive_cell_types():\n",
    "            raise AttributeError('cell types are not mutually exclusive')\n",
    "        if not self.complete_set_of_cell_types():\n",
    "            raise AttributeError('cell types are not complete')\n",
    "\n",
    "        for key in self.defined_thresholds.keys():\n",
    "            self.measure_clones_and_neighbouring_labels_for_ind_thresh(thresh_name=key)\n",
    "\n",
    "    def combine_neighbour_counts_and_measurements(self):\n",
    "        list_df = list(self.results_clones_and_neighbour_counts.values()) + [\n",
    "            self.results_measurements.set_index([\"int_img\", \"label\"])\n",
    "        ]\n",
    "        merged_df = reduce(\n",
    "            lambda left, right: pd.merge(\n",
    "                left,\n",
    "                right,\n",
    "                how=\"left\",\n",
    "                on=[\"int_img\", \"label\"],\n",
    "                suffixes=(None, \"_extra\"),\n",
    "            ),\n",
    "            list_df,\n",
    "        )\n",
    "\n",
    "        merged_df = merged_df.drop(\n",
    "            columns=merged_df.filter(regex=\"extra\").columns\n",
    "        ).reset_index()\n",
    "\n",
    "        # fillna in case merge of results_clones_and_neighbour_counts creates NaN\n",
    "        # when a img_name contains none of a specific cell type\n",
    "        merged_df[merged_df.filter(regex=\"clone$|nc\").columns] = (\n",
    "            merged_df.filter(regex=\"clone$|nc\").fillna(0).astype(np.uint16)\n",
    "        )\n",
    "\n",
    "        return merged_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44d483b9",
   "metadata": {},
   "source": [
    "## CloneCounter subclasses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cefd331",
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class LazyCloneCounter(CloneCounter):\n",
    "    def __init__(\n",
    "        self,\n",
    "        exp_name: str,\n",
    "        img_name_regex: str,\n",
    "        pixel_size: float,\n",
    "        tot_seg_ch: str = \"C0\",\n",
    "    ):\n",
    "        super().__init__(exp_name, img_name_regex, pixel_size)\n",
    "\n",
    "    def add_images(self, **channel_path_globs):\n",
    "        self.image_data = xr.Dataset(\n",
    "            {\"images\": super().add_images(**channel_path_globs)}\n",
    "        )\n",
    "\n",
    "    def add_segmentations(\n",
    "        self,\n",
    "        additional_func_to_map: Callable = None,\n",
    "        ad_func_kwargs: dict = None,\n",
    "        **channel_path_globs,\n",
    "    ):\n",
    "        self.image_data[\"segmentations\"] = super().add_segmentations(\n",
    "            additional_func_to_map, ad_func_kwargs, **channel_path_globs\n",
    "        )\n",
    "\n",
    "    def add_clones_and_neighbouring_labels(\n",
    "        self,\n",
    "        thresholds: list = ['int_img_ch == \"C1\" & mean_intensity > 1000'],\n",
    "        thresh_name: str = \"C1\",\n",
    "        calc_clones: bool = True,\n",
    "    ):\n",
    "        self.image_data[thresh_name] = super().add_clones_and_neighbouring_labels(\n",
    "            thresholds, thresh_name, calc_clones\n",
    "        )\n",
    "        self.defined_thresholds[thresh_name] = thresholds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "579a7aac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class PersistentCloneCounter(CloneCounter):\n",
    "    def __init__(\n",
    "        self,\n",
    "        exp_name: str,\n",
    "        img_name_regex: str,\n",
    "        pixel_size: float,\n",
    "        tot_seg_ch: str = \"C0\",\n",
    "    ):\n",
    "        super().__init__(exp_name, img_name_regex, pixel_size)\n",
    "\n",
    "    def add_images(self, **channel_path_globs):\n",
    "        self.image_data = xr.Dataset(\n",
    "            {\"images\": super().add_images(**channel_path_globs)}\n",
    "        ).persist()\n",
    "\n",
    "    def add_segmentations(\n",
    "        self,\n",
    "        additional_func_to_map: Callable = None,\n",
    "        ad_func_kwargs: dict = None,\n",
    "        **channel_path_globs,\n",
    "    ):\n",
    "        self.image_data[\"segmentations\"] = (\n",
    "            super()\n",
    "            .add_segmentations(\n",
    "                additional_func_to_map, ad_func_kwargs, **channel_path_globs\n",
    "            )\n",
    "            .persist()\n",
    "        )\n",
    "\n",
    "    def add_clones_and_neighbouring_labels(\n",
    "        self,\n",
    "        thresholds: list = ['int_img_ch == \"C1\" & mean_intensity > 1000'],\n",
    "        thresh_name: str = \"C1\",\n",
    "        calc_clones: bool = True,\n",
    "    ):\n",
    "        self.image_data[thresh_name] = (\n",
    "            super()\n",
    "            .add_clones_and_neighbouring_labels(thresholds, thresh_name, calc_clones)\n",
    "            .persist()\n",
    "        )\n",
    "        self.defined_thresholds[thresh_name] = thresholds"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}