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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (43) hide show
  1. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/PKG-INFO +1 -1
  2. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/setup.cfg +1 -1
  3. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/ingest/ingest.py +156 -0
  4. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/ingest/nonst.py +9 -5
  5. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/ingest/st.py +128 -114
  6. vima_spatial-0.2.9/src/vima/ingest/util.py +237 -0
  7. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima_spatial.egg-info/PKG-INFO +1 -1
  8. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima_spatial.egg-info/SOURCES.txt +3 -1
  9. vima_spatial-0.2.9/tests/test_collapse_markers.py +155 -0
  10. vima_spatial-0.2.9/tests/test_st_rasterize.py +363 -0
  11. vima_spatial-0.2.8/src/vima/ingest/util.py +0 -108
  12. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/README.md +0 -0
  13. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/pyproject.toml +0 -0
  14. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/__init__.py +0 -0
  15. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/_settings.py +0 -0
  16. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/cc.py +0 -0
  17. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/data/__init__.py +0 -0
  18. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/data/download.py +0 -0
  19. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/data/patchcollection.py +0 -0
  20. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/data/samples.py +0 -0
  21. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/fingerprints.py +0 -0
  22. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/ingest/__init__.py +0 -0
  23. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/ingest/dimreduce.py +0 -0
  24. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/models/__init__.py +0 -0
  25. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/models/resnet_vae.py +0 -0
  26. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/models/resnetlight_decoder.py +0 -0
  27. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/models/resnetlight_encoder.py +0 -0
  28. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/models/simple_vae.py +0 -0
  29. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/models/vae.py +0 -0
  30. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/patchfeatures.py +0 -0
  31. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/train/__init__.py +0 -0
  32. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/train/logging.py +0 -0
  33. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/train/training.py +0 -0
  34. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/vis/__init__.py +0 -0
  35. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/vis/features.py +0 -0
  36. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/vis/patches.py +0 -0
  37. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/vis/patchexamples.py +0 -0
  38. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/vis/spatial.py +0 -0
  39. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima/vis/umaps.py +0 -0
  40. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima_spatial.egg-info/dependency_links.txt +0 -0
  41. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima_spatial.egg-info/requires.txt +0 -0
  42. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/src/vima_spatial.egg-info/top_level.txt +0 -0
  43. {vima_spatial-0.2.8 → vima_spatial-0.2.9}/tests/test_ra_regression.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: vima-spatial
3
- Version: 0.2.8
3
+ Version: 0.2.9
4
4
  Summary: variational inference-based microniche analysis
5
5
  Home-page: https://github.com/yakirr/vima
6
6
  Author: Yakir Reshef
@@ -1,6 +1,6 @@
1
1
  [metadata]
2
2
  name = vima-spatial
3
- version = 0.2.8
3
+ version = 0.2.9
4
4
  author = Yakir Reshef
5
5
  author_email = yreshef@broadinstitute.org
6
6
  description = variational inference-based microniche analysis
@@ -1,5 +1,6 @@
1
1
  import numpy as numpy
2
2
  import numpy as np
3
+ import pandas as pd
3
4
  import xarray as xr
4
5
  import cv2 as cv2
5
6
  import scanpy as sc
@@ -105,6 +106,161 @@ def add_covs(pca, sid_to_covs):
105
106
  pca[cov_name] = pca['sid'].map(sid_to_covs[cov_name])
106
107
  return ['sid'] + cov_names
107
108
 
