Convergent motifs of early olfactory processing are recapitulated by layer-wise efficient coding.

Code ↔ Paper

3 matches between paragraphs of the paper and lines of its authors' code, computed by the harvester (lexical-v1). Click a colored paragraph or line to see its counterpart.

The 3 matches · 1 of them tie a paragraph to a whole file without given lines: a weak match, whose lines are not tinted

[1] § Setup › Maximizing Mutual Information by Proxy. ↔ model.py, lines 478 to 545 · score 0.59 · mirror descent, Jensen Shannon · Report a mismatch
[2] § Setup › Stimulus and Encoding Model. ↔ sample_generation/generate_binarized_gaussian_samples.py, lines 74 to 131 · score 0.55 · Gaussian copula, binary, Gamma, vector, block, matrix · Report a mismatch
[3] § Results › Receptors. ↔ Figure3/functions/3rdPartyCodes/powerlaws_full_v0.0.10-2012-01-17/randht.m, the whole file · a weak match · score 0.54 · power law, tailed distribution, heavy · Report a mismatch

Paper

Loaded from Europe PMC by your browser; OSCR does not store it: doi.org · Europe PMC

The paper is loaded when this pane is shown.

The paper is loaded by your browser, with JavaScript: read it at doi.org or on Europe PMC.

The authors' code

