ssam-denovo

Author: sebastiantiesmeyer

plankton/utils.py

@@ -1,24 +1,28 @@
+from importlib.resources import path
 from textwrap import fill
 from threading import local
 
 import numpy as np
 # from numpy.ma.extras import _covhelper
-# from numpy.ma.core import dot 
+# from numpy.ma.core import dot
 
 from requests import patch
 from scipy.ndimage import gaussian_filter, maximum_filter
 from plankton.pixelmaps import PixelMap
 
-
+from sklearn.neighbors import NearestNeighbors
+from sklearn.decomposition import FastICA
+from sklearn.cluster import KMeans
 
 import matplotlib.pyplot as plt
 
+from tqdm import tqdm_notebook
 
-def determine_gains(stats1,stats2):
+def determine_gains(stats1, stats2):
 
     norm_counts1 = stats1.counts/stats1.counts.sum()
     norm_counts2 = stats2.counts/stats2.counts.sum()
-    
+
     return norm_counts1/norm_counts2
 
 
@@ -54,6 +58,7 @@ def hbar_compare(stat1, stat2, labels=None, text_display_threshold=0.02, c=None)
     if labels is not None:
         plt.xticks((0, 1), labels)
 
+
 def sorted_bar_compare(stat1, stat2, kwargs1={}, kwargs2={}):
     categories_1 = (stat1.index)
     counts_1 = (np.array(stat1.counts).flatten())
@@ -169,6 +174,7 @@ def sorted_bar_compare(stat1, stat2, kwargs1={}, kwargs2={}):
 #     ax3.xaxis.set_label_position('top')
 #     ax3.set_ylabel('log(count) spatial')
 
+
 def fill_celltypemaps(ct_map, fill_blobs=True, min_blob_area=0, filter_params={}, output_mask=None):
     """
     Post-filter cell type maps created by `map_celltypes`.
@@ -191,25 +197,25 @@ def fill_celltypemaps(ct_map, fill_blobs=True, min_blob_area=0, filter_params={}
     """
 
     from skimage import measure
-    
-    
+
     filtered_ctmaps = np.zeros_like(ct_map) - 1
 
     for cidx in np.unique(ct_map):
-        mask = ct_map==cidx
+        mask = ct_map == cidx
         if min_blob_area > 0 or fill_blobs:
             blob_labels = measure.label(mask, background=0)
             for bp in measure.regionprops(blob_labels):
                 if min_blob_area > 0 and bp.filled_area < min_blob_area:
                     for c in bp.coords:
-                        mask[c[0], c[1],] = 0 
-                        
+                        mask[c[0], c[1], ] = 0
+
                     continue
                 if fill_blobs and bp.area != bp.filled_area:
                     minx, miny,  maxx, maxy, = bp.bbox
-                    mask[minx:maxx, miny:maxy,] |= bp.filled_image
+                    mask[minx:maxx, miny:maxy, ] |= bp.filled_image
 
-        filtered_ctmaps[np.logical_and(mask == 1, np.logical_or(ct_map == -1, ct_map == cidx))] = cidx
+        filtered_ctmaps[np.logical_and(mask == 1, np.logical_or(
+            ct_map == -1, ct_map == cidx))] = cidx
 
     return filtered_ctmaps
 
@@ -240,87 +246,128 @@ def get_histograms(sdata, mins=None, maxs=None, category=None, resolution=5):
 
     return histograms,
 
-def crosscorr(x,y):
-    x -= x.mean(1)[:,None]
-    y -= y.mean(1)[:,None]
 
-    c = (np.dot(x, y.T)/x.shape[1] ).squeeze()
-    return c/x.std(1)[:,None]/y.std(1)
+def crosscorr(x, y):
+    x -= x.mean(1)[:, None]
+    y -= y.mean(1)[:, None]
+    c = (np.dot(x, y.T)/x.shape[1]).squeeze()
+
+    return np.nan_to_num(np.nan_to_num(c/np.array(x.std(1))[:, None])/np.array(y.std(1))[None, :])
 
 
 def ssam(sdata, signatures=None, adata_obs_label='celltype', kernel_bandwidth=2.5, output_um_p_px=5,
-        patch_length=1000, threshold_exp=0.1, threshold_cor=0.1, background_value=-1,
-        fill_blobs=True,min_blob_area=10):
+         patch_length=1000, threshold_exp=0.1, threshold_cor=0.1, background_value=-1,
+         fill_blobs=True, min_blob_area=10):
 
     if (sdata.scanpy is not None) and (signatures is None):
         signatures = sdata.scanpy.generate_signatures(adata_obs_label)
 
