Enable creating graph from datastore YAML file

neural_lam/create_graph.py

@@ -1,5 +1,6 @@
 # Standard library
 import os
+import warnings
 from argparse import ArgumentParser
 
 # Third-party
@@ -12,9 +13,10 @@
 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 import DATASTORES
+from neural_lam.datastore.base import BaseRegularGridDatastore
 
 
 def plot_graph(graph, title=None):
@@ -588,14 +590,25 @@ 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",
+    )
+    parser.add_argument(
+        "--datastore_type",
+        type=str,
+        default="mdp",
+        help="Type of datastore (default: mdp)",
+    )
+
     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 +619,21 @@ 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)
+
+    return datastore
+
+
 if __name__ == "__main__":
     cli()

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/config.yaml"
+    )
+
+    datastore = load_datastore(args)
+    assert isinstance(datastore, DATASTORES["mdp"])