Python · 682 lines · 26 KB · MIT · 1 match

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import flax.linen as nn
import jax.tree_util as jtu
from jax import vmap, lax, jit
from jax.scipy.stats import norm, rankdata
from jax.scipy.special import digamma
from jax.scipy.special import gamma
from jax.scipy.linalg import cholesky
from functools import partial
from typing import NamedTuple
from flax import struct
from functools import partial
plt.rcParams["font.size"] = 20
DATA_PATH = "/n/home10/jfernandezdelcasti/noncanonical-olfaction/rbm/mm_synthetic_data_non_zero_samples_9_April_2025.npy"
class HyperParams(NamedTuple):
N: int = 100 # number of odorants. this needs to match the number of variables (rows) in the data when using the data driven model.
n: int = 2 # sparsity of odor vectors--only applies for the Qin et al 2019 odor model, which we do not use.
M: int = 30 # number of odorant receptors
L: int = 300 # number of olfactory receptor neurons
P: int = 1000 # number of samples when we compute MI
K: int = 30 # number of glomeruli (only relevant if computing with glomerular convergence)
window: int = 32 # this should be set to jnp.round(jnp.sqrt(P) + 0.5).astype(int) but doing this dynamically makes hp not hashable. So change it when you change P!
sigma_0: float = 0.1 # neural noise in Gaussian linear filter model
sigma_c: float = 2.0 # std_dev of log normal Qin et al 2019 odorant model
mu_c: float = 3.0 # mean of log normal odorant distribution. This doesn't matter mathematically, but making the odorants large helps the W optimization.
W_shape: float = 1.0
W_scale: str = "1 / ||c||"
W_log_normal_sigma: float = 1.0
W_log_normal_mu: float = -12.0 # to get to the mean in base 10 of the resulting value, multiply by .43. So this is a mean of around 10^-6.
W_sparsity_mask: float = 1.0
W_init: str = "normal"
kappa_inv_eta_init: str = "normal"
nonlinearity: str = "sigmoid log"
F_max: float = 25.0
hill_exponent: int = 4
odor_model: str = "log normal"
activity_model: str = "glomerular convergence"
loss: str = "jensen_shannon_loss"
phi_w: str = "identity"
psi_w: str = "identity"
phi_e: str = "softplus_phi"
psi_e: str = "softplus_psi"
phi_g: str = "identity"
psi_g: str = "identity"
phi_gain: str = "positive_phi"
psi_gain: str = "positive_psi"
phi_kappa_inv: str = "bounded_phi"
psi_kappa_inv: str = "bounded_psi"
phi_eta: str = "bounded_phi"
psi_eta: str = "bounded_psi"
sigma_kappa_inv: float = 4.0 # this is for Reddy et al. 2018 antagonism model
rho: float = 0.0 # ditto
canonical_E_init: bool = False
balanced_E_init: bool = False
noncanonical_E_init: bool = False
read_in_E: bool = False
canonical_G_init: bool = False
read_in_G: bool = False
initial_gain: float = 1.0
binarize_c_for_MI_computation: bool = False
W_path: str = "/n/home10/jfernandezdelcasti/noncanonical-olfaction/model/results/scratch/W_final_0.npy" # only applies if W_init == "read in"
E_path: str = "/n/home10/jfernandezdelcasti/noncanonical-olfaction/model/results/expression/E_noncanonical_init.npy" # only applies if hp.read_in_E is True
G_path: str = "/n/home10/jfernandezdelcasti/noncanonical-olfaction/model/results/glomerular_convergence/G_final_0.npy" # only applies if read_in_G == True
class TrainingConfig(NamedTuple):
scans: int = 10
epochs_per_scan: int = 100
gamma_W: float = 0.1
gamma_E: float = 0.1
gamma_G: float = 0.1
gamma_gain: float = 0.1
gamma_kappa_inv: float = 0.1
gamma_eta: float = 0.1
gamma_T: float = 0.1
T_hidden_dim: int = 128
T_mode: str = "inner_product"
class LoggingConfig(NamedTuple):
log_interval: int = None
save_model: bool = False
output_dir: str = "."
config_id: str = "0"
class FullConfig(NamedTuple):
hyperparams: HyperParams
training: TrainingConfig
logging: LoggingConfig
seed: int = 0
data_path: str = ''
@struct.dataclass
class Params:
W: jnp.ndarray
E: jnp.ndarray
G: jnp.ndarray
gain: jnp.float32
kappa_inv: jnp.ndarray
eta: jnp.ndarray
class T_estimator(nn.Module):
hidden_dim: int = 128
mode: str = "inner_product"
def g_f(self, v):
return jnp.clip(
jnp.log(2) - jnp.log(1 + jnp.exp(-v)), max=jnp.log(2) - 1e-6
) # see Table 2 of https://arxiv.org/pdf/1606.00709
@nn.compact
def __call__(self, c, r):
# batch should be the first dimension here. this is contra everything else in our pipeline, so be careful! transpose just before passing in c and r
if self.mode == "inner_product":
r = nn.Dense(self.hidden_dim)(r)
r = nn.tanh(r)
r = nn.Dense(self.hidden_dim)(r)
r = nn.tanh(r)
c = nn.Dense(self.hidden_dim)(c)
c = nn.tanh(c)
c = nn.Dense(self.hidden_dim)(c)
c = nn.tanh(c)
v = jnp.einsum("bi,bi->b", c, r) # or use vmap(jnp.dot)(c, r)
return self.g_f(v)
elif self.mode == "concatenate":
# Concatenate c (P, N) and r (P, L) -> (P, N + L)
x = jnp.concatenate([c, r], axis=1)
x = nn.Dense(self.hidden_dim)(x)
x = nn.tanh(x)
x = nn.Dense(self.hidden_dim)(x)
x = nn.tanh(x)
v = nn.Dense(1)(x)
return self.g_f(v)
@struct.dataclass
class TrainingState:
p: any # Your primal parameters
Tp: any # MINE network parameters
key: jax.Array # PRNG key
metrics: dict
def create_one_hot_array(key, L, M):
indices = jax.random.randint(key, (L,), 0, M)
return jax.nn.one_hot(indices, M, dtype=float)
def sigmoid_log(x, hp):
return x / (jnp.abs(x) + 1) # for the cases where x can be negative
def hill_function(x, hp):
x_safe = jnp.clip(x, a_min=1e-10, a_max=None) # Avoid division by zero and extreme values
denom = 1 + (1 / x_safe)**hp.hill_exponent
return 1 / denom
def make_noncanonical_E(key, hp, lower=1, upper=7):
E_rows = []
keys = jax.random.split(key, hp.L)
for i in range(hp.L):
subkey1, subkey2, subkey3 = jax.random.split(keys[i], 3)
k = jax.random.randint(subkey1, shape=(), minval=lower, maxval=upper + 1)
perm = jax.random.permutation(subkey2, hp.M)
indices = perm[:k]
values = jax.random.uniform(subkey3, shape=(k,))
row = jnp.zeros(hp.M)
row = row.at[indices].set(values)
row /= jnp.sum(row)
E_rows.append(row)
E = jnp.stack(E_rows)
return E
def initialize_p(rng, mean_norm_c=None, threshold=None, hp=None) -> Params:
if hp is None:
hp = HyperParams()
W_key, E_key, G_key, z_key, eta_key = jax.random.split(rng, 5)
phi_E = PHI_PSI_REGISTRY[hp.phi_e]
phi_G = PHI_PSI_REGISTRY[hp.phi_g]
phi_W = PHI_PSI_REGISTRY[hp.phi_w]
phi_kappa_inv = PHI_PSI_REGISTRY[hp.phi_kappa_inv]
phi_eta = PHI_PSI_REGISTRY[hp.phi_eta]
if hp.W_init == "gamma":
if isinstance(hp.W_scale, float):
W = float(hp.W_scale) * phi_W(jnp.clip(
jax.random.gamma(W_key, a=hp.W_shape, shape=(hp.M, hp.N)),
min=1e-9,
max=1 - 1e-9,
))
elif hp.W_scale == "1 / ||c||":
W = 1 / mean_norm_c * phi_W(jnp.clip(
jax.random.gamma(W_key, a=hp.W_shape, shape=(hp.M, hp.N)),
min=1e-9,
max=1 - 1e-9,
))
elif hp.W_init == "read in":
W = jnp.clip(jnp.load(hp.W_path), min=1e-16)
elif hp.W_init == "normal":
if hp.W_scale == "1 / ||c||":
W = (1 / mean_norm_c) * phi_W(jax.random.normal(W_key, shape=(hp.M, hp.N)))
if hp.W_scale == "threshold":
W = threshold * phi_W(jax.random.normal(W_key, shape=(hp.M, hp.N)))
elif hp.W_scale == "1 / root(N)":
W = 1 / jnp.sqrt(hp.N) * phi_W(jax.random.normal(W_key, shape=(hp.M, hp.N))) # this is a bad scaling because c is log-normal distributed and sparse
elif hp.W_init == "log normal":
key, W_key, mask_key = jax.random.split(W_key, 3)
Z = jax.random.normal(W_key, shape=(hp.M, hp.N))
mask = jax.random.bernoulli(mask_key, p = hp.W_sparsity_mask, shape=(hp.M, hp.N))
if hp.W_scale == "1 / ||c||":
W = jnp.clip((1 / mean_norm_c) * jnp.exp(hp.W_log_normal_sigma * Z) * mask, min=1e-16, max=1e-1)
if hp.W_scale == "threshold":
W = jnp.clip(threshold * jnp.exp(hp.W_log_normal_sigma * Z) * mask, min=1e-16, max=1e-1)
if hp.W_scale == "custom":
W = jnp.clip(jnp.exp(hp.W_log_normal_sigma * Z + hp.W_log_normal_mu) * mask, min=1e-16, max=1e-1)
if hp.canonical_E_init:
if hp.balanced_E_init:
try:
E = jnp.repeat(
jnp.eye(hp.M), hp.L // hp.M, axis=0
).astype(float) # this is a canonical initialization with exactly equal numbers representing each receptor
except:
raise ValueError(
"For balanced canonical initialization, set L to a postive multiple of M"
)
else:
E = create_one_hot_array(E_key, hp.L, hp.M)
elif hp.read_in_E:
E = jnp.load(hp.E_path)
elif hp.noncanonical_E_init:
E = make_noncanonical_E(E_key, hp) # no softmax here: set lower to upper (3 to 7) entries to be non-zero, and divide by row sum.
else: # random init
E = phi_E(
0.5 + 0.1 * jax.random.normal(E_key, shape=(hp.L, hp.M))
) # this is a random initialization where every neuron expresses roughly the same amount of every receptor.
if hp.canonical_G_init:
G = E.T # this only works if number of receptors (hp.M) == number of glomeruli (hp.K)
elif hp.read_in_G:
G = jnp.load(hp.G_path)
else:
G = phi_G(1/hp.L * jax.random.normal(G_key, shape=(hp.K, hp.L))) # this scaling is key. Otherwise the scale of g is enormous compared to r, and loss = 0 no matter what. the signal is swamped.
gain = hp.initial_gain
if hp.kappa_inv_eta_init == "normal":
kappa_inv = phi_kappa_inv(jax.random.normal(eta_key, shape=(hp.M, hp.N)))
eta = phi_eta(jax.random.normal(z_key, shape=(hp.M, hp.N)))
elif hp.kappa_inv_eta_init == "near zero":
kappa_inv = phi_kappa_inv(jax.random.normal(eta_key, shape=(hp.M, hp.N)) - 5)
eta = phi_eta(jax.random.normal(z_key, shape=(hp.M, hp.N)) - 5)
else: # kappa_inv = jax.random.lognormal(kappa_inv_key, sigma=sigma_kappa_inv, shape=(hp.M, hp.N)) # see "Olfactory Encoding Model" in Reddy and Zak 2018
eta = jax.random.lognormal(eta_key, shape=(hp.M, hp.N))
z = jax.random.normal(z_key, shape=(hp.M, hp.N))
kappa_inv = jnp.exp(
hp.sigma_kappa_inv * (hp.rho * jnp.log(eta) + jnp.sqrt(1 - hp.rho**2) * z)
)
# kappa_inv and eta are now both positive, so no need to pass them through phi
return hp, Params(W, E, G, gain, kappa_inv, eta)
def initialize_training_state(subkey, hp, p_init, training_config):
subkey_cs, subkey_activity, subkey_T, subkey_state = jax.random.split(subkey, 4)
T = T_estimator(
hidden_dim = training_config.T_hidden_dim,
mode = training_config.T_mode
)
draw_cs = ODOR_MODEL_REGISTRY[hp.odor_model]
activity_function = ACTIVITY_FUNCTION_REGISTRY[hp.activity_model]
cs = draw_cs(subkey_cs, hp)
r = activity_function(hp, p_init, cs, subkey_activity)
if hp.loss == "jensen_shannon_loss":
Tp_init = T.init(subkey_T, cs.T, r.T) # this is for variational MI estimation
else:
Tp_init = 0 # this is a dummy for Gaussian MI estimation (not variational, so no T network needed)
init_state = TrainingState(p_init, Tp_init, subkey_state, metrics={"mi": 0.0})
return init_state
def draw_c_sparse_log_normal(subkey, hp):
c = jnp.zeros(hp.N)
sub1, sub2 = jax.random.split(subkey)
non_zero_indices = jax.random.choice(sub1, hp.N, shape=(hp.n,), replace=False)
concentrations = jnp.exp(hp.mu_c) * jax.random.lognormal(sub2, sigma=hp.sigma_c, shape=(hp.n,))
c = c.at[non_zero_indices].set(concentrations)
return c
@partial(jax.jit, static_argnames=['hp'])
def draw_cs_sparse_log_normal(subkey, hp):
keybatch = jax.random.split(subkey, hp.P)
cs = vmap(draw_c_sparse_log_normal, in_axes=(0, None))(keybatch, hp).T
return cs
def closure_draw_cs_data_driven(
path,
): # closures are nice https://apxml.com/courses/advanced-jax/chapter-1-advanced-jax-transformations-control-flow/closures-jax-staging
loaded_data = jnp.load(path)
P_total = loaded_data.shape[1]
def draw_cs_data_driven_binary(subkey, hp):
idx = jax.random.choice(subkey, P_total, shape=(hp.P,), replace=True)
batch = jax.device_put(loaded_data[:, idx]) # move to device only once per call
return batch # shape (N, P), like the log_normal one
def draw_cs_data_driven_log_normal(subkey, hp):
subkey_index, subkey_concentration = jax.random.split(subkey)
idx = jax.random.choice(subkey_index, P_total, shape=(hp.P,), replace=True)
batch = jax.device_put(loaded_data[:, idx]) # move to device only once per call
batch = jnp.where(
batch,
jnp.exp(hp.mu_c) * jax.random.lognormal(
subkey_concentration, sigma=hp.sigma_c, shape=batch.shape
),
jnp.zeros(batch.shape),
)
return batch
return draw_cs_data_driven_binary, draw_cs_data_driven_log_normal
draw_cs_data_driven_binary, draw_cs_data_driven_log_normal = (
closure_draw_cs_data_driven(DATA_PATH)
)
def compute_receptor_activity(subkey, hp, p, c):
pre_activations = p.W @ c
nl = NL_REGISTRY[hp.nonlinearity]
r = nl(pre_activations) + hp.sigma_0 * jax.random.normal(
subkey, shape=pre_activations.shape
)
return r
def compute_linear_filter_activity(hp, p, c, subkey):
receptor_activations = p.W @ c # p.W is receptor x odorant
c_AMP = p.E @ receptor_activations # p.E is neuron x receptor
nl = NL_REGISTRY[hp.nonlinearity]
r = nl(c_AMP, hp) + hp.sigma_0 * jax.random.normal(subkey, shape=c_AMP.shape)
return r
def linear_filter_plus_glomerular_layer(hp, p, c, subkey):
r = compute_linear_filter_activity(hp, p, c, subkey)
g = p.G @ jnp.tanh(p.gain * r)
return g
def compute_osn_firing_rate_with_antagonism(hp, p, c):
"""generalized version of equation 14 in Reddy, Zak et al. 2018. eta and kappa_inv are receptor x odorant, and the gene expression matrix p.E is neuron x receptor"""
receptor_induced_activities = p.E @ (
jnp.matmul((p.eta * p.kappa_inv), c) / (1 + jnp.matmul(p.kappa_inv, c))
)
denominator = 1 + receptor_induced_activities**-hp.hill_exponent
F = hp.F_max / denominator
return F
def compute_osn_firing_with_antagonism(hp, p, c, subkey):
firing_rate = compute_osn_firing_rate_with_antagonism(hp, p, c)
# firing = jax.random.poisson(subkey, lam=firing_rate)
firing = firing_rate + hp.sigma_0 * jax.random.normal(
subkey, shape=firing_rate.shape
)
return firing # this should be neuron x samples, where c is odorants x samples
def compute_entropy(r, hp):
entropy = compute_sum_of_marginal_entropies(r, hp) - compute_information(r)
return entropy
def compute_sum_of_marginal_entropies(r, hp):
compute_entropy_vmap = vmap(vasicek_entropy, in_axes=(0, None))
marginal_entropies = compute_entropy_vmap(r, hp)
return jnp.sum(marginal_entropies)
def vasicek_entropy(
X, hp
): # see https://github.com/scipy/scipy/blob/main/scipy/stats/_entropy.py#L378
n = X.shape[-1]
X = jnp.sort(X, axis=-1)
X = pad_along_last_axis(X, hp.window)
start1 = 2 * hp.window
length = hp.P
differences = lax.dynamic_slice(X, (start1,), (length,)) - lax.dynamic_slice(
X, (0,), (length,)
)
logs = jnp.log(n / (2 * hp.window) * differences)
return jnp.mean(logs, axis=-1)
def pad_along_last_axis(X, window):
first_value = X[0]
last_value = X[-1]
Xl = lax.full_like(x=jnp.empty((window,)), fill_value=first_value)
Xr = lax.full_like(x=jnp.empty((window,)), fill_value=last_value)
return jnp.concatenate((Xl, X, Xr))
def compute_information(r):
M, P = r.shape
G = norm.ppf(
(rankdata(r.T, axis=0) / (P + 1)), loc=0, scale=1
) # this is just ranking the data and making it normally distributed.
bias_correction = 0.5 * jnp.sum(
digamma((P - jnp.arange(1, M + 1) + 1) / 2) - jnp.log(P / 2)
)
cov_matrix = jnp.cov(G, rowvar=False)
chol_decomp = cholesky(cov_matrix)
log_det = jnp.sum(jnp.log(jnp.diag(chol_decomp)))
I = -(
log_det - bias_correction
) # remember: entropy overall is sum of marginals minus information.
return I
@partial(jax.jit, static_argnames=["hp"])
def entropy_loss(p, hp, cs, subkey):
r = compute_linear_filter_activity(subkey, hp, p, cs)
e = compute_entropy(r, hp)
return -e
@partial(jax.jit, static_argnames=["hp"])
def log_det_loss(p, hp, cs, subkey):
r = compute_linear_filter_activity(subkey, hp, p, cs)
cov = jnp.cov(r, rowvar=True)
chol_decomp = cholesky(cov)
log_det = jnp.sum(jnp.log(jnp.diag(chol_decomp)))
return -log_det
def f_star_JSD(t): # see Table 2 of https://arxiv.org/pdf/1606.00709
return -jnp.log(2 - jnp.exp(t))
def jensen_shannon_loss(p, Tp, hp, cs, key):
key, subkey_firing, subkey_shuffle = jax.random.split(key, 3)
activity_function = ACTIVITY_FUNCTION_REGISTRY[hp.activity_model]
r = activity_function(hp, p, cs, subkey_firing)
r_shuffled = jax.random.permutation(subkey_shuffle, r, axis=1)
if hp.binarize_c_for_MI_computation:
cs = (cs != 0).astype(float)
T_joint = T_estimator().apply(Tp, cs.T, r.T)
T_marginal = T_estimator().apply(Tp, cs.T, r_shuffled.T)
mi_estimate = jnp.mean(T_joint) - jnp.mean(f_star_JSD(T_marginal))
return -mi_estimate
def flatten_metrics(metrics_list):
"""Convert list of metric dicts to dict of arrays."""
flattened = {k: [] for k in metrics_list[0].keys()} # Initialize with keys
for metrics in metrics_list:
for k, v in metrics.items():
flattened[k].extend(v)
return {k: jnp.array(v) for k, v in flattened.items()}
@partial(jax.jit, static_argnames=("hp"))
def update_params_T(p, Tp, hp, cs, key, gammas):
# Split keys
key, loss_key = jax.random.split(key)
# parse gammas
gamma_W, gamma_E, gamma_G, gamma_gain, gamma_kappa_inv, gamma_eta, gamma_T = gammas
# Compute gradients for MINE network (T) and model params (E)
if hp.loss != "jensen_shannon_loss":
raise ValueError(
"only the jensen shannon loss requires a T network. For other losses, use update_params"
)
loss = LOSS_REGISTRY[hp.loss]
phi_W = PHI_PSI_REGISTRY[hp.phi_w]
psi_W = PHI_PSI_REGISTRY[hp.psi_w]
phi_E = PHI_PSI_REGISTRY[hp.phi_e]
psi_E = PHI_PSI_REGISTRY[hp.psi_e]
phi_G = PHI_PSI_REGISTRY[hp.phi_g]
psi_G = PHI_PSI_REGISTRY[hp.psi_g]
phi_gain = PHI_PSI_REGISTRY[hp.phi_gain]
psi_gain = PHI_PSI_REGISTRY[hp.psi_gain]
phi_kappa_inv = PHI_PSI_REGISTRY[hp.phi_kappa_inv]
psi_kappa_inv = PHI_PSI_REGISTRY[hp.psi_kappa_inv]
phi_eta = PHI_PSI_REGISTRY[hp.phi_eta]
psi_eta = PHI_PSI_REGISTRY[hp.psi_eta]
mi_estimate, grads = jax.value_and_grad(loss, argnums=(0, 1))(
p, Tp, hp, cs, loss_key
)
grads_p, grads_T = grads
# Update MINE network (T)
Tp = jax.tree.map(lambda param, grad: param - gamma_T * grad, Tp, grads_T)
# natural gradient in dual space is equivalent to mirror descent in primal space. A "straight through" trick.
# phi:unconstrained --> constrained, psi is inverse of phi.
# see https://vene.ro/blog/mirror-descent.html#generalizing-the-projected-gradient-method-with-divergences for a very clear exposition
def update_param(p, param_name, psi, phi, grads, gamma, epsilon=1e-8):
# Map to dual space
U_current = psi(getattr(p, param_name))
# Normalize the gradient
grad = getattr(grads, param_name)
grad_norm = jnp.linalg.norm(grad)
normalized_grad = grad / (grad_norm + epsilon)
# Gradient step in dual space
U_new = U_current - gamma * normalized_grad
# Map back to primal space
new_value = phi(U_new)
# Replace in parameter object
return p.replace(**{param_name: new_value})
# Apply update for each parameter
p = update_param(p, "W", psi_W, phi_W, grads_p, gamma_W)
p = update_param(p, "E", psi_E, phi_E, grads_p, gamma_E)
p = update_param(p, "G", psi_G, phi_G, grads_p, gamma_G)
p = update_param(p, "gain", psi_gain, phi_gain, grads_p, gamma_gain)
p = update_param(p, "kappa_inv", psi_kappa_inv, phi_kappa_inv, grads_p, gamma_kappa_inv)
p = update_param(p, "eta", psi_eta, phi_eta, grads_p, gamma_eta)
return p, Tp, -mi_estimate
def update_params(p, hp, cs, key, gamma_E):
''' this is old--comes from estimating MI directly a la Qin et al. 2019, not using the variational formulation '''
# Split keys
key, loss_key = jax.random.split(key)
loss = LOSS_REGISTRY[hp.loss]
phi = PHI_PSI_REGISTRY[hp.phi]
psi = PHI_PSI_REGISTRY[hp.psi]
# Compute gradients for MINE network (T) and model params (E)
mi_estimate, grads_p = jax.value_and_grad(loss, argnums=(0))(
p, hp, cs, loss_key
) # loss needs to have this signature
# natural gradient in dual space is equivalent to mirror descent in primal space. A "straight through" trick.
# phi:unconstrained --> constrained, psi is inverse of phi.
# see https://vene.ro/blog/mirror-descent.html#generalizing-the-projected-gradient-method-with-divergences for a very clear exposition
U_current = psi(p.E) # map to dual space
U_new = (
U_current - gamma_E * grads_p.E
) # grads_p.E is correctly with respect to the primal parameters and evaluated on the primal parameters
E_new = phi(U_new) # map back to primal space
p = p.replace(E=E_new)
return p, mi_estimate
def scan_step(state, hp, gammas):
key, subkey_odors, subkey_loss = jax.random.split(state.key, 3)
draw_cs = ODOR_MODEL_REGISTRY[hp.odor_model]
cs = draw_cs(subkey_odors, hp)
if hp.loss == "jensen_shannon_loss":
p_new, Tp_new, mi_estimate = update_params_T(
state.p, state.Tp, hp, cs, subkey_loss, gammas
)
else:
p_new, mi_estimate = update_params(state.p, hp, cs, subkey_loss, gammas[0])
Tp_new = 0
new_metrics = {"mi": mi_estimate.astype(float)}
return TrainingState(p_new, Tp_new, key, new_metrics), new_metrics
# @partial(jax.jit, static_argnames=('hp', 'epochs_per_scan'))
def scan_phase(state, hp, gammas, epochs_per_scan):
"""Process multiple epochs in a single scanned phase"""
# Scan over epochs_per_scan steps
def scan_closure(
state, _
): # this is just a trick so you can pass hp into your scan_step.
state, new_metrics = scan_step(state, hp, gammas)
return state, new_metrics
final_state, metrics = jax.lax.scan(
scan_closure, state, length=epochs_per_scan
)
return final_state, metrics
def train_natural_gradient_scan_over_epochs(
initial_state, hp, gammas, scans=2, epochs_per_scan=2
):
num_total_epochs = scans * epochs_per_scan
metrics_store = {
"mi": jnp.zeros(num_total_epochs)
}
epoch_counter = 0
for scan_idx in range(scans):
state, scan_metrics = scan_phase(
initial_state, hp, gammas, epochs_per_scan
)
scan_metrics = jax.device_get(scan_metrics) # ensure it’s on host
mi_vals = scan_metrics["mi"] # shape: (epochs_per_scan,)
metrics_store["mi"] = metrics_store["mi"].at[epoch_counter:epoch_counter + epochs_per_scan].set(mi_vals)
epoch_counter += epochs_per_scan
if scan_idx % 1 == 0:
avg_mi = jnp.mean(jnp.array(scan_metrics["mi"]))
print(f"Scan {scan_idx}: Avg MI {avg_mi:.3f}")
initial_state = state
return initial_state, metrics_store
def make_constant_gammas(epochs_per_scan, scans, gamma_W, gamma_E, gamma_G, gamma_gain, gamma_kappa_inv, gamma_eta, gamma_T):
gamma_Ws = [gamma_W] * (epochs_per_scan * scans)
gamma_Es = [gamma_E] * (epochs_per_scan * scans)
gamma_Gs = [gamma_G] * (epochs_per_scan * scans)
gamma_gains = [gamma_gain] * (epochs_per_scan * scans)
gamma_kappa_invs = [gamma_kappa_inv] * (epochs_per_scan * scans)
gamma_etas = [gamma_eta] * (epochs_per_scan * scans)
gamma_Ts = [gamma_T] * (epochs_per_scan * scans)
return jnp.array([gamma_Ws, gamma_Es, gamma_Gs, gamma_kappa_invs, gamma_etas, gamma_Ts]).T
NL_REGISTRY = {
"sigmoid log": sigmoid_log,
"hill function": hill_function,
"relu": jax.nn.relu,
"tanh": jax.nn.tanh,
"id": lambda x: x,
}
ACTIVITY_FUNCTION_REGISTRY = {
"linear filter": compute_linear_filter_activity,
"antagonism": compute_osn_firing_with_antagonism,
"glomerular convergence": linear_filter_plus_glomerular_layer
}
PHI_PSI_REGISTRY = {
"softplus_phi": lambda u: jax.nn.softmax(
u, axis=1
), # rows of expression matrix should sum to 1.
"softplus_psi": lambda x: jnp.log(jnp.clip(x, min=1e-6, max=1 - 1e-6)),
"bounded_phi": lambda u: 1 / (1 + jnp.exp(-u)),
"bounded_psi": lambda x: jnp.log(
jnp.clip(x, min=1e-6, max=1 - 1e-6) / (1 - jnp.clip(x, 1e-6, 1 - 1e-6))
),
"positive_phi": jnp.exp,
"positive_psi": jnp.log,
"identity": lambda x: x
}
LOSS_REGISTRY = {
"jensen_shannon_loss": jensen_shannon_loss,
"entropy_loss": entropy_loss,
"log_det_loss": log_det_loss,
}
ODOR_MODEL_REGISTRY = {
"data-driven binary": draw_cs_data_driven_binary,
"data-driven log normal": draw_cs_data_driven_log_normal,
"log normal": draw_cs_sparse_log_normal,
}

