nnjai support #129

yuvraajnarula · 2025-01-20T15:46:48Z

Pull Request for issue #128

Description

This code provides a data pipeline for handling and processing satellite data using the NNJAI library. It focuses on creating a PyTorch Dataset class (AMSUDataset) for efficiently loading and accessing Advanced Microwave Sounding Unit (AMSU) data. Below is a detailed breakdown of its functionality:

Authentication Check: Verifies access to NNJAI's DataCatalog using _check_authentication().
Dynamic Dataset Initialization: Loads and filters AMSU data based on specified primary descriptors and additional variables.
PyTorch Integration: Implements Dataset and DataLoader for seamless compatibility with PyTorch-based pipelines.
Flexible Configuration: Users can customize dataset parameters, timestamp filters, and metadata selection.
Fallback for Missing Dependencies: Prompts users to install the NNJAI library if not already available.

This implementation enhances satellite data processing workflows by providing an efficient and modular solution.

Fixes #128

How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce.
Please also list any relevant details for your test configuration.

Unit Tests:
- Tested dataset loading with valid and invalid dataset names.
- Verified filtering based on time, primary_descriptors, and additional_variables.
- Checked the output format of __getitem__ to ensure data integrity.
Integration Tests:
- Used DataLoader to iterate over the dataset and verified the randomness of shuffled outputs.

Yes

If your changes affect data processing, have you plotted any changes? i.e. have you done a quick sanity check?

Yes

Checklist:

My code follows OCF's coding style guidelines
I have performed a self-review of my own code
I have made corresponding changes to the documentation
I have added tests that prove my fix is effective or that my feature works
I have checked my code and corrected any misspellings

jacobbieker

Thank you for this! This is a good start, I think there needs to be a few different changes before this can be merged, but a good start!

jacobbieker · 2025-01-20T16:42:04Z

graph_weather/data/nnjai_wrapp.py

