Feature/issue 62 visualisation

[oleeviyababu]

Describe your changes

Ensemble Plots: Added plot_ensemble_prediction to visualize ground truth, ensemble mean, spread, and individual members.

Latent Inspection: Added plot_latent_samples to compare prior and variational distributions for Graph-EFM latent space.

Modular Refactoring: Switched all functions to the BaseRegularGridDatastore api to ensure compatibility across different regional grids.

Bug Fixes: Resolved a memory leak in plot_spatial_error using .detach() and fixed an alpha-channel dimension mismatch

Motivation

These changes are required to support the transition to probabilistic modeling (Issue #62). By moving these into vis.py we keep the main model classes decoupled from visualization logic.

Dependencies

pytorch, matplotlib, cartopy, and the updated datastore module.

Issue link

#62

Type of change

• 🐛 Bug fix (non-breaking change that fixes an issue)
• ✨ New feature (non-breaking change that adds functionality)
• 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
• 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

• My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
• I have performed a self-review of my code
• For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
• I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
• I have updated the README to cover introduced code changes
• I have added tests that prove my fix is effective or that my feature works
• I have given the PR a name that clearly describes the change, written in imperative form (context).
• I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

• the code is readable
• the code is well tested
• the code is documented (including return types and parameters)
• the code is easy to maintain

Author checklist after completed review

• I have added a line to the CHANGELOG describing this change, in a section
  reflecting type of change (add section where missing):
  • added: when you have added new functionality
  • changed: when default behaviour of the code has been changed
  • fixes: when your contribution fixes a bug
  • maintenance: when your contribution is relates to repo maintenance, e.g. CI/CD or documentation

Checklist for assignee

• PR is up to date with the base branch
• the tests pass
• (if the PR is not just maintenance/bugfix) the PR is assigned to the next milestone. If it is not, propose it for a future milestone.
• author has added an entry to the changelog (and designated the change as added, changed, fixed or maintenance)
• Once the PR is ready to be merged, squash commits and merge the PR.

[kshirajahere]

Hi i went through the code had saw issue in latent-space plotting helper

[kshirajahere]

plot_latent_samples() seems to assume the latent mesh can always be reshaped into a square image via sqrt(num_mesh_nodes). That works for image-like lattices, but for the graph meshes used by the probabilistic models this is not generally true. In those cases this will either fail on the reshape or silently misrepresent the latent field.

Since this PR is for the #62 visualization path, I think this needs either:

• an explicit guard that this helper only supports square regular-grid layouts, or
• a plotting path that uses actual mesh/grid coordinates instead of sqrt(N) x sqrt(N) reshaping.

[joeloskarsson]

Hi, a few high-level points here that you can look at before I give a more detailed review:

1. Please give this PR a more descriptive name.
2. This seems to build on Extract plot_examples from ARmodel.py into stateless function in vis.py #209 if I am not mistaken? Such PR dependencies need to be made VERY explicit, as I then assume that we need to merge Extract plot_examples from ARmodel.py into stateless function in vis.py #209 before this. If this is not the case, there are definitely unrelated changes here.
3. While we could maybe merge ensemble map visualizations before we have any ensemble models, I don't think it makes sense to merge any plotting for latent variables before any model using this exists.
4. Seems like a good idea to coordinate this addition with discussion in [RFC/Design] Standardize probabilistic vs deterministic return contract to unblock evaluation integrations  #335.
5. The memory leak bugfix seems unrelated to this, and should therefore be in its own issue+PR. Try to make each PR do one thing, never a collection of unrelated things.

[sadamov]

Closing this one. The visualisation work here is now covered by #612, which ports plot_ensemble_prediction and plot_latent_samples from prob_model_lam with a cleaner pcolormesh approach via plot_on_axis. @kshirajahere's inline note about the plot_latent_samples reshape was a real concern that #612 sidesteps. Thanks @oleeviyababu for the contribution.