Publish 2026_W39_unit

README.md

@@ -24,12 +24,11 @@ The model has three stages:
 2. a **25-component Gaussian mixture** over the standardised latents, whose
    components read as putative cell types. Component shapes are global; the **mixture weights
    depend on the molecular context** (MERFISH + AGEA PCA volumes) at the unit's position;
-3. a distance-weighted **k=20 nearest-neighbour projection** onto training units maps latents
-   back to 10 interpretable waveform features: `depolarisation_slope`, `recovery_slope`, `repolarisation_slope`, `spatial_spread_um`, `tip_val`, `spike_width_secs`, `predepolarisation_width_secs`, `spike_amplitude`, `peak_to_trough_ratio_log`, `polarity`.
+3. a **context-local exemplar readout** maps the model back to 10 interpretable waveform features: `depolarisation_slope`, `recovery_slope`, `repolarisation_slope`, `spatial_spread_um`, `tip_val`, `spike_width_secs`, `predepolarisation_width_secs`, `spike_amplitude`, `peak_to_trough_ratio_log`, `polarity`. At a position, each component is represented by its training members whose molecular context resembles the position's (the 4096 nearest in a 3-d context key, shrunk toward all of the component's members -- the more so the farther the context lies from the training data, and entirely where there is no context). So the context sets the mixture weights **and** selects which training exemplars represent each component; the latent density depends on it only through the weights.
 
 > **The input is a position.** `predict` takes `x, y, z` (IBL/Allen frame, metres) and returns the
-> expected phenotype there: the local component weights times each component's expected features.
-> `mixture_weights` returns the local putative cell-type composition itself.
+> expected phenotype there: the local component weights times each component's expected features
+> there. `mixture_weights` returns the local putative cell-type composition itself.
 
 ## Quickstart
 
@@ -41,6 +40,7 @@ model = load_pretrained("int-brain-lab/ea-encoder-unit", revision="2026_W39")
 df = pd.read_parquet("example/positions_sample.parquet")   # positions (x, y, z)
 out = model.predict(df)                                     # one pred_<feature> column each
 weights = model.mixture_weights(df)                         # putative cell-type composition
+draws = model.sample(df, n_samples=16)                   # training units drawn from the predictive distribution
 model.selftest()                                            # reproduces the shipped outputs
 ```
 
@@ -52,10 +52,9 @@ as in `config.json`), then assigned to a putative cell type with `model.assign(z
 - `autoencoder.pt`, `shared_latent_scaler.joblib` -- the phenotype encoder and its latent scaler.
 - `global_gmm.joblib` -- the mixture's component means and covariances.
 - `context_transform.joblib`, `context_weight_model_bundle.pt` -- context -> mixture weights.
-- `knn_bank.npz` -- standardised latents and phenotype features of the training units (the kNN
-  exemplars; part of the model, like the stored points of any nearest-neighbour model).
-- `component_feature_expectations.npz` -- each component's expected features, which `predict`
-  mixes with the local weights.
+- `knn_bank.npz` -- standardised latents, phenotype features, GMM components and context keys of the training units (the exemplars; part of the model, like the stored points of any nearest-neighbour model).
+- `component_feature_expectations.npz` -- each component's expected features over the whole
+  brain.
 - `agea_vol_pca.npy`, `merfish_vol_pca.npy` -- the molecular-context volumes, the same frozen
   volumes as the channel-level model `int-brain-lab/ea-encoder-channel` of this vintage.
 - `split.json`, `config.json`, `preprocessing/unit_stats.npz`, `results/` -- the probe split, the

autoencoder.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:df36b2645806cfa5b87bbcdcf40375c597d13917096d4e8f1c68268de54e2f68
-size 22725861
+oid sha256:2346c6d07e78136054f78f08f075108d7a42c2e1beb42b91c6337464b38b686b
+size 22726117

checksums.json

@@ -8,13 +8,13 @@
     },
     {
       "path": "autoencoder.pt",
-      "hash": "3b24b551896a9e3457f2688838b44dde0166cf10",
-      "bytes": 22725861
+      "hash": "184f905a06df05208ea6ef7ed15a0ce96e315189",
+      "bytes": 22726117
     },
     {
       "path": "component_feature_expectations.npz",
-      "hash": "fd157f4216b410d88c65e4dc488b5b68b6fb4da0",
-      "bytes": 1983
+      "hash": "e17db88bed240b4b7bf6907481cb43b9e21de24e",
+      "bytes": 1974
     },
     {
       "path": "conditioning_summary.json",
@@ -23,8 +23,8 @@
     },
     {
       "path": "config.json",
-      "hash": "5b6b1ca80f31d9a047ba365546f4c3d9e3190868",
-      "bytes": 2910
+      "hash": "6edf96900240ca5a309974290b526a98bc79fdbd",
+      "bytes": 3134
     },
     {
       "path": "context_transform.joblib",
@@ -38,8 +38,8 @@
     },
     {
       "path": "ephysatlas_model.json",
-      "hash": "ec11272198df580311b04f780d881cbfdbae690d",
-      "bytes": 2779
+      "hash": "3a4df0fb1d0ed75e50a4be2357de11220c9b0e30",
+      "bytes": 3411
     },
     {
       "path": "example/expected_latents.npy",
@@ -48,8 +48,8 @@
     },
     {
       "path": "example/expected_predictions.parquet",
-      "hash": "dcb3aa478eba00d80f881c8c80697195ada4e46a",
-      "bytes": 32885
+      "hash": "01809e9956fc0f26fec921398e696ab87a9690e4",
+      "bytes": 32897
     },
     {
       "path": "example/positions_sample.parquet",
@@ -73,8 +73,8 @@
     },
     {
       "path": "knn_bank.npz",
-      "hash": "97f669f75d2e8c7e775933d798585abd2501dd5f",
-      "bytes": 14794605
+      "hash": "31801285d6c244e8764e3a467253f4fb452efcef",
+      "bytes": 14966230
     },
     {
       "path": "merfish_vol_pca.npy",
@@ -86,10 +86,15 @@
       "hash": "b2202de2690f78ac77893569b44fcf96640390e3",
       "bytes": 4683
     },