model.py at commit 8c23d3a, under MIT · at the source

Overview

Authors: Juan Carlos Fernández del Castillo^1,2, Farhad Pashakhanloo¹, Venkatesh N Murthy^1,3,4, Jacob A Zavatone-Veth^1,5

ORCID iDs: Juan Carlos Fernández del Castillo, Farhad Pashakhanloo, Venkatesh N Murthy, Jacob A Zavatone-Veth

Center for Brain Science, Harvard University, Cambridge, MA 02138
Program in Biophysics, Harvard University, Cambridge, MA 02138
Department of Molecular and Cellular Biology, Harvard University, Cambridge, MA 02138
Kempner Institute for the Study of Natural and Artificial Intelligence, Harvard University, Cambridge, MA 02138
Society of Fellows, Harvard University, Cambridge, MA 02138

Institutions: Harvard University (United States)

Journal: Proceedings of the National Academy of Sciences of the United States of America, volume 123, issue 13, article e2524661123

Dates: received 4 September 2025; accepted 23 February 2026; published online 24 March 2026; in print 31 March 2026

Type: Research article · Language: English

License: CC BY-NC-ND

Identifiers: DOI 10.1073/pnas.2524661123 · PMID 41875150 · PMCID PMC13037853 · OpenAlex W7140220544