109
+ def collapse_markers(normeddir, markers, pseudomarker, outdir, masksdir=None):
110
+ """
111
+ Keep a subset of markers and fold all the others into one pseudomarker.
112
+
113
+ Reads the normalized pixel matrices written by `prepare_merfish`,
114
+ `prepare_xenium5k`, or `nonst.prepare` (for transcriptomic data the markers
115
+ are genes), retains `markers`, and replaces every other marker with a single
116
+ channel named `pseudomarker` holding their combined signal. Because the
117
+ stored values are log-normalized, they are exponentiated, summed, and
118
+ log-normalized again, so the pseudomarker is on the same scale as a real
119
+ marker. The dataset-wide means and stds stored on each file are rewritten to
120
+ match the new marker set: retained markers keep their existing values and the
121
+ pseudomarker's are computed from the collapsed data.
122
+
123
+ Parameters
124
+ ----------
125
+ normeddir
126
+ Directory of normalized ``.nc`` files, e.g. ``{outdir}/normalized``.
127
+ markers
128
+ Markers to retain, in the order they should appear; the pseudomarker is
129
+ appended after them. Markers absent from the data are skipped with a
130
+ warning.
131
+ pseudomarker
132
+ Name for the new channel, e.g. ``'nonimmune'``.
133
+ outdir
134
+ Directory to write the collapsed ``.nc`` files to. This is a
135
+ ``normalized``-style directory, so to build a dataset that
136
+ `pca_pixels` can be pointed at, pass ``f'{newroot}/normalized'`` and copy
137
+ the masks to ``f'{newroot}/masks'``.
138
+ masksdir
139
+ Directory of tissue masks, used to compute the pseudomarker's moments
140
+ over non-empty pixels only. Defaults to ``masks`` alongside `normeddir`.
141
+ """
142
+ import netCDF4
143
+
144
+ if masksdir is None:
145
+ masksdir = os.path.join(os.path.dirname(os.path.normpath(normeddir)), 'masks')
146
+
147
+ sids = [os.path.splitext(f)[0]
148
+ for f in os.listdir(normeddir) if f.endswith('.nc') and not f.startswith('.')]
149
+ if len(sids) == 0:
150
+ logger.warning(f'No .nc files found in {normeddir}. Check your path and try again.')
151
+ return
152
+ os.makedirs(outdir, exist_ok=True)
153
+
154
+ # decide on the new marker set, using the first sample as the reference
155
+ da = xr.open_dataarray(f'{normeddir}/{sids[0]}.nc')
156
+ ref_markers = list(da.marker.values)
157
+ ref_means, ref_stds = da.attrs['means'], da.attrs['stds']
158
+ da.close(); del da
159
+
160
+ if pseudomarker in ref_markers:
161
+ raise ValueError(f"'{pseudomarker}' is already a marker; pick a name for the "
162
+ f"pseudomarker that isn't in the data.")
163
+ missing = [m for m in markers if m not in ref_markers]
164
+ if missing:
165
+ logger.warning(f'{len(missing)} of the {len(markers)} requested markers are not in the '
166
+ f'data and will be skipped: {missing}')
167
+ keep = [m for m in markers if m in ref_markers]
168
+ if len(keep) == 0:
169
+ raise ValueError('None of the requested markers are in the data.')
170
+ if len(keep) == len(ref_markers):
171
+ logger.warning(f'All {len(ref_markers)} markers were retained, so {pseudomarker} will be '
172
+ f'empty.')
173
+ logger.info(f'Keeping {len(keep)} markers and collapsing the other '
174
+ f'{len(ref_markers) - len(keep)} into {pseudomarker}.')
175
+
176
+ # collapse each sample, accumulating the pseudomarker's per-sample moments as we go
177
+ sample_sids, sample_means, sample_stds, sample_npixels = [], [], [], []
178
+ for sid in settings.progress(sids, name='collapsing markers'):
179
+ da = xr.open_dataarray(f'{normeddir}/{sid}.nc')
180
+ mask_da = xr.open_dataarray(f'{masksdir}/{sid}.nc')
181
+
182
+ sid_markers = list(da.marker.values)
183
+ if sid_markers != ref_markers:
184
+ logger.warning(f'{sid} has different markers ({len(sid_markers)}) than {sids[0]} '
185
+ f'({len(ref_markers)}); matching by name.')
186
+ absent = [m for m in keep if m not in sid_markers]
187
+ if absent:
188
+ raise ValueError(f'{sid} is missing {len(absent)} of the markers to retain: {absent}')
189
+ ix = {m: i for i, m in enumerate(sid_markers)}
190
+ keep_ix = [ix[m] for m in keep]
191
+ keep_ix_set = set(keep_ix)
192
+ drop_ix = [i for i in range(len(sid_markers)) if i not in keep_ix_set]
193
+
194
+ # load raw arrays and close before doing any work, so we never hold two
195
+ # full (H x W x n_markers) copies simultaneously
196
+ x = da.x.values; y = da.y.values
197
+ data = da.values
198
+ mask = mask_da.values
199
+ da.close(); mask_da.close()
200
+ del da, mask_da; gc.collect()
201
+
202
+ # accumulate the dropped markers one at a time: fancy-indexing them all at
203
+ # once would materialize the (y, x, n_dropped) copy this function exists to
204
+ # avoid. log-normalized values are exponentiated before summing so the
205
+ # pseudomarker ends up on the same scale as the markers we kept.
206
+ acc = np.zeros(data.shape[:2], dtype=np.float64)
207
+ for j in drop_ix:
208
+ acc += np.expm1(data[..., j].astype(np.float64))
209
+ other = np.log1p(acc)
210
+ del acc
211
+
212
+ # empty pixels are zero in the input, and log1p(sum(expm1(0))) is zero, so
213
+ # they stay empty without any special handling
214
+ s = xr.DataArray(
215
+ np.concatenate([data[..., keep_ix], other[..., None].astype(np.float32)], axis=-1),
216
+ dims=['y', 'x', 'marker'],
217
+ coords={'x': x, 'y': y, 'marker': keep + [pseudomarker]})
218
+ s.name = sid
219
+ util.write_xarray(s, f'{outdir}/{sid}.nc')
220
+ del data, s; gc.collect()
221
+
222
+ # a sample with an empty mask contributes nothing rather than a nan that
223
+ # would poison the pooled moments
224
+ vals = other[mask]
225
+ if len(vals) == 0:
226
+ logger.warning(f'{sid} has no non-empty pixels; excluding it from the '
227
+ f'{pseudomarker} moments.')
228
+ else:
229
+ sample_sids.append(sid)
230
+ sample_means.append(vals.mean(dtype=np.float64))
231
+ sample_stds.append(vals.std(dtype=np.float64))
232
+ sample_npixels.append(len(vals))
233
+ del other, mask, vals; gc.collect()
234
+
235
+ if len(sample_sids) == 0:
236
+ raise ValueError(f'No non-empty pixels in any sample, so {pseudomarker} has no moments. '
237
+ f'Check that {masksdir} holds the masks for {normeddir}.')
238
+
239
+ # pool the pseudomarker's moments across samples the same way get_sumstats does
240
+ pseudo_mean, pseudo_std = util.pool_moments(
241
+ pd.DataFrame([sample_means], index=[pseudomarker], columns=sample_sids),
242
+ pd.DataFrame([sample_stds], index=[pseudomarker], columns=sample_sids),
243
+ sample_npixels)
244
+ pseudo_mean, pseudo_std = pseudo_mean.iloc[0], pseudo_std.iloc[0]
245
+ if pseudo_std == 0:
246
+ # a constant channel stays exactly 0 after standardization rather than
247
+ # dividing by zero downstream
248
+ logger.warning(f'{pseudomarker} has zero variance; setting its std to 1.')
249
+ pseudo_std = 1.
250
+ logger.info(f'{pseudomarker}: mean {pseudo_mean:.3f}, std {pseudo_std:.3f}')
251
+
252
+ # the pseudomarker's moments aren't known until every sample has been read, so
253
+ # the files are written above without them and stamped here; this rewrites
254
+ # metadata only, rather than re-collapsing every sample a second time
255
+ keep_ix_ref = [ref_markers.index(m) for m in keep]
256
+ means = np.append(np.asarray(ref_means)[keep_ix_ref], pseudo_mean).astype(np.float32)
257
+ stds = np.append(np.asarray(ref_stds)[keep_ix_ref], pseudo_std).astype(np.float32)
258
+ for sid in sids:
259
+ with netCDF4.Dataset(f'{outdir}/{sid}.nc', 'a') as ds:
260
+ v = ds.variables[sid]
261
+ v.setncattr('means', means)
262
+ v.setncattr('stds', stds)
263
+
108
264
  def pca_pixels(outdir, repname, nmetamarkers=10, npixels_to_plot=50000,
109
265
  total_n_metapixels=2_000_000, sid_to_covs=None):
