ebin 1.0__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.
ebin-1.0/PKG-INFO ADDED
@@ -0,0 +1,12 @@
1
+ Metadata-Version: 2.4
2
+ Name: ebin
3
+ Version: 1.0
4
+ Summary: Posterior bin-expectation normalization for sorting-based MPRA
5
+ Requires-Python: >=3.9
6
+ Requires-Dist: jax>=0.4
7
+ Requires-Dist: numpy
8
+ Requires-Dist: scipy
9
+ Requires-Dist: pandas
10
+ Provides-Extra: netcdf
11
+ Requires-Dist: xarray; extra == "netcdf"
12
+ Requires-Dist: h5netcdf; extra == "netcdf"
ebin-1.0/README.md ADDED
@@ -0,0 +1,316 @@
1
+ # EBin: Normalization of MPRA data (FACS bins × replicates × sequences)
2
+
3
+ Raw per-bin read counts are turned into one number per sequence per cell line —
4
+ the **posterior-mean expected bin**, `E[bin]` — together with the posterior SD
5
+ that states how far it can be trusted. The counts enter a compound NB–Poisson
6
+ likelihood; what is estimated is the sequence's latent expression distribution.
7
+
8
+
9
+
10
+ <img src="http://data.georgy.top/media/ebin_model_layers.png" alt="EBin model"/>
11
+
12
+
13
+ ```MPRA
14
+ X[n,s,b] | T[n,s,b] ~ NB(R[s,b] · T[n,s,b], P[s,b]) reads from cells
15
+ T[n,s,b] ~ Poisson(a_n · Pi[n,b] · lambda_s) cells in bin b
16
+ Pi[n,b] = Phi((q_b − mu_n)/sigma_n) − Phi((q_{b−1} − mu_n)/sigma_n)
17
+ ```
18
+ Each sequence `n` carries a latent Gaussian effect law `N(mu_n, sigma_n²)` — its
19
+ expression distribution over cells — and an abundance `a_n`. The library's
20
+ marginal effect law is the **`N`-component mixture** of those per-sequence
21
+ normals, `G(x) = Σ_n w_n Phi((x − mu_n)/sigma_n)` with `w_n = 1/N`, and the `B`
22
+ sorting bins are its quantile cells: the cut points `q_1 < … < q_{B−1}` solve
23
+ `G(q_j) = j/B`. They are shared, and re-solved at every step, so every sequence
24
+ is tied to the marginal through them. This approach has 2 benefits: first one is that
25
+ it allows us to cut the number of params down to 2 instead of `B-1`, and the second one is that those params are further regularized, both by the normality assumption and the fact that they are all should form the global marginal distribution.
26
+
27
+ The activity that ships is the posterior **mean** of `E[bin] = Σ_b b·Pi[n,b]`,
28
+ integrated over the posterior of `(mu_n, sigma_n)` with the abundance profiled
29
+ per grid point. The plug-in point estimate is a posterior *mode*, and for a
30
+ low-count sequence the mode sits at a simplex vertex — activity exactly 1 or
31
+ `B`, regardless of depth.
32
+
33
+ ## Install
34
+ ```
35
+ pip install ebin
36
+ ```
37
+ or
38
+
39
+ ```bash
40
+ pip install -e /path/to/package # or place this directory on PYTHONPATH
41
+ ```
42
+
43
+ `jax` (GPU-version strongly recommended: `pip install jax[cuda]`), `numpy`, `scipy` and `pandas` are required.
44
+ `xarray` + `h5netcdf` are needed only for the netCDF export and `matplotlib`
45
+ only for the diagnostics. Everything runs in float64 as some parts of the algorithm require high precision, namely the tridiagonal eigenvalue solver for orthogonal polynomials.
46
+
47
+ ## Input
48
+
49
+ One row per sequence, a 3-level column header `(cell line, replicate, bin)`:
50
+
51
+ ```
52
+ seq, A549,A549,A549,A549,A549,...
53
+ rep1,rep1,rep1,rep1,rep2,...
54
+ bin1,bin2,bin3,bin4,bin1,...
55
+ ACGT..., 112, 340, 88, 15, 97, ...
56
+ ```
57
+ (in other words, the tabular should be readable with `pd.read_csv(..., header=[0,1,2])`)
58
+
59
+
60
+ Bins are ordered by the number in their name. Cell lines are fitted
61
+ independently. **Empty fields are read as zero counts, not as missing data** —
62
+ some exports write every zero as a blank, and treating those as missin
63
+
64
+ ## Install
65
+
66
+ ```bash
67
+ pip install -e /path/to/package # or place this directory on PYTHONPATH
68
+ ```
69
+
70
+ `jax` (GPU strongly recommended), `numpy`, `scipy` and `pandas` are required.
71
+ `xarray` + `h5netcdf` are needed only for the netCDF export and `matplotlib`
72
+ only for the diagnostics. Everything runs in float64.
73
+
74
+ ## Input
75
+
76
+ One row per sequence, a 3-level column header `(cell line, replicate, bin)`:
77
+
78
+ ```
79
+ seq, A549,A549,A549,A549,A549,...
80
+ rep1,rep1,rep1,rep1,rep2,...
81
+ bin1,bin2,bin3,bin4,bin1,...
82
+ ACGT..., 112, 340, 88, 15, 97, ...
83
+ ```
84
+
85
+ Bins are ordered by the number in their name. Cell lines are fitted
86
+ independently. **Empty fields are read as zero counts, not as missing data** —
87
+ some exports write every zero as a blank, and treating those as missing deletes
88
+ the informative zeros of the low-abundance sequences.
89
+
90
+ ## Use
91
+
92
+ ```bash
93
+ python -m ebin counts.csv -o activity.csv
94
+ ```
95
+
96
+ ```python
97
+ from ebin import activity_table
98
+
99
+ table, results = activity_table("counts.csv", out="activity.csv")
100
+ table[("A549", "activity")] # posterior-mean E[bin] per sequence
101
+ ```
102
+
103
+ The CSV has two header rows, `(cell line, field)`, one row per sequence, and
104
+ carries everything the model estimates per sequence:
105
+
106
+ | field | meaning |
107
+ | --- | --- |
108
+ | `activity` | posterior-mean `E[bin]` — the normalized value intended for downstream use |
109
+ | `activity_sd` | its posterior SD — per-sequence confidence |
110
+ | `activity_map` | plug-in mode, kept for reference only |
111
+ | `mu`, `mu_sd` | posterior-mean effect location and its standard error |
112
+ | `sigma`, `sigma_sd` | posterior-mean effect scale and its standard error |
113
+ | `abundance` | fitted size factor `a_n` |
114
+ | `n_eff` | posterior grid support — large means the estimate is diffuse |
115
+ | `tot` | total observed reads |
116
+
117
+ All of it comes from the same posterior weights, so the columns are mutually
118
+ consistent: `activity` is `Σ_b b·bin_prob`. `mu` and `sigma` are in gauge units
119
+ (median cut 0, first cut −1), so they are comparable across sequences within a
120
+ cell line. A subset is selected with `--fields activity,activity_sd,tot` (or
121
+ `fields=[...]` in the API).
122
+
123
+ The `B−1` cut points are a property of the cell line, not of a sequence, so
124
+ they are not columns: `cuts_table(results)` returns them as one row per cell
125
+ line, and the CLI writes them next to the activity table as
126
+ `<out>_cuts.csv`. The gauge fixes the first two at −1 and 0, so what varies
127
+ between lines is where the upper cuts land (A549 0.990, MDA-MB-231 0.964).
128
+
129
+ Sequences with no reads at all in a cell line keep an activity (they are fitted
130
+ as structural zeros and land near the neutral prior mean, ~2.5 for `B = 4`);
131
+ sequences with no observed cell at all are `NaN`.
132
+
133
+ Useful CLI options:
134
+
135
+ ```
136
+ --groups A549,HepG2 cell lines to run, in output order
137
+ --abundance gamma integrate a_n out instead of fitting it
138
+ --zero-truncation drop all-zero rows instead of fitting them
139
+ --grid 61x25 readout grid (n_mu x n_log_sigma)
140
+ --save-fits fits.pkl keep the fitted parameters
141
+ --warm-from fits.pkl start from an earlier fit of the same data
142
+ --netcdf activity.nc also write the full posterior (see below)
143
+ --plots DIR also write the diagnostic plots
144
+ ```
145
+
146
+ ### Per-bin distribution
147
+
148
+ The one field too wide for a flat table is `bin_prob`, the posterior-mean
149
+ probability of each bin — an `(N, B)` array per cell line. It is returned in
150
+ `results`, and `--netcdf` writes it alongside every scalar field:
151
+
152
+ ```python
153
+ from ebin import activity_table, write_netcdf
154
+
155
+ table, results = activity_table("counts.csv")
156
+ results["A549"]["bin_prob"] # (N, B), rows sum to 1
157
+ write_netcdf("activity.nc", results, table.index)
158
+ ```
159
+
160
+ Readouts can be redone from saved fits without refitting, which is the cheap
161
+ way to try a different grid or prior:
162
+
163
+ ```python
164
+ from ebin import readout_table
165
+ table, results = readout_table("counts.csv", "fits.pkl")
166
+ ```
167
+
168
+ ## Diagnostics
169
+
170
+ Three plots, one panel per cell line, drawn from the table alone (`--plots DIR`
171
+ on the CLI writes all three):
172
+
173
+ ```python
174
+ from ebin import (plot_marginal_distribution, plot_activity_vs_raw,
175
+ plot_gc_vs_abundance)
176
+
177
+ plot_marginal_distribution(table, out="marginal.png")
178
+ plot_activity_vs_raw(table, "counts.csv", out="vs_raw.png")
179
+ plot_gc_vs_abundance(table, out="gc.png")
180
+ ```
181
+
182
+ **`plot_marginal_distribution`** — the marginal effect law: the `N`-component
183
+ mixture `G(x) = Σ_n w_n Phi((x − mu_n)/sigma_n)`, one normal per sequence,
184
+ drawn as the density it defines and sliced into the sorting bins, each in its
185
+ own colour. The x ticks are the cut points with their values. Weights are
186
+ uniform by default, as in the fit; `weights=` takes a per-sequence weight (the
187
+ fitted abundances, say) to see the marginal over cells rather than over
188
+ sequences. Every slice holds exactly `1/B` of the mass by construction, so the
189
+ gauge (cuts at −1 and 0) and any unusual tightness or skew in a cell line's
190
+ effect law are visible here.
191
+
192
+ **`plot_activity_vs_raw`** — fitted activity against the depth-normalized raw
193
+ mass centre (each replicate×bin column divided by its own library size, then
194
+ `Σ_b b·p_b`), coloured by fitted abundance. Well-measured sequences sit on the
195
+ diagonal; the horizontal band at the centre is the low-abundance tail, where the
196
+ raw estimate is noise and the model shrinks toward neutral. A cell line whose
197
+ *high*-abundance sequences leave the diagonal is a real warning sign.
198
+
199
+ **`plot_gc_vs_abundance`** — GC content against the fitted size factor, with the
200
+ binned median trend. A strong trend means `a_n` is absorbing a
201
+ cloning/sequencing bias rather than biological abundance alone. GC is taken from
202
+ the table index when the index is the sequence; otherwise it is supplied through
203
+ `sequences=`.
204
+
205
+ `raw_mass_center(counts)` and `gc_content(seqs)` are exported separately for
206
+ cases where the raw numbers are more appropriate than a picture.
207
+
208
+ ## Configuration
209
+
210
+ Defaults are the shipped configuration; they are not arbitrary and are worth
211
+ understanding before they are changed.
212
+
213
+ | option | default | why |
214
+ | --- | --- | --- |
215
+ | `conditional` | `False` | Zero **inflation**: all-zero rows are fitted with a structural-zero probability `phi`, so they enter the mixture that sets the cuts. Zero **truncation** (`True`) drops them, which sets the cuts on a zero-excluded distribution — measurably worse, both in-library and on transfer to an independent library. |
216
+ | `abundance` | `True` | A free size factor `a_n` per sequence. Without one, every sequence has the same expected latent total, so real depth variation has nowhere to go but `Pi` — which tilts toward the bins with the largest channel scale, and `Pi` stops being a profile. |
217
+ | `abundance_prior` | `None` | `'gamma'` integrates `a_n` out under `Gamma(kappa, mean=1)` instead: 2 parameters per sequence instead of 3, low-count sequences regularized, and the abundance uncertainty propagated into the effect. The activity profile is essentially unchanged (ρ ≈ 0.9996 against free `a_n`), so this is a parameter-economy. Costs ~2× runtime. |
218
+ | `mu_prior` | `1.3` | SD of the prior on the effect location. Bounding the effect *spread* alone still lets a single-bin low-count sequence run `mu_n` large and walk `Pi` to a vertex. The gauge pins Q50 = 0 and Q25 = −1, so the population SD is ≈ 1.5 and the between-sequence SD ≈ 1.1 — the gauge-consistent `tau` is ~1.3. A larger value (2–3) is inconsistent with the gauge and under-regularizes. |
219
+ | `sigma_prior` | `(0.0, 0.5)` | Prior on `log sigma_n`. The per-sequence MLE is unbounded: a sequence whose reads land only in the outer bins wants a two-point mass, reachable only as `sigma → ∞`. Without this, `sigma` hits its bound in several cell lines; those are bad optima, so the prior often *improves* the data likelihood. `(0, 1)` fits marginally better on well-behaved lines but leaves stuck ones stuck. |
220
+ | `lambda_fix` | `50.0` | The latent Poisson level is not identified upward — it rides the `R·lambda` ridge, and only products like `R(1−P)/P · Pi · lambda` are identified. It is held fixed and `R` absorbs the scale. |
221
+ | `a_max` | `15.0` | Bound on `a_n`; also sizes the truncated latent-count grid, so raising it costs runtime and memory. |
222
+ | readout grid | `121 × 33` | `(mu, log sigma)` grid for the posterior mean. Converged in the bulk: doubling to `181 × 49` moves the median activity by 2e-6 (a handful of extreme low-count sequences move by up to 0.04), while halving to `61 × 25` moves it by 5e-3. |
223
+
224
+ The readout re-uses the fit's own priors by default, so the posterior mean is
225
+ coherent with the fitted model. A flat effect location is obtained with
226
+ `readout=dict(tau_within=None)`, and `mu_center=<array>` shrinks toward a
227
+ per-sequence target (for instance the same sequence's consensus across cell
228
+ lines).
229
+
230
+ ## Reproducing the shipped tables
231
+
232
+ The shipped `activity_zi_full.csv` (30k sequences × 20 cell lines) is the
233
+ default configuration:
234
+
235
+ ```bash
236
+ python -m ebin data.csv -o activity.csv --save-fits fits.pkl
237
+ ```
238
+
239
+ Cell lines come out in the order they appear in the count table, and every field
240
+ is written. The shipped table carried only `activity, activity_sd,
241
+ activity_map, tot` in an explicit cell-line order, so
242
+
243
+ ```bash
244
+ python -m ebin data.csv -o activity.csv \
245
+ --fields activity,activity_sd,activity_map,tot --groups A549,HepG2,MCF7,...
246
+ ```
247
+
248
+ reproduces that exact layout. The extra columns are additional output, not a
249
+ change to the numbers.
250
+
251
+ Checked against the shipped table three ways:
252
+
253
+ * **Readout.** Given the shipped fitted parameters, `readout_table` reproduces
254
+ `activity`, `activity_sd`, `activity_map` and `tot` to `4e-16` — exact up to
255
+ CSV round-off.
256
+ * **Objective.** The log-likelihood and its gradient agree **bit for bit** with
257
+ the original implementation at the same parameter point, for zero inflation
258
+ and zero truncation, free and Gamma abundance — including the gradient
259
+ through the implicit cut-point solve.
260
+ * **Full refit.** Most cell lines land on the shipped optimum (`max|Δ|` on
261
+ activity ≤ 1e-6). A few of the hardest lines do not: the objective is
262
+ multi-modal there, and the optimizer can settle in a neighbouring basin,
263
+ which moves individual sequences by up to ~1e-1 (Pearson r ≥ 0.9998 against
264
+ the shipped column). This is a property of the problem, not of the
265
+ repackaging — the original code has the same instability on the same lines
266
+ (re-fitting HCT116 with it lands 13 nats from its own shipped value), and
267
+ both implementations reproduce those lines when fitted on their own.
268
+
269
+ The external-library (lib2) target used the same fit with a smaller readout
270
+ grid, selected with `--grid 61x25`.
271
+
272
+ The Gamma-abundance variant was warm-started from the shipped fits, which is
273
+ what makes it affordable, and used a shorter finishing schedule:
274
+
275
+ ```python
276
+ from ebin import activity_table
277
+ activity_table("data.csv", out="activity_gamma.csv", warm_from="fits.pkl",
278
+ abundance=False, abundance_prior="gamma", n_abund_nodes=30,
279
+ max_outer=1, finish_maxiter=2000, finish_rounds=4)
280
+ ```
281
+
282
+ The CLI runs the same variant on the default schedule:
283
+ `--abundance gamma --nodes 30 --warm-from fits.pkl`.
284
+
285
+ ## Runtime
286
+
287
+ One cell line of 30k sequences × 4 bins takes ~2–4 min on a GPU (fit + readout)
288
+ and roughly 20× longer on CPU; the Gamma-abundance variant costs ~2× the
289
+ free-`a_n` fit. Memory is dominated by the latent-count grid
290
+ (`a_max · lambda_fix` wide); the likelihood scans over abundance/rate components
291
+ so only one is live at a time. `XLA_PYTHON_CLIENT_PREALLOCATE=false` is set on
292
+ import so JAX does not claim 75% of VRAM.
293
+
294
+
295
+ ## Parameters
296
+
297
+
298
+ | symbol | shape | role | what it is | example (A549) |
299
+ | --- | --- | --- | --- | --- |
300
+ | `X[n,s,b]` | N×S×B | observed | reads counted for sequence n in channel (s,b) | 1686 median total |
301
+ | `T[n,s,b]` | N×S×B | latent | cells of n that landed in bin b — never observed, summed out | ~63 cells |
302
+ | `mu_n`, `sigma_n` | N each | per sequence | the sequence's expression law over cells, in gauge units | σ median 1.46 |
303
+ | `a_n` | N | per sequence | size factor: how many cells n contributes at all. Geometric mean pinned to 1, capped at `a_max` | 1.42 median |
304
+ | `Pi[n,b]` | N×B | derived | bin profile, rows sum to 1. Not free — read off `(mu_n, sigma_n)` and the cuts | — |
305
+ | `q_j` | B−1 | shared | cut points = quantiles of the mixture `G` at `j/B`. Gauge pins `q_1 = −1`, `q_2 = 0` | −1, 0, 0.99 |
306
+ | `lambda_s` | S | fixed | cell level of replicate s. Not identified — rides the `R·lambda` ridge, so held fixed while `R` absorbs the scale | 50 |
307
+ | `R[s,b]` | S×B | per channel | NB size per latent cell: `R·T` is the shape of the read distribution | 12.7 11.4 11.6 20.1 |
308
+ | `P[s,b]` | S×B | per channel | NB probability. Small `P` = a noisier channel at the same depth | .33 .32 .31 .52 |
309
+ | `phi` | 1 | shared | probability a sequence is a structural zero (dropout, never a cell) | 8e−7 |
310
+ | `kappa` | 1 | optional | Gamma shape when `a_n` is integrated out instead of fitted; CV = 1/√κ | 0.9 (variant) |
311
+
312
+ so `X, T, R, P` are 3D arrays:
313
+ <center>
314
+ <img src="http://data.georgy.top/media/ebin_index_cube.png" alt="Index cube: n sequences, s replicates, b bins" width="196">
315
+ </center>
316
+
@@ -0,0 +1,32 @@
1
+ """EBin -- posterior bin-expectation normalization for sorting-based MPRA.
2
+
3
+ Reads raw per-bin read counts and returns, for every sequence in every cell
4
+ line, the posterior-mean expected bin E[bin] (the activity) together with its
5
+ posterior SD. The model is a compound NB-Poisson likelihood over a latent
6
+ Gaussian effect law per sequence; see model.py, effects.py and activity.py.
7
+
8
+ from ebin import activity_table
9
+ table, _ = activity_table("counts.csv", out="activity.csv")
10
+ """
11
+
12
+ import os
13
+
14
+ # jax grabs 75% of VRAM otherwise; float64 throughout (the likelihood sums
15
+ # thousands of log terms per object)
16
+ os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
17
+
18
+ import jax # noqa: E402
19
+
20
+ jax.config.update("jax_enable_x64", True)
21
+
22
+ from .data import GroupData, read_counts, load_groups, prep # noqa: E402,F401
23
+ from .initialize import initialize, init_from_fit, InitResult # noqa: E402,F401
24
+ from .fit import fit_effects, FitResult # noqa: E402,F401
25
+ from .activity import posterior_activity # noqa: E402,F401
26
+ from .pipeline import (activity_table, readout_table, save_fits, # noqa: E402,F401
27
+ load_fits, write_netcdf)
28
+ from .plots import (plot_marginal_distribution, plot_activity_vs_raw, # noqa: E402,F401
29
+ plot_gc_vs_abundance, plot_all, raw_mass_center,
30
+ gc_content, marginal_density)
31
+
32
+ __version__ = "1.0"
@@ -0,0 +1,76 @@
1
+ """Command line: python -m ebin counts.csv -o activity.csv"""
2
+
3
+ import argparse
4
+
5
+ from .pipeline import activity_table, write_netcdf, TABLE_FIELDS
6
+
7
+
8
+ def _parse_args(argv=None):
9
+ p = argparse.ArgumentParser(
10
+ prog="python -m ebin",
11
+ description="Fit the EBin model to a bin-count table and write the "
12
+ "per-cell-line activity table. Defaults are the shipped "
13
+ "configuration.")
14
+ p.add_argument("data", help="count table, 3-level header "
15
+ "(cell line, replicate, bin)")
16
+ p.add_argument("-o", "--out", default="activity.csv", help="output CSV")
17
+ p.add_argument("--groups", help="comma-separated cell lines to run, in "
18
+ "output order (default: all)")
19
+ p.add_argument("--zero-truncation", action="store_true",
20
+ help="drop all-zero rows instead of fitting them as "
21
+ "structural zeros")
22
+ p.add_argument("--abundance", choices=("free", "gamma"), default="free",
23
+ help="per-object abundance: a free a_n, or integrated out "
24
+ "under a Gamma prior")
25
+ p.add_argument("--nodes", type=int, default=24,
26
+ help="quadrature nodes for the Gamma abundance integral")
27
+ p.add_argument("--a-max", type=float, default=15.0)
28
+ p.add_argument("--lambda-fix", type=float, default=50.0)
29
+ p.add_argument("--mu-prior", type=float, default=1.3,
30
+ help="sd of the prior on the effect location (0 disables)")
31
+ p.add_argument("--sigma-prior", default="0,0.5",
32
+ help="mean,sd of the prior on log sigma ('none' disables)")
33
+ p.add_argument("--grid", default="121x33",
34
+ help="readout grid, n_mu x n_log_sigma")
35
+ p.add_argument("--fields", help="comma-separated subset of the output "
36
+ "fields (default: all of them)")
37
+ p.add_argument("--save-fits", help="pickle the fitted parameters here")
38
+ p.add_argument("--warm-from", help="warm-start from a fits pickle")
39
+ p.add_argument("--netcdf", help="also write the full posterior to netCDF")
40
+ p.add_argument("--plots", metavar="DIR",
41
+ help="also write the diagnostic plots into DIR")
42
+ p.add_argument("-q", "--quiet", action="store_true")
43
+ return p.parse_args(argv)
44
+
45
+
46
+ def main(argv=None):
47
+ a = _parse_args(argv)
48
+ n_mu, n_ls = (int(v) for v in a.grid.lower().split("x"))
49
+ sigma_prior = None if a.sigma_prior.lower() == "none" else \
50
+ tuple(float(v) for v in a.sigma_prior.split(","))
51
+ fit_kw = dict(conditional=a.zero_truncation, a_max=a.a_max,
52
+ lambda_fix=a.lambda_fix, sigma_prior=sigma_prior,
53
+ mu_prior=a.mu_prior or None)
54
+ if a.abundance == "gamma":
55
+ fit_kw.update(abundance=False, abundance_prior="gamma",
56
+ n_abund_nodes=a.nodes)
57
+
58
+ table, results = activity_table(
59
+ a.data, groups=a.groups.split(",") if a.groups else None,
60
+ out=a.out, fits_out=a.save_fits, warm_from=a.warm_from,
61
+ readout=dict(n_mu=n_mu, n_ls=n_ls),
62
+ fields=tuple(a.fields.split(",")) if a.fields else TABLE_FIELDS,
63
+ verbose=not a.quiet, **fit_kw)
64
+ print(f"[saved] {a.out} {table.shape[0]} sequences x "
65
+ f"{len(results)} cell lines")
66
+ if a.netcdf:
67
+ write_netcdf(a.netcdf, results, table.index, order=list(results),
68
+ attrs=dict(model="EBin", readout=f"{n_mu}x{n_ls} grid"))
69
+ print(f"[saved] {a.netcdf}")
70
+ if a.plots:
71
+ from .plots import plot_all
72
+ plot_all(table, a.data, a.plots, groups=list(results))
73
+
74
+
75
+ if __name__ == "__main__":
76
+ main()
@@ -0,0 +1,158 @@
1
+ """Posterior-mean activity E[bin] and its uncertainty.
2
+
3
+ The plug-in readout E[bin] = sum_b b * Pi[n,b] at the fitted (mu_n, sigma_n) is
4
+ a MODE. For a low-information object the posterior over (mu_n, sigma_n) is
5
+ diffuse and skewed, so its mode sits near a simplex vertex (activity 1 or B,
6
+ regardless of depth) while the posterior MEAN is a mild tilt with a large SD --
7
+ the honest summary, and the one that ships.
8
+
9
+ E[bin] is a functional of theta = (mu_n, sigma_n), so its posterior mean is
10
+
11
+ activity_n = sum_j w_{n,j} ebin(theta_j), w_{n,j} ~ p(X_n | theta_j) p(theta_j)
12
+
13
+ on a shared (mu, log sigma) grid. The abundance is PROFILED per grid point (a
14
+ size factor a = total_n / E[total | a = 1, theta_j]), so a deep object's
15
+ likelihood is sharply peaked and a 3-read object's is flat: depth enters
16
+ honestly instead of being frozen at the joint mode. The prior is the same
17
+ penalty the fit used, so the posterior mean is coherent with the fitted model.
18
+ """
19
+
20
+ import numpy as np
21
+ import jax
22
+ import jax.numpy as jnp
23
+
24
+ from .model import LoglikBuilder
25
+ from .effects import normal_bin_probs
26
+ from .truncation import log_zero_prob, log1mexp
27
+
28
+ jax.config.update("jax_enable_x64", True)
29
+
30
+
31
+ def posterior_activity(res, X, mask, *, mu_center=0.0, tau_within="auto",
32
+ sigma_prior="auto", conditional="auto",
33
+ n_mu=121, n_ls=33,
34
+ mu_lim=12.0, ls_lim=(0.2, 5.0)):
35
+ """Posterior summary of every object of one cell line.
36
+
37
+ res : a fitted ``FitResult`` (supplies R, P, rates, cuts, config).
38
+ X, mask : the same (N, S, B) arrays the fit saw.
39
+ mu_center : scalar or (N,) shrinkage target of the mu prior (0 = neutral).
40
+ tau_within : sd of the mu prior; "auto" takes the fit's ``mu_prior``,
41
+ None gives a flat effect location.
42
+ sigma_prior : (m, s) log-sigma prior; "auto" takes the fit's.
43
+ conditional : whether the per-object likelihood is zero-truncated; "auto"
44
+ takes the fit's. Under zero inflation phi is a global
45
+ constant and drops out of the per-object weights.
46
+
47
+ Everything comes from the SAME posterior weights, so the summaries are
48
+ mutually consistent (activity == bin_prob @ [1..B], q2 == mu). Returns a
49
+ dict of (N,) arrays, NaN off ``res.observed``:
50
+
51
+ activity, activity_sd posterior mean and SD of E[bin]
52
+ activity_map the plug-in mode, for reference
53
+ mu, sigma posterior-mean effect location and scale
54
+ mu_sd, sigma_sd their posterior SDs (standard errors)
55
+ q1, q2, q3 quartiles of the posterior-mean effect law
56
+ abundance posterior-mean profiled size factor a_n
57
+ n_eff posterior grid support (larger = more diffuse)
58
+ tot observed reads
59
+ observed (N,) bool
60
+
61
+ plus bin_prob, the (N, B) posterior-mean distribution over bins.
62
+ """
63
+ cfg = res.config
64
+ X = np.asarray(X)
65
+ N, S, B = X.shape
66
+ bvec = np.arange(1, B + 1, dtype=float)
67
+ lambda_fix = cfg["lambda_fix"]
68
+ # profile the abundance up to a_max whenever the fit modelled one, whether
69
+ # as a free a_n or integrated out under the Gamma prior
70
+ a_max = (cfg["a_max"]
71
+ if (cfg["abundance"] or cfg.get("abundance_prior")) else 1.0)
72
+ rate_max = lambda_fix * a_max
73
+ if tau_within == "auto":
74
+ tau_within = cfg.get("mu_prior")
75
+ if isinstance(sigma_prior, str):
76
+ sigma_prior = cfg.get("sigma_prior")
77
+ if conditional == "auto":
78
+ conditional = cfg.get("conditional", True)
79
+
80
+ cuts = jnp.asarray(res.cuts)
81
+ R, P = jnp.asarray(res.R), jnp.asarray(res.P)
82
+ rates, log_w = jnp.asarray(res.rates), jnp.asarray(res.log_w)
83
+ lam = np.asarray(res.rates)[:, 0] # (S,) K = 1
84
+ c = np.asarray(res.R) * (1 - np.asarray(res.P)) / np.asarray(res.P)
85
+
86
+ builder = LoglikBuilder(X, mask=mask, rate_max=rate_max)
87
+ kept = np.asarray(res.observed)
88
+ Xz = np.where(np.asarray(builder.mask), np.nan_to_num(X, nan=0.0), 0.0)
89
+ total = Xz.sum((1, 2)) # observed reads
90
+
91
+ # shared (mu, log sigma) grid, its activity readout and expected total
92
+ mu_g = np.linspace(-mu_lim, mu_lim, n_mu)
93
+ ls_g = np.linspace(np.log(ls_lim[0]), np.log(ls_lim[1]), n_ls)
94
+ MU, LS = np.meshgrid(mu_g, ls_g, indexing="ij")
95
+ MU, LS = MU.ravel(), LS.ravel()
96
+ G = MU.size
97
+ Pi_g = np.asarray(normal_bin_probs(jnp.asarray(MU), jnp.exp(jnp.asarray(LS)),
98
+ cuts))
99
+ ebin_g = Pi_g @ bvec # (G,)
100
+ denom_g = (c[None] * Pi_g[:, None, :]).sum(2) @ lam # E[tot | a = 1]
101
+
102
+ @jax.jit
103
+ def ll_grid(pi_row, a_col):
104
+ M = a_col[:, None] * pi_row[None, :]
105
+ obj = builder.object_loglik(R, P, M, rates, log_w, 0.0)
106
+ if not conditional:
107
+ return obj
108
+ lz = jnp.minimum(log_zero_prob(R, P, M, rates, log_w, builder.mask),
109
+ -1e-12)
110
+ return obj - log1mexp(lz)
111
+
112
+ ll = np.empty((N, G))
113
+ a_grid = np.empty((N, G))
114
+ for j in range(G):
115
+ a_col = np.clip(total / max(denom_g[j], 1e-12), 1e-4, a_max)
116
+ a_grid[:, j] = a_col
117
+ ll[:, j] = np.asarray(ll_grid(jnp.asarray(Pi_g[j]), jnp.asarray(a_col)))
118
+
119
+ # per-object prior over the grid: N(mu; center_n, tau^2) * N(ls; m_s, s_s)
120
+ logprior = np.zeros((N, G))
121
+ if tau_within is not None:
122
+ center = np.broadcast_to(np.asarray(mu_center, float), (N,))
123
+ logprior = logprior - (MU[None, :] - center[:, None]) ** 2 \
124
+ / (2.0 * tau_within ** 2)
125
+ if sigma_prior is not None:
126
+ m_s, s_s = sigma_prior
127
+ logprior = logprior - ((LS - m_s) ** 2 / (2.0 * s_s ** 2))[None, :]
128
+
129
+ logpost = ll + logprior
130
+ logpost -= logpost.max(1, keepdims=True)
131
+ w = np.exp(logpost)
132
+ w /= w.sum(1, keepdims=True)
133
+ ebin_pm = w @ ebin_g
134
+ ebin_sd = np.sqrt(np.clip(w @ (ebin_g ** 2) - ebin_pm ** 2, 0, None))
135
+
136
+ sig_g = np.exp(LS)
137
+ mu_pm = w @ MU
138
+ sigma_pm = w @ sig_g
139
+ z25 = 0.6744897501960817 # N(0,1) third quartile
140
+ return dict(
141
+ activity=np.where(kept, ebin_pm, np.nan),
142
+ activity_sd=np.where(kept, ebin_sd, np.nan),
143
+ activity_map=np.where(kept, np.asarray(res.Pi) @ bvec, np.nan),
144
+ mu=np.where(kept, mu_pm, np.nan),
145
+ mu_sd=np.where(kept, np.sqrt(np.clip(w @ (MU ** 2) - mu_pm ** 2,
146
+ 0, None)), np.nan),
147
+ sigma=np.where(kept, sigma_pm, np.nan),
148
+ sigma_sd=np.where(kept, np.sqrt(np.clip(w @ (sig_g ** 2)
149
+ - sigma_pm ** 2, 0, None)),
150
+ np.nan),
151
+ q1=np.where(kept, w @ (MU - z25 * sig_g), np.nan),
152
+ q2=np.where(kept, mu_pm, np.nan),
153
+ q3=np.where(kept, w @ (MU + z25 * sig_g), np.nan),
154
+ abundance=np.where(kept, (w * a_grid).sum(1), np.nan),
155
+ n_eff=np.where(kept, 1.0 / (w ** 2).sum(1), np.nan),
156
+ tot=np.where(kept, total, np.nan),
157
+ bin_prob=np.where(kept[:, None], w @ Pi_g, np.nan),
158
+ observed=kept)