Open access: hybrid, a free copy (OpenAlex)

Status: code verified (how the code is found and verified)

Categories: mouse (organism), cellular / molecular (subfield)

Keywords: olfaction, efficient coding, convergent evolution

MeSH: Olfactory Pathways*, Olfactory Receptor Neurons*, Smell*, Animals, Biological Evolution, Mice, Models, Neurological, Receptors, Odorant (* major topic)

Topic: Olfactory and Sensory Function Studies (Sensory Systems, Neuroscience), according to OpenAlex

Funding: Harvard | Society of Fellows, Harvard University (Society of Fellows) (N/A); NTT Research (N/A, A47994); NIGMS NIH HHS (T32 GM158477); NIH HHS (DP5 OD037354); HHS | National Institutes of Health (DP5OD037354, T32GM158477)

Citations: cited by 2 papers (Europe PMC); 103 references in the paper

Abstract

The abstract is not reproduced here: the paper's license (CC BY-NC-ND) does not allow it. Read it in the paper, at the publisher or on Europe PMC.

Repositories

Its files are read in the Code ↔ Paper reader above, with 3 matches between paragraphs and lines of code.

VNMurthyLab/olfactory-ec

License: MIT

State: the link answers, verified on 30 September 2026

Evidence: files inventoried

Commit: 8c23d3af55f7e7573bb34483ab8ddbb229477c1f, 24 March 2026

