def test_splinelg_mle_small_rf():
w_true, X, y, dims, dt = generate_2d_rf_data(noise='white')
df = [3, 4]
model = splineLG(X, y, dims=dims, dt=dt, df=df, compute_mle=True)
assert mse(uvec(model.w_mle), uvec(w_true.flatten())) < 1e-1
def fit_nonparametric_nonlinearity(self, nbins=50, w=None):
if w is None:
if self.w_spl is not None:
w = self.w_spl.flatten()
elif self.w_mle is not None:
w = self.w_mle.flatten()
elif self.w_sta is not None:
w = self.w_sta.flatten()
else:
w = jnp.array(w)
X = self.X
X = X.reshape(X.shape[0], -1)
y = self.y
output_raw = X @ uvec(w)
output_spk = X[y != 0] @ uvec(w)
hist_raw, bins = jnp.histogram(output_raw, bins=nbins, density=True)
hist_spk, _ = jnp.histogram(output_spk, bins=bins, density=True)
mask = ~(hist_raw == 0)
yy0 = hist_spk[mask] / hist_raw[mask]
self.nl_bins = bins[1:]
self.fnl_nonparametric = interp1d(bins[1:][mask], yy0)
def test_splinelg_small_rf():
w_true, X, y, dims, dt = generate_2d_rf_data(noise='white')
df = [3, 4]
model = splineLG(X, y, dims=dims, dt=dt, df=df)
model.fit(metric='corrcoef', num_iters=100, verbose=0, tolerance=10, beta=0.01)
assert mse(uvec(model.w_opt), uvec(w_true.flatten())) < 1e-1
def test_glm_3d_outputnl():
w_true, X, y, dims, dt = generate_3d_rf_data(noise='white')
df = tuple([int(np.maximum(dim / 3, 3)) for dim in dims])
model = GLM(distr='gaussian', output_nonlinearity='exponential')
model.add_design_matrix(X, dims=dims, df=df, smooth='cr', kind='train', filter_nonlinearity='none',
name='stimulus')
model.initialize(num_subunits=1, dt=dt, method='mle', random_seed=42, compute_ci=False, y=y)
model.fit(y={'train': y}, num_iters=300, verbose=0, step_size=0.03, beta=0.001, metric='corrcoef')
assert model.score(X, y, metric='corrcoef') > 0.4
assert mse(uvec(model.w['opt']['stimulus']), uvec(w_true.flatten())) < 0.01
def test_asd_small_rf():
w_true, X, y, dims, dt = generate_2d_rf_data(noise='white')
model = ASD(X, y, dims=dims)
model.fit(p0=[
1.,
1.,
6.,
6.,
], num_iters=10, verbose=10)
w_fit = model.optimized_C_post @ X.T @ y / model.optimized_params[0]**2
assert mse(uvec(w_fit), uvec(w_true.flatten())) < 1e-1
def V1complex_2d(dims=(30, 40), scale=(.025, .03)):
dt = 1 / 60 # time bin size
nt = dims[0]
nx = dims[1]
tt = np.arange(-nt * dt, 0, dt)
kt1 = scipy.stats.gamma.pdf(-tt, dims[0] / 7.5, scale=scale[0])
kt2 = scipy.stats.gamma.pdf(-tt, dims[1] / 6, scale=scale[1])
kt1 /= np.linalg.norm(kt1)
kt2 /= -np.linalg.norm(kt2)
kt = np.vstack([kt1, kt2]).T
xx = np.linspace(-2, 2, nx)
kx1 = np.cos(2 * np.pi * xx / 2 + np.pi / 5) * np.exp(
-1 / (2 * 0.35**2) * xx**2)
kx2 = np.sin(2 * np.pi * xx / 2 + np.pi / 5) * np.exp(
-1 / (2 * 0.35**2) * xx**2)
kx1 /= np.linalg.norm(kx1)
kx2 /= np.linalg.norm(kx2)
kx = np.vstack([kx1, kx2])
k = kt @ kx
return uvec(k)
def bs(x, df, degree=3):
"""
B-spline basis. Knots placed equally by percentile.
Simplified from `patsy.bs`:
https://github.com/pydata/patsy/blob/master/patsy/splines.py
"""
from scipy.interpolate import BSpline
def _get_all_sorted_knots(_x, _df, _degree):
order = _degree + 1
n_inner_knots = _df - order
knot_quantiles = np.linspace(0, 1, n_inner_knots + 2)[1:-1] * 100
inner_knots = np.percentile(_x, knot_quantiles)
all_knots = np.hstack(([np.min(_x), np.max(_x)] * order, inner_knots))
all_knots = np.sort(all_knots)
return all_knots
x = np.asarray(x)
df = np.asarray(df)
knots = _get_all_sorted_knots(x, df, degree)
n_bases = len(knots) - (degree + 1)
coeff = np.eye(n_bases)
basis = np.vstack(
[BSpline(knots, coeff[i], degree)(x) for i in range(n_bases)]).T
return uvec(basis)
def test_ald_small_rf():
w_true, X, y, dims, dt = generate_2d_rf_data(noise='white')
sigma0 = [1.3]
rho0 = [0.8]
params_t0 = [3., 20., 3., 20.9] # taus, nus, tauf, nuf
params_y0 = [3., 20., 3., 20.9]
p0 = sigma0 + rho0 + params_t0 + params_y0
model = ALD(X, y, dims=dims)
model.fit(p0=p0, num_iters=30, verbose=10)
w_fit = model.optimized_C_post @ X.T @ y / model.optimized_params[0]**2
assert mse(uvec(w_fit), uvec(w_true.flatten())) < 1e-1
def test_glm_3d_history():
w_true, X, y, dims, dt = generate_3d_rf_data(noise='white')
(X_train, y_train), (X_dev, y_dev), (_, _) = split_data(X, y, dt, frac_train=0.8, frac_dev=0.2)
df = tuple([int(np.maximum(dim / 3, 3)) for dim in dims])
model = GLM(distr='gaussian', output_nonlinearity='none')
model.add_design_matrix(X_train, dims=dims, df=df, smooth='cr', kind='train', filter_nonlinearity='none',
name='stimulus')
model.add_design_matrix(y_train, dims=[5], df=[3], smooth='cr', kind='train', filter_nonlinearity='none',
name='history')
model.add_design_matrix(X_dev, dims=dims, df=df, name='stimulus', kind='dev')
model.add_design_matrix(y_dev, dims=[5], df=[3], kind='dev', name='history')
model.initialize(num_subunits=1, dt=dt, method='mle', random_seed=42, compute_ci=False, y=y_train)
model.fit(
y={'train': y_train, 'dev': y_dev}, num_iters=200, verbose=100, step_size=0.03, beta=0.001, metric='corrcoef')
assert model.score({"stimulus": X_train, 'history': y_train}, y_train, metric='corrcoef') > 0.6
assert model.score({"stimulus": X_dev, 'history': y_dev}, y_dev, metric='corrcoef') > 0.4
assert mse(uvec(model.w['opt']['stimulus']), uvec(w_true.flatten())) < 0.01
def gabor2d(dims=(25, 25), omega=0.5, theta=np.pi / 6, func=np.cos, K=1.):
radius = (int(dims[1] / 2.0), int(dims[0] / 2.0))
[x, y] = np.meshgrid(range(-radius[0], radius[0] + 1),
range(-radius[1], radius[1] + 1))
x1 = x * np.cos(theta) + y * np.sin(theta)
y1 = -x * np.sin(theta) + y * np.cos(theta)
gauss = omega**2 / (4 * np.pi * K**2) * np.exp(-omega**2 / (8 * K**2) *
(4 * x1**2 + y1**2))
sinusoid = func(omega * x1) * np.exp(K**2 / 2)
gabor = gauss * sinusoid
return uvec(gabor)[:dims[0], :dims[1]]
def tp(x, df):
"""
Simplified implementation of the truncated Thin Plate (TP) regression spline.
See Wood, S. (2017) p.216-217
"""
def eta(r):
return r**2 * np.log(r + 1e-10)
E = eta(np.abs(x.ravel() - x.ravel().reshape(-1, 1)))
U, _, _ = np.linalg.svd(E)
basis = U[:, :int(df)]
return uvec(basis)
def cr(x, df):
"""
Natural cubic regression splines. Knots placed equally by percentile.
Simplified from `patsy.cr`:
https://github.com/pydata/patsy/blob/master/patsy/mgcv_cubic_splines.py
"""
def _get_all_sorted_knots(_x, _df):
n_inner_knots = _df - 2
knot_quantiles = np.linspace(0, 1, n_inner_knots + 2)[1:-1] * 100
inner_knots = np.percentile(np.unique(_x), knot_quantiles)
all_knots = np.concatenate(([np.min(_x), np.max(_x)], inner_knots))
all_knots = np.unique(all_knots)
return all_knots
knots = _get_all_sorted_knots(x, df)
j = np.maximum(np.searchsorted(knots, x) - 1, 0)
h = np.mean(knots[1:] - knots[:-1]) # constant
ajm = (knots[j + 1] - x) / h
ajp = (x - knots[j]) / h
cjm = ((knots[j + 1] - x)**3 / h - h * (knots[j + 1] - x)) / 6
cjp = ((x - knots[j])**3 / h - h * (x - knots[j])) / 6
B = np.sum([
np.diag([1, 2, 1][i] * h * np.ones(df - [3, 2, 3][i]), [-1, 0, 1][i]) /
[6, 3, 6][i] for i in range(3)
], 0)
D = np.sum([[1, -2, 1][i] * np.pad(
np.eye(df - 2) / h, pad_width=((0, 0), (i, 2 - i)), mode='constant')
for i in range(3)], 0)
f = np.vstack([np.zeros(df), np.linalg.solve(B, D), np.zeros(df)])
i = np.eye(df)
basis = ajm * i[j, :].T + ajp * i[j + 1, :].T + cjm * f[j, :].T + cjp * f[
j + 1, :].T
return uvec(basis.T)
def subunits2d(num_subunits=5,
dims=(25, 25),
std=(3, 3),
offset=(3, 3),
kind='gaussian',
random_seed=2046):
if kind == 'gaussian':
filter2d = gaussian2d
elif kind == 'mexicanhat':
filter2d = mexicanhat2d
else:
raise NotImplementedError(kind)
w = np.zeros(list(dims) + [num_subunits])
f = filter2d(dims, std)
for i in range(num_subunits):
if random_seed is not None:
np.random.seed(random_seed + i + 5)
h, v = np.random.randint(low=-offset[0], high=offset[1], size=2)
w[:, :, i] = uvec(np.roll(np.roll(f, h, axis=0), v, axis=1))
return w
def fit_STC(self,
prewhiten=False,
n_repeats=10,
percentile=100.,
random_seed=2046,
verbose=5):
"""
Spike-triggered Covariance Analysis.
Parameters
==========
prewhiten: bool
n_repeats: int
Number of repeats for STC significance test.
percentile: float
Valid range of STC significance test.
verbose: int
random_seed: int
"""
def get_stc(_X, _y, _w):
n = len(_X)
ste = _X[_y != 0]
proj = ste - ste * _w * _w.T
stc = proj.T @ proj / (n - 1)
_eigvec, _eigval, _ = jnp.linalg.svd(stc)
return _eigvec, _eigval
key = random.PRNGKey(random_seed)
y = self.y
if prewhiten:
if self.compute_mle is False:
self.XtX = self.X.T @ self.X
self.w_mle = jnp.linalg.solve(self.XtX, self.XtY)
X = jnp.linalg.solve(self.XtX, self.X.T).T
w = uvec(self.w_mle)
else:
X = self.X
w = uvec(self.w_sta)
eigvec, eigval = get_stc(X, y, w)
self.w_stc = dict()
if n_repeats:
print('STC significance test: ')
eigval_null = []
for counter in range(n_repeats):
if verbose:
if counter % int(verbose) == 0:
print(f' {counter + 1:}/{n_repeats}')
y_randomize = random.permutation(key, y)
_, eigval_randomize = get_stc(X, y_randomize, w)
eigval_null.append(eigval_randomize)
else:
if verbose:
print(f'Done.')
eigval_null = jnp.vstack(eigval_null)
max_null, min_null = jnp.percentile(eigval_null,
percentile), jnp.percentile(
eigval_null,
100 - percentile)
mask_sig_pos = eigval > max_null
mask_sig_neg = eigval < min_null
mask_sig = jnp.logical_or(mask_sig_pos, mask_sig_neg)
self.w_stc['eigvec'] = eigvec
self.w_stc['pos'] = eigvec[:, mask_sig_pos]
self.w_stc['neg'] = eigvec[:, mask_sig_neg]
self.w_stc['eigval'] = eigval
self.w_stc['eigval_mask'] = mask_sig
self.w_stc['eigval_pos_mask'] = mask_sig_pos
self.w_stc['eigval_neg_mask'] = mask_sig_neg
self.w_stc['max_null'] = max_null
self.w_stc['min_null'] = min_null
else:
self.w_stc['eigvec'] = eigvec
self.w_stc['eigval'] = eigval
self.w_stc['eigval_mask'] = jnp.ones_like(eigval).astype(bool)
def mexicanhat1d(dims=200, std=15., a=0.8):
g0 = gaussian1d(dims, std)
g1 = gaussian1d(dims, std * a)
m = g1 - 0.65 * g0
return uvec(m)
def mexicanhat2d(dims=(25, 25), std=(3., 3.), a=0.55):
g0 = gaussian2d(dims, std)
g1 = gaussian2d(dims, np.array(std) * a)
m = g1 - 0.65 * g0
return uvec(m)
def gaussian1d(dim=200, std=15.):
return uvec(scipy.signal.gaussian(dim, std=std))
def cc(x, df):
"""
Cyclic cubic regression splines. Knots placed equally by percentile.
Simplified from `patsy.cc`:
https://github.com/pydata/patsy/blob/master/patsy/mgcv_cubic_splines.py
"""
def _map_cyclic(_x, lbound, ubound):
_x = np.copy(_x)
_x[_x > ubound] = lbound + (_x[_x > ubound] - ubound) % (ubound -
lbound)
_x[_x < lbound] = ubound - (lbound - _x[_x < lbound]) % (ubound -
lbound)
return _x
def _get_all_sorted_knots(_x, _df):
n_inner_knots = _df - 2
knot_quantiles = np.linspace(0, 1, n_inner_knots + 2)[1:-1] * 100
inner_knots = np.percentile(np.unique(_x), knot_quantiles)
all_knots = np.concatenate(([np.min(_x), np.max(_x)], inner_knots))
all_knots = np.unique(all_knots)
return all_knots
def _get_cyclic_f(_knots):
kd = _knots[1:] - _knots[:-1]
n = _knots.size - 1
b = np.zeros((n, n))
d = np.zeros((n, n))
b[0, 0] = (kd[n - 1] + kd[0]) / 3.
b[0, n - 1] = kd[n - 1] / 6.
b[n - 1, 0] = kd[n - 1] / 6.
d[0, 0] = -1. / kd[0] - 1. / kd[n - 1]
d[0, n - 1] = 1. / kd[n - 1]
d[n - 1, 0] = 1. / kd[n - 1]
for i in range(1, n):
b[i, i] = (kd[i - 1] + kd[i]) / 3.
b[i, i - 1] = kd[i - 1] / 6.
b[i - 1, i] = kd[i - 1] / 6.
d[i, i] = -1. / kd[i - 1] - 1. / kd[i]
d[i, i - 1] = 1. / kd[i - 1]
d[i - 1, i] = 1. / kd[i - 1]
return np.linalg.solve(b, d)
knots = _get_all_sorted_knots(x, df) # length = df
x = _map_cyclic(x, min(knots), max(knots))
df -= 1
j = np.maximum(np.searchsorted(knots, x) - 1, 0)
h = np.mean(knots[1:] - knots[:-1]) # constant
ajm = (knots[j + 1] - x) / h
ajp = (x - knots[j]) / h
cjm = ((knots[j + 1] - x)**3 / h - h * (knots[j + 1] - x)) / 6
cjp = ((x - knots[j])**3 / h - h * (x - knots[j])) / 6
f = _get_cyclic_f(knots)
i = np.eye(df)
j1 = j + 1
j1[j1 == df] = 0
basis = ajm * i[j, :].T + ajp * i[j1, :].T + cjm * f[j, :].T + cjp * f[
j1, :].T
return uvec(basis.T)
def test_lnp_mle_small_rf():
w_true, X, y, dims, dt = generate_2d_rf_data(noise='white')
model = LNP(X, y, dims=dims, dt=dt, compute_mle=True)
assert mse(uvec(model.w_mle), uvec(w_true.flatten())) < 1e-1
def gabor3d(dims, std, omega, theta, func=np.cos, K=np.pi):
g_t = np.gradient(gaussian1d(dims[0], std))
g_s = gabor2d(dims[1:], omega, theta, func, K).flatten()
g = np.kron(g_t, g_s)
return uvec(g).reshape(dims)
def gaussian3d(dims, std):
gaussian_t = np.gradient(gaussian1d(dims[0], std[0]))
gaussian_s = gaussian2d(dims[1:], std[1:]).flatten()
return uvec(np.kron(gaussian_t, gaussian_s)).reshape(dims)
def mexicanhat3d(dims, std, a=0.3):
g_t = np.gradient(gaussian1d(dims[0], std[0]))
m_s = mexicanhat2d(dims[1:], std[1:], a).flatten()
m = np.kron(g_t, m_s)
return uvec(m).reshape(dims)
def build_spline_matrix(dims, df, smooth, dtype=np.float64):
"""
Building spline matrix for n-dimensional RF (n=[1,2,3]) with tensor product smooth.
Parameters
==========
dims : list or array_like, shape (d, )
Dimensions or shape of the RF to estimate. Assumed order [t, sx, sy]
df : list or array_like, shape (d,)
Degree of freedom for spline / smooth basis. Same length as dims.
smooth : str
Spline or smooth to be used. Current supported methods include:
* `bs`: B-spline
* `cr`: Cubic Regression spline
* `tp`: (Simplified) Thin Plate regression spline
dtype : dtype
Data type S will be cast to before returning
Return
======
S : array_like, shape (n_features, n_spline_coef)
Spline matrix. Each column is one basis.
Note
====
---outdated
A mesh-free (actually simpler) way to do this is to get the spline bases for each dimension,
then calculate the kronecker product of them, for example:
>>> dims, df = (10, 10, 10), (3, 3, 3)
>>> St, Sy, Sx = [basis(np.arange(d), f), for d, f in zip(dims, df)]
>>> S = np.kron(St, np.kron(Sy, Sx))
Here we use a mesh-based `te` approach to keep consistent with the Patsy inplementation / Wood, S. (2017).
----
Now we switched to the mesh-free inmplemtation.
"""
ndim = len(dims) # store RF dimemsion
# initialize list of degree of freedom for each dimension
if len(df) != ndim:
raise ValueError("`df` must have the same length as `dims`")
if smooth == 'cr':
basis = cr # Natural cubic regression spline
elif smooth == 'cc':
basis = cc # cyclic cubic regression spline
elif smooth == 'bs':
basis = bs # b-spline
elif smooth == 'tp':
basis = tp # thin-plate spline
else:
raise ValueError("Input method `{}` is not supported.".format(smooth))
# build spline matrix
if ndim == 1:
g0 = np.arange(dims[0])
S = basis(g0.ravel(), df[0])
elif ndim == 2:
g0 = np.arange(dims[0])
g1 = np.arange(dims[1])
St = basis(g0.ravel(), df[0])
Sx = basis(g1.ravel(), df[1])
S = np.kron(St, Sx)
elif ndim == 3:
g0 = np.arange(dims[0])
g1 = np.arange(dims[1])
g2 = np.arange(dims[2])
St = basis(g0.ravel(), df[0])
Sx = basis(g1.ravel(), df[1])
Sy = basis(g2.ravel(), df[2])
S = np.kron(St, np.kron(Sx, Sy))
else:
raise NotImplementedError(ndim)
return uvec(S).astype(dtype)
def gaussian2d(dims=(25, 25), std=(2., 2.)):
gaussian_x = gaussian1d(dims[0], std=std[0])
gaussian_y = gaussian1d(dims[1], std=std[1])
return uvec(np.kron(gaussian_x, gaussian_y)).reshape(dims)
