vima-spatial 0.2.7__tar.gz → 0.2.8__tar.gz

{vima_spatial-0.2.7 → vima_spatial-0.2.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vima-spatial
-Version: 0.2.7
+Version: 0.2.8
 Summary: variational inference-based microniche analysis
 Home-page: https://github.com/yakirr/vima
 Author: Yakir Reshef

{vima_spatial-0.2.7 → vima_spatial-0.2.8}/setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = vima-spatial
-version = 0.2.7
+version = 0.2.8
 author = Yakir Reshef
 author_email = yreshef@broadinstitute.org
 description = variational inference-based microniche analysis

{vima_spatial-0.2.7 → vima_spatial-0.2.8}/src/vima/ingest/dimreduce.py

@@ -16,10 +16,10 @@ def metapixels_allsamples(normedpixelsdir, masksdir, sids, total_n_metapixels):
     """
     Pool metapixels across all samples for a more robust PCA fit.
 
-    Loads each sample's normalized pixels, standardizes them with the stored
-    per-marker means/stds, builds metapixels, and randomly downsamples each
-    sample to roughly ``total_n_metapixels // len(sids)`` metapixels. Warns if a
-    sample's markers differ from the first sample's.
+    Loads each sample's normalized pixels, builds roughly
+    ``total_n_metapixels // len(sids)`` randomly chosen metapixels from it, and
+    standardizes them with the stored per-marker means/stds. Warns if a sample's
+    markers differ from the first sample's.
 
     Parameters
     ----------
@@ -55,7 +55,7 @@ def metapixels_allsamples(normedpixelsdir, masksdir, sids, total_n_metapixels):
     for i, sid in enumerate(settings.progress(sids, name='creating metapixels')):
         da = xr.open_dataarray(f'{normedpixelsdir}/{sid}.nc')
         mask_da = xr.open_dataarray(f'{masksdir}/{sid}.nc')
-        
+
         # ensure same markers in same order in all files
         markers = list(da.marker.values)
         if ref_markers is None:
@@ -68,18 +68,21 @@ def metapixels_allsamples(normedpixelsdir, masksdir, sids, total_n_metapixels):
                       f'{len(missing)} missing, {len(extra)} extra vs {ref_sid}')
             logger.warning(f'{sid} has different markers ({len(markers)}) '
                            f'than {ref_sid} ({len(ref_markers)}): {detail}')
-        
-        means = xr.DataArray(da.attrs['means'], dims='marker')
-        stds = xr.DataArray(da.attrs['stds'], dims='marker')
-        da = ((da - means) / stds).where(mask_da, 0)
 
-        all_metapixels[sid], all_npixels[sid] = metapixels(da, mask_da)
+        means, stds = da.attrs['means'], da.attrs['stds']
+
+        # metapixels are built from the un-standardized pixels and standardized
+        # afterward: averaging over a window and the affine (x-mean)/std commute,
+        # so this is identical to standardizing the full array first but avoids
+        # materializing several (y, x, marker)-sized temporaries.
+        mp, all_npixels[sid] = metapixels(da, mask_da, n_metapixels=nmp_per_sample)
         da.close(); mask_da.close()
         del da, mask_da
-        if len(all_metapixels[sid]) > nmp_per_sample:
-            ix = np.random.choice(len(all_metapixels[sid]), nmp_per_sample, replace=False)
-            all_metapixels[sid] = all_metapixels[sid].iloc[ix]
-            all_npixels[sid] = all_npixels[sid][ix]
+
+        mp -= means
+        mp /= stds
+        all_metapixels[sid] = pd.DataFrame(data=mp, columns=markers)
+        del mp
 
         # visualize distribution of num non-empty pixels per metapixel in this sample
         if settings.show_plots():
@@ -95,33 +98,61 @@ def metapixels_allsamples(normedpixelsdir, masksdir, sids, total_n_metapixels):
 
     return all_metapixels, all_npixels
 
-def metapixels(s, mask, npixels_thresh=0):
+def metapixels(s, mask, npixels_thresh=0, n_metapixels=None, window=5):
     """
     Pool each pixel with its neighbors into a metapixel.
 
-    Sums each marker over a 5x5 window centered on every pixel and divides by
-    the number of non-empty (masked) pixels contributing to that window, so each
+    Averages each marker over a ``window``-by-``window`` window centered on a
+    pixel, using only the non-empty (masked) pixels in that window, so each
     metapixel is the average over its non-empty neighbors. Metapixels with at
     most ``npixels_thresh`` contributing pixels are dropped.
 
+    Parameters
+    ----------
+    n_metapixels
+        If given, uniformly sample at most this many metapixel centers and
+        compute only those. Since the caller typically keeps a small random
+        subset anyway, this avoids convolving the whole (y, x, marker) array,
+        which dominates the cost for large marker panels.
+
     Returns
     -------
     tuple
-        ``(metapixel_df, npixels)``: a marker-columned DataFrame of retained
-        metapixels and the per-metapixel count of contributing non-empty pixels.
+        ``(metapixels, npixels)``: a ``(metapixel, marker)`` float32 array and
+        the per-metapixel count of contributing non-empty pixels.
     """
-    markers = s.marker.values
-
-    # make metapixels and compute how many non-empty pixels and transcripts are in each metapixel
-    kernel = np.ones((5, 5), np.float32)
-    mp = convolve(s.data, kernel[:, :, None], mode="constant")
-    npixels = convolve(mask.data.astype('float32'), kernel, mode="constant")
-
-    # filter out metapixels with few non-empty pixels
-    metapixels_mask = npixels > npixels_thresh
+    mask = mask.data
+    H, W = mask.shape
+
+    # how many non-empty pixels contribute to each candidate metapixel (cheap: 2D only)
+    kernel = np.ones((window, window), np.float32)
+    npixels = convolve(mask.astype(np.float32), kernel, mode="constant")
+
+    # pick the metapixel centers, sampling before doing any work over markers
+    centers = np.flatnonzero(npixels.ravel() > npixels_thresh)
+    if n_metapixels is not None and len(centers) > n_metapixels:
+        centers = centers[np.random.choice(len(centers), n_metapixels, replace=False)]
+    npixels = npixels.ravel()[centers]
+
+    # sum each window by gathering its non-empty pixels, one neighbor offset at a time
+    data = s.data.reshape(H * W, -1)
+    mask = mask.ravel()
+    r, c = np.divmod(centers, W)
+    mp = np.zeros((len(centers), data.shape[1]), np.float32)
+    rad = window // 2
+    for dr in range(-rad, rad + 1):
+        rr = r + dr
+        for dc in range(-rad, rad + 1):
+            cc = c + dc
+            neighbor = rr * W + cc
+            contributes = (rr >= 0) & (rr < H) & (cc >= 0) & (cc < W)
+            contributes &= mask[np.where(contributes, neighbor, 0)]
+            i = np.flatnonzero(contributes)
+            mp[i] += data[neighbor[i]]
 
     # divide each metapixel by the # of non-empty pixels that contributed to it and return
-    return pd.DataFrame(data=mp[metapixels_mask] / npixels[metapixels_mask][:,None], columns=markers), npixels[metapixels_mask]
+    mp /= npixels[:, None]
+    return mp, npixels
 
 # mps should be an array of dataframes containing metapixels
 def pca_metapixels(mps, k):
@@ -142,17 +173,33 @@ def pca_metapixels(mps, k):
     Returns
     -------
     tuple
-        ``(loadings, C, allmp)``: the gene-by-component loading matrix, the
-        feature correlation matrix, and the standardized metapixel AnnData.
+        ``(loadings, allmp)``: the gene-by-component loading matrix and the
+        standardized metapixel AnnData.
     """
     logger.info('merging and standardizing metapixels')
-    allmp = pd.concat(mps)
-    allmp -= allmp.values.mean(axis=0, dtype=np.float64)
-    allmp /= allmp.values.std(axis=0, dtype=np.float64)
-    allmp = allmp.fillna(0)
-    allmp.index = np.arange(len(allmp)).astype(str)
-    allmp = ad.AnnData(X=allmp)
-    C = np.corrcoef(allmp.X[::max(1,(len(allmp)//50000))].T)
+    mps = list(mps)
+    markers = mps[0].columns
+    n = sum(len(mp) for mp in mps)
+
+    # merge into one preallocated float32 matrix and standardize it in place,
+    # accumulating the moments in float64. Doing this in pandas instead upcasts
+    # the whole matrix to float64 and copies it once per operation, which at
+    # these sizes costs more than the PCA itself.
+    allmp = np.empty((n, len(markers)), np.float32)
+    i = 0
+    for mp in mps:
+        allmp[i:i+len(mp)] = mp.to_numpy(np.float32, copy=False)
+        i += len(mp)
+    del mps
+
+    allmp -= (allmp.sum(axis=0, dtype=np.float64) / n).astype(np.float32)
+    stds = np.sqrt(np.einsum('ij,ij->j', allmp, allmp, dtype=np.float64) / n)
+    stds[stds == 0] = 1  # constant features stay exactly 0, as the old fillna(0) left them
+    allmp /= stds.astype(np.float32)
+
+    allmp = ad.AnnData(X=allmp,
+                       obs=pd.DataFrame(index=np.arange(n).astype(str)),
+                       var=pd.DataFrame(index=markers))
     logger.info(f'Metapixel matrix: {allmp.shape[0]:,} pixels × {allmp.shape[1]} features')
 
     logger.info('performing PCA...')
@@ -176,7 +223,7 @@ def pca_metapixels(mps, k):
         plt.xticks(range(len(loadings.columns)), loadings.columns, rotation=90)
         settings.show('pc_loadings')
 
-    return loadings, C, allmp
+    return loadings, allmp
 
 def pca_pixels(normedpixelsdir, masksdir, pcloadings, sids):
     """
@@ -193,16 +240,18 @@ def pca_pixels(normedpixelsdir, masksdir, pcloadings, sids):
         column giving the source sample.
     """
     pcs = []
-    sid_labels = []
+    sid_codes = []
+    # project in float32: a DataFrame (or float64) right-hand side silently
+    # promotes the result, doubling both the projection cost and the size of
+    # the returned table, which has a row per pixel in the whole dataset
+    loadings = np.ascontiguousarray(np.asarray(pcloadings), dtype=np.float32)
 
     logger.info('Applying PCA projection to each sample')
-    for sid in settings.progress(sids, name='pixels -> PCA space'):
+    for code, sid in enumerate(settings.progress(sids, name='pixels -> PCA space')):
         da = xr.open_dataarray(f'{normedpixelsdir}/{sid}.nc')
         mask_da = xr.open_dataarray(f'{masksdir}/{sid}.nc')
         
-        means = xr.DataArray(da.attrs['means'], dims='marker')
-        stds = xr.DataArray(da.attrs['stds'], dims='marker')
-        da = ((da - means) / stds).where(mask_da, 0)
+        means, stds = da.attrs['means'], da.attrs['stds']
 
         # load raw arrays and close before dtype conversion so we never hold
         # two full (H × W × n_genes) copies simultaneously
@@ -213,19 +262,29 @@ def pca_pixels(normedpixelsdir, masksdir, pcloadings, sids):
         pl = data.astype(np.float32, copy=False)[mask]
         del data, mask; gc.collect()
 
-        pl_pca = pl.dot(pcloadings)
+        # standardize the non-empty pixels only, rather than the full
+        # (y, x, marker) array; empty pixels are dropped by the mask anyway
+        pl -= means
+        pl /= stds
+
+        pl_pca = pl.dot(loadings)
         pcs.append(pl_pca)
-        sid_labels.append(np.full(pl_pca.shape[0], sid, dtype=object))
+        sid_codes.append(np.full(pl_pca.shape[0], code, dtype=np.int32))
         del pl; gc.collect()
 
     # concatenate
     pcs = np.vstack(pcs)
-    sid_labels = np.concatenate(sid_labels)
+    sid_codes = np.concatenate(sid_codes)
 
     allpixels_pca = pd.DataFrame(
         pcs,
-        columns=[f'PC{i}' for i in range(1, pcloadings.shape[1] + 1)]
+        columns=[f'PC{i}' for i in range(1, loadings.shape[1] + 1)]
     )
-    allpixels_pca['sid'] = sid_labels
+    # categorical rather than an object column: one code per pixel instead of
+    # one pointer, over tens of millions of rows
+    # drop categories for samples that contributed no pixels, so downstream
+    # get_dummies (Harmony) never sees an all-zero batch column
+    allpixels_pca['sid'] = pd.Categorical.from_codes(
+        sid_codes, categories=list(sids)).remove_unused_categories()
 
     return allpixels_pca

{vima_spatial-0.2.7 → vima_spatial-0.2.8}/src/vima/ingest/ingest.py

@@ -149,7 +149,7 @@ def pca_pixels(outdir, repname, nmetamarkers=10, npixels_to_plot=50000,
                                                           total_n_metapixels=total_n_metapixels)
 
     # PCA the metapixels
-    loadings, C, allmp = dimreduce.pca_metapixels(metapixels.values(), nmetamarkers)
+    loadings, allmp = dimreduce.pca_metapixels(metapixels.values(), nmetamarkers)
     loadings.to_feather(f'{processeddir}/_pcloadings.feather')
     del metapixels, allmp; gc.collect()
 

{vima_spatial-0.2.7 → vima_spatial-0.2.8}/src/vima_spatial.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vima-spatial
-Version: 0.2.7
+Version: 0.2.8
 Summary: variational inference-based microniche analysis
 Home-page: https://github.com/yakirr/vima
 Author: Yakir Reshef