+"""
+
+import numpy as np
+from nnja.io import _check_authentication


I would suggest only checking for the library when importing the package. Then the AMSUDataset doesn't have to be in a giant if/else.

Suggested change

from nnja.io import _check_authentication

try:

from nnja.io import _check_authentication

from nnja import DataCatalog

Except ModuleImportError:

print("NNJA-AI library not installed. Please install with `pip install git+https://github.com/brightbandtech/nnja-ai.git`")

graph_weather/data/nnjai_wrapp.py

jacobbieker · 2025-01-20T16:50:53Z

graph_weather/data/nnjai_wrapp.py

+            latitude = row["LAT"]
+            longitude = row["LON"]
+            metadata = np.array([row[col] for col in self.metadata_columns], dtype=np.float32)
+            return time, latitude, longitude, metadata


I would suggest returning this as a dictionary, so its easier to track.

Suggested change

return time, latitude, longitude, metadata

return {"timestamp": time, "latitude": latitude, "longitude": longitude, "metadata": metadata}

graph_weather/data/nnjai_wrapp.py

for more information, see https://pre-commit.ci

yuvraajnarula · 2025-01-20T17:21:39Z

Here is the output for the suggestions you requested.

(py310env) PS D:\code\os_contri\graph_weather> C:\Users\DELL\miniconda3\envs\py310env\python.exe graph_weather/data/nnjai_wrapp.py
Batch size: 4
Timestamps shape: torch.Size([4])
Latitudes shape: torch.Size([4])
Longitudes shape: torch.Size([4])
Metadata shape: torch.Size([4, 1])

All items in batch:

Item 0:
Time: 1609524352.0
Latitude: -31.94420051574707
Longitude: 60.85639953613281
Metadata: tensor([180.2300])

Item 1:
Time: 1609531264.0
Item 1:
Time: 1609531264.0
Latitude: -68.8927993774414
Item 1:
Time: 1609531264.0
Latitude: -68.8927993774414
Longitude: 57.41469955444336
Metadata: tensor([209.7500])
Item 1:
Time: 1609531264.0
Latitude: -68.8927993774414
Longitude: 57.41469955444336
Metadata: tensor([209.7500])

Item 1:
Time: 1609531264.0
Latitude: -68.8927993774414
Longitude: 57.41469955444336
Metadata: tensor([209.7500])

Item 2:
Latitude: -68.8927993774414
Longitude: 57.41469955444336
Metadata: tensor([209.7500])

Item 2:
Time: 1609520640.0
Latitude: -79.92410278320312
Longitude: 97.82360076904297
Metadata: tensor([164.9300])
Latitude: -68.8927993774414
Longitude: 57.41469955444336
Metadata: tensor([209.7500])

Item 2:
Time: 1609520640.0
Latitude: -79.92410278320312
Longitude: 97.82360076904297
Metadata: tensor([164.9300])

Item 3:
Time: 1609537536.0
Latitude: -14.215499877929688
Longitude: 179.46139526367188
Metadata: tensor([223.6900])
(py310env) PS D:\code\os_contri\graph_weather>
Latitude: -68.8927993774414
Longitude: 57.41469955444336
Metadata: tensor([209.7500])

Item 2:
Time: 1609520640.0
Latitude: -79.92410278320312
Longitude: 97.82360076904297
Metadata: tensor([164.9300])


Item 2:
Time: 1609520640.0
Latitude: -79.92410278320312
Longitude: 97.82360076904297
Metadata: tensor([164.9300])

Item 3:
Time: 1609537536.0
Latitude: -14.215499877929688
Longitude: 179.46139526367188
Metadata: tensor([223.6900])

jacobbieker · 2025-01-20T17:38:00Z

Here is the output for the suggestions you requested.

(py310env) PS D:\code\os_contri\graph_weather> C:\Users\DELL\miniconda3\envs\py310env\python.exe graph_weather/data/nnjai_wrapp.py
Batch size: 4
Timestamps shape: torch.Size([4])
Latitudes shape: torch.Size([4])
Longitudes shape: torch.Size([4])
Metadata shape: torch.Size([4, 1])

All items in batch:

Item 0:
Time: 1609524352.0
Latitude: -31.94420051574707
Longitude: 60.85639953613281
Metadata: tensor([180.2300])

Item 1:
Time: 1609531264.0
Item 1:
Time: 1609531264.0
Latitude: -68.8927993774414
Item 1:
Time: 1609531264.0
Latitude: -68.8927993774414
Longitude: 57.41469955444336
Metadata: tensor([209.7500])
Item 1:
Time: 1609531264.0
Latitude: -68.8927993774414
Longitude: 57.41469955444336
Metadata: tensor([209.7500])

Item 1:
Time: 1609531264.0
Latitude: -68.8927993774414
Longitude: 57.41469955444336
Metadata: tensor([209.7500])

Item 2:
Latitude: -68.8927993774414
Longitude: 57.41469955444336
Metadata: tensor([209.7500])

Item 2:
Time: 1609520640.0
Latitude: -79.92410278320312
Longitude: 97.82360076904297
Metadata: tensor([164.9300])
Latitude: -68.8927993774414
Longitude: 57.41469955444336
Metadata: tensor([209.7500])

Item 2:
Time: 1609520640.0
Latitude: -79.92410278320312
Longitude: 97.82360076904297
Metadata: tensor([164.9300])

Item 3:
Time: 1609537536.0
Latitude: -14.215499877929688
Longitude: 179.46139526367188
Metadata: tensor([223.6900])
(py310env) PS D:\code\os_contri\graph_weather>
Latitude: -68.8927993774414
Longitude: 57.41469955444336
Metadata: tensor([209.7500])

Item 2:
Time: 1609520640.0
Latitude: -79.92410278320312
Longitude: 97.82360076904297
Metadata: tensor([164.9300])


Item 2:
Time: 1609520640.0
Latitude: -79.92410278320312
Longitude: 97.82360076904297
Metadata: tensor([164.9300])

Item 3:
Time: 1609537536.0
Latitude: -14.215499877929688
Longitude: 179.46139526367188
Metadata: tensor([223.6900])

Hi, could you push the changes to this branch? I don't see any differences in the code?

for more information, see https://pre-commit.ci

yuvraajnarula · 2025-01-21T07:49:04Z

@jacobbieker any updates?

jacobbieker

Its getting there! The only thing I would go with changing is to add a unit test in the tests/ directory that checks the output shapes from the dataset to ensure it is what it expects. You basically have the test here, under the main bit, just change it to check for the output shapes and put that part in a unit test, then it can be merged!

jacobbieker · 2025-01-21T08:55:30Z

graph_weather/data/nnjai_wrapp.py

+    return {key: torch.stack([item[key] for item in batch]) for key in batch[0].keys()}
+
+
+if __name__ == "__main__":


Could you add a unit test for this dataset under the tests/? Basically, have it do this, but instead of printing it out, check for the expected shapes and such.

for more information, see https://pre-commit.ci

jacobbieker

This looks good to me! Thanks for making the changes.

yuvraajnarula · 2025-01-22T02:56:50Z

Thank you for guiding me. If there's anything more I could do for this issue then do let me know.

yuvraajnarula added 2 commits January 20, 2025 21:11

nnjai support

2a8c6e8

ruff format

c70d3c1

yuvraajnarula force-pushed the nnjai_support branch from bb78c36 to c70d3c1 Compare January 20, 2025 16:26

jacobbieker requested changes Jan 20, 2025

View reviewed changes

yuvraajnarula force-pushed the nnjai_support branch from a0358a2 to c70d3c1 Compare January 20, 2025 17:21

[pre-commit.ci] auto fixes from pre-commit.com hooks

9041657

for more information, see https://pre-commit.ci

yuvraajnarula and others added 2 commits January 20, 2025 23:22

changes as per requested

999eafb

[pre-commit.ci] auto fixes from pre-commit.com hooks

7440174

for more information, see https://pre-commit.ci

jacobbieker requested changes Jan 21, 2025

View reviewed changes

yuvraajnarula and others added 2 commits January 21, 2025 22:57

test_nnjai.py

7759818

[pre-commit.ci] auto fixes from pre-commit.com hooks

c400f0c

for more information, see https://pre-commit.ci

jacobbieker approved these changes Jan 21, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

nnjai support #129

nnjai support #129

yuvraajnarula commented Jan 20, 2025

jacobbieker left a comment

jacobbieker Jan 20, 2025

jacobbieker Jan 20, 2025

yuvraajnarula commented Jan 20, 2025

jacobbieker commented Jan 20, 2025

yuvraajnarula commented Jan 21, 2025

jacobbieker left a comment

jacobbieker Jan 21, 2025

jacobbieker left a comment

yuvraajnarula commented Jan 22, 2025

-from nnja.io import _check_authentication
+try:
+    from nnja.io import _check_authentication
+    from nnja import DataCatalog
+Except ModuleImportError:
+   print("NNJA-AI library not installed. Please install with `pip install git+https://github.com/brightbandtech/nnja-ai.git`")

	return time, latitude, longitude, metadata
	return {"timestamp": time, "latitude": latitude, "longitude": longitude, "metadata": metadata}

		return {key: torch.stack([item[key] for item in batch]) for key in batch[0].keys()}


		if __name__ == "__main__":

nnjai support #129

Are you sure you want to change the base?

nnjai support #129

Conversation

yuvraajnarula commented Jan 20, 2025

Pull Request for issue #128

Description

How Has This Been Tested?

Checklist:

jacobbieker left a comment

Choose a reason for hiding this comment

jacobbieker Jan 20, 2025

Choose a reason for hiding this comment

jacobbieker Jan 20, 2025

Choose a reason for hiding this comment

yuvraajnarula commented Jan 20, 2025

jacobbieker commented Jan 20, 2025

yuvraajnarula commented Jan 21, 2025

jacobbieker left a comment

Choose a reason for hiding this comment

jacobbieker Jan 21, 2025

Choose a reason for hiding this comment

jacobbieker left a comment

Choose a reason for hiding this comment

yuvraajnarula commented Jan 22, 2025