refactor: add ForecastBatch contract

[kshirajahere]

Refs #613

Describe your changes

Adds a named ForecastBatch contract for WeatherDataset outputs and migrates the current post-#208 internal consumers away from positional batch access.

Changes

• Add ForecastBatch(NamedTuple) in neural_lam/batch.py.
• Return ForecastBatch from WeatherDataset.__getitem__.
• Update ForecasterModule, compute_standardization_stats.py, and tests to use named fields.
• Keep native PyTorch default_collate support without a custom collate function.
• Add regression coverage for DataLoader collation and named batch access.

Issue Link

Closes #613

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]

@sadamov @leifdenby please have a look 🙏 thank you

[sadamov]

Nice PR @kshirajahere, thanks! One ask before merge: please add boundary as a 4th field of ForecastBatch now, so this lands cleanly alongside #635 in v0.8.0 without a re-rebase of every consumer. Relabelled as enhancement, proposed for v0.8.0.

[sadamov]

Add boundary and propagate through .to():

Suggested change

	init_states: torch.Tensor
	target_states: torch.Tensor
	forcing: torch.Tensor
	target_times: torch.Tensor
	def to(self, device: torch.device | str) -> "ForecastBatch":
	"""Move all tensor fields to the same device."""
	return ForecastBatch(
	init_states=self.init_states.to(device),
	target_states=self.target_states.to(device),
	forcing=self.forcing.to(device),
	target_times=self.target_times.to(device),
	)
	init_states: torch.Tensor
	target_states: torch.Tensor
	forcing: torch.Tensor
	boundary: torch.Tensor
	target_times: torch.Tensor
	def to(self, device: torch.device | str) -> "ForecastBatch":
	"""Move all tensor fields to the same device."""
	return ForecastBatch(
	init_states=self.init_states.to(device),
	target_states=self.target_states.to(device),
	forcing=self.forcing.to(device),
	boundary=self.boundary.to(device),
	target_times=self.target_times.to(device),
	)

[sadamov]

Suggested change

	return ForecastBatch(
	init_states=batch[0],
	target_states=batch[1],
	forcing=batch[2],
	target_times=batch[3],
	)
	return ForecastBatch(
	init_states=batch[0],
	target_states=batch[1],
	forcing=batch[2],
	boundary=batch[3],
	target_times=batch[4],
	)

[sadamov]

Pass boundary through here too (zero tensor of matching shape is fine until #635 wires up real loading).

[sadamov]

Nit: training_step / validation_step / test_step each coerce, then common_step coerces again. Could lift to one place. Not blocking.

[sadamov]

Link the PR, not the issue.

Suggested change

	- Add a named `ForecastBatch` contract for `WeatherDataset` batches and update
	internal consumers to use field access [\#613](https://github.com/mllam/neural-lam/issues/613) @kshirajahere
	- Add a named `ForecastBatch` contract for `WeatherDataset` batches and update
	internal consumers to use field access [\#637](https://github.com/mllam/neural-lam/pull/637) @kshirajahere

[leifdenby]

I really appreciate the effort to construct a data-structure for the entire training sample, rather than just a tuple of tensors, this will be fantastic 🌟 . But I think having a key in the named-tuple called boundary as well as keys with names state, static, forcing doesn't quite make sense relative to the discussions we've been having on #635. This is because the state,static,forcing labels denote what these tensors are used for (i.e. are they input/output, input but needs broadcasting across the batch and time, or only input) whereas boundary denotes what spatial domain does the data come from. I think it is a bit confusing to mix the two. Or said different, from just the name boundary it isn't obvious to me if this boundary-forcing data or boundary-static data.

What we probably need is either:

1. Data-structure where both spatial domain and use of data is explicit through a nested structure, i.e. batch.{domain}.{data_category}, which means there would be tensors like batch.interior.state and batch.boundary.forcing, or
2. We coerce the spatial domain and data category into a single key, i.e. batch.{domain}__{data_category} with examples like batch.interior__state and batch.boundary__forcing

I would prefer 1. above, coercing multiple things into a single string is always brittle. In addition, I would actually advocate for the {domain} identifier above we actuallyuse the datastore names that are being introduced #635 (based on discussions in #652), so that the data-structure becomes batch.{datastore_name}.{data_category} with examples batch.danra.state and batch.era5.forcing.

Thoughts? :)