Migration to torch distributions and scoringrules integration

[MicheleCattaneo]

Proposal

• Remove the dependency of TFP in order to move to Keras3 and rely on torch as backend and torch.distributions for the probabilistic layers.
• Rely on scoringrules for probabilistic loss functions.

Current progress

• Update poetry dependencies.
• scoringrules has been tested on normal and exponential distributions with crps.
• Update code to remove tf dependencies.
• Update all models to the new structure
  • TCN and phyisical layers are temporarily not supported.
• Update tests to the new structure.
• Create the equivalent of DistributionLossWrapper for sample losses.
• Ensure configuration files are compatible with the updates
• Optional torch CPU/GPU
• Add distributions:
  • Censored normal (custom)
  • Truncated normal (custom)
  • Multivariate normal
  • Gamma
  • Beta
  • LogNormal
• Update README.md (via notebook)
• Add test for SHAP (include as dev dependency)
• Extend support to python 3.12
• Format with black
• ...

Some implementation details:

• Every distribution will implement BaseParametricDistribution
• A probabilistic layer will implement BaseDistributionLayer, which contains an underlying distribution and is responsible for mapping its input into the correct number of parameters.
• DistributionLossWrapper allows to wrap a function from scoringrules such that the loss can be computed between a tensor of observations and a predicted distribution.
• There is a clear separation between models and layers.

Useful links:

• Table with common tf operations and their corresponding keras3 variant: here

[MicheleCattaneo]

test inference code:

        # make predictions
        out_batch = [model(ds.x).sample(tfd_samples).numpy() for _ in range(mc_samples)]
        out_batch = np.concatenate(out_batch, axis=0).astype("float32")
        assert out_batch.shape[0] == tfd_samples * mc_samples