move_to method #320

kyralianaka · 2024-04-08T11:34:15Z

This PR adds a method to base.py that allows one to move cells or networks to specific x, y, z coordinates.

It takes a float, such that the first compartment of the first branch of the first cell is moved to that float coordinate, and everything else is shifted accordingly.

However, it can also take an array where the x, y, z dimensions all have to be the same (and the same as the number of cells), and this allows moving cells to specific x, y, z locations all at once. I wanted to write this so that one doesn't have to loop over all of the x, y, z coordinates of each cell to assign their locations, which would be very cumbersome for a large network.

I also wrote a file with tests for move() (which were missing before as far as I could tell) and move_to().

I look forward to hearing feedback and thoughts, because I'm sure there is some room for improvement.

This reverts commit a43de93.

…es and black formatted" This reverts commit 5ebf460.

…h xyzr dims

kyralianaka · 2024-04-08T11:41:33Z

I just realized that I vectorized the part that takes the coordinate arrays but not the part that takes floats, so I will vectorize that "half" of the method too now:)

…not implemented yet

michaeldeistler

Awesome, thanks!

Please write a docstring that explains exactly what this method is doing.

jaxley/modules/base.py

michaeldeistler · 2024-04-08T14:11:47Z

jaxley/modules/base.py

+            # Compute the within cell branch offsets
+            tup_indices = np.array([view.cell_index, view.branch_index])
+            branches_per_cell = np.unique(tup_indices, axis=1)[0]
+            _, first_branches = np.unique(branches_per_cell, return_index=True)


to get the first branches, can you not do net[:,0,0].view.index?

Not really because what I need is the index into xyzr which has length equal to the number of branches, not the number of compartments like view

tests/test_moving.py

michaeldeistler · 2024-04-08T14:15:50Z

I think you can almost exactly copy the PR description into the docstring :)

kyralianaka · 2024-04-08T15:24:15Z

Thanks for the feedback @michaeldeistler!

I added lots of comments, so you can let me know now if it's too much or too little.

Unfortunately I still cannot use move_to() with array inputs and groups, because group.xyzr doesn't exist (like network.xyzr), but we can talk about whether that's worth implementing.

michaeldeistler

Thanks for pointing out that this is not yet possible :) But I think this can be done without much code, see me comment below. Let me know if you want to have a quick zoom.

tests/test_moving.py

michaeldeistler · 2024-04-08T15:55:50Z

jaxley/modules/base.py

@@ -1021,6 +1020,73 @@ def _move(self, x: float, y: float, z: float, view):
            self.xyzr[i][:, 1] += y
            self.xyzr[i][:, 2] += z

+    def move_to(


Okay I just saw that this cannot move submodules, right? I.e.

net.cell(0).move_to(x, y, z)

does not work, right?

I think this should definitely be possible. The way to do this is to pass self.nodes to the ._move_to() method in Module and to add a move_to method to View and pass self.view to _move_to() in the View. You can take inspiration from other methods: For example the .show() method passes self.nodes here and self.view in the View here

That actually does work, I wrote the function just like move() so it is also in view (these lines). Hope this is alright :)

Well sorry, this works to update the coordinates in the view, but then I guess it doesn't update the module's xyzr, so I will try to fix that

you probably need self.pointer.xyzr

Ok I made it so that it can take cell view input now. It is not very pretty underneath, but when one does net.cell(0).move_to(x, y, z), one is really operating on net.xyzr, and it makes sense that this is distinct from operating on a cell.xyzr, which would only be the case if one does cell.move_to(x, y, z). I am totally open to prettier ways of doing this, but for now it works, and it is even possible to do net.cell([0, 1, ...]).move_to(x, y, z) where x, y, and z are arrays the length of the number of cells!

I also added ._update_nodes_with_xyz() to both the move and move_to methods at the end so that the results are shown in the view.

Another thing I added is that if someone tries to call move() or move_to() with float input and xyz is not calculated yet, then it will print something and run .compute_xyz(). If this isn't done, then xyzr will just stay entirely nans, and the user might be confused as to why.

jnsbck · 2024-04-09T09:30:07Z

This adds xyzr to Views.

@property
def xyzr(self):
    idxs = self.view.global_comp_index
    return np.array(self.pointer.xyzr)[idxs]

Works until .comp but for .loc one would need some kind of interpolation.

…d z in view

jnsbck · 2024-04-09T09:58:35Z

jaxley/modules/base.py

+
+    @property
+    def xyzr(self):
+        idxs = self.view.global_comp_index


