Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Get correct centers of mass from rois.com even when swap_dim is True in get_contours #1370

Merged
merged 7 commits into from
Jul 9, 2024
Merged

Get correct centers of mass from rois.com even when swap_dim is True in get_contours #1370

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
522d456
Fix CoM calculation when swap_dim is true
ethanbb Jun 28, 2024
1f88e49
Fix plot_contours to use CoMs from updated get_contours
ethanbb Jul 2, 2024
b38f0fa
Add rewritten com function with C/F order and add test
ethanbb Jul 2, 2024
467c591
Restore individual dimension arguments for com
ethanbb Jul 5, 2024
d883598
Add slice_dim option for get_contours and 3D test for swap_dim
ethanbb Jul 5, 2024
830967c
Minor changes to make contours more consistent/predictable
ethanbb Jul 5, 2024
04d1ca2
Use None instead of 'auto'
ethanbb Jul 8, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 15 additions & 29 deletions caiman/base/rois.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,46 +28,32 @@
pass


def com(A: np.ndarray, d1: int, d2: int, d3: Optional[int] = None) -> np.array:
def com(A, *dims: int, order='F') -> np.ndarray:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Provided people were only doing positional argument passing before, this is okay (and there's no in-repo code that does otherwise), although people will also have to remember that because dims is thirsty, order can only be passed as a keyword.

>>> def foo(alpha, *beta, gamma):
        print(f"{alpha=}, {beta=}, {gamma=}")

>>>
>>> foo('a', 'b1', 'b2','b3','b4','g')
alpha='a', beta=('b1', 'b2', 'b3', 'b4', 'g'), gamma=None

I hope this doesn't trip anyone up.

I wonder if we might want to keep the type definition for A but expand it to be a Union of all those types.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Provided people were only doing positional argument passing before, this is okay (and there's no in-repo code that does otherwise)

Right, I doubt anyone was passing the dimensions as keywords, but just in case we could add keyword arguments for them as well (although then we'd have to deal with some edge cases if both positional and keyword dimensions are passed).

There's also no strong reason to modify the signature at all except for adding order. Do you think we should just use com(A, d1, d2, d3=None, order='F')?

I wonder if we might want to keep the type definition for A but expand it to be a Union of all those types.

Can definitely add that if you want to err on the side of annotating every argument, although I don't feel strongly about it.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just adding order would eliminate the possibility of disruption; probably better to do that.

WRT A, in this case, maybe that type definition would be so bulky that it wouldn't be worth it; it's fine to remove it. I'll add the full form back later if it gets me enough the next time I go tackling bugs with mypy.

"""Calculation of the center of mass for spatial components

Args:
A: np.ndarray
matrix of spatial components (d x K)
A: np.ndarray or scipy.sparse array or matrix
matrix of spatial components (d x K); column indices map to dimensions according to F-order.

d1: int
number of pixels in x-direction

d2: int
number of pixels in y-direction

d3: int
number of pixels in z-direction
*dims: D ints
each positional argument after A defines the extent of one of the dimensions of the data.

order: 'C' or 'F'
how each column of A should be reshaped to match dims.

Returns:
cm: np.ndarray
center of mass for spatial components (K x 2 or 3)
center of mass for spatial components (K x D)
"""

if 'csc_matrix' not in str(type(A)):
A = scipy.sparse.csc_matrix(A)

if d3 is None:
Coor = np.matrix([np.outer(np.ones(d2), np.arange(d1)).ravel(),
np.outer(np.arange(d2), np.ones(d1)).ravel()],
dtype=A.dtype)
else:
Coor = np.matrix([
np.outer(np.ones(d3),
np.outer(np.ones(d2), np.arange(d1)).ravel()).ravel(),
np.outer(np.ones(d3),
np.outer(np.arange(d2), np.ones(d1)).ravel()).ravel(),
np.outer(np.arange(d3),
np.outer(np.ones(d2), np.ones(d1)).ravel()).ravel()
],
dtype=A.dtype)

cm = (Coor * A / A.sum(axis=0)).T
# make coordinate arrays where coor[d] increases from 0 to npixels[d]-1 along the dth axis
coors = np.meshgrid(*[range(d) for d in dims], indexing='ij')
coor = np.stack([c.ravel(order=order) for c in coors])

# take weighted sum of pixel positions along each coordinate
cm = (coor @ A / A.sum(axis=0)).T
return np.array(cm)


Expand Down
17 changes: 17 additions & 0 deletions caiman/tests/test_toydata.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

import caiman.source_extraction.cnmf.params
from caiman.source_extraction import cnmf as cnmf
from caiman.utils.visualization import get_contours


#%%
Expand Down Expand Up @@ -64,6 +65,22 @@ def pipeline(D):
]
npt.assert_allclose(corr, 1, .05)

# Check that get_contours works regardless of swap_dim
if D == 2: # Can't expect swap_dim to work for 3D data in the same way
coor_normal = get_contours(cnm.estimates.A, dims, swap_dim=False)
coor_swapped = get_contours(cnm.estimates.A, dims[::-1], swap_dim=True)
for c_normal, c_swapped in zip(coor_normal, coor_swapped):
# have to sort coordinates b/c starting point is unimportant & depend on orientation
# also, the first point is repeated and this may be a different point depending on orientation.
# to get around this, compare differences instead (have to take absolute value b/c direction may be opposite)
def normalize_coords(coords):
abs_diffs = np.abs(np.diff(coords, axis=0))
sort_order = np.lexsort(abs_diffs.T)
return abs_diffs[sort_order, :]

npt.assert_allclose(normalize_coords(c_normal['coordinates']), normalize_coords(c_swapped['coordinates'][:, ::-1]))
npt.assert_allclose(c_normal['CoM'], c_swapped['CoM'][::-1])


def test_2D():
pipeline(2)
Expand Down
23 changes: 6 additions & 17 deletions caiman/utils/visualization.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -392,18 +392,12 @@ def get_contours(A, dims, thr=0.9, thr_method='nrg', swap_dim=False):
if 'csc_matrix' not in str(type(A)):
A = csc_matrix(A)
d, nr = np.shape(A)
# if we are on a 3D video
if len(dims) == 3:
d1, d2, d3 = dims
x, y = np.mgrid[0:d2:1, 0:d3:1]
else:
d1, d2 = dims
x, y = np.mgrid[0:d1:1, 0:d2:1]
d1, d2 = dims[:2]

coordinates = []

# get the center of mass of neurons( patches )
cm = caiman.base.rois.com(A, *dims)
cm = caiman.base.rois.com(A, *dims, order='C' if swap_dim else 'F')

# for each patches
for i in range(nr):
Expand Down Expand Up @@ -1098,16 +1092,11 @@ def plot_contours(A, Cn, thr=None, thr_method='max', maxthr=0.2, nrgthr=0.9, dis
plt.plot(*v.T, c=colors, **contour_args)

if display_numbers:
d1, d2 = np.shape(Cn)
d, nr = np.shape(A)
cm = caiman.base.rois.com(A, d1, d2)
nr = A.shape[1]
if max_number is None:
max_number = A.shape[1]
for i in range(np.minimum(nr, max_number)):
if swap_dim:
ax.text(cm[i, 0], cm[i, 1], str(i + 1), color=colors, **number_args)
else:
ax.text(cm[i, 1], cm[i, 0], str(i + 1), color=colors, **number_args)
max_number = nr
for i, c in zip(range(np.minimum(nr, max_number)), coordinates):
ax.text(c['CoM'][1], c['CoM'][0], str(i + 1), color=colors, **number_args)
return coordinates

def plot_shapes(Ab, dims, num_comps=15, size=(15, 15), comps_per_row=None,
Expand Down