Model request: non-negative spatial factorization (NSF)

[aribenjamin]

Is anyone currently working on implementing NSF? It's a nice probabilistic method designed with spatial transcriptomics in mind. https://www.nature.com/articles/s41592-022-01687-w

It'd be nice to work with this here, within scvi-tools. The authors' implementation (@ https://github.com/willtownes/nsf-paper) is a little hard to parse and implemented with Tensorflow as a backend (which I don't work in).

[adamgayoso]

I think this would be a nice addition. It might be easiest to implement using tensorflow probability and jax as the kernel functions should be the same; unless of course it's also straightforward in pytorch. Are you willing to contribute this?

[aribenjamin]

Yes, I'm happy to go for it. It'd be good practice contributing. I can't guarantee any timeline, though, so if someone wants this sooner please ping me.

[adamgayoso]

Looking into this more, it should be possible with just PyTorch, or just jax+numpyro (which we already depend on), without the introduction of tensorflow probability. @martinkim0 do you want to try this model? It looks both fairly straightforward and fun.

[aribenjamin]

fyi I have back-burnered this to the point of inaction. If you want to work on this, know that I'm not working in parallel. I'll update this when and if I begin.