from typing import cast
import numpy as np
from matplotlib.axes import Axes
from neuralib.suite2p import Suite2PResult, Suite2pRoiStat
__all__ = ['get_soma_pixel',
'plot_soma_center']
[docs]
def get_soma_pixel(s2p: Suite2PResult,
neuron_ids: np.ndarray | None = None,
*,
color_diff: bool = True,
include_overlap_pixel: bool = True) -> np.ndarray:
"""
Get the image mask of the registered neuronal soma shape.
The mask value equals to the neuron's id + 1.
:param s2p: :class:`~neuralib.imaging.suite2p.core.Suite2PResult`
:param neuron_ids: index ROIs array or bool mask array. `Array[int|bool, N]`
:param color_diff: whether show color difference across neurons
:param include_overlap_pixel: if taking overlap area into account
:return: 2d imaging array, shape: (xpix, ypix)
"""
neuron_pix = np.zeros((s2p.image_width, s2p.image_height))
if neuron_ids is None:
neuron_ids = np.arange(s2p.n_neurons)
elif neuron_ids.dtype == bool:
if len(neuron_ids) != s2p.stat.shape[0]:
raise ValueError('invalid shape for the mask array')
neuron_ids = np.nonzero(neuron_ids)[0]
for i in neuron_ids:
stat = cast(Suite2pRoiStat, s2p.stat[i])
ypix = stat['ypix']
xpix = stat['xpix']
if not include_overlap_pixel:
overlap = stat['overlap']
ypix = ypix[~overlap]
xpix = xpix[~overlap]
neuron_pix[xpix, ypix] = i + 1 if color_diff else 1
return neuron_pix
[docs]
def plot_soma_center(ax: Axes,
s2p: Suite2PResult,
neuron_ids: np.ndarray | None = None,
*,
invert_xy: bool = True,
with_index: bool = True,
font_size: float = 5,
**kwargs) -> None:
"""Plot center of roi soma and its corresponding index
:param ax: :class:`matplotlib.axes.Axes`
:param s2p: :class:`~neuralib.imaging.suite2p.core.Suite2PResult`
:param neuron_ids: index ROIs array or bool mask array. `Array[int|bool, N]`
:param invert_xy: If invert the FOV xy pixel
:param with_index: whether plot the index ROI index nearby
:param font_size: font size of the text
:param kwargs: additional args to ``plt.scatter()``
"""
if neuron_ids is None:
n_neuron = s2p.stat.shape[0]
neuron_ids = np.arange(n_neuron)
elif neuron_ids.dtype == bool:
if len(neuron_ids) != s2p.stat.shape[0]:
raise ValueError('invalid shape for the mask array')
neuron_ids = np.nonzero(neuron_ids)[0]
assert neuron_ids is not None
coords = []
for i in neuron_ids:
stat = cast(Suite2pRoiStat, s2p.stat[i])
coords.append((
np.mean(stat['ypix' if invert_xy else 'xpix']),
np.mean(stat['xpix' if invert_xy else 'ypix'])
))
coords = np.array(coords)
ax.scatter(coords[:, 0], coords[:, 1], **kwargs)
if with_index:
for (x, y), idx in zip(coords, neuron_ids, strict=False):
ax.text(x, y, str(idx), fontsize=font_size)