We might wanna raise on .loc or add `interpolation of coordinates like in

jaxley/jaxley/utils/cell_utils.py

Line 195 in 08d5538

def interpolate_xyz(loc: float, coords: np.ndarray):

kyralianaka · 2024-04-10T08:20:17Z

I was getting some weird behavior trying to use the xyzr property before, so I removed it, but I will add it back in when I make sure that it works alright. Then I'll see if it can be used to clean up the method a bit. Thanks for your patience :)

kyralianaka · 2024-04-16T09:01:18Z

A couple of updates on this pull request:

I added the xyzr property to View, but this cannot then be called within the module methods, because they don't actually seem to get the cell view passed to them. One could access the full module view in these methods with self.show(), and then index into it using the "view" dataframe as I am already doing, but it doesn't seem like one can simply call view.xyzr in these module methods (because view is just a dataframe then as created by set_global_index_and_index).
The move() method can take inhomogeneously shaped xyzr input created by the functions that read swc files. Allowing these inhomogeneously shaped coordinate arrays means that the xyzr property of the network cannot be operated on as an array, and this is a big limitation in terms of computational time. I adjusted the code so that the xyzr property is agnostic to the form of the self.xyzr entries, but move_to() requires the form of xyzr to be [# branches, [branch start, branch end], [x, y, z, r]] (this is checked). I have some ideas about how to deal with this situation moving forward, but for now, move_to() does what I want it to do and works on cell views.

Looking forward to hearing your thoughts :)

Edit: I guess I will work next on getting move_to to work with the swc coordinate arrays, hopefully vectorized but we'll see (right now it just raises and says to use move().

michaeldeistler

Thanks a lot!

I think we definitely have to support SWC files though. I think an easy way (that will also save a ton of code) would be to have a function that modifies one coordinate and jit that. Would that not work? We can also discuss on Thu.

michaeldeistler · 2024-04-16T10:05:40Z

jaxley/modules/base.py

+
+        NOTE: This method works when module.xyzr is not of the form
+        (n branches, [start_branch, end_branch], 4) but rather constructed by the code
+        that reads swc files."""


So it does not work if the xzy is constructed with self.compute_xyz()?

Okay I only now (after having seen code below) understand this sentence. So, IIUC, it does work for coordinates generated from self.compute_xyz() and it does not work for those generated from read_swc, right?

Anyways, I think this is a pretty big limitation. Can we not jit-compile the for loop across branches and make it fast without having to rely on matrix operations?

Yes, sorry that was not worded correctly.

I should definitely get something working for xyzrs from SWCs -- the only reason I didn't before is just because I didn't realize they were different in shape from xyzr generated by compute_xyz().

The one thing that I find a bit strange about xyzr from SWC files though is that some branches have a (4, ) shape array as the xyzr, so not even the start and end of the branch. I guess this then means that those branches have length zero or something. I will have to look into this more though.

As for the jit compiled loop, jitting would require static argument shapes, and different networks will have different numbers of branches. If we jit the moving of one branch, I'm not terribly convinced there would be a great speed-up because the function is then just one variable assignment, but I can test it. vmap might be really helpful though.

michaeldeistler · 2024-04-16T10:06:02Z

jaxley/modules/base.py

@@ -1024,6 +1028,102 @@ def _move(self, x: float, y: float, z: float, view):
            self.xyzr[i][:, 0] += x
            self.xyzr[i][:, 1] += y
            self.xyzr[i][:, 2] += z
+        self._update_nodes_with_xyz()


why are you modifying ._move()?

Without this line, the xyz coords don't update in the view, so calling network.show() after calling .move() would not update the xyz. This is run at the end of move_to() as well. We could require that the user call _update_notes_with_xyz() after calling move() and move_to(), but I thought that was a bit of a hassle

michaeldeistler · 2024-04-16T10:07:06Z

jaxley/modules/base.py

+        If x, y, and z are arrays, then they must each have a length equal to the number
+        of cells being moved. Then the first compartment of the first branch of each
+        cell is moved to the specified location. This avoids the need to move cells in
+        loops.


michaeldeistler · 2024-04-16T10:07:47Z

jaxley/modules/base.py

+        x: Union[float, np.ndarray] = 0.0,
+        y: Union[float, np.ndarray] = 0.0,
+        z: Union[float, np.ndarray] = 0.0,
+    ):


type hint for return type: -> None.

This is not done for any of the other methods in base.py though as far as I see, should I add that everywhere? :)

Or I could just do it here and then I guess that would be a separate PR

jaxley/modules/base.py

michaeldeistler · 2024-04-16T10:23:16Z

jaxley/modules/base.py

+            # Find the indices in xyzr of the first branch of each cell
+            _, first_branches = np.unique(cell_inds, return_index=True)
+            # Select the xyz coordinates of the first branch of each cell
+            xyz_first_branches = xyzr_arr[first_branches, :, :3]


can you not do pointer[cell_inds, 0, 0].xyzr?

michaeldeistler · 2024-04-16T10:24:38Z

jaxley/modules/base.py

+            # Compute the distance between first branches and their connected branches
+            branch_offsets = xyzr_arr[:, :, :3] - xyz_firsts_expanded
+
+            # Compute new coord arr [# branches, [branch start, branch end], [x, y, z]]


I think if we just use a jitted for-loop we can get rid of almost all of this, no?

michaeldeistler

This is looking great now! Thanks a ton!

kyralianaka added 11 commits March 22, 2024 16:58

Added move_to() and tests

37bdd17

Switched from commenting out to removing the include_groups arg

a43de93

Revert "Switched from commenting out to removing the include_groups arg"

7b06f1c

This reverts commit a43de93.

Revert "Removed invalid include_groups argument from _get_local_indic…

d7d2a1b

…es and black formatted" This reverts commit 5ebf460.

Minor formatting

aae78d9

Some progress on setting cell coords differently but still issues wit…

ab49f45

…h xyzr dims

Merge remote-tracking branch 'origin/main' into move_to

7c99451

Merge branch 'main' into move_to

26faa7c

Fixed move_to for networks

70fb9d7

move_to now working with numpy array coordinate input and test written

801c3f0

Merge branch 'main' into move_to

13e0a06

kyralianaka requested a review from michaeldeistler April 8, 2024 11:36

kyralianaka added 3 commits April 8, 2024 14:36

Vectorized and improved move_to with float input

caaa823

Edited docstring because tests only work for networks, moving groups …

e345ea9

…not implemented yet

Removed incorrect docstring note

8406f97

michaeldeistler reviewed Apr 8, 2024

View reviewed changes

Added docstrings and comments

87c423d

Small comment edits

da0f322

michaeldeistler requested changes Apr 8, 2024

View reviewed changes

Added xyzr as property, warning in move, and autoupdating of x, y, an…

b0351f7

…d z in view

jnsbck reviewed Apr 9, 2024

View reviewed changes

kyralianaka added 5 commits April 9, 2024 16:43

Enabled move_to() operation on cell view

0232654

Enabled move_to() operation on cell view

e2ac9b0

Enabled array input for move_to() operation on cell view

fba12ed

Removed xyzr property for now because not working yet

c9d7dd1

Merge branch 'main' into move_to

f5e327e

michaeldeistler force-pushed the main branch from fac8c1a to 536098f Compare April 11, 2024 09:57

kyralianaka added 3 commits April 12, 2024 15:04

Merge branch 'main' into move_to

6196917

Added xyzr property to view

77b2fee

Dealt with alternate xyzr shapes created when reading SWC files

881d5e9

michaeldeistler requested changes Apr 16, 2024

View reviewed changes

kyralianaka added 7 commits April 17, 2024 18:28

xyzr property cleanup

49d10bf

Rewritten move_to(), to be cleaned

bd65c7c

Lots of restructuring

d06fecd

Added test and bug fixes for SWC cell moving

a668798

Merge branch 'main' into move_to

23dbb69

Merge branch 'main' into move_to

d5e2de6

Method simplifications

5268960

michaeldeistler approved these changes Apr 24, 2024

View reviewed changes

jnsbck mentioned this pull request Apr 24, 2024

Make View behave more like Module #351

Closed

9 tasks

kyralianaka merged commit c6d81cc into main Apr 24, 2024
1 check passed

kyralianaka deleted the move_to branch April 24, 2024 13:37

jnsbck added a commit that referenced this pull request Apr 24, 2024

add/rm: add groups attr and remove xyzr since #320 was merged

c2da3af

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

move_to method #320

move_to method #320

kyralianaka commented Apr 8, 2024

kyralianaka commented Apr 8, 2024

michaeldeistler left a comment

michaeldeistler Apr 8, 2024

kyralianaka Apr 8, 2024

michaeldeistler commented Apr 8, 2024

kyralianaka commented Apr 8, 2024

michaeldeistler left a comment

michaeldeistler Apr 8, 2024

kyralianaka Apr 9, 2024 •

edited

Loading

kyralianaka Apr 9, 2024

michaeldeistler Apr 9, 2024

kyralianaka Apr 9, 2024

jnsbck commented Apr 9, 2024 •

edited

Loading

jnsbck Apr 9, 2024

kyralianaka commented Apr 10, 2024

kyralianaka commented Apr 16, 2024 •

edited

Loading

michaeldeistler left a comment

michaeldeistler Apr 16, 2024

michaeldeistler Apr 16, 2024

kyralianaka Apr 17, 2024

michaeldeistler Apr 16, 2024

kyralianaka Apr 17, 2024

michaeldeistler Apr 16, 2024

michaeldeistler Apr 16, 2024

kyralianaka Apr 17, 2024

kyralianaka Apr 17, 2024

michaeldeistler Apr 16, 2024

michaeldeistler Apr 16, 2024

michaeldeistler left a comment

move_to method #320

move_to method #320

Conversation

kyralianaka commented Apr 8, 2024

kyralianaka commented Apr 8, 2024

michaeldeistler left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

michaeldeistler commented Apr 8, 2024

kyralianaka commented Apr 8, 2024

michaeldeistler left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kyralianaka Apr 9, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jnsbck commented Apr 9, 2024 • edited Loading

Choose a reason for hiding this comment

kyralianaka commented Apr 10, 2024

kyralianaka commented Apr 16, 2024 • edited Loading

michaeldeistler left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

michaeldeistler left a comment

Choose a reason for hiding this comment

kyralianaka Apr 9, 2024 •

edited

Loading

jnsbck commented Apr 9, 2024 •

edited

Loading

kyralianaka commented Apr 16, 2024 •

edited

Loading