110
266
  """
@@ -1,5 +1,6 @@
1
1
  import os, glob, gc
2
2
  import numpy as np
3
+ import pandas as pd
3
4
  import cv2 as cv2
4
5
  import xarray as xr
5
6
  from . import util
@@ -134,12 +135,14 @@ def prepare(load, filepaths, orig_pixel_size, markers, get_foreground, norm_by_b
134
135
  )
135
136
  for sid in settings.progress(sids)])
136
137
  gc.collect()
137
- _, pixels = norm_by_background(pixels)
138
+ goodmarkers, pixels = norm_by_background(pixels)
138
139
  ntranscripts = pixels.sum(axis=1, dtype=np.float64)
139
140
  med_ntranscripts = np.median(ntranscripts)
140
141
  pixels = np.log1p(med_ntranscripts * pixels / (ntranscripts[:,None] + 1e-6)) # adding to denominator in case pixel is all 0s
141
- means = pixels.mean(axis=0, dtype=np.float64)
142
- stds = pixels.std(axis=0, dtype=np.float64)
142
+ # indexed by marker name so each sample's moments are aligned by name below,
143
+ # rather than positionally against whatever markers that sample kept
144
+ means = pd.Series(pixels.mean(axis=0, dtype=np.float64), index=goodmarkers)
145
+ stds = pd.Series(pixels.std(axis=0, dtype=np.float64), index=goodmarkers)
143
146
  del pixels; gc.collect()
144
147
 
145
148
  logger.info('Normalizing and writing')
@@ -153,6 +156,7 @@ def prepare(load, filepaths, orig_pixel_size, markers, get_foreground, norm_by_b
153
156
  pl = np.log1p(med_ntranscripts * pl / (pl.sum(axis=1)[:,None] + 1e-6)) # adding to denominator in case pixel is all 0s
154
157
  s = s.sel(marker=goodmarkers)
155
158
  util.set_pixels(s, mask, pl)
156
- s.attrs['means'] = means.reindex(s.marker.values, fill_value=1).values.astype(np.float32)
157
- s.attrs['stds'] = stds.reindex(s.marker.values, fill_value=0).values.astype(np.float32)
159
+ # a marker with no moments standardizes to a no-op rather than dividing by zero
160
+ s.attrs['means'] = means.reindex(s.marker.values, fill_value=0).values.astype(np.float32)
161
+ s.attrs['stds'] = stds.reindex(s.marker.values, fill_value=1).values.astype(np.float32)
158
162
  util.write_xarray(s, f'{normeddir}/{sid}.nc')
@@ -6,50 +6,72 @@ from matplotlib import pyplot as plt
6
6
  from . import util
7
7
  from .._settings import logger, settings
8
8
 
9
- def med_ntranscripts(load, filepaths, x_col, y_col, pixel_size=10):
9
+ def rasterize_all(load, filepaths, x_col, y_col, gene_col, pixel_size=10):
10
+ """
11
+ Bin every sample's transcripts to a pixel grid, once.
12
+
13
+ Ingest makes three passes over the samples -- normalization factor, then
14
+ HVGs and dataset-wide moments, then the final matrices -- and each pass
15
+ needs the same binned counts. Binning once and holding the sparse result
16
+ trades memory for two fewer reads and two fewer binning passes per sample;
17
+ the counts are ~1-2% dense, so a 34-sample Xenium 5K cohort costs a few GB.
18
+
19
+ Returns
20
+ -------
21
+ list
22
+ ``(sample_id, SparsePixels)`` per file, in ``filepaths`` order.
23
+ """
24
+ rasterized = []
25
+ for filepath in settings.progress(filepaths, name='binning transcripts'):
26
+ sid, data = load(filepath)
27
+ pixels = util.transcriptlist_to_sparsepixels(data, x_col, y_col, gene_col,
28
+ pixel_size=pixel_size)
29
+ del data; gc.collect()
30
+ logger.info(f'\t{sid}: {pixels.ntranscripts/1e6:.2f}M tx → '
31
+ f'{len(pixels.nonempty())} non-empty pixels '
32
+ f'({pixels.nbytes()/1e9:.2f} GB sparse)')
33
+ rasterized.append((sid, pixels))
34
+
35
+ held = sum(p.nbytes() for _, p in rasterized)
36
+ logger.info(f'\tHolding {held/1e9:.2f} GB of binned counts for reuse across passes.')
37
+ return rasterized
38
+
39
+ def med_ntranscripts(rasterized):
10
40
  """