Languages: Python (27), Shell (1)

Size: 2,472 files, 28 scripts

Software Heritage: not archived

Found in: “Data, Materials, and Software Availability”

Holds: README, license file

Not found: CITATION.cff, environment file, tests, continuous integration, documentation

Tools: JAX (27 files), Matplotlib (27 files), NumPy (7 files), SciPy (5 files)

Availability: 1 check, the latest on 30 September 2026: the link answers

30 September 2026: the link answers

30 files

configs/glomerular_convergence/update.sh, Shell, 37 lines
figures/E_vs_W_opt_or_shuffle.py, Python, 416 lines
figures/E_vs_environment_phase_diagram.py, Python, 292 lines
figures/W_distributions.py, Python, 120 lines
figures/W_histograms_vs_M.py, Python, 310 lines
figures/adult_fly_receptor_heatmap.py, Python, 218 lines
figures/downsample_data.py, Python, 53 lines
figures/fly_larva_expression.py, Python, 132 lines
figures/fly_larva_receptor_heatmap.py, Python, 222 lines
figures/glomerular_convergence.py, Python, 273 lines
figures/inspect_G.py, Python, 161 lines
figures/inspect_W_opt.py, Python, 235 lines
figures/inspect_r.py, Python, 138 lines
figures/inspect_receptors.py, Python, 284 lines
figures/neural_noise_sweep.py, Python, 97 lines
figures/odorant_model_schematic.py, Python, 80 lines
figures/plot_utils.py, Python, 10 lines
figures/receptor_figure.py, Python, 511 lines
figures/receptor_softrank_vs_M.py, Python, 330 lines
figures/summary_stats.py, Python, 66 lines
figures/varying_W_model_figure.py, Python, 133 lines
figures/varying_hill_coef.py, Python, 141 lines
model.py, Python, 682 lines, 1 match
plotting.py, Python, 390 lines
run_optimization.py, Python, 167 lines
sample_generation/generate_binarized_gaussian_samples.py, Python, 171 lines, 1 match
shuffle_and_fit_W.py, Python, 159 lines
vary_W_model.py, Python, 334 lines
LICENSE, License, 21 lines
README.md, Text, 17 lines

