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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vima-spatial
-Version: 0.2.7
+Version: 0.2.9
 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.9}/setup.cfg

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

{vima_spatial-0.2.7 → vima_spatial-0.2.9}/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.9}/src/vima/ingest/ingest.py

@@ -1,5 +1,6 @@
 import numpy as numpy
 import numpy as np
+import pandas as pd
 import xarray as xr
 import cv2 as cv2
 import scanpy as sc
@@ -105,6 +106,161 @@ def add_covs(pca, sid_to_covs):
         pca[cov_name] = pca['sid'].map(sid_to_covs[cov_name])
     return ['sid'] + cov_names
 
+def collapse_markers(normeddir, markers, pseudomarker, outdir, masksdir=None):
+    """
+    Keep a subset of markers and fold all the others into one pseudomarker.
+
+    Reads the normalized pixel matrices written by `prepare_merfish`,
+    `prepare_xenium5k`, or `nonst.prepare` (for transcriptomic data the markers
+    are genes), retains `markers`, and replaces every other marker with a single
+    channel named `pseudomarker` holding their combined signal. Because the
+    stored values are log-normalized, they are exponentiated, summed, and
+    log-normalized again, so the pseudomarker is on the same scale as a real
+    marker. The dataset-wide means and stds stored on each file are rewritten to
+    match the new marker set: retained markers keep their existing values and the
+    pseudomarker's are computed from the collapsed data.
+
+    Parameters
+    ----------
+    normeddir
+        Directory of normalized ``.nc`` files, e.g. ``{outdir}/normalized``.
+    markers
+        Markers to retain, in the order they should appear; the pseudomarker is
+        appended after them. Markers absent from the data are skipped with a
+        warning.
+    pseudomarker
+        Name for the new channel, e.g. ``'nonimmune'``.
+    outdir
+        Directory to write the collapsed ``.nc`` files to. This is a
+        ``normalized``-style directory, so to build a dataset that
+        `pca_pixels` can be pointed at, pass ``f'{newroot}/normalized'`` and copy
+        the masks to ``f'{newroot}/masks'``.
+    masksdir
+        Directory of tissue masks, used to compute the pseudomarker's moments
+        over non-empty pixels only. Defaults to ``masks`` alongside `normeddir`.
+    """
+    import netCDF4
+
+    if masksdir is None:
+        masksdir = os.path.join(os.path.dirname(os.path.normpath(normeddir)), 'masks')
+
+    sids = [os.path.splitext(f)[0]
+        for f in os.listdir(normeddir) if f.endswith('.nc') and not f.startswith('.')]
+    if len(sids) == 0:
+        logger.warning(f'No .nc files found in {normeddir}. Check your path and try again.')
+        return
+    os.makedirs(outdir, exist_ok=True)
+
+    # decide on the new marker set, using the first sample as the reference
+    da = xr.open_dataarray(f'{normeddir}/{sids[0]}.nc')
+    ref_markers = list(da.marker.values)
+    ref_means, ref_stds = da.attrs['means'], da.attrs['stds']
+    da.close(); del da
+
+    if pseudomarker in ref_markers:
+        raise ValueError(f"'{pseudomarker}' is already a marker; pick a name for the "
+                         f"pseudomarker that isn't in the data.")
+    missing = [m for m in markers if m not in ref_markers]
+    if missing:
+        logger.warning(f'{len(missing)} of the {len(markers)} requested markers are not in the '
+                       f'data and will be skipped: {missing}')
+    keep = [m for m in markers if m in ref_markers]
+    if len(keep) == 0:
+        raise ValueError('None of the requested markers are in the data.')
+    if len(keep) == len(ref_markers):
+        logger.warning(f'All {len(ref_markers)} markers were retained, so {pseudomarker} will be '
+                       f'empty.')
+    logger.info(f'Keeping {len(keep)} markers and collapsing the other '
+                f'{len(ref_markers) - len(keep)} into {pseudomarker}.')
+
+    # collapse each sample, accumulating the pseudomarker's per-sample moments as we go
+    sample_sids, sample_means, sample_stds, sample_npixels = [], [], [], []
+    for sid in settings.progress(sids, name='collapsing markers'):
+        da = xr.open_dataarray(f'{normeddir}/{sid}.nc')
+        mask_da = xr.open_dataarray(f'{masksdir}/{sid}.nc')
+
+        sid_markers = list(da.marker.values)
+        if sid_markers != ref_markers:
+            logger.warning(f'{sid} has different markers ({len(sid_markers)}) than {sids[0]} '
+                           f'({len(ref_markers)}); matching by name.')
+        absent = [m for m in keep if m not in sid_markers]
+        if absent:
+            raise ValueError(f'{sid} is missing {len(absent)} of the markers to retain: {absent}')
+        ix = {m: i for i, m in enumerate(sid_markers)}
+        keep_ix = [ix[m] for m in keep]
+        keep_ix_set = set(keep_ix)
+        drop_ix = [i for i in range(len(sid_markers)) if i not in keep_ix_set]
+
+        # load raw arrays and close before doing any work, so we never hold two
+        # full (H x W x n_markers) copies simultaneously
+        x = da.x.values; y = da.y.values
+        data = da.values
+        mask = mask_da.values
+        da.close(); mask_da.close()
+        del da, mask_da; gc.collect()
+
+        # accumulate the dropped markers one at a time: fancy-indexing them all at
+        # once would materialize the (y, x, n_dropped) copy this function exists to
+        # avoid. log-normalized values are exponentiated before summing so the
+        # pseudomarker ends up on the same scale as the markers we kept.
+        acc = np.zeros(data.shape[:2], dtype=np.float64)
+        for j in drop_ix:
+            acc += np.expm1(data[..., j].astype(np.float64))
+        other = np.log1p(acc)
+        del acc
+
+        # empty pixels are zero in the input, and log1p(sum(expm1(0))) is zero, so
+        # they stay empty without any special handling
+        s = xr.DataArray(
+            np.concatenate([data[..., keep_ix], other[..., None].astype(np.float32)], axis=-1),
+            dims=['y', 'x', 'marker'],
+            coords={'x': x, 'y': y, 'marker': keep + [pseudomarker]})
+        s.name = sid
+        util.write_xarray(s, f'{outdir}/{sid}.nc')
+        del data, s; gc.collect()
+
+        # a sample with an empty mask contributes nothing rather than a nan that
+        # would poison the pooled moments
+        vals = other[mask]
+        if len(vals) == 0:
+            logger.warning(f'{sid} has no non-empty pixels; excluding it from the '
+                           f'{pseudomarker} moments.')
+        else:
+            sample_sids.append(sid)
+            sample_means.append(vals.mean(dtype=np.float64))
+            sample_stds.append(vals.std(dtype=np.float64))
+            sample_npixels.append(len(vals))
+        del other, mask, vals; gc.collect()
+
+    if len(sample_sids) == 0:
+        raise ValueError(f'No non-empty pixels in any sample, so {pseudomarker} has no moments. '
+                         f'Check that {masksdir} holds the masks for {normeddir}.')
+
+    # pool the pseudomarker's moments across samples the same way get_sumstats does
+    pseudo_mean, pseudo_std = util.pool_moments(
+        pd.DataFrame([sample_means], index=[pseudomarker], columns=sample_sids),
+        pd.DataFrame([sample_stds],  index=[pseudomarker], columns=sample_sids),
+        sample_npixels)
+    pseudo_mean, pseudo_std = pseudo_mean.iloc[0], pseudo_std.iloc[0]
+    if pseudo_std == 0:
+        # a constant channel stays exactly 0 after standardization rather than
+        # dividing by zero downstream
+        logger.warning(f'{pseudomarker} has zero variance; setting its std to 1.')
+        pseudo_std = 1.
+    logger.info(f'{pseudomarker}: mean {pseudo_mean:.3f}, std {pseudo_std:.3f}')
+
+    # the pseudomarker's moments aren't known until every sample has been read, so
+    # the files are written above without them and stamped here; this rewrites
+    # metadata only, rather than re-collapsing every sample a second time
+    keep_ix_ref = [ref_markers.index(m) for m in keep]
+    means = np.append(np.asarray(ref_means)[keep_ix_ref], pseudo_mean).astype(np.float32)
+    stds  = np.append(np.asarray(ref_stds)[keep_ix_ref],  pseudo_std).astype(np.float32)
+    for sid in sids:
+        with netCDF4.Dataset(f'{outdir}/{sid}.nc', 'a') as ds:
+            v = ds.variables[sid]
+            v.setncattr('means', means)
+            v.setncattr('stds', stds)
+
 def pca_pixels(outdir, repname, nmetamarkers=10, npixels_to_plot=50000,
                total_n_metapixels=2_000_000, sid_to_covs=None):
     """
@@ -149,7 +305,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.9}/src/vima/ingest/nonst.py

@@ -1,5 +1,6 @@
 import os, glob, gc
 import numpy as np
+import pandas as pd
 import cv2 as cv2
 import xarray as xr
 from . import util
@@ -134,12 +135,14 @@ def prepare(load, filepaths, orig_pixel_size, markers, get_foreground, norm_by_b
             )
         for sid in settings.progress(sids)])
     gc.collect()
-    _, pixels = norm_by_background(pixels)
+    goodmarkers, pixels = norm_by_background(pixels)
     ntranscripts = pixels.sum(axis=1, dtype=np.float64)
     med_ntranscripts = np.median(ntranscripts)
     pixels = np.log1p(med_ntranscripts * pixels / (ntranscripts[:,None] + 1e-6)) # adding to denominator in case pixel is all 0s
-    means = pixels.mean(axis=0, dtype=np.float64)
-    stds = pixels.std(axis=0, dtype=np.float64)
+    # indexed by marker name so each sample's moments are aligned by name below,
+    # rather than positionally against whatever markers that sample kept
+    means = pd.Series(pixels.mean(axis=0, dtype=np.float64), index=goodmarkers)
+    stds = pd.Series(pixels.std(axis=0, dtype=np.float64), index=goodmarkers)
     del pixels; gc.collect()
     
     logger.info('Normalizing and writing')
@@ -153,6 +156,7 @@ def prepare(load, filepaths, orig_pixel_size, markers, get_foreground, norm_by_b
         pl = np.log1p(med_ntranscripts * pl / (pl.sum(axis=1)[:,None] + 1e-6)) # adding to denominator in case pixel is all 0s
         s = s.sel(marker=goodmarkers)
         util.set_pixels(s, mask, pl)
-        s.attrs['means'] = means.reindex(s.marker.values, fill_value=1).values.astype(np.float32)
-        s.attrs['stds'] = stds.reindex(s.marker.values, fill_value=0).values.astype(np.float32)
+        # a marker with no moments standardizes to a no-op rather than dividing by zero
+        s.attrs['means'] = means.reindex(s.marker.values, fill_value=0).values.astype(np.float32)
+        s.attrs['stds'] = stds.reindex(s.marker.values, fill_value=1).values.astype(np.float32)
         util.write_xarray(s, f'{normeddir}/{sid}.nc')