11
41
  Compute the normalization target: median transcripts per dense pixel.
12
42
 
13
- Bins each sample's transcripts into ``pixel_size`` square pixels and takes
14
- the median transcript count over pixels with at least 10 transcripts; returns
15
- the mean of these per-sample medians, used as the total-count normalization
16
- target.
43
+ Takes the median transcript count over pixels with at least 10 transcripts
44
+ in each sample, and returns the mean of those per-sample medians; this is
45
+ the total-count normalization target. Counts come from the same binning the
46
+ rest of ingest uses, so any transcript with no gene assigned is excluded.
17
47
  """
18
48
  medians = []
19
- for i, filepath in enumerate(settings.progress(filepaths, name='median #transcripts')):
20
- sid, df = load(filepath)
21
-
22
- df["px"] = (df[x_col] // pixel_size).astype(np.int32)
23
- df["py"] = (df[y_col] // pixel_size).astype(np.int32)
24
- pg = (
25
- df
26
- .groupby(["px", "py"])
27
- .size()
28
- .reset_index(name="txcount")
29
- )
30
- med = pg[pg.txcount > 10].txcount.median()
31
- logger.info(f"\tSample {i+1}/{len(filepaths)}: {sid} ({len(df)/1e6:.2f}M tx) "
49
+ for i, (sid, pixels) in enumerate(settings.progress(rasterized, name='median #transcripts')):
50
+ # float64 so the median (and the normalization factor built from it)
51
+ # carries full precision, as it did when counts were an int64 table
52
+ totals = np.asarray(pixels.counts.sum(axis=1), dtype=np.float64).ravel()
53
+ med = np.median(totals[totals > 10])
54
+ logger.info(f"\tSample {i+1}/{len(rasterized)}: {sid} ({pixels.ntranscripts/1e6:.2f}M tx) "
32
55
  f"→ median {med} transcripts per pixel with >=10 transcripts")
33
56
  medians.append(med)
34
- del df, pg; gc.collect()
35
57
  return np.mean(np.array(medians))
36
58
 
37
- def get_sumstats(load, filepaths, target_sum, x_col, y_col, gene_col, n_top_genes_per_sample=200, genes_to_add=[],
38
- pixel_size=10, min_mean=0.01, min_ntranscripts_per_pixel=10, min_ngenes_per_pixel=1,
59
+ def get_sumstats(rasterized, target_sum, n_top_genes_per_sample=200, genes_to_add=[],
60
+ min_mean=0.01, min_ntranscripts_per_pixel=10, min_ngenes_per_pixel=1,
39
61
  min_npixels=20, min_totalcounts=500):
40
62
  """