At the source: github.com/VNMurthyLab/olfactory-ec

samuellab/larval-orn

License: MIT

State: the link answers, verified on 30 September 2026

Evidence: files inventoried

Commit: 1618850795214b9d1649caccc5ab96fbc9ada7ba, 31 March 2020

Languages: MATLAB (39)

Size: 102 files, 39 scripts

Software Heritage: not archived

Found in: “Data, Materials, and Software Availability”

Holds: README, license file

Not found: CITATION.cff, environment file, tests, continuous integration, documentation

Tools: Curve Fitting Toolbox (4 files), Statistics and Machine Learning Toolbox (4 files)

Availability: 1 check, the latest on 30 September 2026: the link answers

30 September 2026: the link answers

41 files

Figure2/Figure2A_FigureS3C.m, MATLAB, 94 lines
Figure2/Figure2B.m, MATLAB, 45 lines
Figure2/FigureS3A.m, MATLAB, 178 lines
Figure2/FigureS3B.m, MATLAB, 39 lines
Figure2/FigureS3D.m, MATLAB, 90 lines
Figure2/functions/FindColorMap.m, MATLAB, 15 lines
Figure2/functions/GetBestOrder.m, MATLAB, 138 lines
Figure2/functions/LoadData.m, MATLAB, 86 lines
Figure3/Figure3A_S4A.m, MATLAB, 254 lines
Figure3/Figure3B_3C_3D.m, MATLAB, 97 lines
Figure3/Figure3E.m, MATLAB, 65 lines
Figure3/FigureS4B_S4C.m, MATLAB, 119 lines
Figure3/FigureS4D_G.m, MATLAB, 14 lines
Figure3/FigureS4H.m, MATLAB, 52 lines
Figure3/FigureS4I.m, MATLAB, 92 lines
Figure3/functions/3rdPartyCodes/powerlaws_full_v0.0.10-2012-01-17/deta.m, MATLAB, 198 lines
Figure3/functions/3rdPartyCodes/powerlaws_full_v0.0.10-2012-01-17/plfit.m, MATLAB, 284 lines
Figure3/functions/3rdPartyCodes/powerlaws_full_v0.0.10-2012-01-17/plplot.m, MATLAB, 100 lines
Figure3/functions/3rdPartyCodes/powerlaws_full_v0.0.10-2012-01-17/plpva.m, MATLAB, 318 lines
Figure3/functions/3rdPartyCodes/powerlaws_full_v0.0.10-2012-01-17/plvar.m, MATLAB, 286 lines
Figure3/functions/3rdPartyCodes/powerlaws_full_v0.0.10-2012-01-17/randht.m, MATLAB, 111 lines, 1 match
Figure3/functions/3rdPartyCodes/powerlaws_full_v0.0.10-2012-01-17/zeta.m, MATLAB, 42 lines
Figure3/functions/EnsembleMiniSearch.m, MATLAB, 42 lines
Figure3/functions/FitMLE_global.m, MATLAB, 67 lines
Figure3/functions/FitMLE_local.m, MATLAB, 53 lines
Figure3/functions/MLEScript.m, MATLAB, 46 lines
Figure3/functions/SetCurveType.m, MATLAB, 351 lines
Figure3/functions/invMLEFit.m, MATLAB, 47 lines
Figure3/functions/likelihood_fun_global.m, MATLAB, 32 lines
Figure3/functions/likelihood_fun_inv_curve.m, MATLAB, 13 lines
Figure3/functions/likelihood_fun_local.m, MATLAB, 17 lines
Figure3/functions/plotInvDistributions.m, MATLAB, 19 lines
Figure3/functions/prep_ds_data.m, MATLAB, 51 lines
Figure3/functions/prob_den.m, MATLAB, 76 lines
Figure4/Figure4.m, MATLAB, 134 lines
Figure5/Figure5A_5B_S5A_S5B.m, MATLAB, 198 lines
Figure5/functions/getfilter2.m, MATLAB, 31 lines
Figure5/functions/myRescale.m, MATLAB, 29 lines
Figure5/functions/objfunFigS5B.m, MATLAB, 7 lines
LICENSE, License, 21 lines
README.md, Text, 49 lines

