docs(dataset): migrate WeatherDataset to Google-style and add type hints by yukthagangadhari5 · Pull Request #505 · mllam/neural-lam

yukthagangadhari5 · 2026-03-24T05:14:40Z

Description

This PR refactors WeatherDataset and WeatherDataModule in neural_lam/weather_dataset.py to use Google-style docstrings. This ensures compatibility with mkdocstrings and improves clarity for researchers.

Key Changes

Docstring Migration: Converted NumPy-style to Google-style (Args:, Returns:, Raises:).
Tensor Shapes: Added explicit dimensions like (B, T, L, C) to docstrings.
Type Hinting: Added PEP 484 type hints throughout the file.
Refactoring: Introduced _get_dataloader helper in WeatherDataModule to reduce code duplication.

Type of change

Documentation update
Refactoring (non-breaking change which improves code quality)

How Has This Been Tested?

Verified dataset initialization: from neural_lam.weather_dataset import WeatherDataset.
Confirmed tensor shapes in code match the new documentation.
Ran local linting checks to ensure no syntax errors were introduced.

Checklist:

My code follows the style guidelines of this project.
I have performed a self-review of my own code.
I have commented my code, particularly in hard-to-understand areas.
I have made corresponding changes to the documentation.
My changes generate no new warnings.

…t-google-style

kshirajahere · 2026-03-24T08:30:30Z

@yukthagangadhari5 please follow PR template

kshirajahere

Thanks for the docstring/type-hint cleanup effort here. I think there is useful work in this PR, but I don’t think it is mergeable in the current form because it is carrying some non-doc behavior changes as well.

I reproduced this locally with:
pytest tests/test_datasets.py -k "ensemble or out_of_bounds or forecast_len or no_forcing or shorter_forcing" -q
which fails multiple current ensemble tests, including:

test_ensemble_len_scales_with_default_all_members
test_ensemble_index_mapping_is_time_major
test_ensemble_forcing_uses_same_member_when_available

My suggestion would be:

rebase onto current main
keep only pure docstring/type-hint edits in this PR
move the interrogate hook to a separate PR if that is still wanted

I’ve left a few inline comments on the concrete spots where behavior changed.

kshirajahere · 2026-03-24T08:40:27Z

-
-            # check that there are enough forecast steps available to create
-            # samples given the number of autoregressive steps requested
+            if self.datastore.is_ensemble:


This seems to reintroduce pre-#332 behavior. WeatherDataset on main now iterates all ensemble members by default, whereas this branch warns that only the first member is used and later hardcodes ensemble_member=0 again below. I can reproduce current test failures from this revert (test_ensemble_len_scales_with_default_all_members, test_ensemble_index_mapping_is_time_major, test_ensemble_forcing_uses_same_member_when_available).

kshirajahere · 2026-03-24T08:40:49Z

-        else:
-            da_forcing = None
+        da_forcing = self.da_forcing
+        if da_forcing is not None and "ensemble_member" in da_forcing.dims:


This also looks like a revert of merged #332 behavior. Main now supports forcing data with its own ensemble_member dimension by selecting the matching member; raising NotImplementedError here breaks the current contract again.

kshirajahere · 2026-03-24T08:41:13Z

