Task refactor

[RasmusOrsoe]

This PR is the first in a series of smaller PRs that introduces normalizing flows to graphnet.

In this first PR, the aim is to introduce new flexibility in the generic Task to allow for easier domain/ML paradigm -specific implementations.

The main constraints of our current Task is:

• The loss computation is fixed; some ML paradigms require a different set of inputs for loss computation.
• All Tasks contain affine/learnable embeddings; some ML paradigms is incompatible with this.

This PR does the following:

• Remove _forward, forward and compute_loss from generic Task. Domain / ML Paradigm specific implementation of Task are required to implement these according to their needs. Keeping them as @abstractmethod in Task is not possible in this case because of the inputs might be different.
• Introduce LearnedTask: A generic Task with a learned embedding - forward is implemented with @final but _forward and compute_loss is left as @abstractmethod for implementation by subclasses. E.g. autoencoding tasks would want to subclass this.
• Introduce StandardLearnedTask: Essentially our old Task class; compute_loss is implemented with @final but _forward is left as @abstractmethod. All of our current physics tasks now inherits from this.
• Introduce FlowTask: A problem-specific implementation of Task that is aimed directly for normalizing flows.

I've included a quick illustration that shows this new subclassing structure.

[RasmusOrsoe]

This looks good.

I'm not sure if we will be able to guarantee that StandardModel will be compatible with all tasks in the future. In particular, I could imagine that the predict/predict_as_dataframe, might not work with all imagineable tasks in the future. Perhaps we should restrict the tasks provided to StandardModel to be StandardLearnedTask?

You are absolutely right. The next few PRs in the normalizing flow series introduce a few extra lines of code that makes .fit and .predict compatible with this new FlowTask. I have also just now introduced a check in this PR to make sure all tasks are StandardLearnedTask.

Now that we are going about refactoring tasks I think that the aggregation parts of the GNN's might be better placed in the task section. The reason would be that the GNN's would now only be responsible for feature learning, and it would be easier to do transfer-learning between various tasks, where you might have one tasks which predicts on an entire pulsemap and another task, which does regression/classification.

In particular I think we could introduce an Aggregation module, which will take in a variable length pulsemap and return a "fixed length" tensor. This Module would then act as the "link" between the feature-pulsemap in a GNN and the MLP in the StandardLearnedTask, but would be omitted when doing per-node reconstructions.

Perhaps - but I do think it makes a lot of sense to allow models to have aggregation in them. To me it sounds like you're describing ensemble functionality, but perhaps we can discuss this on Tuesday's meeting.