vima-spatial 0.2.0__tar.gz → 0.2.2__tar.gz

{vima_spatial-0.2.0/src/vima_spatial.egg-info → vima_spatial-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vima-spatial
-Version: 0.2.0
+Version: 0.2.2
 Summary: variational inference-based microniche analysis
 Home-page: https://github.com/yakirr/vima
 Author: Yakir Reshef
@@ -25,7 +25,7 @@ Requires-Dist: netcdf4
 Requires-Dist: seaborn
 Requires-Dist: pandas>=2.2.3
 Requires-Dist: scipy
-Requires-Dist: cna>=0.2.3
+Requires-Dist: cna>=0.2.4
 Requires-Dist: tqdm
 Requires-Dist: pyarrow
 Requires-Dist: scikit-image

{vima_spatial-0.2.0 → vima_spatial-0.2.2}/setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = vima-spatial
-version = 0.2.0
+version = 0.2.2
 author = Yakir Reshef
 author_email = yreshef@broadinstitute.org
 description = variational inference-based microniche analysis
@@ -35,7 +35,7 @@ install_requires =
 	seaborn
 	pandas>=2.2.3
 	scipy
-	cna>=0.2.3
+	cna>=0.2.4
 	tqdm
 	pyarrow
 	scikit-image

{vima_spatial-0.2.0 → vima_spatial-0.2.2}/src/vima/cc.py

@@ -28,32 +28,52 @@ def anndata(patchmeta, Z, var_names=None, use_rep='X', n_comps=10, **kwargs):
 
     return d
 
-def apply(models, P, batch_size=1000):
+def apply(models, P, batch_size=1000, with_mse=False):
     P.pytorch_mode()
     P.augmentation_off()
     for model in models:
         model.eval()
-    
+
     eval_loader = DataLoader(
         dataset=P,
         batch_size=batch_size,
         shuffle=False)
 
     Zs = {modelid: [] for modelid in range(len(models))}
+    if with_mse:
+        MSEs = {modelid: [] for modelid in range(len(models))}
     with torch.no_grad():
         for batch in pb(eval_loader):
             for modelid, model in enumerate(models):
-                Zs[modelid].append(model.embedding(batch).detach().cpu().numpy())
-
-    return np.array([
-        np.concatenate(Z) for Z in Zs.values()])
-
-def latentreps(models, P, use_rep='X', n_comps=100, **kwargs):
+                if with_mse:
+                    x_recon, mean, _ = model(batch, sample_from_latent=False)
+                    Zs[modelid].append(mean.reshape(len(batch[0]), -1).detach().cpu().numpy())
+                    MSEs[modelid].append(((x_recon - batch[0]) ** 2).mean(dim=(2, 3)).detach().cpu().numpy())
+                else:
+                    Zs[modelid].append(model.embedding(batch).detach().cpu().numpy())
+
+    Zs_out = np.array([np.concatenate(Z) for Z in Zs.values()])
+    if with_mse:
+        MSEs_out = np.array([np.concatenate(M) for M in MSEs.values()])
+        return Zs_out, MSEs_out
+    return Zs_out
+
+def latentreps(models, P, use_rep='X', n_comps=100, with_mse=True, **kwargs):
     print('applying models')
-    Zs = apply(models, P)
+    result = apply(models, P, with_mse=with_mse)
+    Zs, MSEs = result if with_mse else (result, None)
+
     print('computing nearest-neighbor graphs')
     ds = [anndata(P.meta, Z, use_rep=use_rep, n_comps=n_comps, **kwargs) for Z in pb(Zs)]
-    return Fingerprints.from_list(ds)
+    fp = Fingerprints.from_list(ds)
+    
+    if MSEs is not None:
+        mean_mse = MSEs.mean(axis=0)
+        marker_names = next(iter(P.samples.values())).coords['marker'].values
+        per_channel = pd.DataFrame(mean_mse, index=fp.obs.index, columns=marker_names)
+        fp.obsm['per_channel_mse'] = per_channel
+        fp.obs['mse'] = per_channel.mean(axis=1)
+    return fp
 
 def _association(MAMresid, M, Nmodels, y, batches, donorids, ks=None, Nnull=1000, show_progress=False):
     # prep data

{vima_spatial-0.2.0 → vima_spatial-0.2.2}/src/vima/fingerprints.py

@@ -120,7 +120,7 @@ class Fingerprints:
             self._adata.obsp[f'connectivities_{i}'] = d.obsp['connectivities']
             self._adata.obsp[f'distances_{i}'] = d.obsp['distances']
             self._adata.uns[f'neighbors_{i}'] = d.uns['neighbors']
-
+        
     def sample_pcs(self, sid_name='sid'):
         D = self.avg_graph(make_umap=False)
         NAM, _ = cna.tl.nam(D, sid_name)
@@ -129,6 +129,16 @@ class Fingerprints:
         U, _, _ = np.linalg.svd(NAM, full_matrices=False)
         return pd.DataFrame(U, index=NAM.index,
                             columns=[f'PC{i+1}' for i in range(U.shape[1])])
+    
+    def mn_pcs(self, sid_name='sid'):
+        D = self.avg_graph(make_umap=False)
+        NAM, _ = cna.tl.nam(D, sid_name)
+        NAM -= NAM.mean(axis=0)
+        NAM /= NAM.std(axis=0)
+        _, _, VT = np.linalg.svd(NAM, full_matrices=False)
+        print(VT.shape)
+        return pd.DataFrame(VT.T, index=NAM.columns,
+                            columns=[f'PC{i+1}' for i in range(VT.shape[0])])
 
     def to_anndata(self):
         X = np.hstack([self._adata.obsm[f'X_{i}'] for i in range(self.nmodels)])

{vima_spatial-0.2.0 → vima_spatial-0.2.2}/src/vima/ingest/dimreduce.py

@@ -25,16 +25,31 @@ def metapixels_allsamples(normedpixelsdir, masksdir, sids, total_n_metapixels, p
     # figure out how many metapixels to store per sample
     nsamples = len(sids)
     nmp_per_sample = total_n_metapixels // nsamples
-    
 
     if plot:
         fig = plt.figure(figsize=(7,5))
 
     print('Creating metapixels prior to PCA')
     print(f'\t(will randomly downsample to {nmp_per_sample} metapixels per sample if needed.)')
+    ref_markers = None
+    ref_sid = None
     for i, sid in enumerate(pb(sids)):
         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:
+            ref_markers, ref_sid = markers, sid
+        elif markers != ref_markers:
+            missing = set(ref_markers) - set(markers)
+            extra   = set(markers) - set(ref_markers)
+            order_mismatch = not missing and not extra
+            detail = (f'order differs' if order_mismatch else
+                      f'{len(missing)} missing, {len(extra)} extra vs {ref_sid}')
+            print(f'\033[93mWARNING: {sid} has different markers ({len(markers)}) '
+                  f'than {ref_sid} ({len(ref_markers)}): {detail}\033[0m')
+        
         all_metapixels[sid], all_npixels[sid] = metapixels(da.astype(np.float32), mask_da)
         da.close(); mask_da.close()
         del da, mask_da
@@ -118,10 +133,14 @@ def pca_pixels(normedpixelsdir, masksdir, pcloadings, sids):
     for sid in pb(sids):
         da = xr.open_dataarray(f'{normedpixelsdir}/{sid}.nc')
         mask_da = xr.open_dataarray(f'{masksdir}/{sid}.nc')
-        pl = util.xr_to_pixellist(da.astype(np.float32), mask_da)
-        # close file handles before the dot product so HDF5 cache is released
+        # load raw arrays and close before dtype conversion so we never hold
+        # two full (H × W × n_genes) copies simultaneously
+        data = da.values
+        mask = mask_da.values
         da.close(); mask_da.close()
         del da, mask_da; gc.collect()
+        pl = data.astype(np.float32, copy=False)[mask]
+        del data, mask; gc.collect()
 
         pl_pca = pl.dot(pcloadings)
         pcs.append(pl_pca)

{vima_spatial-0.2.0 → vima_spatial-0.2.2}/src/vima/ingest/ingest.py

@@ -100,7 +100,7 @@ def pca_pixels(outdir, repname, nmetamarkers=10, plot=True, npixels_to_plot=5000
 
     if plot:
         visualize_pixels(pca, npixels_to_plot, 'metamarkers', cov_names)
-    return pca
+    return pca, loadings
 
 def harmonize(allpixels_pca, sid_to_covs=None, npixels_to_plot=50000, plot=True):
     import harmonypy as hm
@@ -129,12 +129,12 @@ def write_harmonized(outdir, repname, harmpixels):
         pl = harmpixels[harmpixels.sid == sid]
         s_ = np.zeros((*mask.shape, len(hpcs)))
         s_[mask.data] = pl[pcs].values
-        mask.close(); del mask
         s = xr.DataArray(s_,
              dims=['y', 'x', 'marker'],
              coords={'x': mask.x, 'y': mask.y, 'marker': hpcs})
         s.name = sid
         s.to_netcdf(f'{processeddir}/{sid}.nc', encoding={s.name: util.compression}, engine="netcdf4")
+        mask.close(); del mask
         gc.collect()
 
 def sanity_checks(outdir, repname, npcs=1, nskip=3):

{vima_spatial-0.2.0 → vima_spatial-0.2.2}/src/vima/ingest/st.py

@@ -171,7 +171,7 @@ def transcriptlist_to_normedpixelmatrix(sid, data, x_col, y_col, gene_col, pixel
 
 def rasterize_and_normalize_generic(load, filepaths, x_col, y_col, gene_col, n_top_genes_per_sample, pixel_size, outdir,
                                     min_ntranscripts_per_pixel, min_ngenes_per_pixel,
-                                    genes_to_add=[], plot_mean_var=True, plot_spatial_hvgs=False):
+                                    genes_to_add=[], basic_plots=True, plot_spatial_hvgs=False):
     if len(filepaths) == 0:
         print('No files found. Check your filepaths and try again.')
         return
@@ -181,7 +181,7 @@ def rasterize_and_normalize_generic(load, filepaths, x_col, y_col, gene_col, n_t
     print('Finding HVGs and dataset-wide mean and variance per gene...')
     hvgs, means, stds = get_sumstats(load, filepaths, normfactor, x_col, y_col, gene_col,
                                      n_top_genes_per_sample=n_top_genes_per_sample,
-                                     genes_to_add=genes_to_add, pixel_size=pixel_size, plot_mean_var=plot_mean_var,
+                                     genes_to_add=genes_to_add, pixel_size=pixel_size, plot_mean_var=basic_plots,
                                      plot_spatial_hvgs=plot_spatial_hvgs,
                                      min_ntranscripts=min_ntranscripts_per_pixel)
     print('Final number of genes used =', len(hvgs))
@@ -195,7 +195,7 @@ def rasterize_and_normalize_generic(load, filepaths, x_col, y_col, gene_col, n_t
         sid, data = load(filepath)
         print(f'Processing sample {i+1}/{len(filepaths)}: {sid}')
         mask, pm = transcriptlist_to_normedpixelmatrix(sid, data, x_col, y_col, gene_col, pixel_size,
-                                                       normfactor, means, stds, genes=hvgs, plots=plot_mean_var,
+                                                       normfactor, means, stds, genes=hvgs, plots=basic_plots,
                                                        min_ntranscripts_per_pixel=min_ntranscripts_per_pixel,
                                                        min_ngenes_per_pixel=min_ngenes_per_pixel)
         del data; gc.collect()
@@ -203,25 +203,25 @@ def rasterize_and_normalize_generic(load, filepaths, x_col, y_col, gene_col, n_t
         util.write_xarray(pm, f'{normdir}/{pm.name}.nc')
 
 def prepare_xenium5k(load, filepaths, x_col, y_col, gene_col, n_top_genes_per_sample, outdir,
-                     pixel_size=10, genes_to_add=[], plot_mean_var=True, plot_spatial_hvgs=False,
+                     pixel_size=10, genes_to_add=[], basic_plots=True, plot_spatial_hvgs=False,
                      min_ntranscripts_per_pixel=11, min_ngenes_per_pixel=5):
     rasterize_and_normalize_generic(load, filepaths, x_col, y_col, gene_col,
                                   n_top_genes_per_sample,
                                   pixel_size=pixel_size,
                                   outdir=outdir,
                                   genes_to_add=genes_to_add,
-                                  plot_mean_var=plot_mean_var,
+                                  basic_plots=basic_plots,
                                   plot_spatial_hvgs=plot_spatial_hvgs,
                                   min_ntranscripts_per_pixel=min_ntranscripts_per_pixel,
                                   min_ngenes_per_pixel=min_ngenes_per_pixel)
 
 def prepare_merfish(load, filepaths, x_col, y_col, gene_col, outdir,
-                    pixel_size=10, plot_mean_var=True,
+                    pixel_size=10, basic_plots=True,
                     min_ntranscripts_per_pixel=11, min_ngenes_per_pixel=1):
     rasterize_and_normalize_generic(load, filepaths, x_col, y_col, gene_col,
                                   None,
                                   pixel_size=pixel_size,
                                   outdir=outdir,
-                                  plot_mean_var=plot_mean_var,
+                                  basic_plots=basic_plots,
                                   min_ntranscripts_per_pixel=min_ntranscripts_per_pixel,
                                   min_ngenes_per_pixel=min_ngenes_per_pixel)

{vima_spatial-0.2.0 → vima_spatial-0.2.2}/src/vima/patchfeatures.py

@@ -4,7 +4,7 @@ import pandas as pd
 import xarray as xr
 from scipy.stats import rankdata
 from tqdm import tqdm
-pb = lambda x: tqdm(x, ncols=100)
+pb = lambda x, desc: tqdm(x, ncols=100, desc=desc)
 
 def cell_type_counts(
     cells,
@@ -18,6 +18,7 @@ def cell_type_counts(
     patch_x_microns_col='x_microns',
     patch_y_microns_col='y_microns',
     patch_size_in_pixels_col='patchsize',
+    include_totalcells=False,
     pixel_size_microns=10
 ):
     """Return per-patch cell type counts.
@@ -46,7 +47,7 @@ def cell_type_counts(
     cell_types = sorted(cells[celltype_col].unique())
     counts = pd.DataFrame(0, index=patch_meta.index, columns=cell_types, dtype=int)
 
-    for sid, sid_cells in pb(cells.groupby(sid_col)):
+    for sid, sid_cells in pb(cells.groupby(sid_col), 'cell_type_counts'):
         sid_patches = patch_meta[patch_meta[patch_sid_col] == sid]
         if len(sid_patches) == 0 or len(sid_cells) == 0:
             continue
@@ -70,6 +71,7 @@ def cell_type_counts(
     if normalized:
         totals = counts.sum(axis=1)
         counts = counts.div(totals, axis=0).fillna(0)
+    if include_totalcells:
         counts['totalcells'] = totals
 
     return counts
@@ -107,7 +109,7 @@ def expression_profiles(
 
     result = None
 
-    for sid in pb(sids):
+    for sid in pb(sids, 'expression_profiles'):
         sid_patches = patch_meta[patch_meta[patch_sid_col] == sid]
         sample = xr.open_dataarray(os.path.join(directory, f'{sid}.nc')).load()
         marker_names = sample.coords['marker'].values.tolist()
@@ -137,45 +139,94 @@ def expression_profiles(
     return result
 
 
+def _permutation_pvals(X_ranked, group_a, group_b, donors, n_perms, rng):
+    sum_a    = X_ranked[group_a].sum(axis=0)
+    sum_b    = X_ranked[group_b].sum(axis=0)
+    count_a  = float(group_a.sum())
+    count_b  = float(group_b.sum())
+    obs_diff = sum_a / count_a - sum_b / count_b
+
+    unique_donors = np.unique(donors)
+    n_donors      = len(unique_donors)
+    da_sum   = np.zeros((n_donors, X_ranked.shape[1]))
+    db_sum   = np.zeros((n_donors, X_ranked.shape[1]))
+    da_count = np.zeros(n_donors)
+    db_count = np.zeros(n_donors)
+    for i, d in enumerate(unique_donors):
+        in_d        = donors == d
+        da_sum[i]   = X_ranked[in_d & group_a].sum(axis=0)
+        db_sum[i]   = X_ranked[in_d & group_b].sum(axis=0)
+        da_count[i] = (in_d & group_a).sum()
+        db_count[i] = (in_d & group_b).sum()
+
+    delta_sum   = db_sum - da_sum
+    delta_count = db_count - da_count
+    flip = (rng.random((n_perms, n_donors)) < 0.5).astype(float)
+
+    sum_a_null   = sum_a   + flip @ delta_sum
+    count_a_null = count_a + flip @ delta_count
+    sum_b_null   = (sum_a + sum_b) - sum_a_null
+    count_b_null = (count_a + count_b) - count_a_null
+    count_a_null = np.maximum(count_a_null, 1.0)
+    count_b_null = np.maximum(count_b_null, 1.0)
+    null_diff = sum_a_null / count_a_null[:, None] - sum_b_null / count_b_null[:, None]
+
+    return ((np.abs(null_diff) >= np.abs(obs_diff)).sum(axis=0) + 1) / (n_perms + 1)
+
+
+def _ttest_pvals(X_ranked, group_a, group_b, donors):
+    from scipy import stats
+    global_mean_a = X_ranked[group_a].mean(axis=0)
+    global_mean_b = X_ranked[group_b].mean(axis=0)
+    diffs = []
+    for d in np.unique(donors):
+        in_d   = donors == d
+        a_rows = X_ranked[in_d & group_a]
+        b_rows = X_ranked[in_d & group_b]
+        mean_a = a_rows.mean(axis=0) if len(a_rows) > 0 else global_mean_a
+        mean_b = b_rows.mean(axis=0) if len(b_rows) > 0 else global_mean_b
+        diffs.append(mean_a - mean_b)
+    if len(diffs) < 2:
+        raise ValueError(f'T-test requires at least 2 units; got {len(diffs)}')
+    _, pvals = stats.ttest_1samp(np.array(diffs), 0, axis=0)
+    return pvals
+
+
 def test_features(
     features,
     group_a,
     group_b=None,
     *,
-    perm_key,
-    method='mean_of_ranks',
+    unit_of_analysis,
     n_perms=100000,
     seed=None,
     corr_method='benjamini-hochberg',
+    Ttest=False,
 ):
-    """Compare feature distributions between two patch groups via donor-level permutation.
+    """Compare feature distributions between two patch groups.
 
-    Group labels are flipped at the donor level so the null distribution respects
-    within-donor correlations.
+    Uses mean rank difference (Wilcoxon/AUC equivalent) as the test statistic.
+    By default, significance is assessed via donor-level permutation (group labels
+    are flipped at the unit_of_analysis level). Pass Ttest=True to instead run a
+    paired T-test on per-unit mean rank differences.
 
     Args:
         features: DataFrame (n_patches × n_features), e.g. from cell_type_counts or
                   expression_profiles.
         group_a: boolean array, length n_patches — first group (e.g. associated patches).
         group_b: boolean array or None — second group; defaults to ~group_a.
-        perm_key: array-like of donor/sample IDs aligned with features rows. Permutations
-                  flip labels at the level of unique values in perm_key.
-        method: test statistic to use —
-                'mean_of_ranks'    mean rank difference (Wilcoxon/AUC equivalent, default),
-                'mean_of_medians'  mean of per-donor median differences; only donors
-                                   with patches in both groups contribute.
-        n_perms: number of permutations (default 100000).
-        seed: random seed for reproducibility.
+        unit_of_analysis: array-like of donor/sample IDs aligned with features rows.
+                  Permutations (or T-test pairing) operate at the level of unique values.
+        n_perms: number of permutations (default 100000); ignored when Ttest=True.
+        seed: random seed for reproducibility; ignored when Ttest=True.
         corr_method: multiple-testing correction — 'benjamini-hochberg' (default) or 'bonferroni'.
+        Ttest: if True, use a one-sample two-sided T-test on per-unit mean rank differences
+               instead of permutation testing (default False).
 
     Returns:
         DataFrame indexed by feature name, columns: median_a, median_b, diff, pvals,
         pvals_adj. diff = median_a - median_b. Sorted by pvals ascending.
     """
-    if method not in ('mean_of_medians', 'mean_of_ranks'):
-        raise ValueError(
-            f'method must be "mean_of_medians", or "mean_of_ranks"; got {method!r}'
-        )
     if corr_method not in ('benjamini-hochberg', 'bonferroni'):
         raise ValueError(
             f'corr_method must be "benjamini-hochberg" or "bonferroni"; got {corr_method!r}'
@@ -196,7 +247,7 @@ def test_features(
 
     group_a  = _align(group_a, 'group_a').astype(bool)
     group_b  = ~group_a if group_b is None else _align(group_b, 'group_b').astype(bool)
-    donors   = _align(perm_key, 'perm_key')
+    donors   = _align(unit_of_analysis, 'unit_of_analysis')
     X = features.values.astype(float)
 
     print(f'Comparing {group_a.sum()} patches (Group A) to {group_b.sum()} patches (Group B).')
@@ -206,41 +257,12 @@ def test_features(
     median_a_raw = np.median(X[group_a], axis=0)
     median_b_raw = np.median(X[group_b], axis=0)
 
-    if method == 'mean_of_ranks':
-        X = rankdata(X, axis=0, nan_policy='raise')
-
-    sum_a    = X[group_a].sum(axis=0)
-    sum_b    = X[group_b].sum(axis=0)
-    count_a  = float(group_a.sum())
-    count_b  = float(group_b.sum())
-    obs_diff = sum_a / count_a - sum_b / count_b
-
-    unique_donors = np.unique(donors)
-    n_donors      = len(unique_donors)
-    da_sum   = np.zeros((n_donors, X.shape[1]))
-    db_sum   = np.zeros((n_donors, X.shape[1]))
-    da_count = np.zeros(n_donors)
-    db_count = np.zeros(n_donors)
-    for i, d in enumerate(unique_donors):
-        in_d        = donors == d
-        da_sum[i]   = X[in_d & group_a].sum(axis=0)
-        db_sum[i]   = X[in_d & group_b].sum(axis=0)
-        da_count[i] = (in_d & group_a).sum()
-        db_count[i] = (in_d & group_b).sum()
+    X = rankdata(X, axis=0, nan_policy='raise')
 
-    delta_sum   = db_sum - da_sum
-    delta_count = db_count - da_count
-    flip = (rng.random((n_perms, n_donors)) < 0.5).astype(float)
-
-    sum_a_null   = sum_a   + flip @ delta_sum
-    count_a_null = count_a + flip @ delta_count
-    sum_b_null   = (sum_a + sum_b) - sum_a_null
-    count_b_null = (count_a + count_b) - count_a_null
-    count_a_null = np.maximum(count_a_null, 1.0)
-    count_b_null = np.maximum(count_b_null, 1.0)
-    null_diff = sum_a_null / count_a_null[:, None] - sum_b_null / count_b_null[:, None]
-
-    pvals = ((np.abs(null_diff) >= np.abs(obs_diff)).sum(axis=0) + 1) / (n_perms + 1)
+    if Ttest:
+        pvals = _ttest_pvals(X, group_a, group_b, donors)
+    else:
+        pvals = _permutation_pvals(X, group_a, group_b, donors, n_perms, rng)
 
     if corr_method == 'benjamini-hochberg':
         from statsmodels.stats.multitest import multipletests
@@ -260,6 +282,7 @@ def test_features(
             'mean_b':      mean_b,
             'log2fc':            np.log2((median_a_raw + 1e-9) / (median_b_raw + 1e-9)),
             'log2fc_means': np.log2((mean_a + 1e-9) / (mean_b + 1e-9)),
+            'diff_median':     median_a_raw - median_b_raw,
         }, index=features.columns)
     else:
         result = pd.DataFrame({

{vima_spatial-0.2.0 → vima_spatial-0.2.2}/src/vima/vis/__init__.py

@@ -5,4 +5,4 @@ from .patchexamples import (scaler, apply_colormap, plot_with_reconstruction,
                       plot_patches_separatechannels, plot_patches_overlaychannels,
                       plot_patches_overlaychannels_linsum,
                       plot_patches_overlaychannels_sorted, plot_patches_fourcolors)
-from .spatial import spatialplot, annotate_spatialplot
+from .spatial import spatialplot, annotate_spatialplot, plot_npatches_per_sample, plot_samples_with_patches, plot_sample_with_patches

{vima_spatial-0.2.0 → vima_spatial-0.2.2}/src/vima/vis/patches.py

@@ -1,3 +1,4 @@
+import os
 import warnings
 import matplotlib.pyplot as plt
 import matplotlib.patches as mpatches
@@ -34,7 +35,8 @@ def _plot_separate(patches, markers, vmin, vmax, cmap='seismic', show=True):
     return fig
 
 
-def _plot_composite(patches, markers, colors, vmin, vmax, features=None, nx=5, ny=5, show=True):
+def _plot_composite(patches, markers, colors, vmin, vmax, features=None, nx=5, ny=5, show=True,
+                    subfig=None):
     N, ps, K = patches.shape[0], patches.shape[1], len(markers)
 
     rgb = np.zeros((N, ps, ps, 3))
@@ -58,7 +60,8 @@ def _plot_composite(patches, markers, colors, vmin, vmax, features=None, nx=5, n
     else:
         cell_to_patch = {i: i for i in range(N)}
 
-    fig, axs = plt.subplots(ny, nx, figsize=(nx, ny))
+    fig = subfig if subfig is not None else plt.figure(figsize=(nx, ny))
+    axs = fig.subplots(ny, nx)
     for ax in axs.flatten():
         ax.axis('off')
     for cell_i, patch_i in cell_to_patch.items():
@@ -68,7 +71,8 @@ def _plot_composite(patches, markers, colors, vmin, vmax, features=None, nx=5, n
     legend_handles = [mpatches.Patch(facecolor=colors[k], label=markers[k]) for k in range(K)]
     fig.legend(handles=legend_handles, loc='lower center', ncol=K, frameon=False,
                fontsize=8, bbox_to_anchor=(0.5, 0), bbox_transform=fig.transFigure)
-    plt.tight_layout(rect=[0, 0.08, 1, 1])
+    if subfig is None:
+        plt.tight_layout(rect=[0, 0.08, 1, 1])
     if show:
         plt.show()
     return fig
@@ -102,6 +106,10 @@ class MarkersInSpace:
         self._arrays = {}   # {sid: np.ndarray(H, W, K)}
         self.vmin = {}      # {marker: float}
         self.vmax = {}      # {marker: float}
+        if samples is not None:
+            self._all_sids = set(samples.keys())
+        else:
+            self._all_sids = {f[:-3] for f in os.listdir(directory) if f.endswith('.nc')}
         if markers:
             self.add_markers(markers)
 
@@ -144,6 +152,7 @@ class MarkersInSpace:
         for sid, arr in pb(self._arrays.items(), f'Adding {len(new_markers)} markers'):
             new_data = self._read_sid(sid, new_markers)
             self._arrays[sid] = np.concatenate([arr, new_data], axis=-1)
+        self._ensure_sids_loaded(self._all_sids)  # load remaining sids for dataset-wide stats
         self._update_stats()
 
     def _ensure_sids_loaded(self, sids):
@@ -203,7 +212,7 @@ class MarkersInSpace:
         return _plot_separate(patches, markers, vmin, vmax, cmap, show=show)
 
     def show_composite(self, patchmeta, markers=None, features=None, colors=None,
-                       n=25, nx=5, ny=5, seed=None, vmin=None, vmax=None, show=True):
+                       n=25, nx=5, ny=5, seed=None, vmin=None, vmax=None, show=True, subfig=None):
         """Show patches as additive RGB composites in an nx × ny grid.
 
         Args:
@@ -244,7 +253,8 @@ class MarkersInSpace:
         marker_indices = [self._marker_to_idx[m] for m in markers]
         patches = self._extract_patches(patchmeta, marker_indices)
         vmin, vmax = self._resolve_scale(markers, vmin, vmax)
-        return _plot_composite(patches, markers, colors, vmin, vmax, features, nx, ny, show=show)
+        return _plot_composite(patches, markers, colors, vmin, vmax, features, nx, ny, show=show,
+                               subfig=subfig)
 
 
 # ── standalone convenience functions (backed by a global MarkersInSpace) ─────
@@ -273,7 +283,8 @@ def show_patches_separate(patchmeta, markers, directory, samples=None,
 
 def show_patches_composite(patchmeta, markers, directory, samples=None,
                             features=None, colors=None,
-                            n=25, nx=5, ny=5, seed=None, vmin=None, vmax=None, show=True):
+                            n=25, nx=5, ny=5, seed=None, vmin=None, vmax=None, show=True,
+                            subfig=None):
     """Convenience wrapper around MarkersInSpace.show_composite using a global cache.
 
     On the first call (or when directory changes) a new MarkersInSpace instance is
@@ -282,12 +293,13 @@ def show_patches_composite(patchmeta, markers, directory, samples=None,
     """
     return _get_default_mis(directory, samples).show_composite(
         patchmeta, markers, features=features, colors=colors,
-        n=n, nx=nx, ny=ny, seed=seed, vmin=vmin, vmax=vmax, show=show)
+        n=n, nx=nx, ny=ny, seed=seed, vmin=vmin, vmax=vmax, show=show, subfig=subfig)
 
 
 def show_patches_cells(patchmeta, cells, x_col, y_col, celltype_col,
                        pixelsize_microns, nx=5, ny=5, sid_col='sid',
-                       colors=None, seed=None, s=8, show=True):
+                       colors=None, seed=None, s=8, show=True, subfig=None,
+                       include_only=None):
     """Show an nx×ny grid of randomly chosen patches with cells overlaid as colored dots.
 
     Args:
@@ -312,13 +324,17 @@ def show_patches_cells(patchmeta, cells, x_col, y_col, celltype_col,
 
     extent = int(patchmeta['patchsize'].iloc[0]) * pixelsize_microns
 
+    if include_only is not None:
+        cells = cells[cells[celltype_col].isin(include_only)]
+
     # Build color map for cell types
     all_types = sorted(cells[celltype_col].dropna().unique())
     if colors is None:
         palette = plt.cm.tab20(np.linspace(0, 1, max(len(all_types), 1)))
         colors = {ct: palette[i] for i, ct in enumerate(all_types)}
 
-    fig, axs = plt.subplots(ny, nx, figsize=(nx * 2, ny * 2), squeeze=False)
+    fig = subfig if subfig is not None else plt.figure(figsize=(nx * 2, ny * 2))
+    axs = fig.subplots(ny, nx, squeeze=False)
     for ax in axs.flatten():
         ax.set_visible(False)
 
@@ -352,7 +368,8 @@ def show_patches_cells(patchmeta, cells, x_col, y_col, celltype_col,
     fig.legend(handles=handles, loc='lower center', ncol=ncol_legend,
                frameon=False, fontsize=7,
                bbox_to_anchor=(0.5, 0), bbox_transform=fig.transFigure)
-    plt.tight_layout(rect=[0, bottom_margin, 1, 1])
+    if subfig is None:
+        plt.tight_layout(rect=[0, bottom_margin, 1, 1])
 
     if show:
         plt.show()

{vima_spatial-0.2.0 → vima_spatial-0.2.2}/src/vima/vis/spatial.py

@@ -129,14 +129,17 @@ def spatialplot(patchmeta, values, sids=None, cmap='viridis', vmin=None, vmax=No
     for ax in axs[len(sids):]:
         ax.axis('off')
 
+    fig._vima_sid_to_ax = sid_to_ax
     fig.tight_layout()
     if show:
         plt.show()
-    else:
-        return sid_to_ax
+    return fig
 
 
-def annotate_spatialplot(sid_to_ax, patchmeta, highlight, color, thickness=3, show=True):
+def annotate_spatialplot(patchmeta, highlight, color, thickness=3, show=True, fig=None):
+    if fig is None:
+        fig = plt.gcf()
+    sid_to_ax = fig._vima_sid_to_ax
     for sid, ax in sid_to_ax.items():
         mypatches = patchmeta[patchmeta.sid == sid]
         if len(mypatches) == 0:
@@ -161,9 +164,8 @@ def annotate_spatialplot(sid_to_ax, patchmeta, highlight, color, thickness=3, sh
             cnt = cnt.squeeze()
             if cnt.ndim == 1:
                 continue
-            ax.plot(cnt[:, 0], cnt[:, 1], color=color, linewidth=thickness)
+            ax.plot(np.append(cnt[:, 0], cnt[0, 0]), np.append(cnt[:, 1], cnt[0, 1]), color=color, linewidth=thickness)
 
     if show:
         plt.show()
-    else:
-        return sid_to_ax
+    return fig

{vima_spatial-0.2.0 → vima_spatial-0.2.2/src/vima_spatial.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vima-spatial
-Version: 0.2.0
+Version: 0.2.2
 Summary: variational inference-based microniche analysis
 Home-page: https://github.com/yakirr/vima
 Author: Yakir Reshef
@@ -25,7 +25,7 @@ Requires-Dist: netcdf4
 Requires-Dist: seaborn
 Requires-Dist: pandas>=2.2.3
 Requires-Dist: scipy
-Requires-Dist: cna>=0.2.3
+Requires-Dist: cna>=0.2.4
 Requires-Dist: tqdm
 Requires-Dist: pyarrow
 Requires-Dist: scikit-image

{vima_spatial-0.2.0 → vima_spatial-0.2.2}/src/vima_spatial.egg-info/requires.txt

@@ -11,7 +11,7 @@ netcdf4
 seaborn
 pandas>=2.2.3
 scipy
-cna>=0.2.3
+cna>=0.2.4
 tqdm
 pyarrow
 scikit-image