+        Args:
+            stage (str, optional): "fit", "test", or None.
+        """
+        common_kwargs = {


WeatherDataModule still stores load_single_member, but setup() no longer forwards it into WeatherDataset(...). That makes the flag a no-op through the main DataModule path, which seems like an accidental functional change rather than part of the docstring migration.

kshirajahere · 2026-03-24T08:41:33Z

      - id: flake8
        description: Check Python code for correctness, consistency and adherence to best practices
        additional_dependencies: [Flake8-pyproject]
+  - repo: https://github.com/econchick/interrogate


This feels out of scope for this PR. Adding a repo-wide interrogate --fail-under=100 hook changes contributor workflow for the entire package, not just weather_dataset.py or the files touched here. If we want this policy, I think it should be proposed and landed separately.

kshirajahere · 2026-03-24T08:43:00Z

@@ -1,7 +1,8 @@
 # Standard library
 import dataclasses
+import argparse


This import looks unused, and the branch currently fails flake8 on touched files. I think the type-hint/docstring cleanup should probably be rebased and kept lint-clean before review, especially since the PR already has larger scope drift elsewhere.

yukthagangadhari5 · 2026-03-24T15:54:07Z

Hi @kshirajahere, I have updated the PR description to follow the repository's official template. Thanks for pointing that out!

kshirajahere · 2026-03-24T16:12:32Z

@yukthagangadhari5 i reckon its still not in the "correct" PR template

kshirajahere · 2026-03-24T16:14:18Z

Follow this mate

## Describe your changes
<FILL THIS WITH DESCRIPTION>
### Dependencies
No new external dependencies.

## Issue Link
<CLOSES #?>

## 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](README.MD) 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](https://www.gitkraken.com/learn/git/best-practices/git-commit-message#using-imperative-verb-form)).
- [ ] 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.

sadamov · 2026-03-25T02:04:10Z

@yukthagangadhari5 I looked through your changes here and there seem to be changes to the type hints, the docstrings and the code. For everything related to docstrings please directly contribute to #196. I you still think there is a need to improve the type hints specifically, you can reduce this PR in scope and rebase it on the main branch.

Mohit-Lakra and others added 8 commits March 18, 2026 20:38

Datastore docstring Updated

73897c5

Precommit Updated for 100 percent docs coverage check

6ddc4ae

fix: improve datastore error message clarity

43e55fe

fix: improve datastore error message clarity

e5dc2d6

docs: improve CLI description for graph visualization

3230b35

docs(dataset): migrate WeatherDataset to Google-style and add type hints

208ffc5

Merge remote-tracking branch 'upstream/main' into docs/weather-datase…

6edeb4a

…t-google-style

chore: resolve merge conflicts and update documentation

6cd6115

kshirajahere suggested changes Mar 24, 2026

View reviewed changes

sadamov closed this Mar 25, 2026

Navigation Menu

Uh oh!

docs(dataset): migrate WeatherDataset to Google-style and add type hints - #505

docs(dataset): migrate WeatherDataset to Google-style and add type hints#505
yukthagangadhari5 wants to merge 8 commits into
mllam:mainfrom
yukthagangadhari5:docs/weather-dataset-google-style

yukthagangadhari5 commented Mar 24, 2026 •

edited

Loading

Uh oh!

kshirajahere commented Mar 24, 2026

Uh oh!

kshirajahere left a comment

Uh oh!

kshirajahere Mar 24, 2026

Uh oh!

kshirajahere Mar 24, 2026

Uh oh!

kshirajahere Mar 24, 2026

Uh oh!

kshirajahere Mar 24, 2026

Uh oh!

kshirajahere Mar 24, 2026

Uh oh!

yukthagangadhari5 commented Mar 24, 2026

Uh oh!

kshirajahere commented Mar 24, 2026

Uh oh!

kshirajahere commented Mar 24, 2026 •

edited

Loading

Uh oh!

sadamov commented Mar 25, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

yukthagangadhari5 commented Mar 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Key Changes

Type of change

How Has This Been Tested?

Checklist:

Uh oh!

kshirajahere commented Mar 24, 2026

Uh oh!

kshirajahere left a comment

Choose a reason for hiding this comment

Uh oh!

kshirajahere Mar 24, 2026

Choose a reason for hiding this comment

Uh oh!

kshirajahere Mar 24, 2026

Choose a reason for hiding this comment

Uh oh!

kshirajahere Mar 24, 2026

Choose a reason for hiding this comment

Uh oh!

kshirajahere Mar 24, 2026

Choose a reason for hiding this comment

Uh oh!

kshirajahere Mar 24, 2026

Choose a reason for hiding this comment

Uh oh!

yukthagangadhari5 commented Mar 24, 2026

Uh oh!

kshirajahere commented Mar 24, 2026

Uh oh!

kshirajahere commented Mar 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sadamov commented Mar 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

yukthagangadhari5 commented Mar 24, 2026 •

edited

Loading

kshirajahere commented Mar 24, 2026 •

edited

Loading

sadamov commented Mar 25, 2026 •

edited

Loading