41
63
  Select highly variable genes and compute dataset-wide per-gene moments.
42
64
 
43
- For each sample, rasterizes transcripts to pixels, QC-filters pixels,
44
- log-normalizes to ``target_sum``, and selects the top highly variable genes;
45
- the union of HVGs across samples (plus any ``genes_to_add``) is returned.
46
- Also returns the pixel-count-weighted grand mean and grand std per gene
47
- across the whole dataset (pooling within- and between-sample variance).
65
+ For each sample, QC-filters the binned pixels, log-normalizes them to
66
+ ``target_sum``, and selects the top highly variable genes; the union of HVGs
67
+ across samples (plus any ``genes_to_add``) is returned. Also returns the
68
+ pixel-count-weighted grand mean and grand std per gene across the whole
69
+ dataset (pooling within- and between-sample variance).
48
70
 
49
71
  Parameters
50
72
  ----------
51
- load
52
- Callable mapping a file path to ``(sample_id, transcript_dataframe)``.
73
+ rasterized
74
+ ``(sample_id, SparsePixels)`` per sample, from `rasterize_all`.
53
75
  target_sum
54
76
  Per-pixel total-count normalization target (from `med_ntranscripts`).
55
77
  n_top_genes_per_sample
@@ -71,29 +93,31 @@ def get_sumstats(load, filepaths, target_sum, x_col, y_col, gene_col, n_top_gene
71
93
  stds = []
72
94
  npixels = []
73
95
 
74
- for i, filepath in enumerate(settings.progress(filepaths, name='HVGs & moments')):
75
- sid, transcripts = load(filepath)
76
- ntranscripts = len(transcripts)
77
-
78
- # create pixellist
79
- pl = util.transcriptlist_to_pixellist(transcripts, x_col, y_col, gene_col, pixel_size=pixel_size)
80
- del transcripts; gc.collect()
81
-
82
- # create scanpy object and filter empty/near-empty pixels
83
- genes = pl.columns[2:]
84
- obs = pl[['pixel_x', 'pixel_y']].copy()
85
- obs.index = obs.index.astype(str)
86
- pl = sc.AnnData(pl[genes].values.astype(np.float32), var=pd.DataFrame(index=genes), obs=obs)
87
- sc.pp.filter_cells(pl, min_counts=min_ntranscripts_per_pixel, inplace=True)
88
- sc.pp.filter_cells(pl, min_genes=min_ngenes_per_pixel, inplace=True)
89
-
96
+ for i, (sid, pixels) in enumerate(settings.progress(rasterized, name='HVGs & moments')):
97
+ ntranscripts = pixels.ntranscripts
98
+
99
+ # filter empty/near-empty pixels; both thresholds are per-pixel, so
100
+ # applying them together matches filtering on one and then the other
101
+ genes = pixels.genes
102
+ keep, totals = util.qc_pixels(pixels.counts, min_ntranscripts_per_pixel, min_ngenes_per_pixel)
103
+ rows = np.flatnonzero(keep)
104
+ counts = pixels.counts[keep]
105
+
90
106
  # compute moments
91
- X = pl.X / pl.X.sum(axis=1, keepdims=True) * target_sum
92
- X = np.log1p(X)
93
- means.append(pd.Series(np.array(X.mean(axis=0, dtype=np.float64)).squeeze(), index=genes))
94
- stds.append(pd.Series(np.array(X.std(axis=0, dtype=np.float64)).squeeze(), index=genes))
95
- npixels.append(pl.n_obs)
107
+ X = util.lognormalize(counts, totals[keep], target_sum)
108
+ mean, std = util.sparse_moments(X)
109
+ means.append(pd.Series(mean, index=genes))
110
+ stds.append(pd.Series(std, index=genes))
111
+ npixels.append(len(rows))
96
112
  union_allgenes.update(genes)
113
+ del X; gc.collect()
114
+
115
+ # create scanpy object over the surviving pixels
116
+ pixel_x, pixel_y = pixels.coords_of(rows)
117
+ obs = pd.DataFrame({'pixel_x': pixel_x, 'pixel_y': pixel_y})
118
+ obs.index = obs.index.astype(str)
119
+ pl = sc.AnnData(counts, var=pd.DataFrame(index=genes), obs=obs)
120
+ del counts; gc.collect()
97
121
 
98
122
  # QC genes and compute HVGs for this sample
99
123
  sc.pp.filter_genes(pl, min_cells=min_npixels, inplace=True)
@@ -149,43 +173,36 @@ def get_sumstats(load, filepaths, target_sum, x_col, y_col, gene_col, n_top_gene
149
173
  plt.close(fig)
150
174
 
151
175
  union_hvgs.update(hvgs)
152
- logger.info(f'\tSample {i+1}/{len(filepaths)}: {sid} ({ntranscripts/1e6:.2f}M tx) → {len(union_hvgs)} HVGs across all samples so far.')
176
+ logger.info(f'\tSample {i+1}/{len(rasterized)}: {sid} ({ntranscripts/1e6:.2f}M tx) → {len(union_hvgs)} HVGs across all samples so far.')
153
177
  else:
154
178
  union_hvgs.update(genes)
155
- logger.info(f'\tSample {i+1}/{len(filepaths)}: {sid} ({ntranscripts/1e6:.2f}M tx) → Using all {len(genes)} genes.')
179
+ logger.info(f'\tSample {i+1}/{len(rasterized)}: {sid} ({ntranscripts/1e6:.2f}M tx) → Using all {len(genes)} genes.')
156
180
 
157
- pl.X = None; del pl; del X; gc.collect()
181
+ pl.X = None; del pl; gc.collect()
158
182
 
159
183
  means_df = pd.concat([m.reindex(index=union_allgenes, fill_value=0) for m in means], axis=1)
160
184
  stds_df = pd.concat([s.reindex(index=union_allgenes, fill_value=0) for s in stds], axis=1)
161
- w = np.array(npixels, dtype=np.float64)
162
- W = w.sum()
163
-
164
- grand_mean = np.sum((means_df * w).values, axis=1, dtype=np.float64) / W
165
- mean_of_vars = np.sum((stds_df ** 2 * w).values, axis=1, dtype=np.float64) / W
166
- var_of_means = ((means_df.subtract(grand_mean, axis=0).values.astype(np.float64) ** 2) * w).sum(axis=1) / W
167
- grand_mean = pd.Series(grand_mean, index=means_df.index)
168
- grand_std = pd.Series(np.sqrt(mean_of_vars + var_of_means), index=means_df.index)
185
+ grand_mean, grand_std = util.pool_moments(means_df, stds_df, npixels)
169
186
 
170
187
  return list(union_hvgs), grand_mean, grand_std
171
188
 
172
- def transcriptlist_to_normedpixelmatrix(sid, data, x_col, y_col, gene_col, pixel_size, target_sum, means, stds,
173
- genes, min_ngenes_per_pixel, min_ntranscripts_per_pixel):
189
+ def pixels_to_normedpixelmatrix(sid, pixels, target_sum, means, stds,
190
+ genes, min_ngenes_per_pixel, min_ntranscripts_per_pixel):
174
191
  """
175
- Rasterize one sample's transcripts to a log-normalized ``(y, x, gene)`` matrix.
192
+ Turn one sample's binned pixels into a log-normalized ``(y, x, gene)`` matrix.
176
193
 
177
- Bins transcripts into square pixels, drops pixels failing the QC thresholds,
178
- log-normalizes surviving pixels to ``target_sum``, restricts to ``genes``
179
- (filling absent genes with zeros), and returns the pixel matrix alongside its
180
- tissue mask. The dataset-wide ``means`` and ``stds`` are stored as attributes
181
- on the matrix for later standardization (they are not applied here).
194
+ Drops pixels failing the QC thresholds, log-normalizes the survivors to
195
+ ``target_sum``, restricts to ``genes`` (filling absent genes with zeros), and
196
+ returns the pixel matrix alongside its tissue mask. The dataset-wide
197
+ ``means`` and ``stds`` are stored as attributes on the matrix for later
198
+ standardization (they are not applied here).
182
199
 
183
200
  Parameters
184
201
  ----------
185
202
  sid
186
203
  Sample ID; used to name the output arrays.
187
- data
188
- Transcript table for this sample.
204
+ pixels
205
+ This sample's `util.SparsePixels`, from `rasterize_all`.
189
206
  target_sum
190
207
  Per-pixel total-count normalization target.
191
208
  means, stds
@@ -199,49 +216,39 @@ def transcriptlist_to_normedpixelmatrix(sid, data, x_col, y_col, gene_col, pixel
199
216
  ``(mask, matrix)``: the boolean tissue mask and the normalized
200
217
  ``(y, x, gene)`` DataArray.
201
218
  """
202
- logger.info(f'\tNumber of transcripts: {len(data)/1e6:.2f}M')
203
-
204
- # process data
205
- pl = util.transcriptlist_to_pixellist(
206
- data,
207
- x_col,
208
- y_col,
209
- gene_col,
210
- pixel_size=pixel_size
211
- )
212
- markers = pl.columns[2:]
213
- all_x = np.sort(pl['pixel_x'].unique())
214
- all_y = np.sort(pl['pixel_y'].unique())
215
- logger.info(f'\tMaking pixel list... {len(pl)} pixels.')
219
+ logger.info(f'\tNumber of transcripts: {pixels.ntranscripts/1e6:.2f}M')
220
+ logger.info(f'\t{len(pixels.nonempty())} non-empty pixels.')
216
221
 
217
222
  if settings.show_plots():
223
+ occupied_x, occupied_y = pixels.coords_of(pixels.nonempty())
218
224
  plt.figure(figsize=(5,5))
219
- plt.scatter(pl.pixel_x, pl.pixel_y, c='gray', s=0.1, alpha=0.2)
220
- pl = pl[(pl[markers] != 0).sum(axis=1) >= min_ngenes_per_pixel]
221
- pl = pl[(pl[markers].sum(axis=1) >= min_ntranscripts_per_pixel) ]#& (pl[list(set(markers) & set(genes))].sum(axis=1) > 0)]
225
+ plt.scatter(occupied_x, occupied_y, c='gray', s=0.1, alpha=0.2)
226
+ # both thresholds are per-pixel, so applying them together matches
227
+ # filtering on one and then the other
228
+ keep, totals = util.qc_pixels(pixels.counts, min_ntranscripts_per_pixel, min_ngenes_per_pixel)
229
+ rows = np.flatnonzero(keep)
222
230
  if settings.show_plots():
223
- plt.scatter(pl.pixel_x, pl.pixel_y, c=pl[markers].sum(axis=1), s=0.1, alpha=0.8, vmin=0, vmax=100)
231
+ kept_x, kept_y = pixels.coords_of(rows)
232
+ plt.scatter(kept_x, kept_y, c=totals[rows], s=0.1, alpha=0.8, vmin=0, vmax=100)
224
233
  plt.gca().set_aspect('equal'); plt.title('transcript density (gray = failed qc)'); plt.axis('off')
225
234
  settings.show(f'transcript_density_{sid}')
226
- logger.info(f'\t{len(pl)} pixels after QC.')
235
+ logger.info(f'\t{len(rows)} pixels after QC.')
227
236
 
228
237
  logger.info('\tLog-normalizing...')
229
- pl[markers] = pl[markers].div(pl[markers].sum(axis=1), axis=0).fillna(0) * target_sum
230
- pl[markers] = np.log1p(pl[markers])
238
+ X = util.lognormalize(pixels.counts[keep], totals[keep], target_sum)
231
239
 
232
240
  logger.info(f'\trestricting to {len(genes)} genes')
233
- pl = pl.reindex(columns=['pixel_x', 'pixel_y'] + genes, fill_value=0) # zeros for genes not present in this sample
234
- mask_pl = pl[['pixel_x', 'pixel_y']].copy()
235
- mask_pl['nonempty'] = 1
236
-
237
- s = util.pixellist_to_pixelmatrix(pl, genes).reindex({'x': all_x, 'y': all_y}, fill_value=0)
241
+ # sorted to match the marker order the pandas pivot used to impose; genes
242
+ # absent from this sample are scattered in as zeros
243
+ s = util.sparsepixels_to_pixelmatrix(X, pixels.genes, rows, pixels.x, pixels.y, sorted(genes))
244
+ del X; gc.collect()
238
245
  s.attrs['means'] = means.reindex(s.marker.values).values.astype(np.float32)
239
246
  s.attrs['stds'] = stds.reindex(s.marker.values).values.astype(np.float32)
240
- mask = util.pixellist_to_pixelmatrix(mask_pl, ['nonempty']).squeeze().astype(bool).reindex({'x': all_x, 'y': all_y}, fill_value=False)
247
+ mask = util.sparsepixels_to_mask(rows, pixels.x, pixels.y)
241
248
  s.name = sid; mask.name = sid
242
249
  gc.collect()
243
250
  logger.info(f'\tMaking pixel matrix... done. shape: {s.shape}')
244
-
251
+
245
252
  return mask, s.astype(np.float32)
246
253
 
247
254
  def rasterize_and_normalize_generic(load, filepaths, x_col, y_col, gene_col, n_top_genes_per_sample, pixel_size, outdir,
@@ -250,10 +257,12 @@ def rasterize_and_normalize_generic(load, filepaths, x_col, y_col, gene_col, n_t
250
257
  """
251
258
  Shared driver that rasterizes and normalizes every sample to disk.
252
259
 
253
- Backs `prepare_merfish` and `prepare_xenium5k`: computes the normalization
254
- factor and dataset-wide gene set/moments once, then rasterizes and
255
- normalizes each sample, writing its normalized matrix and mask under
256
- ``outdir/normalized`` and ``outdir/masks``.
260
+ Backs `prepare_merfish` and `prepare_xenium5k`: bins every sample's
261
+ transcripts once, computes the normalization factor and dataset-wide gene
262
+ set/moments from those bins, then normalizes each sample, writing its matrix
263
+ and mask under ``outdir/normalized`` and ``outdir/masks``. The binned counts
264
+ are held in memory across all three passes; each sample's are released as
265
+ soon as its matrix has been written.
257
266
 
258
267
  Parameters
259
268
  ----------
@@ -268,31 +277,36 @@ def rasterize_and_normalize_generic(load, filepaths, x_col, y_col, gene_col, n_t
268
277
  logger.warning('No files found. Check your filepaths and try again.')
269
278
  return
270
279
 
280
+ logger.info('Binning transcripts to pixels...')
281
+ rasterized = rasterize_all(load, filepaths, x_col, y_col, gene_col, pixel_size=pixel_size)
282
+
271
283
  logger.info('Computing normalization factor...')
272
- normfactor = med_ntranscripts(load, filepaths, x_col, y_col, pixel_size=pixel_size)
284
+ normfactor = med_ntranscripts(rasterized)
273
285
  logger.info('Finding HVGs and dataset-wide mean and variance per gene...')
274
- hvgs, means, stds = get_sumstats(load, filepaths, normfactor, x_col, y_col, gene_col,
286
+ hvgs, means, stds = get_sumstats(rasterized, normfactor,
275
287
  n_top_genes_per_sample=n_top_genes_per_sample,
276
- genes_to_add=genes_to_add, pixel_size=pixel_size,
288
+ genes_to_add=genes_to_add,
277
289
  min_ntranscripts_per_pixel=min_ntranscripts_per_pixel,
278
290
  min_ngenes_per_pixel=min_ngenes_per_pixel)
279
291
  logger.info(f'Final number of genes used = {len(hvgs)}')
280
292
 
281
- logger.info('Rasterizing and normalizing...')
293
+ logger.info('Normalizing...')
282
294
  normdir = f'{outdir}/normalized'
283
295
  masksdir = f'{outdir}/masks'
284
296
  os.makedirs(normdir, exist_ok=True)
285
297
  os.makedirs(masksdir, exist_ok=True)
286
- for i, filepath in enumerate(settings.progress(filepaths, name='rasterize & normalize')):
287
- sid, data = load(filepath)
288
- logger.info(f'Processing sample {i+1}/{len(filepaths)}: {sid}')
289
- mask, pm = transcriptlist_to_normedpixelmatrix(sid, data, x_col, y_col, gene_col, pixel_size,
290
- normfactor, means=means, stds=stds, genes=hvgs,
291
- min_ntranscripts_per_pixel=min_ntranscripts_per_pixel,
292
- min_ngenes_per_pixel=min_ngenes_per_pixel)
293
- del data; gc.collect()
298
+ for i in settings.progress(range(len(rasterized)), name='normalize & write'):
299
+ sid, pixels = rasterized[i]
300
+ rasterized[i] = None # this sample's bins are no longer needed after here
301
+ logger.info(f'Processing sample {i+1}/{len(rasterized)}: {sid}')
302
+ mask, pm = pixels_to_normedpixelmatrix(sid, pixels, normfactor, means=means, stds=stds,
303
+ genes=hvgs,
304
+ min_ntranscripts_per_pixel=min_ntranscripts_per_pixel,
305
+ min_ngenes_per_pixel=min_ngenes_per_pixel)
306
+ del pixels; gc.collect()
294
307
  util.write_xarray(mask, f'{masksdir}/{pm.name}.nc')
295
308
  util.write_xarray(pm, f'{normdir}/{pm.name}.nc')
309
+ del mask, pm; gc.collect()
296
310
 
297
311
  def prepare_xenium5k(load, filepaths, x_col, y_col, gene_col, n_top_genes_per_sample, outdir,
298
312
  pixel_size=10, genes_to_add=[],