[Feature request]: Reconciling multi task models with ase Calculator interface.

[laserkelvin]

Feature/behavior summary

Per #215 MultiTaskLitModules are a supported type of task to be passed into the calculator, but the calculate step doesn't treat it any differently even when it is actually required.

Request attributes

• Would this be a refactor of existing code?
• Does this proposal require new package dependencies?
• Would this change break backwards compatibility?
• Does this proposal include a new model?
• Does this proposal include a new dataset?
• Does this proposal include a new task/workflow?

Related issues

No response

Solution description

We need a way to:

1. identify when inference results come from a MultiTaskLitModule
2. In the case where only one subtask is relevant, figure out how to extract those results and map them into the Calculator
3. In the case where multiple subtasks are relevant, figure out how to aggregate the result and report a scalar as expected by ase

Additional notes

No response

[laserkelvin]

@melo-gonzo please add scenarios I might have missed!

[melo-gonzo]

Some more notes for consideration:

Related to 3, two subtasks may output relevant information but not complete information. For example using a ScalarRegressionTask for energy prediction, and GradFreeForceRegression for forces.

Initially I think doing basic combinations of output heads with the same task keys would be best, and future aggregations could be PR'd later as features. Perhaps starting with mean, and weighted mean (weigh by subtask) would be a good starting point.

Finally, I think related to #104, the datasets which are used with MultiTaskModules may have slightly varying keys (energy, energy_total, corrected_total_energy, ...) so introducing a mapping to proper/uniform outputs could be important as well.