From ae0f945b4380b39633099073e779604a075d02af Mon Sep 17 00:00:00 2001 From: Jacob Mathias Schreiner Date: Thu, 13 Feb 2025 16:59:43 +0100 Subject: [PATCH 01/10] use explicit imports instead of relative --- neural_lam/create_graph.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/neural_lam/create_graph.py b/neural_lam/create_graph.py index ef979be3c..610178665 100644 --- a/neural_lam/create_graph.py +++ b/neural_lam/create_graph.py @@ -12,9 +12,9 @@ import torch_geometric as pyg from torch_geometric.utils.convert import from_networkx -# Local -from .config import load_config_and_datastore -from .datastore.base import BaseRegularGridDatastore +# First-party +from neural_lam.config import load_config_and_datastore +from neural_lam.datastore.base import BaseRegularGridDatastore, MDPDatastore def plot_graph(graph, title=None): From dd0aec134f7d55ac4dd914f5f2f213d71726457e Mon Sep 17 00:00:00 2001 From: Jacob Mathias Schreiner Date: Thu, 13 Feb 2025 16:59:59 +0100 Subject: [PATCH 02/10] enable using a datastore arg --- neural_lam/create_graph.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/neural_lam/create_graph.py b/neural_lam/create_graph.py index 610178665..6be06cbdb 100644 --- a/neural_lam/create_graph.py +++ b/neural_lam/create_graph.py @@ -588,14 +588,20 @@ def cli(input_args=None): action="store_true", help="Generate hierarchical mesh graph (default: False)", ) + parser.add_argument( + "--datastore", + type=str, + help="Path to datastore", + ) + args = parser.parse_args(input_args) assert ( - args.config_path is not None - ), "Specify your config with --config_path" + args.config_path or args.datastore + ), "Specify your config with --config_path or datastore with --datastore" # Load neural-lam configuration and datastore to use - _, datastore = load_config_and_datastore(config_path=args.config_path) + datastore = load_datastore(args) create_graph_from_datastore( datastore=datastore, @@ -606,5 +612,14 @@ def cli(input_args=None): ) +def load_datastore(args): + if args.config_path: + _, datastore = load_config_and_datastore(config_path=args.config_path) + else: + datastore = MDPDatastore.from_path(args.datastore) + + return datastore + + if __name__ == "__main__": cli() From ab9ed6858fe3ae80e55a245d17e6e1d8bd8bd5bb Mon Sep 17 00:00:00 2001 From: Jacob Mathias Schreiner Date: Thu, 13 Feb 2025 17:50:41 +0100 Subject: [PATCH 03/10] add type argument with mdp default --- neural_lam/create_graph.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/neural_lam/create_graph.py b/neural_lam/create_graph.py index 6be06cbdb..b30d7779c 100644 --- a/neural_lam/create_graph.py +++ b/neural_lam/create_graph.py @@ -14,7 +14,8 @@ # First-party from neural_lam.config import load_config_and_datastore -from neural_lam.datastore.base import BaseRegularGridDatastore, MDPDatastore +from neural_lam.datastore import DATASTORES +from neural_lam.datastore.base import BaseRegularGridDatastore def plot_graph(graph, title=None): @@ -593,6 +594,12 @@ def cli(input_args=None): type=str, help="Path to datastore", ) + parser.add_argument( + "--datastore_type", + type=str, + default="mdp", + help="Type of datastore (default: mdp)", + ) args = parser.parse_args(input_args) @@ -616,7 +623,8 @@ def load_datastore(args): if args.config_path: _, datastore = load_config_and_datastore(config_path=args.config_path) else: - datastore = MDPDatastore.from_path(args.datastore) + assert args.datastore_type in DATASTORES + DATASTORES[args.datastore_type](args.datastore) return datastore From f9813d6a2e644cb77ecb4d3744dd4ea8e2c8f4f4 Mon Sep 17 00:00:00 2001 From: Jacob Mathias Schreiner Date: Thu, 13 Feb 2025 17:52:43 +0100 Subject: [PATCH 04/10] named arg --- neural_lam/create_graph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/neural_lam/create_graph.py b/neural_lam/create_graph.py index b30d7779c..3dc2c8bae 100644 --- a/neural_lam/create_graph.py +++ b/neural_lam/create_graph.py @@ -624,7 +624,7 @@ def load_datastore(args): _, datastore = load_config_and_datastore(config_path=args.config_path) else: assert args.datastore_type in DATASTORES - DATASTORES[args.datastore_type](args.datastore) + DATASTORES[args.datastore_type](config_path=args.datastore) return datastore From c71a0867c915606f7d26bdd7e09dc1e0d34259e3 Mon Sep 17 00:00:00 2001 From: Jacob Mathias Schreiner Date: Thu, 13 Feb 2025 18:49:39 +0100 Subject: [PATCH 05/10] define the datastore --- neural_lam/create_graph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/neural_lam/create_graph.py b/neural_lam/create_graph.py index 3dc2c8bae..b5536a3ee 100644 --- a/neural_lam/create_graph.py +++ b/neural_lam/create_graph.py @@ -624,7 +624,7 @@ def load_datastore(args): _, datastore = load_config_and_datastore(config_path=args.config_path) else: assert args.datastore_type in DATASTORES - DATASTORES[args.datastore_type](config_path=args.datastore) + datastore = DATASTORES[args.datastore_type](config_path=args.datastore) return datastore From 94fd63820beb208ec3cb3768507581653562cbb1 Mon Sep 17 00:00:00 2001 From: Jacob Mathias Schreiner Date: Thu, 13 Feb 2025 19:09:15 +0100 Subject: [PATCH 06/10] add tests --- tests/test_graph_creation.py | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/tests/test_graph_creation.py b/tests/test_graph_creation.py index 93a7a55f4..e41124dfa 100644 --- a/tests/test_graph_creation.py +++ b/tests/test_graph_creation.py @@ -1,13 +1,14 @@ # Standard library import tempfile from pathlib import Path +from unittest import mock # Third-party import pytest import torch # First-party -from neural_lam.create_graph import create_graph_from_datastore +from neural_lam.create_graph import create_graph_from_datastore, load_datastore from neural_lam.datastore import DATASTORES from neural_lam.datastore.base import BaseRegularGridDatastore from tests.conftest import init_datastore_example @@ -117,3 +118,25 @@ def test_graph_creation(datastore_name, graph_name): assert r.shape[0] == 2 # adjacency matrix uses two rows elif file_id.endswith("_features"): assert r.shape[1] == d_features + + +def test_load_datastore_works_for_datastore_path(): + args = mock.MagicMock() + args.config_path = None + args.datastore_type = "mdp" + args.datastore = ( + "tests/datastore_examples/mdp/danra_100m_winds/danra.datastore.yaml" + ) + + datastore = load_datastore(args) + assert isinstance(datastore, DATASTORES[args.datastore_type]) + + +def test_load_datastore_works_for_config_path(): + args = mock.MagicMock() + args.config_path = ( + "tests/datastore_examples/mdp/danra_100m_winds/danra.datastore.yaml" + ) + + datastore = load_datastore(args) + assert isinstance(datastore, DATASTORES["mdp"]) From 35bf9faa5d85f8c0ebfd0b3277ccdcb46da07f3a Mon Sep 17 00:00:00 2001 From: Jacob Mathias Schreiner Date: Thu, 13 Feb 2025 19:14:46 +0100 Subject: [PATCH 07/10] path to config --- tests/test_graph_creation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_graph_creation.py b/tests/test_graph_creation.py index e41124dfa..fc274dba1 100644 --- a/tests/test_graph_creation.py +++ b/tests/test_graph_creation.py @@ -135,7 +135,7 @@ def test_load_datastore_works_for_datastore_path(): def test_load_datastore_works_for_config_path(): args = mock.MagicMock() args.config_path = ( - "tests/datastore_examples/mdp/danra_100m_winds/danra.datastore.yaml" + "tests/datastore_examples/mdp/danra_100m_winds/config.yaml" ) datastore = load_datastore(args) From 74c7dc9e6f84ffb69b836982ce05286ade1d428b Mon Sep 17 00:00:00 2001 From: Jacob Mathias Schreiner Date: Fri, 7 Mar 2025 10:47:31 +0100 Subject: [PATCH 08/10] redundant comment --- neural_lam/create_graph.py | 1 - 1 file changed, 1 deletion(-) diff --git a/neural_lam/create_graph.py b/neural_lam/create_graph.py index b5536a3ee..d34dbfc06 100644 --- a/neural_lam/create_graph.py +++ b/neural_lam/create_graph.py @@ -607,7 +607,6 @@ def cli(input_args=None): args.config_path or args.datastore ), "Specify your config with --config_path or datastore with --datastore" - # Load neural-lam configuration and datastore to use datastore = load_datastore(args) create_graph_from_datastore( From c9c88ffa03a05eddb0578e73de03b454fbdaf803 Mon Sep 17 00:00:00 2001 From: Jacob Mathias Schreiner Date: Fri, 7 Mar 2025 10:52:39 +0100 Subject: [PATCH 09/10] deprecation warning --- neural_lam/create_graph.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/neural_lam/create_graph.py b/neural_lam/create_graph.py index d34dbfc06..de25efa78 100644 --- a/neural_lam/create_graph.py +++ b/neural_lam/create_graph.py @@ -1,5 +1,6 @@ # Standard library import os +import warnings from argparse import ArgumentParser # Third-party @@ -621,6 +622,11 @@ def cli(input_args=None): def load_datastore(args): if args.config_path: _, datastore = load_config_and_datastore(config_path=args.config_path) + warnings.warn( + "Using config_path is deprecated, use datastore instead, instead use --datastore and --datastore_type (default: to mdp)", + DeprecationWarning, + ) + else: assert args.datastore_type in DATASTORES datastore = DATASTORES[args.datastore_type](config_path=args.datastore) From d7c20bf696f86938a697b7f585255a7fd045a561 Mon Sep 17 00:00:00 2001 From: Jacob Mathias Schreiner Date: Fri, 7 Mar 2025 11:08:10 +0100 Subject: [PATCH 10/10] break string --- neural_lam/create_graph.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/neural_lam/create_graph.py b/neural_lam/create_graph.py index de25efa78..093fa3352 100644 --- a/neural_lam/create_graph.py +++ b/neural_lam/create_graph.py @@ -623,7 +623,8 @@ def load_datastore(args): if args.config_path: _, datastore = load_config_and_datastore(config_path=args.config_path) warnings.warn( - "Using config_path is deprecated, use datastore instead, instead use --datastore and --datastore_type (default: to mdp)", + """Using config_path is deprecated, use datastore instead, + instead use --datastore and --datastore_type (default: to mdp)""", DeprecationWarning, )