-    kernel_bandwidth_px=kernel_bandwidth/output_um_p_px
+    kernel_bandwidth_px = kernel_bandwidth/output_um_p_px
 
     out_shape = (np.ceil(sdata.x.max()/output_um_p_px+kernel_bandwidth_px*3).astype(int),
                  np.ceil(sdata.y.max()/output_um_p_px+kernel_bandwidth_px*3).astype(int))
 
     ct_map = np.zeros((out_shape), dtype=int)+background_value
     vf_norm = np.zeros_like(ct_map)
 
-    range_x = np.floor(sdata.x.min()).astype(int),np.ceil(sdata.x.max()).astype(int)
-    patch_delimiters_x = list(range(range_x[0],range_x[1],patch_length))+[range_x[1]]
+    range_x = np.floor(sdata.x.min()).astype(
+        int), np.ceil(sdata.x.max()).astype(int)
+    patch_delimiters_x = list(
+        range(range_x[0], range_x[1], patch_length))+[range_x[1]]
+
+    range_y = np.floor(sdata.y.min()).astype(
+        int), np.ceil(sdata.y.max()).astype(int)
+    patch_delimiters_y = list(
+        range(range_y[0], range_y[1], patch_length))+[range_y[1]]
+
+    print(list(patch_delimiters_x), list(patch_delimiters_y))
+
+    with tqdm_notebook(total=(len(patch_delimiters_x)-1)*(len(patch_delimiters_y)-1)) as pbar:
+        for i, x in enumerate(patch_delimiters_x[:-1]):
+            for j, y in enumerate(patch_delimiters_y[:-1]):
+                sdata_patch = sdata.raw().spatial[x:patch_delimiters_x[i+1],
+                                                y:patch_delimiters_y[j+1]]
+
+                if not len(sdata_patch):
+                    break
+
+                mins = sdata_patch.coordinates.min(0).astype(int)
+                maxs = sdata_patch.coordinates.max(
+                    0).astype(int)+kernel_bandwidth*3
+                hists = get_histograms(
+                    sdata_patch, mins=mins, maxs=maxs, resolution=output_um_p_px)
+                hists = np.concatenate(
+                    [gaussian_filter(h, kernel_bandwidth_px) for h in hists])
+
+                # print(hists.shape)
+
+                norm = np.sum(hists, axis=0)
+
+                vf_norm[(x+mins[0])//output_um_p_px:(x+mins[0])//output_um_p_px+norm.shape[0],
+                        (y+mins[1])//output_um_p_px:(y+mins[1])//output_um_p_px+norm.shape[1], ] = norm
+
+                mask = norm > threshold_exp
+
+                exps = np.zeros((len(sdata.genes), mask.sum()))
+
+                # print(exps.shape,signatures.shape)
+                exps[sdata.stats.loc[sdata_patch.genes].gene_ids.values,
+                    :] = hists[:, mask]
+
+                local_ct_map = mask.astype(int)-1
+
+                corrs = crosscorr(exps.T, signatures)
+                # print(corrs.min())
+                corrs_winners = corrs.argmax(1)
+                corrs_winners[corrs.max(1) < threshold_cor] = background_value
+                local_ct_map[mask] = corrs_winners
 
-    range_y = np.floor(sdata.y.min()).astype(int),np.ceil(sdata.y.max()).astype(int)
-    patch_delimiters_y = list(range(range_y[0],range_y[1],patch_length))+[range_y[1]]
+                # idcs=sdata_patch.index
 
-    print(list(patch_delimiters_x),list(patch_delimiters_y))
+                x_ = int(max(0, (x+mins[0])-kernel_bandwidth_px*3))//output_um_p_px
+                y_ = int(max(0, (y+mins[1])-kernel_bandwidth_px*3))//output_um_p_px
 
-    for i,x in enumerate(patch_delimiters_x[:-1]):
-        for j,y in enumerate(patch_delimiters_y[:-1]):
-            sdata_patch = sdata.raw().spatial[x:patch_delimiters_x[i+1],
-                                        y:patch_delimiters_y[j+1]]   
+                ct_map[x_:x_+local_ct_map.shape[0],
+                    y_:y_+local_ct_map.shape[1]] = local_ct_map
 
-            if not len(sdata_patch): break
-            
-            mins = sdata_patch.coordinates.min(0).astype(int)
-            maxs = sdata_patch.coordinates.max(0).astype(int)+kernel_bandwidth*3
-            hists = get_histograms(sdata_patch,mins=mins,maxs=maxs, resolution=output_um_p_px)
-            hists = np.concatenate([gaussian_filter(h, kernel_bandwidth_px) for h in hists])
+                pbar.update(1)
 
-            print(hists.shape)
+    ct_map = fill_celltypemaps(
+        ct_map, min_blob_area=min_blob_area, fill_blobs=fill_blobs)
+    return PixelMap(ct_map.T, px_p_um=1/output_um_p_px)
 
-            norm = np.sum(hists, axis=0)
 
-            vf_norm[(x+mins[0])//output_um_p_px:(x+mins[0])//output_um_p_px+norm.shape[0],
-                    (y+mins[1])//output_um_p_px:(y+mins[1])//output_um_p_px+norm.shape[1],]=norm
+def localmax_sampling(sdata, n_clusters=10, min_distance=3, bandwidth=4):
+    n_bins = np.array(sdata.spatial.shape)
 
-            mask = norm > threshold_exp
+    vf = (gaussian_filter(np.histogram2d(
+        *sdata.coordinates.T, bins=n_bins)[0], 2))
+    localmaxs = np.where((maximum_filter(vf, min_distance) == vf) & (vf > 0.2))
+    knn = NearestNeighbors(n_neighbors=150)
+    knn.fit(sdata.coordinates)
+    dists, nbrs = knn.kneighbors(np.array(localmaxs).T)
+    neighbor_types = np.array(sdata.gene_ids)[nbrs]
 
-            exps = np.zeros((len(sdata.genes),mask.sum()))
+    counts = np.zeros((dists.shape[0], len(sdata.genes)))
 
-            print(exps.shape,signatures.shape)
-            exps[sdata.stats.loc[sdata_patch.genes].gene_ids.values,:]=hists[:, mask]
+    bandwidth = bandwidth
+    def kernel(x): return np.exp(-x**2/(2*bandwidth**2))
 
-            # signatures -= signatures.mean(0)
-            # signatures /= signatures.std(0)
+    for i in range(0, dists.shape[1]):
+        counts[np.arange(dists.shape[0]), neighbor_types[:, i]
+               ] += kernel(dists[:, i])
 
-            # exps -= exps.mean(0)
-            # exps -= exps.std(0)
+    assert (all(counts.sum(1)) > 0)
 
-            local_ct_map = mask.astype(int)-1
-            # corrs = np.inner(exps.T, signatures)
+    counts=counts/counts.sum(1)[:,None]
 
-            corrs = crosscorr(exps.T,signatures)
-            print(corrs.min())
-            corrs_winners = corrs.argmax(1)
-            corrs_winners[corrs.max(1)<threshold_cor]=background_value
-            local_ct_map[mask] = corrs_winners
+    ica = FastICA(n_components=30)
+    facs = ica.fit_transform(counts)
 
-            # idcs=sdata_patch.index
+    km = KMeans(
+        n_clusters=n_clusters,
+        max_iter=100,
+        n_init=1,).fit(counts[:])
 
+    return km.cluster_centers_
 
-            x_=int(max(0,(x+mins[0])-kernel_bandwidth_px*3))//output_um_p_px
-            y_=int(max(0,(y+mins[1])-kernel_bandwidth_px*3))//output_um_p_px
-            
-            
-            ct_map[x_:x_+local_ct_map.shape[0],
-                   y_:y_+local_ct_map.shape[1]]=local_ct_map
-    
-    ct_map = fill_celltypemaps(ct_map,min_blob_area=min_blob_area, fill_blobs=fill_blobs)
-    return PixelMap(ct_map.T,px_p_um=1/output_um_p_px)