ENH: sparse: New nD sparse array functions

[dschult]

Is your feature request related to a problem? Please describe.

CVXPY is looking for a sparse COO-like nD type. Functionality needed seems close to what is provided by the new COO nD format here.

Describe the solution you'd like.

Some helpful NumPy functions wanted for the sparse array are:

• expands_dims
• swapaxes, (or transpose?)
• nd-kron
• slicing along one axis. (either with slices, or arrays, e.g. A[:, [4,8,12], :])

Describe alternatives you've considered.

No response

Additional context (e.g. screenshots, GIFs)

This stems from a scipy-Discord discussion about 11/18/2024

[dschult]

From @Transurgeon on discord:

Yes that would be really helpful. we have a function select_rows that is called quite often actually.
and it just does the following (this is for the numpy backend, hence why its 3d):

  def func(x):
      return x[:, rows, :]

and if you are curious, our current implementation for the same behavior in scipy is as follows (where p is the size of the third dimension)
and just a reminder, the trick we use is to stack the third dimension vertically to create a strided array (with many more rows than cols):

  def func(x, p):
      if p == 1:
          return x[rows, :]
      else:
          m = x.shape[0] // p
          return x[np.tile(rows, p) + np.repeat(np.arange(p) * m, len(rows)), :]

[dschult]

Could you take a look at our current kron. I think we just need to switch from handling row and col to handling the tuples of coords A.coords in a tuple comprehension way. I envision a PR to make kron handle nD (including 1D) and add extend_dims and swapaxes. The code for extend_dims involves creating a new array of zeros put into coords and updating self._shape. But first check whether a reshape with the new shape adding a 1 is sufficient. That would be a really short function and would mimic the numpy function. The swapaxes function would need to rearrange the tuple self.coords and self._shape. It might be able to use transpose, but I haven't checked that.

[Transurgeon]

Hey @dschult thanks for opening this issue!
Sorry for the confusion but this is actually my github account. (If you remember, I had contributed to networkx also last year :), on adding the broadcasting module, thanks again for your tremendous review!).
I am just commenting to have this issue in the tracker. Also gonna reference #21435 to also keep track of it.

[dschult]

Nice! Hello William! I had not made that connection.
:)
Can we tag Parth too? Not sure of their github username.

[rgommers]

These additions seem reasonable to me, and having a real-world need for them is certainly helpful.

Re swapaxes: use permute_dims I'd say (xref permute_dims in the array API standard and data-apis/array-api#483 (comment)).

[izaid]

Hi all,

Can I just point out that there are open (and even approved!!) PRs that generalise SciPy sparse arrays to N dimensions? Some of that work has already been merged, and some of it still needs someone to just merge it. These were done by @anushkasuyal with guidance from myself and @rgommers over the summer. See:

#21435
#21613

The work is already done for months now, so it would really be great if we can work out how to proceed. cc @rgommers @steppi @dschult @perimosocordiae

[lucascolley]

Can I just point out that there are open (and even approved!!) PRs that generalise SciPy sparse arrays to N dimensions?

Yep, I think @PTNobel said on discord that they'd drop a review on gh-21435.

[PTNobel]

Yep, I think @PTNobel said on discord that they'd drop a review on gh-21435.

Yes, it's on my to-do list for the weekend. William and I (well, almost all William) went through the operations in the PRs last weekend to make the list of additional operations he requested on Discord.

I also am shopping around a design doc on our end to get permission to put coo_array in our public API.

[izaid]

That's fantastic, thanks @PTNobel! Please do ping @anushkasuyal and I as you go.

[anushkasuyal]

Could you take a look at our current kron. I think we just need to switch from handling row and col to handling the tuples of coords A.coords in a tuple comprehension way. I envision a PR to make kron handle nD (including 1D) and add extend_dims and swapaxes. The code for extend_dims involves creating a new array of zeros put into coords and updating self._shape. But first check whether a reshape with the new shape adding a 1 is sufficient. That would be a really short function and would mimic the numpy function. The swapaxes function would need to rearrange the tuple self.coords and self._shape. It might be able to use transpose, but I haven't checked that.

Hi, just wanted to mention that an implementation of kron for nD COO arrays does exist here - e18c383. I've not put up this PR yet as I was waiting for the previous two to get merged first (so I could rebase etc.). The branch that the above commit belongs to also implements tensorsolve, in case that's of interest too!

[anushkasuyal]

Yes, it's on my to-do list for the weekend. William and I (well, almost all William) went through the operations in the PRs last weekend to make the list of additional operations he requested on Discord.

If this is of any help, here is the link to a tracker that I'm maintaining to track how much functionality for COO arrays has been extended to nD. As far as the work done is concerned, all the methods in this tracker (apart from resize) have a working implementation. A few of them have been implemented in #21435 and #21613, whereas some are in a branch consisting of the last chunk of work I did - extend-coo-nd-kron-tensorsolve.

[Transurgeon]

Hi, just wanted to mention that an implementation of kron for nD COO arrays does exist here - e18c383. I've not put up this PR yet as I was waiting for the previous two to get merged first (so I could rebase etc.). The branch that the above commit belongs to also implements tensorsolve, in case that's of interest too!

That looks really good, I think as long as it matches NumPy's behavior it will be good for our use-case. I was just wondering if this implementation would support having one of the kron operands be dense aswell?

[Transurgeon]

If this is of any help, here is the link to a tracker that I'm maintaining to track how much functionality for COO arrays has been extended to nD. As far as the work done is concerned, all the methods in this tracker (apart from resize) have a working implementation. A few of them have been implemented in #21435 and #21613, whereas some are in a branch consisting of the last chunk of work I did - extend-coo-nd-kron-tensorsolve.

Thanks for your great work @anushkasuyal this is a tremendous contribution!
I guess we are just missing indexing as mentioned above? Let us know if we could be of any help for that!
I would be really happy to work with you on this after my exams :).