+    {
+      "path": "readout_summary.json",
+      "hash": "c3a1c76392649d00fab09c8c7f4463569dcd271b",
+      "bytes": 1209
+    },
     {
       "path": "results/summary.json",
-      "hash": "ec12ba0eb657743eecea98569ba977c0e4fa4664",
-      "bytes": 2210
+      "hash": "f8a976e54314e3b152c86bce68815b10bc797707",
+      "bytes": 2793
     },
     {
       "path": "shared_latent_scaler.joblib",

component_feature_expectations.npz

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3b38760939c6be13d8ce001d0df1665a39413602ec846b50a27158016db814e9
-size 1983
+oid sha256:a52e08355cd58ce06fe7a02a82f3966769253668006f5b5dda9a2f0729122dda
+size 1974

config.json

@@ -55,6 +55,17 @@
   "rare_component_power": 0.5,
   "rare_component_weight_cap": 10.0,
   "knn_decoder_k": 20,
+  "readout_key_dim": 3,
+  "readout_key_alphas": [
+    100.0,
+    1000.0,
+    10000.0,
+    30000.0,
+    100000.0
+  ],
+  "readout_neighbours": 4096,
+  "readout_shrinkage": 10.0,
+  "readout_off_data_quantile": 0.99,
   "mirror_x_to_single_hemisphere": true,
   "mirror_x_sign": -1.0,
   "region_gaussian_variance_floor": 0.001,

ephysatlas_model.json

@@ -92,7 +92,7 @@
     "context": {
       "n_cell_pcs": 50,
       "n_gene_pcs": 50,
-      "conditions": "mixture weights only; component means/covariances are global",
+      "conditions": "mixture weights, and which training exemplars represent each component in the phenotype readout; component means/covariances are global",
       "source_model": "int-brain-lab/ea-encoder-channel"
     },
     "preprocessing": {
@@ -100,6 +100,16 @@
       "stats_file": "preprocessing/unit_stats.npz",
       "mirror_x_to_single_hemisphere": true,
       "mirror_x_sign": -1.0
+    },
+    "readout": {
+      "method": "context_local_members",
+      "description": "the molecular context sets the mixture weights gamma_k(x) and selects which training exemplars represent each component: its members nearest to x in a context key, shrunk toward all of its members -- further as the key leaves the training data, and entirely without context",
+      "key_dim": 3,
+      "key_ridge_alpha": 30000.0,
+      "neighbours": 4096,
+      "shrinkage": 10.0,
+      "off_data_quantile": 0.99,
+      "absent_context": "global member means"
     }
   },
   "data_source": {

example/expected_predictions.parquet

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:51da16f91e81be1d49c2a879c8ea96a1ee5e1c6ac48134546da78b0ec66e6258
-size 32885
+oid sha256:1956eff735fde1ad163940ffdce8b3134a17bee774ebf7b5d2b9ef5ec455f428
+size 32897

knn_bank.npz

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bfa52dac944ae2233a65c7e1d71ea79e0c61f412bb55cfcbc4ce2a3e6e1896ba
-size 14794605
+oid sha256:a8e49ee3ecf897fc1296d1d297fff06af8746a9fe3ef1d0ff9c0847a6b1986d2
+size 14966230

readout_summary.json

@@ -0,0 +1,64 @@
+{
+  "method": "context_local_members",
+  "key_features": [
+    "depolarisation_slope",
+    "recovery_slope",
+    "repolarisation_slope",
+    "spatial_spread_um",
+    "tip_val",
+    "spike_width_secs",
+    "predepolarisation_width_secs",
+    "spike_amplitude",
+    "peak_to_trough_ratio_log"
+  ],
+  "key_alpha": 30000.0,
+  "key_alpha_validation_mse": {
+    "100.0": 0.8998134946502477,
+    "1000.0": 0.8993491239299802,
+    "10000.0": 0.8983127253407119,
+    "30000.0": 0.8982345222443193,
+    "100000.0": 0.9005961719454074
+  },
+  "key_dim": 3,
+  "key_singular_values": [
+    80.27060833307382,
+    48.52855142725006,
+    43.25424518501049,
+    27.807332889109052,
+    16.724424931676705,
+    11.457588376291048,
+    6.943266663920836,
+    6.228926358749439,
+    2.4549947322639083
+  ],
+  "neighbours": 4096,
+  "shrinkage": 10.0,
+  "off_data_quantile": 0.99,
+  "train_members_per_component": [
+    1672,
+    4436,
+    3821,
+    1876,
+    1507,
+    2522,
+    2396,
+    5296,
+    2604,
+    5307,
+    463,
+    2184,
+    3571,
+    2285,
+    2279,
+    1403,
+    1539,
+    1210,
+    1981,
+    3729,
+    1566,
+    1131,
+    969,
+    1045,
+    1366
+  ]
+}

results/summary.json

@@ -12,10 +12,19 @@
     "gmm_components": 25,
     "gmm_covariance_type": "full",
     "gmm_geometry": "global means and covariance matrices",
-    "context_conditioning": "mixture weights gamma only",
+    "context_conditioning": "mixture weights gamma, and the TRAIN exemplars representing each component",
     "context_definition": "50 MERFISH PCs + 50 AGEA PCs",
     "knn_k": 20,
-    "knn_policy": "distance-weighted empirical TRAIN-exemplar projection in standardized latent space"
+    "knn_policy": "distance-weighted empirical TRAIN-exemplar projection in standardized latent space",
+    "phenotype_readout": {
+      "method": "context-local members: gamma_k(x) weights the mean of component k's TRAIN members nearest to x in a context key, shrunk toward all its members",
+      "key_dim": 3,
+      "key_ridge_alpha": 30000.0,
+      "neighbours": 4096,
+      "shrinkage": 10.0,
+      "off_data_quantile": 0.99,
+      "absent_context": "global member means"
+    }
   },
   "split_counts": {
     "train_units": 58158,
@@ -57,7 +66,7 @@
   "invariants": [
     "TEST PIDs are never used for AE/GMM/context training or as kNN retrieval exemplars.",
     "The GMM means and covariance matrices are global.",
-    "Molecular context changes only the GMM mixture weights gamma_k(x).",
+    "Molecular context sets the GMM mixture weights gamma_k(x) and selects which TRAIN exemplars represent each component in the phenotype readout; the conditional latent density depends on it only through gamma_k(x).",
     "The kNN stage changes phenotype projection only; it does not change the conditional latent density.",
     "All spatial modeling uses the configured single-hemisphere mirroring convention."
   ]