"""
Module that hold techniques to project the brain volume onto 2D images for visualisation purposes
"""
from functools import lru_cache
import logging
import numpy as np
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt
import matplotlib.colors
from iblutil.numerical import ismember
from iblutil.util import Bunch
from iblutil.io.hashfile import md5
import one.remote.aws as aws
from ibllib.atlas.atlas import AllenAtlas, BrainRegions
_logger = logging.getLogger(__name__)
[docs]@lru_cache(maxsize=1, typed=False)
def circles(N=5, atlas=None, display='flat'):
"""
:param N: number of circles
:param atlas: brain atlas at 25 m
:param display: "flat" or "pyramid"
:return: 2D map of indices, ap_coordinate, ml_coordinate
"""
atlas = atlas if atlas else AllenAtlas()
sz = np.array([])
level = np.array([])
for k in np.arange(N):
nlast = 2000 # 25 um for 5mm diameter
n = int((k + 1) * nlast / N)
r = .4 * (k + 1) / N
theta = (np.linspace(0, 2 * np.pi, n) + np.pi / 2)
sz = np.r_[sz, r * np.exp(1j * theta)]
level = np.r_[level, theta * 0 + k]
atlas.compute_surface()
iy, ix = np.where(~np.isnan(atlas.top))
centroid = np.array([np.mean(iy), np.mean(ix)])
xlim = np.array([np.min(ix), np.max(ix)])
ylim = np.array([np.min(iy), np.max(iy)])
s = Bunch(
x=np.real(sz) * np.diff(xlim) + centroid[1],
y=np.imag(sz) * np.diff(ylim) + centroid[0],
level=level,
distance=level * 0,
)
# compute the overall linear distance for each circle
d0 = 0
for lev in np.unique(s['level']):
ind = s['level'] == lev
diff = np.abs(np.diff(s['x'][ind] + 1j * s['y'][ind]))
s['distance'][ind] = np.cumsum(np.r_[0, diff]) + d0
d0 = s['distance'][ind][-1]
fcn = interp1d(s['distance'], s['x'] + 1j * s['y'], fill_value='extrap')
d = np.arange(0, np.ceil(s['distance'][-1]))
s_ = Bunch({
'x': np.real(fcn(d)),
'y': np.imag(fcn(d)),
'level': interp1d(s['distance'], level, kind='nearest')(d),
'distance': d
})
if display == 'flat':
ih = np.arange(atlas.bc.nz)
iw = np.arange(s_['distance'].size)
image_map = np.zeros((ih.size, iw.size), dtype=np.int32)
iw, ih = np.meshgrid(iw, ih)
# i2d = np.ravel_multi_index((ih[:], iw[:]), image_map.shape)
iml, _ = np.meshgrid(np.round(s_.x).astype(np.int32), np.arange(atlas.bc.nz))
iap, idv = np.meshgrid(np.round(s_.y).astype(np.int32), np.arange(atlas.bc.nz))
i3d = atlas._lookup_inds(np.c_[iml.flat, iap.flat, idv.flat])
i3d = np.reshape(i3d, [atlas.bc.nz, s_['x'].size])
image_map[ih, iw] = i3d
elif display == 'pyramid':
for i in np.flipud(np.arange(N)):
ind = s_['level'] == i
dtot = s_['distance'][ind]
dtot = dtot - np.mean(dtot)
if i == N - 1:
ipx = np.arange(np.floor(dtot[0]), np.ceil(dtot[-1]) + 1)
nh = atlas.bc.nz * N
X0 = int(ipx[-1])
image_map = np.zeros((nh, ipx.size), dtype=np.int32)
iw = np.arange(np.sum(ind))
iw = np.int32(iw - np.mean(iw) + X0)
ih = atlas.bc.nz * i + np.arange(atlas.bc.nz)
iw, ih = np.meshgrid(iw, ih)
iml, _ = np.meshgrid(np.round(s_.x[ind]).astype(np.int32), np.arange(atlas.bc.nz))
iap, idv = np.meshgrid(np.round(s_.y[ind]).astype(np.int32), np.arange(atlas.bc.nz))
i3d = atlas._lookup_inds(np.c_[iml.flat, iap.flat, idv.flat])
i3d = np.reshape(i3d, [atlas.bc.nz, s_['x'][ind].size])
image_map[ih, iw] = i3d
x, y = (atlas.bc.i2x(s.x), atlas.bc.i2y(s.y))
return image_map, x, y
# if display == 'flat':
# fig, ax = plt.subplots(2, 1, figsize=(16, 5))
# elif display == 'pyramid':
# fig, ax = plt.subplots(1, 2, figsize=(14, 12))
# ax[0].imshow(ba._label2rgb(ba.label.flat[image_map]), origin='upper')
# ax[1].imshow(ba.top)
# ax[1].plot(centroid[1], centroid[0], '*')
# ax[1].plot(s.x, s.y)
[docs]def swanson(filename="swanson2allen.npz"):
# filename could be "swanson2allen_original.npz", or "swanson2allen.npz" for remapped indices to match
# existing labels in the brain atlas
OLD_MD5 = [
'bb0554ecc704dd4b540151ab57f73822', # version 2022-05-02 (remapped)
'7722c1307cf9a6f291ad7632e5dcc88b', # version 2022-05-09 (removed wolf pixels and 2 artefact regions)
]
npz_file = AllenAtlas._get_cache_dir().joinpath(filename)
if not npz_file.exists() or md5(npz_file) in OLD_MD5:
npz_file.parent.mkdir(exist_ok=True, parents=True)
_logger.info(f'downloading swanson image from {aws.S3_BUCKET_IBL} s3 bucket...')
aws.s3_download_file(f'atlas/{npz_file.name}', npz_file)
s2a = np.load(npz_file)['swanson2allen'] # inds contains regions ids
return s2a
[docs]def plot_swanson(acronyms=None, values=None, ax=None, hemisphere=None, br=None,
orientation='landscape', annotate=False, empty_color='silver', **kwargs):
"""
Displays the 2D image corresponding to the swanson flatmap.
This case is different from the others in the sense that only a region maps to another regions, there
is no correspondency from the spatial 3D coordinates.
:param acronyms:
:param values:
:param hemisphere: hemisphere to display, options are 'left', 'right', 'both' or 'mirror'
:param br: ibllib.atlas.BrainRegions object
:param ax: matplotlib axis object to plot onto
:param orientation: 'landscape' (default) or 'portrait'
:param annotate: (False) if True, labels regions with acronyms
:param empty_color: (grey) matplotlib color code or rgb_a int8 tuple defining the filling
of brain regions not provided. Defaults to 'silver'
:param kwargs: arguments for imshow
:return:
"""
mapping = 'Swanson'
br = BrainRegions() if br is None else br
br.compute_hierarchy()
s2a = swanson()
# both hemishpere
if hemisphere == 'both':
_s2a = s2a + np.sum(br.id > 0)
_s2a[s2a == 0] = 0
_s2a[s2a == 1] = 1
s2a = np.r_[s2a, np.flipud(_s2a)]
mapping = 'Swanson-lr'
elif hemisphere == 'mirror':
s2a = np.r_[s2a, np.flipud(s2a)]
if orientation == 'portrait':
s2a = np.transpose(s2a)
if acronyms is None:
regions = br.mappings[mapping][s2a]
im = br.rgba[regions]
iswan = None
else:
ibr, vals = br.propagate_down(acronyms, values)
# we now have the mapped regions and aggregated values, map values onto swanson map
iswan, iv = ismember(s2a, ibr)
im = np.zeros_like(s2a, dtype=np.float32)
im[iswan] = vals[iv]
im[~iswan] = np.nan
if not ax:
ax = plt.gca()
ax.set_axis_off() # unless provided we don't need scales here
ax.imshow(im, **kwargs)
# overlay the boundaries if value plot
imb = np.zeros((*s2a.shape[:2], 4), dtype=np.uint8)
# fill in the empty regions with the blank regions colours if necessary
if iswan is not None:
imb[~iswan] = (np.array(matplotlib.colors.to_rgba(empty_color)) * 255).astype('uint8')
imb[s2a == 0] = 255
# imb[s2a == 1] = np.array([167, 169, 172, 255])
imb[s2a == 1] = np.array([0, 0, 0, 255])
ax.imshow(imb)
if annotate:
annotate_swanson(ax=ax, orientation=orientation, br=br)
# provides the mean to see the region on axis
def format_coord(x, y):
ind = s2a[int(y), int(x)]
ancestors = br.ancestors(br.id[ind])['acronym']
return f'x={x:1.4f}, y={x:1.4f}, {br.acronym[ind]} \n {ancestors}'
ax.format_coord = format_coord
return ax
@lru_cache(maxsize=None)
def _swanson_labels_positions():
"""
This functions computes label positions to overlay on the Swanson flatmap
:return: dictionary where keys are acronyms
"""
NPIX_THRESH = 20000 # number of pixels above which region is labeled
s2a = swanson()
iw, ih = np.meshgrid(np.arange(s2a.shape[1]), np.arange(s2a.shape[0]))
# compute the center of mass of all regions (fast enough to do on the fly)
bc = np.maximum(1, np.bincount(s2a.flatten()))
cmw = np.bincount(s2a.flatten(), weights=iw.flatten()) / bc
cmh = np.bincount(s2a.flatten(), weights=ih.flatten()) / bc
bc[0] = 1
NWH, NWW = (200, 600)
h, w = s2a.shape
labels = {}
for ilabel in np.where(bc > NPIX_THRESH)[0]:
x, y = (cmw[ilabel], cmh[ilabel])
# the polygon is convex and the label is outside. Dammit !!!
if s2a[int(y), int(x)] != ilabel:
# find the nearest point to the center of mass
ih, iw = np.where(s2a == ilabel)
iimin = np.argmin(np.abs((x - iw) + 1j * (y - ih)))
# get the center of mass of a window around this point
sh = np.arange(np.maximum(0, ih[iimin] - NWH), np.minimum(ih[iimin] + NWH, h))
sw = np.arange(np.maximum(0, iw[iimin] - NWW), np.minimum(iw[iimin] + NWW, w))
roi = s2a[sh][:, sw] == ilabel
roi = roi / np.sum(roi)
# ax.plot(x, y, 'k+')
# ax.plot(iw[iimin], ih[iimin], '*k')
x = sw[np.searchsorted(np.cumsum(np.sum(roi, axis=0)), .5) - 1]
y = sh[np.searchsorted(np.cumsum(np.sum(roi, axis=1)), .5) - 1]
# ax.plot(x, y, 'r+')
labels[ilabel] = (x, y)
return labels
[docs]def annotate_swanson(ax, acronyms=None, orientation='landscape', br=None, **kwargs):
"""
Display annotations on the flatmap
:param ax:
:param acronyms: (None) list or np.array of acronyms or allen region ids. If None plot all.
:param orientation:
:param br: BrainRegions object
:param kwargs: arguments for the annotate function
:return:
"""
br = br or BrainRegions()
if acronyms is None:
indices = np.arange(br.id.size)
else: # tech debt: here in fact we should remap and compute labels for hierarchical regions
aids = br.parse_acronyms_argument(acronyms)
_, indices, _ = np.intersect1d(br.id, br.remap(aids, 'Swanson-lr'), return_indices=True)
labels = _swanson_labels_positions()
for ilabel in labels:
# do not display uwanted labels
if ilabel not in indices:
continue
# rotate the labels if the dislay is in portrait mode
xy = reversed(labels[ilabel]) if orientation == 'portrait' else labels[ilabel]
ax.annotate(br.acronym[ilabel], xy=xy, ha='center', va='center', **kwargs)