feat: Generalize NIRData API and Conversions

[benkroehs]

With this PR I propose adding the following features to NIRData:

• Frameworks do not need to check in their from-NIRData version for the type of the observable data (TimeGriddedData or (Valued)EventData) , they can just call data.get_event()
• The continuous time_shift parameter is substituted by a boolean dynamic_before_transition to represent the two existing implementations for time-gridded simulators: propagating the the dynamics first and checking if the threshold has been exceeded after or on the contrary first checking for the threshold-crossing and then calculating the dynamics accordingly
• The dimension order for TimeGriddedData is no longer fixed but stored in the dimension_order attribute. The conversion happens lazily if a different format is requested. The __getitem__ and __setitem__ functions use per default the most popular ordering (time, batch, neuron), but this is flexible.

An exemplary implementation of these changes can be seen in the most recent NIRTorch PR.

[Jegp]

Very very nice. I have one question around the use of coordinate-based representations in the EventData class. And there's a merge conflict after the Ruff branch merged.

[Jegp]

Is it correct that EventData is collecting a series (what you call sample) of time-stepped events? So, that's basically a "batched" coordinate AER view, correct?

And is it correct that n_events cap each sample to n_events? If so,

1. why not ask for a lower cap as well? Would you want to "slice" the samples, like you'd slice a normal list?
2. and if you are more or less slicing, what's the logic below good for? Why not just do idx[:, start:end]? I can see you're doing checks for invalid events, but I'm puzzled why that has to happen in the getter. Shouldn't that be a constructor check to make sure the data itself is well-formed upon construction?

[benkroehs]

So EventData is not time-stepped data but continuous/event-based. For each spike you have a time when it occured and and index of the spiking neuron idx.
Maybe the phrase "each event is discrete" in the docstring is misleading. What I meant by that is that an event does not has a length but happens at one point in time.

And yes, n_events cap the events per sample to this number. One could say that n_events is the number of events per sample. If a sample has 5 spikes but n_events is 10, the last 5 events would be "empty/invalid events" with idx=-1`. But if there were 15 spikes, the last 5 would be dropped. I implemented this for a better data handling.

Does this clear things up on your side?

[Jegp]

Same comment as for L233

[Jegp]

Ruff complains :P

[Jegp]

Additionally, I wonder whether it would make sense to actually model the relationship between the NIRNode and the NIRData. Currently, if I'm understanding it correctly, you're simply checking that the dimensionality of the data fits with the shape of the node. Wouldn't it be nicer to have a direct link, so you can access the node in the data format? Something like an augmented NIRNode where you add a data entry, or similar? Like

@dataclass
class NIRNodeData(NIRNode):
    observables: dict[str, EventData | TimeGriddedData]
    
    def check_observables(self):
        """
        Check that the shapes of the observables match the node's output shapes
        """
        output_shape = self.output_type["output"]
        return all(obs.n_neurons == output_shape for obs in self.observables.values())

Wouldn't that make the relationship unambiguous?

[muffgaga]

(Ben and I are just sitting together and thinking about this :) .)

Regarding your last comment @Jegp — you suggest NIRNodeData to inherit from NIRNode and construct with an NIRNode instance reference to have an explicit link to a existing node in the NIRGraph?
This does not look correct to me — primitives are NIRNodes (…and NIRGraph, nesting is possible which is nice :D); however, the existing NIRNodeData doesn't "process" anything, it's not a vertex in the signal-flow graph.
If you suggest the inverse, having "data" in the NIRNodes → yes, that does not sound wrong — however, in this case I don't think it should be inheritance ("is-a") but rather composition ("has-a").
In that case we could have a ternary .data member that provides None | ObservableData — None would provide the sentinel value for "wasn't recorded/observed" and 0-length EventData would mean "no spike" observed.

Advantage: no separate observable data structure that has a "string-based" link to the topology data structure (and no possibility to have a mismatch in "keys", i.e. no need for a "key-match" check).
Disadvantage: one more member variable in the previously-only-topology data structure and the NIRGraph would mix topology and observables; from/to_nir_data() would be weird and NIR implementations for individual backends have to provide the variable (so topology im/export-only backends might be more "confusing" to the user as you would get None by default (which might be ok)).

We don't have a strong opinion ;) …

[Jegp]

Thank you for your thoughts @muffgaga and @benkroehs. Hmmm, yes I agree, the relationship is more a "has-a" rather than an "is-a".

I found that check_observables raises on any node whose output isn't 1-d. This is a problem on a Conv2d, for instance. I also noted that the check_observables assumes every observable spans the whole population. On hardware, I assume you typically record all spikes from a population but a membrane trace from a single neuron, and that combination is rejected, is that on purpose?

NIRNodeData({"spikes": spikes_512}).check_observables(node)               # True
NIRNodeData({"spikes": spikes_512, "v": v_rec}).check_observables(node)  # False

I'm not saying that the string-keyed link is bad. But note that check_nodes can't simply move into __post_init__ as things stand. NIRGraphData holds no reference to the NIRGraph. So, if we want the keys checked at construction, that link has to be made somewhere.

What about putting it on the data side? The extension should know about NIR, not the other way around. And it would keep the topology IR free of observables, as you say, @muffgaga. As an example:

class NIRNodeData:
    observables: Dict[str, ObservableData]
    node: NIRNode | None = None

This doesn't change the file format (HDF5 will store a name). Here, __post_init__ can then check the shapes etc.
Does that approach seem reasonable? Is there a reason .data on NIRNode would be easier for your import/export paths?

Separately, on n_events: I found two problems with making it optional. The default only exists on TimeGriddedData. On EventData, ValuedEventData, and the ObservableData, n_events is still a required positional argument.

>>> EventData(...).get_event()
    TypeError: EventData.get_event() missing 1 required positional argument: 'n_events'

And where the default does apply it's n_time_steps, while a single sample can hold up to n_time_steps × n_neurons events. Everything past the cap is dropped in the getter with no warning. On a 100-step, 256-neuron tensor:

 rate=0.001  events=  24  kept= 24  lost=   0
 rate=0.005  events= 109  kept=100  lost=   9
 rate=0.01   events= 267  kept=100  lost= 167
 rate=0.05   events=1298  kept=100  lost=1198

Truncating starts as soon as the man rate exceeds 1/n_neurons. Is that intentional?