At the source: github.com/samuellab/Larval-ORN/blob/master/Figure3/results/MLEFit.mat

The paper's code and data availability statement is in the Data section.

Tracing map

Proposed by the machine: these links were found in the paper and verified at the source, without human review. The map will receive a Zenodo DOI once one of the paper's authors has validated it with their ORCID.

What the map holds:

2 repositories of the authors' code, each at its verified commit, with its license and how the link was found in the paper;
67 scripts, each with its path and the digest of its content;
3 matches between paragraphs of the paper and lines of the code (method lexical-v1);
neither the text of the paper nor the code itself.

Its JSON (tracing-map.json) is deposited on Zenodo with its DOI once the map is validated.

Data

No dataset and no data link were found in the paper.

Code and data availability statement

The paper has a code and data availability statement. Its license (CC BY-NC-ND) does not allow reproducing it here; in short, from what the harvester recognized in it:

it points to the authors' code: samuellab/larval-orn, VNMurthyLab/olfactory-ec

Read it in the paper: doi.org/10.1073/pnas.2524661123.

Versions

The history of this record: each version stored by the harvester or made by a correction of its authors or of the maintainers of its code, and what changed in its facts. The texts of the paper (its abstract, its availability statements) are not part of it; versions that changed only those are not listed.

Version 1, 30 September 2026: the first record

