Benchmark Making SimState.system_idx a sparse tensor

[curtischong]

Correct me if I'm wrong - but I recall that @thomasloux performed some benchmarking and concluded that storing the system_idx as a sparse tensor performs just as well as storing system_idx as a dense tensor (I think on scatter_sum operations).

We should probably make it a sparse tensor because:

1. It saves memory
2. It might make some operations cheaper to perform

[orionarcher]

By "sparse tensor" do you mean storing n_atoms_per_system instead of system_idx?

The two are fungible and imo it depends which is being used more. system_idx is used as a dense tensor for many scatter operations, which would be one reason to keep as is. Would need to do a bit of code review to determine whether system_idx -> n_atoms_per_system or n_atoms_per_system -> system_idx is more common. That'd be a good task for Claude Code.

If memory is slightly more efficient and usage is equal, I'd be fine making the switch, we could just make system_idx an @property calculated from n_atoms_per_system.

[curtischong]

Yes I mean it stores n_atoms_per_system instead of system_idx.

I am aware that the scatter operations rely on it. But I thought someone said (in our constraints call) that there are similar operations where we can pass in a sparse tensor? (or maybe they meant to construct a dense tensor on the fly)

Using @property makes sense.

[orionarcher]

@thomasloux anything to add here? I remember something about this too.

[thomasloux]

What I briefly tried is getting value through a mask or index:

import torch
import timeit

n = 10_0000
frac = 0.95
indices = torch.arange(int(n*frac))
mask = torch.full((int(n*frac),), True)
array = torch.randn(n)

t_mask = timeit.timeit("array[mask]", number=1000, globals=globals())
t_idx  = timeit.timeit("array[indices]", number=1000, globals=globals())

print("mask indexing:", t_mask) # 0.16
print("index indexing:", t_idx) # 0.06

Index is faster, even is it's full: meaning here you select a vast majority of you array.

But, we can try to benchmark scatter_reduce versus segment_reduce operation: https://docs.pytorch.org/docs/stable/generated/torch.segment_reduce.html

My guess is that segment_reduce is faster, but I didn't try. We should check Pytorch Geometric code to know what they use.

Btw we may want to decorate more properties with @cached_decorator (https://docs.python.org/dev/library/functools.html#functools.cached_property) for properties that don't change.

[thomasloux]

We may want to use this as well: https://pytorch-scatter.readthedocs.io/en/latest/functions/segment_coo.html
No idea what is the different the best option. When reading this thread, it seems like torch.segment_reduce is netter when compiled, otherwise torch_scatter is better: rusty1s/pytorch_scatter#404