Recorded: type, language, journal, volume, issue, pages, dates, 4 authors, 3 keywords, 8 MeSH terms, 5 funders, 91 references.

Cite

This paper

del Castillo, J. C. F., Pashakhanloo, F., Murthy, V. N., & Zavatone-Veth, J. A. (2026). Convergent motifs of early olfactory processing are recapitulated by layer-wise efficient coding. Proceedings of the National Academy of Sciences of the United States of America, 123(13), e2524661123. https://doi.org/10.1073/pnas.2524661123

BibTeX

@article{delcastillo2026convergent, author = {del Castillo, Juan Carlos Fernández and Pashakhanloo, Farhad and Murthy, Venkatesh N and Zavatone-Veth, Jacob A}, title = {{Convergent motifs of early olfactory processing are recapitulated by layer-wise efficient coding}}, journal = {Proceedings of the National Academy of Sciences of the United States of America}, year = {2026}, month = mar, volume = {123}, number = {13}, pages = {e2524661123}, publisher = {National Academy of Sciences}, issn = {0027-8424}, doi = {10.1073/pnas.2524661123}, url = {https://doi.org/10.1073/pnas.2524661123}, pmid = {41875150}, pmcid = {PMC13037853} }

RIS

TY - JOUR AU - del Castillo, Juan Carlos Fernández AU - Pashakhanloo, Farhad AU - Murthy, Venkatesh N AU - Zavatone-Veth, Jacob A TI - Convergent motifs of early olfactory processing are recapitulated by layer-wise efficient coding T2 - Proceedings of the National Academy of Sciences of the United States of America J2 - Proc Natl Acad Sci U S A PY - 2026 DA - 2026/03/24 VL - 123 IS - 13 SP - e2524661123 SN - 0027-8424 PB - National Academy of Sciences DO - 10.1073/pnas.2524661123 UR - https://doi.org/10.1073/pnas.2524661123 LA - en ER -

CSL-JSON

{ "id": "10.1073/pnas.2524661123", "type": "article-journal", "title": "Convergent motifs of early olfactory processing are recapitulated by layer-wise efficient coding", "container-title": "Proceedings of the National Academy of Sciences of the United States of America", "author": [ { "family": "del Castillo", "given": "Juan Carlos Fernández" }, { "family": "Pashakhanloo", "given": "Farhad" }, { "family": "Murthy", "given": "Venkatesh N" }, { "family": "Zavatone-Veth", "given": "Jacob A" } ], "container-title-short": "Proc Natl Acad Sci U S A", "volume": "123", "issue": "13", "page": "e2524661123", "DOI": "10.1073/pnas.2524661123", "PMID": "41875150", "PMCID": "PMC13037853", "ISSN": "0027-8424", "publisher": "National Academy of Sciences", "URL": "https://doi.org/10.1073/pnas.2524661123", "language": "en", "issued": { "date-parts": [ [ 2026, 3, 24 ] ] } }

The tracing map gets a citation of its own once an author has validated it and it has a DOI.

Contribute

The authors of this paper can claim it, correct its record and validate its tracing map, and the maintainers of its code (its owner, or a public member of its organization) correct what it says of their repository; anyone signed in can ask for its removal. Every request goes to OSCR's own machine, which answers it; your account page follows them.

Sign in with ORCID to claim this paper as one of its authors, correct its record or validate its tracing map: when the paper's metadata lists your ORCID iD, you are recognized at once. Maintainers of its code: sign in with GitHub, then claim the repository on your account page.

Add the badge to the code's README

If you are an author or a maintainer of this paper's code, you can add this badge. It says what OSCR knows of the paper (its code, its validated tracing map) and links the code to this page. Paste one of these snippets under the README's title. Keep the two marks around it, so that the badge can be replaced in place and is never added twice. OSCR changes nothing in your repository.

Markdown

<!-- oscr-badge:start -->
[![OSCR](https://openscicode.org/badge/doi/10.1073/pnas.2524661123.svg)](https://openscicode.org/badge/doi/10.1073/pnas.2524661123)
<!-- oscr-badge:end -->

reStructuredText

.. oscr-badge:start

.. image:: https://openscicode.org/badge/doi/10.1073/pnas.2524661123.svg
   :target: https://openscicode.org/badge/doi/10.1073/pnas.2524661123
   :alt: OSCR

.. oscr-badge:end

HTML

<!-- oscr-badge:start -->
<a href="https://openscicode.org/badge/doi/10.1073/pnas.2524661123"><img src="https://openscicode.org/badge/doi/10.1073/pnas.2524661123.svg" alt="OSCR"></a>
<!-- oscr-badge:end -->

On GitHub, open the README, choose Edit (the pencil), paste the snippet under the title, then "Commit changes".

If you choose to, you can have the badge kept up to date for you, with oscr badge --write from the command line or with the GitHub Action in your repository's workflows. Every other form and option (a card, AsciiDoc, LaTeX, shields.io) is on the badge page.

Request its removal

To ask OSCR to remove this record, the copies of its authors' scripts or its tracing map, use the removal request page: signed in, you say who you are, what to remove and why, then review and confirm the request. Published rules decide every request (how).

Discussion, reproductions, activity

Discussion: code errors and mismatches between the code and the paper about this paper and its code, from readers and its authors, beside the paper and its tracing map.

Report a mismatch between the code and the paper · Report a code error (signed in)

Reproductions: reports from readers who ran the authors' code and did not get the paper's result, what they ran, with which environment, commit and data.

Report a reproduction failure (signed in)

Versions of its code: the releases of its code tied to a version of this paper (the preprint, the accepted manuscript, the version of record…), each with the tracing map versioned with it.

Activity: what happens around this paper: new versions of its record, its map's validation, discussions and reproductions. It opens with sign-in.