diff --git a/CHANGELOG.md b/CHANGELOG.md index 6cf3ab0..5223180 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Add `__version__` attribute to the package init [\#56](https://github.com/mllam/weather-model-graphs/pull/56) @AdMub +### Changed + +- Make `test_create_decode_mask` more specific and future-proof by checking only m2g (decoding) edges instead of total edges. [\#104](https://github.com/mllam/weather-model-graphs/pull/104) @zweihuehner + ## [v0.3.0](https://github.com/mllam/weather-model-graphs/releases/tag/v0.3.0) diff --git a/README.md b/README.md index e22fd8d..41d2265 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,8 @@ -# weather-model-graphs +# poseidon-graphprep [![tests](https://github.com/mllam/weather-model-graphs/actions/workflows/ci-tests.yml/badge.svg)](https://github.com/mllam/weather-model-graphs/actions/workflows/ci-tests.yml) [![linting](https://github.com/mllam/weather-model-graphs/actions/workflows/pre-commit.yml/badge.svg)](https://github.com/mllam/weather-model-graphs/actions/workflows/pre-commit.yml) [![Jupyter Book Badge](https://jupyterbook.org/badge.svg)](https://mllam.github.io/weather-model-graphs) -`weather-model-graphs` is a package for creating, visualising and storing graphs used in message-passing graph-based data-driven weather models. +`poseidon-graphprep` is a package for creating, visualising and storing graphs used in message-passing graph-based data-driven weather models. The package is designed to use `networkx.DiGraph` objects as the primary data structure for the graph representation right until the graph is to be stored on disk into a specific format. This makes the graph generation process modular (every step outputs a `networkx.DiGraph`), easy to debug (visualise the graph at any step) and allows output to different file-formats and file-structures to be easily implemented. More details are given in the [background](https://mllam.github.io/weather-model-graphs/background.html) and [design](https://mllam.github.io/weather-model-graphs/design.html) section of the online [documentation](https://mllam.github.io/weather-model-graphs/). @@ -10,34 +10,34 @@ This makes the graph generation process modular (every step outputs a `networkx. ## Installation -If you simply want to install and use `weather-model-graphs` as-is you can install the most recent release directly from pypi with pip +If you simply want to install and use `poseidon-graphprep` as-is you can install the most recent release directly from pypi with pip ```bash -python -m pip install weather-model-graphs +python -m pip install poseidon-graphprep ``` If you want to be able to save to pytorch-geometric data-structure used in [neural-lam](https://github.com/mllam/neural-lam) then you will need to install pytorch and pytorch-geometric too. This can be done by with the `pytorch` -optional extra in `weather-model-graphs`: +optional extra in `poseidon-graphprep`: ```bash -python -m pip install weather-model-graphs[pytorch] +python -m pip install poseidon-graphprep[pytorch] ``` This will install the CPU version of pytorch by default. If you want to install a GPU variant you should [install that first](https://pytorch.org/get-started/locally/) before installing -`weather-model-graphs`. +`poseidon-graphprep`. -## Developing `weather-model-graphs` +## Developing `poseidon-graphprep` -The easiest way to work on developing `weather-model-graphs` is to fork the [main repo](https://github.com/mllam/weather-model-graphs) under your github account, clone this repo locally, install [pdm](https://pdm-project.org/en/latest/), create a venv with pdm and then install `weather-model-graphs` (and all development dependencies): +The easiest way to work on developing `poseidon-graphprep` is to fork the [main repo](https://github.com/mllam/weather-model-graphs) under your github account, clone this repo locally, install [pdm](https://pdm-project.org/en/latest/), create a venv with pdm and then install `poseidon-graphprep` (and all development dependencies): ``` -git clone https://github.com//weather-model-graphs -cd weather-model-graphs +git clone https://github.com//poseidon-graphprep +cd poseidon-graphprep pdm venv create pdm use --venv in-project pdm install --dev @@ -69,13 +69,13 @@ pdm install --group pytorch # Usage -The best way to understand how to use `weather-model-graphs` is to look at the [documentation](https://mllam.github.io/weather-model-graphs) (which are executable Jupyter notebooks!), to have look at the tests in [tests/](tests/) or simply to read through the source code. +The best way to understand how to use `poseidon-graphprep` is to look at the [documentation](https://mllam.github.io/weather-model-graphs) (which are executable Jupyter notebooks!), to have look at the tests in [tests/](tests/) or simply to read through the source code. ## Example, Keisler 2021 flat graph architecture ```python import numpy as np -import weather_model_graphs as wmg +import poseidon_graphprep as wmg # define your (x,y) grid coodinates xs, ys = np.meshgrid(np.linspace(0, 1, 32), np.linspace(0, 1, 32)) @@ -97,4 +97,4 @@ for component, graph in graph_components.items(): # Documentation -The documentation is built using [Jupyter Book](https://jupyterbook.org/intro.html) and can be found at [https://mllam.github.io/weather-model-graphs](https://mllam.github.io/weather-model-graphs). This includes background on graph-based weather models, the design principles of `weather-model-graphs` and how to use it to create your own graph architectures. +The documentation is built using [Jupyter Book](https://jupyterbook.org/intro.html) and can be found at [https://mllam.github.io/weather-model-graphs](https://mllam.github.io/weather-model-graphs). This includes background on graph-based weather models, the design principles of `poseidon-graphprep` and how to use it to create your own graph architectures. diff --git a/docs/background.ipynb b/docs/background.ipynb index d5df8f1..943e8a3 100644 --- a/docs/background.ipynb +++ b/docs/background.ipynb @@ -6,7 +6,7 @@ "source": [ "# Background\n", "\n", - "> *this section provides an introduction to what graph-based data-driven weather models are and how they work. Inline code using `weather-model-graphs` is used to demonstrate different parts of the graph*\n", + "> *this section provides an introduction to what graph-based data-driven weather models are and how they work. Inline code using `poseidon-graphprep` is used to demonstrate different parts of the graph*\n", "\n", "Current graph-based weather models use the [encode-process-decode paradigm](https://arxiv.org/abs/1806.01261) on [message-passing graphs](https://arxiv.org/abs/1704.01212) to do the auto-regressive temporal prediction of the atmospheric weather state to produce a weather forecast. \n", "The graphs are directed acyclic graphs (DAGs) with the nodes representing features (physical variables) at a given location in space and the edges representing flow of information.\n", @@ -14,7 +14,7 @@ "\n", "## The graph nodes\n", "\n", - "Using the nomenclature of [Lam et al 2022](https://arxiv.org/abs/2212.12794) the nodes in `weather-model-graphs` are split into two types:\n", + "Using the nomenclature of [Lam et al 2022](https://arxiv.org/abs/2212.12794) the nodes in `poseidon-graphprep` are split into two types:\n", "\n", "- **grid nodes**: representing the physical variables of the atmospheric state at a specific `(x,y)` coordinate in the (input) initial state to the model and the (output) prediction of the model\n", "\n", @@ -28,7 +28,7 @@ "outputs": [], "source": [ "import numpy as np\n", - "import weather_model_graphs as wmg\n", + "import poseidon_graphprep as wmg\n", "\n", "# create some fake cartesian coordinates\n", "def _create_fake_xy(N=10):\n", @@ -53,7 +53,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "As you can see in the code snippet above, the only input that graph generation in `weather-model-graphs` requires are the static `(x,y)` *grid* coordinates of the atmospheric state as the state changes over time. These coordinates are used to create the **grid nodes** nodes of the graph, with a node for each `(x,y)` coordinate.\n", + "As you can see in the code snippet above, the only input that graph generation in `poseidon-graphprep` requires are the static `(x,y)` *grid* coordinates of the atmospheric state as the state changes over time. These coordinates are used to create the **grid nodes** nodes of the graph, with a node for each `(x,y)` coordinate.\n", "\n", "In addition to grid nodes the graph also contains **mesh nodes** that represent the latent space of the model at a set of `(x,y)` coordinates (this is in general a different set of coordinates to the **grid nodes** coordinates)." ] diff --git a/docs/creating_the_graph.ipynb b/docs/creating_the_graph.ipynb index 8cbc913..37fc9d4 100644 --- a/docs/creating_the_graph.ipynb +++ b/docs/creating_the_graph.ipynb @@ -9,7 +9,7 @@ "\n", "Within graph-based weather models one single graph is used to represent the encode-process-decode operations of the data-driven weather model. The graph is a directed acyclic graph (DAG) with the nodes representing features at a given location in space and the edges representing flow of information.\n", "\n", - "`weather-model-graphs` provides a framework for creating and visualising graphs for data-driven weather models. The framework is designed to be flexible and allow for the creation of a wide range of graph architectures.\n", + "`poseidon-graphprep` provides a framework for creating and visualising graphs for data-driven weather models. The framework is designed to be flexible and allow for the creation of a wide range of graph architectures.\n", "\n", "The graph is comprised of three components that represent the three encode-process-decode operations:\n", "\n", @@ -26,7 +26,7 @@ "source": [ "# The grid nodes\n", "\n", - "To get started we will create a set of fake grid nodes, which represent the geographical locations where we have values for the physical fields. We will here work with cartesian x/y coordinates. See [this page](./lat_lons.ipynb) for how to use lat/lon coordinates in weather-model-graphs." + "To get started we will create a set of fake grid nodes, which represent the geographical locations where we have values for the physical fields. We will here work with cartesian x/y coordinates. See [this page](./lat_lons.ipynb) for how to use lat/lon coordinates in poseidon-graphprep." ] }, { @@ -37,7 +37,7 @@ "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt\n", - "import weather_model_graphs as wmg" + "import poseidon_graphprep as wmg" ] }, { @@ -80,7 +80,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Lets start with a simple mesh which only has nearest neighbour connections. At the moment `weather-model-graphs` creates a rectangular mesh that sits within the spatial domain spanned by the grid nodes (specifically within the axis-aligned bounding box of the grid nodes). Techniques for adding non-square meshes are in development." + "Lets start with a simple mesh which only has nearest neighbour connections. At the moment `poseidon-graphprep` creates a rectangular mesh that sits within the spatial domain spanned by the grid nodes (specifically within the axis-aligned bounding box of the grid nodes). Techniques for adding non-square meshes are in development." ] }, { @@ -101,7 +101,7 @@ "source": [ "# Graph archetypes\n", "\n", - "To simplify the creation of the full *encode-process-decode* graph, the `weather-model-graphs` package contains implementations of a number of graph archetypes.\n", + "To simplify the creation of the full *encode-process-decode* graph, the `poseidon-graphprep` package contains implementations of a number of graph archetypes.\n", "These architypes principally differ in the way the mesh component of the graph is constructed, but also in the way the grid and mesh components are connected. As more approaches are developed they will be added to the library.\n" ] }, diff --git a/docs/decoding_mask.ipynb b/docs/decoding_mask.ipynb index 7a2d5f4..bcc1379 100644 --- a/docs/decoding_mask.ipynb +++ b/docs/decoding_mask.ipynb @@ -15,7 +15,7 @@ "metadata": {}, "outputs": [], "source": [ - "import weather_model_graphs as wmg\n", + "import poseidon_graphprep as wmg\n", "import numpy as np\n", "import matplotlib.pyplot as plt" ] diff --git a/docs/design.ipynb b/docs/design.ipynb index 7ff8b12..64520b4 100644 --- a/docs/design.ipynb +++ b/docs/design.ipynb @@ -6,14 +6,14 @@ "source": [ "# Design principles\n", "\n", - "The key design principle of `weather-model-graphs` is to work with `networkx.DiGraph` objects as the primary data structure for the graph representation right until the graph is to be stored on disk into a specific format.\n", + "The key design principle of `poseidon-graphprep` is to work with `networkx.DiGraph` objects as the primary data structure for the graph representation right until the graph is to be stored on disk into a specific format.\n", "Using only `networkx.DiGraph` objects as the intermediate representations makes it possible to\n", "\n", "1) easily modularise the whole generation process, with every step outputting a `networkx.DiGraph` object,\n", "2) easily visualise the graph resulting from any step and\n", "3) easily connect graph nodes across graph components, combine graphs and split graphs based on node and edge attributes.\n", "\n", - "The graph generation in `weather-model-graphs` is split into to the following steps:\n", + "The graph generation in `poseidon-graphprep` is split into to the following steps:\n", "\n", "1. Create the three graph components of the message-passing graph that constitute the auto-regressive atmospheric flow model, all represented by `networkx.DiGraph` objects:\n", "\n", @@ -27,9 +27,9 @@ "3. Split the combined graph into the three *output* graph components again (or more if the specific graph architecture requires it).\n", "4. Store each of the *output* graph components in the desired format, for example:\n", "\n", - " - **networkx** `.pickle` file: save `networkx.DiGraph` objects using `pickle` to disk (`weather_model_graphs.save.to_pickle(...)`)\n", + " - **networkx** `.pickle` file: save `networkx.DiGraph` objects using `pickle` to disk (`poseidon_graphprep.save.to_pickle(...)`)\n", "\n", - " - [pytorch-geometric](https://github.com/pyg-team/pytorch_geometric) for [neural-lam](https://github.com/mllam/neural-lam): edges indexes and features are stored in separate `torch.Tensor` objects serialised to disk that can then be loaded into `torch_geometric.data.Data` objects (`weather_model_graphs.save.to_pyg(...)`)" + " - [pytorch-geometric](https://github.com/pyg-team/pytorch_geometric) for [neural-lam](https://github.com/mllam/neural-lam): edges indexes and features are stored in separate `torch.Tensor` objects serialised to disk that can then be loaded into `torch_geometric.data.Data` objects (`poseidon_graphprep.save.to_pyg(...)`)" ] }, { @@ -42,7 +42,7 @@ "source": [ "## Diagram of the graph generation process:\n", "\n", - "Below, the graph generation process is visualised in `weather-model-graphs` for the example given above:\n", + "Below, the graph generation process is visualised in `poseidon-graphprep` for the example given above:\n", "\n", "\n", "```{mermaid}\n", @@ -50,7 +50,7 @@ "\n", "xy_grid[\"(x,y) grid coordinates [numpy.ndarray]\"]\n", "\n", - "subgraph weather_model_graphs[\"weather-model-graphs\"]\n", + "subgraph poseidon_graphprep[\"poseidon-graphprep\"]\n", " subgraph graph_components[\"Graph components\"]\n", " G_g2m[\"G_g2m[networkx.DiGraph]\"]\n", " G_m2m[\"G_m2m[networkx.DiGraph]\"]\n", @@ -103,7 +103,7 @@ "\n", "class hidden_graph1 hidden\n", "class hidden_graph2 hidden\n", - "class weather_model_graphs title_left\n", + "class poseidon_graphprep title_left\n", "class graph_components title_left\n", "class stored_on_disk title_left\n", "class pyg_loaded title_left\n", @@ -111,7 +111,7 @@ "\n", "### Node and edge attributes\n", "\n", - "There are a number of node and edge attributes with special meanings in `weather-model-graphs` which enable the splitting and visualisation of the graph components.\n", + "There are a number of node and edge attributes with special meanings in `poseidon-graphprep` which enable the splitting and visualisation of the graph components.\n", "\n", "#### Node attributes\n", "\n", @@ -138,7 +138,7 @@ "source": [ "### Splitting graphs\n", "\n", - "The splitting of the graph is done with by utilising the edge attributes, and thus it is easy to split the complete graph by either which component the edge belongs to, or by the level of the edge in the graph. This is done using the `weather_model_graphs.split_graph_by_edge_attribute(...)` function.\n" + "The splitting of the graph is done with by utilising the edge attributes, and thus it is easy to split the complete graph by either which component the edge belongs to, or by the level of the edge in the graph. This is done using the `poseidon_graphprep.split_graph_by_edge_attribute(...)` function.\n" ] }, { @@ -151,10 +151,10 @@ "source": [ "## Code layout\n", "\n", - "The code layout of `weather-model-graphs` is organised into submodules by the functionality they provide. The main submodules are:\n", + "The code layout of `poseidon-graphprep` is organised into submodules by the functionality they provide. The main submodules are:\n", "\n", "```\n", - "weather_model_graphs\n", + "poseidon_graphprep\n", " .create\n", " .archetype:\n", " for creating specific archetype graph\n", diff --git a/docs/intro.md b/docs/intro.md index 6cadd2d..4e845ab 100644 --- a/docs/intro.md +++ b/docs/intro.md @@ -1,11 +1,11 @@ # Intro -> *`weather-model-graphs`: "Graphs for data-driven weather models"* +> *`poseidon-graphprep`: "Graphs for data-driven weather models"* -You are currently looking at a Jupyter Book for the [weather-model-graphs](https://github.com/mllam/weather-model-graphs) package. +You are currently looking at a Jupyter Book for the [poseidon-graphprep](https://github.com/mllam/weather-model-graphs) package. This package aims to be a tool for constructing, visualizing and storing graphs for data-driven weather models. -To get started you can have a look at the [background material](background) to get an overview of the concepts of graph-based data-driven weather models, check the [design principles](design) of `weather-model-graphs` or jump right into learning [how to create the graph](creating_the_graph). +To get started you can have a look at the [background material](background) to get an overview of the concepts of graph-based data-driven weather models, check the [design principles](design) of `poseidon-graphprep` or jump right into learning [how to create the graph](creating_the_graph). ## Table of Contents diff --git a/docs/lat_lons.ipynb b/docs/lat_lons.ipynb index fe807e2..623f7e2 100644 --- a/docs/lat_lons.ipynb +++ b/docs/lat_lons.ipynb @@ -66,7 +66,7 @@ "metadata": {}, "outputs": [], "source": [ - "import weather_model_graphs as wmg\n", + "import poseidon_graphprep as wmg\n", "\n", "graph = wmg.create.archetype.create_keisler_graph(coords, mesh_node_distance=10)\n", "fig, ax = plt.subplots(figsize=(15, 9), subplot_kw={\"projection\": ccrs.PlateCarree()})\n", diff --git a/pyproject.toml b/pyproject.toml index 01eb298..6a8967e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,10 +1,14 @@ [project] -name = "weather-model-graphs" +name = "poseidon-graphprep" version = "0.3.0" description = "Tools for creating and visualising data-driven weather model graphs" authors = [ {name = "Leif Denby", email = "lcd@dmi.dk"}, ] + +[tool.setuptools] +packages = ["poseidon_graphprep"] + dependencies = [ "numpy>=1.26.4", "loguru>=0.7.2", diff --git a/src/weather_model_graphs/__init__.py b/src/poseidon_graphprep/__init__.py similarity index 78% rename from src/weather_model_graphs/__init__.py rename to src/poseidon_graphprep/__init__.py index 374bbdb..3806e29 100644 --- a/src/weather_model_graphs/__init__.py +++ b/src/poseidon_graphprep/__init__.py @@ -1,7 +1,7 @@ import importlib.metadata try: - __version__ = importlib.metadata.version("weather-model-graphs") + __version__ = importlib.metadata.version("poseidon-graphprep") except importlib.metadata.PackageNotFoundError: __version__ = "unknown" diff --git a/src/weather_model_graphs/create/__init__.py b/src/poseidon_graphprep/create/__init__.py similarity index 100% rename from src/weather_model_graphs/create/__init__.py rename to src/poseidon_graphprep/create/__init__.py diff --git a/src/weather_model_graphs/create/archetype.py b/src/poseidon_graphprep/create/archetype.py similarity index 72% rename from src/weather_model_graphs/create/archetype.py rename to src/poseidon_graphprep/create/archetype.py index b4d717e..1de5955 100644 --- a/src/weather_model_graphs/create/archetype.py +++ b/src/poseidon_graphprep/create/archetype.py @@ -4,6 +4,77 @@ from .base import create_all_graph_components +def create_mike_mesh_graph( + coords, + mesh_node_distance=3, + coords_crs: pyproj.crs.CRS | None = None, + graph_crs: pyproj.crs.CRS | None = None, + decode_mask: Iterable[bool] | None = None, + return_components: bool = False, +): + """ + Create a flat LAM graph from Oskarsson et al (2023, https://arxiv.org/abs/2309.17370) + This graph setup is inspired by the global graph used by Keisler (2022, https://arxiv.org/abs/2202.07575). + + This graph is a flat single scale graph with nearest neighbour connectivity + (8 neighbours) within the mesh. + + The grid to mesh connectivity connects each mesh node to grid nodes withing + distance 0.51d, where d is the length of diagonal edges between neighbouring + mesh nodes. The choice of 0.51 makes sure that all grid node positions will + be connected to at least one mesh node (see + https://www.desmos.com/calculator/sqqz0ka4ho for a visualization). + The mesh to grid connectivity connects each grid point to the 4 nearest mesh nodes. + + Parameters + ---------- + coords: np.ndarray + 2D array of grid point positions, in coordinate CRS + mesh_node_distance: float + Distance (in x- and y-direction) between created mesh nodes, + in coordinate system of coords + coords_crs: pyproj.crs.CRS or None + Coordinate Reference System (CRS) of the input `coords`. Must be a + pyproj.crs.CRS or cartopy.crs.CRS. + graph_crs: pyproj.crs.CRS or None + CRS where the graph creation should take place. If both `coords_crs` + and `graph_crs` are given, coords are transformed from `coords_crs` to + `graph_crs`. If any one of them is None, the original coords are used. + decode_mask: Iterable[bool] or None + An Iterable of booleans, masking which grid positions should be decoded + to (included in the m2g subgraph), i.e., which positions should be output. + Must have the same length as `coords`. If None (default), all grid nodes + are included. + return_components: bool, default False + If True, the function returns a dict with m2g, m2m, and g2m as separate + graphs. If False, returns one combined graph. + + Returns + ------- + networkx.DiGraph or dict[networkx.DiGraph] + The graph or graph components. + """ + return create_all_graph_components( + coords=coords, + m2m_connectivity="flat", + m2m_connectivity_kwargs=dict(mesh_node_distance=mesh_node_distance), + g2m_connectivity="within_radius", + non_decode_g2m_connectivity="within_radius", + m2g_connectivity="containing_rectangle", + g2m_connectivity_kwargs=dict( + rel_max_dist=0.6, + ), + non_decode_g2m_connectivity_kwargs=dict( + rel_max_dist=1.6, + ), + m2g_connectivity_kwargs=dict(), + coords_crs=coords_crs, + graph_crs=graph_crs, + decode_mask=decode_mask, + return_components=return_components, + allow_zero_degree=False, + remove_zero_degree_mesh=True + ) def create_keisler_graph( coords, @@ -60,13 +131,15 @@ def create_keisler_graph( m2m_connectivity="flat", m2m_connectivity_kwargs=dict(mesh_node_distance=mesh_node_distance), g2m_connectivity="within_radius", - m2g_connectivity="nearest_neighbours", + non_decode_g2m_connectivity="within_radius", + m2g_connectivity="containing_rectangle", g2m_connectivity_kwargs=dict( rel_max_dist=0.51, ), - m2g_connectivity_kwargs=dict( - max_num_neighbours=4, + non_decode_g2m_connectivity_kwargs=dict( + rel_max_dist=1.6, ), + m2g_connectivity_kwargs=dict(), coords_crs=coords_crs, graph_crs=graph_crs, decode_mask=decode_mask, @@ -139,13 +212,15 @@ def create_graphcast_graph( max_num_levels=max_num_levels, ), g2m_connectivity="within_radius", - m2g_connectivity="nearest_neighbours", + non_decode_g2m_connectivity="within_radius", + m2g_connectivity="containing_rectangle", g2m_connectivity_kwargs=dict( rel_max_dist=0.51, ), - m2g_connectivity_kwargs=dict( - max_num_neighbours=4, + non_decode_g2m_connectivity_kwargs=dict( + rel_max_dist=1.6, ), + m2g_connectivity_kwargs=dict(), coords_crs=coords_crs, graph_crs=graph_crs, decode_mask=decode_mask, @@ -223,13 +298,15 @@ def create_oskarsson_hierarchical_graph( max_num_levels=max_num_levels, ), g2m_connectivity="within_radius", - m2g_connectivity="nearest_neighbours", + non_decode_g2m_connectivity="within_radius", + m2g_connectivity="containing_rectangle", g2m_connectivity_kwargs=dict( rel_max_dist=0.51, ), - m2g_connectivity_kwargs=dict( - max_num_neighbours=4, + non_decode_g2m_connectivity_kwargs=dict( + rel_max_dist=1.6, ), + m2g_connectivity_kwargs=dict(), coords_crs=coords_crs, graph_crs=graph_crs, decode_mask=decode_mask, diff --git a/src/weather_model_graphs/create/base.py b/src/poseidon_graphprep/create/base.py similarity index 75% rename from src/weather_model_graphs/create/base.py rename to src/poseidon_graphprep/create/base.py index f922b2d..266bfbe 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/poseidon_graphprep/create/base.py @@ -18,11 +18,6 @@ import scipy.spatial from loguru import logger -from ..networkx_utils import ( - replace_node_labels_with_unique_ids, - split_graph_by_edge_attribute, - split_on_edge_attribute_existance, -) from .grid import create_grid_graph_nodes from .mesh.kinds.flat import ( create_flat_multiscale_mesh_graph, @@ -31,18 +26,98 @@ from .mesh.kinds.hierarchical import create_hierarchical_multiscale_mesh_graph +def _remove_zero_degree_mesh_nodes(G_g2m, G_grid, m2m_graph, grid_connect_graph): + """ + Safely remove mesh nodes that have zero degree in g2m, as long as removing + them does not disconnect the mesh. A node is removed only if ALL of its + m2m neighbors also have zero g2m degree. The check is applied iteratively + so that cascading removals from the boundary of dead zones are handled. + + The functino is not tested, so it could be that it does not safely prohibit + unconnceted islands + + Returns updated (G_g2m, m2m_graph, grid_connect_graph). + """ + grid_nodes_set = set(G_grid.nodes()) + g2m_degrees = dict(G_g2m.degree()) + + zero_degree_mesh_nodes = set( + node for node, deg in g2m_degrees.items() + if deg == 0 and node not in grid_nodes_set + ) + + if not zero_degree_mesh_nodes: + return G_g2m, m2m_graph, grid_connect_graph + + # Build undirected m2m view restricted to mesh nodes + m2m_undirected = m2m_graph.to_undirected() + mesh_nodes_in_m2m = [n for n in m2m_undirected.nodes() if n not in grid_nodes_set] + m2m_mesh_only = m2m_undirected.subgraph(mesh_nodes_in_m2m) + + # Iteratively mark nodes for removal + nodes_to_remove = set() + changed = True + while changed: + changed = False + for node in list(zero_degree_mesh_nodes - nodes_to_remove): + active_neighbors = [ + n for n in m2m_mesh_only.neighbors(node) + if n not in nodes_to_remove + ] + if len(active_neighbors) == 0 or all( + g2m_degrees.get(n, 0) == 0 for n in active_neighbors + ): + nodes_to_remove.add(node) + changed = True + + if nodes_to_remove: + logger.warning( + f"Removing {len(nodes_to_remove)} zero-degree mesh nodes " + f"from g2m/m2m (out of {len(zero_degree_mesh_nodes)} candidates)" + ) + # Rebuild graphs without removed nodes (handles frozen graphs) + keep = lambda graph: [n for n in graph.nodes() if n not in nodes_to_remove] + m2m_graph = m2m_graph.subgraph(keep(m2m_graph)).copy() + grid_connect_graph = grid_connect_graph.subgraph(keep(grid_connect_graph)).copy() + G_g2m = G_g2m.subgraph(keep(G_g2m)).copy() + + kept = len(zero_degree_mesh_nodes) - len(nodes_to_remove) + if kept > 0: + logger.info( + f"Kept {kept} zero-degree mesh nodes in g2m " + f"(at least one m2m neighbor is connected to the grid)" + ) + + # Log remaining zero-degree grid nodes + remaining_zero = [ + node for node, deg in dict(G_g2m.degree()).items() if deg == 0 + ] + remaining_grid = [n for n in remaining_zero if n in grid_nodes_set] + if remaining_grid: + logger.warning( + f"{len(remaining_grid)} grid nodes have zero degree in g2m " + f"(kept): {remaining_grid[:10]}..." + ) + + return G_g2m, m2m_graph, grid_connect_graph + + def create_all_graph_components( coords: np.ndarray, m2m_connectivity: str, m2g_connectivity: str, g2m_connectivity: str, + non_decode_g2m_connectivity: str = None, m2m_connectivity_kwargs={}, m2g_connectivity_kwargs={}, g2m_connectivity_kwargs={}, + non_decode_g2m_connectivity_kwargs={}, coords_crs: pyproj.crs.CRS | None = None, graph_crs: pyproj.crs.CRS | None = None, - decode_mask: Iterable[bool] | None = None, + decode_mask: Iterable | None = None, return_components: bool = False, + allow_zero_degree: bool = True, + remove_zero_degree_mesh: bool = False, ): """ Create all graph components used in creating the message-passing graph, @@ -71,7 +146,7 @@ def create_all_graph_components( - "hierarchical": Create a hierarchical mesh graph with `max_num_levels`, `mesh_node_distance` and `level_refinement_factor`, similar to Oskarsson et al. (2023) - + m2g_connectivity: - "nearest_neighbour": Find the nearest neighbour in mesh for each node in grid - "nearest_neighbours": Find the `max_num_neighbours` nearest neighbours in mesh for each node in grid @@ -94,9 +169,18 @@ def create_all_graph_components( `return_components` is a boolean flag, if True the function returns a dict with m2g, m2m and g2m as separate graphs. If false returns one combined graph. + + `remove_zero_degree_mesh` (bool): If True, safely removes mesh nodes that have zero degree + in the g2m graph (not connected to any grid node), as long as removing them does not create + disconnected islands in the m2m mesh graph. Grid nodes with zero degree are always kept. """ graph_components: dict[networkx.DiGraph] = {} + from ..networkx_utils import ( + replace_node_labels_with_unique_ids, + split_graph_by_edge_attribute, + ) + assert ( len(coords.shape) == 2 and coords.shape[1] == 2 ), "Grid node coordinates should be given as an array of shape [num_grid_nodes, 2]." @@ -153,23 +237,64 @@ def create_all_graph_components( G_grid = create_grid_graph_nodes(xy=xy) - G_g2m = connect_nodes_across_graphs( - G_source=G_grid, - G_target=grid_connect_graph, - method=g2m_connectivity, - **g2m_connectivity_kwargs, - ) - graph_components["g2m"] = G_g2m - if decode_mask is None: # decode to all grid nodes decode_grid = G_grid + + # Same g2m method for all grid nodes + G_g2m = connect_nodes_across_graphs( + G_source=G_grid, + G_target=grid_connect_graph, + method=g2m_connectivity, + **g2m_connectivity_kwargs, + ) else: # Select subset of grid nodes to decode to, where m2g should connect - filter_nodes = [ + decode_filter_nodes = [ n for n, include in zip(G_grid.nodes, decode_mask, strict=True) if include ] - decode_grid = G_grid.subgraph(filter_nodes) + decode_grid = G_grid.subgraph(decode_filter_nodes) + + non_decode_filter_nodes = [ + n for n, decode in zip(G_grid.nodes, decode_mask, strict=True) if not decode + ] + non_decode_grid = G_grid.subgraph(non_decode_filter_nodes) + + # Different g2m connectivity methods + g2m_decode = connect_nodes_across_graphs( + G_source=decode_grid, + G_target=grid_connect_graph, + method=g2m_connectivity, + **g2m_connectivity_kwargs, + ) + assert non_decode_g2m_connectivity is not None + g2m_non_decode = connect_nodes_across_graphs( + G_source=non_decode_grid, + G_target=grid_connect_graph, + method=non_decode_g2m_connectivity, + **non_decode_g2m_connectivity_kwargs, + ) + + # Take union of graphs + G_g2m = nx.compose(g2m_decode, g2m_non_decode) + + graph_components["g2m"] = G_g2m + + if not allow_zero_degree: + # Find all nodes in g2m with degree == 0 + zero_degree_nodes = [ + node for node, degree in dict(G_g2m.degree()).items() if degree == 0 + ] + + if remove_zero_degree_mesh and len(zero_degree_nodes) > 0: + G_g2m, graph_components["m2m"], grid_connect_graph = ( + _remove_zero_degree_mesh_nodes( + G_g2m, G_grid, graph_components["m2m"], grid_connect_graph + ) + ) + graph_components["g2m"] = G_g2m + elif len(zero_degree_nodes) > 0: + assert False, f"Zero-degree nodes in g2m: {zero_degree_nodes}" G_m2g = connect_nodes_across_graphs( G_source=grid_connect_graph, @@ -185,7 +310,6 @@ def create_all_graph_components( graph.edges[edge]["component"] = name if return_components: - # Because merging to a single graph and then splitting again leads to changes in node indexing when converting to `pyg.Data` objects (this in part is due to the to `m2g` and `g2m` having a different set of grid nodes) the ability to return the graph components (`g2m`, `m2m` and `m2g`) has been added here. See https://github.com/mllam/weather-model-graphs/pull/34#issuecomment-2507980752 for details # Give each component unique ids graph_components = { comp_name: replace_node_labels_with_unique_ids(subgraph) @@ -256,11 +380,18 @@ def connect_nodes_across_graphs( Graph containing the nodes in `G_source` and `G_target` and directed edges from nodes in `G_source` to nodes in `G_target` """ + from ..networkx_utils import ( + split_graph_by_edge_attribute, + split_on_edge_attribute_existance, + ) source_nodes_list = list(G_source.nodes) target_nodes_list = list(G_target.nodes) # build kd tree for source nodes (e.g. the mesh nodes when constructing m2g) - xy_source = np.array([G_source.nodes[node]["pos"] for node in G_source.nodes]) + # TODO: Why do we keep sorting nodes in this method? + xy_source = np.array( + [G_source.nodes[node]["pos"] for node in sorted(G_source.nodes)] + ) kdt_s = scipy.spatial.KDTree(xy_source) # Determine method and perform checks once @@ -384,6 +515,7 @@ def _find_neighbour_node_idxs_in_source_mesh(xy_target): key=lambda x: x[2].get("len", 0), )[2]["len"] longest_edge = max(longest_edge, longest_graph_edge) + logger.debug(f"Longest distance in bottom mesh graph: {longest_edge}") query_dist = longest_edge * rel_max_dist else: raise Exception( diff --git a/src/poseidon_graphprep/create/grid/__init__.py b/src/poseidon_graphprep/create/grid/__init__.py new file mode 100644 index 0000000..e4f6de6 --- /dev/null +++ b/src/poseidon_graphprep/create/grid/__init__.py @@ -0,0 +1 @@ +from .grid import create_grid_graph_nodes \ No newline at end of file diff --git a/src/weather_model_graphs/create/grid/grid.py b/src/poseidon_graphprep/create/grid/grid.py similarity index 100% rename from src/weather_model_graphs/create/grid/grid.py rename to src/poseidon_graphprep/create/grid/grid.py diff --git a/src/weather_model_graphs/create/mesh/__init__.py b/src/poseidon_graphprep/create/mesh/__init__.py similarity index 100% rename from src/weather_model_graphs/create/mesh/__init__.py rename to src/poseidon_graphprep/create/mesh/__init__.py diff --git a/src/weather_model_graphs/create/mesh/kinds/__init__.py b/src/poseidon_graphprep/create/mesh/kinds/__init__.py similarity index 100% rename from src/weather_model_graphs/create/mesh/kinds/__init__.py rename to src/poseidon_graphprep/create/mesh/kinds/__init__.py diff --git a/src/weather_model_graphs/create/mesh/kinds/flat.py b/src/poseidon_graphprep/create/mesh/kinds/flat.py similarity index 80% rename from src/weather_model_graphs/create/mesh/kinds/flat.py rename to src/poseidon_graphprep/create/mesh/kinds/flat.py index 92f47b0..47a18ec 100644 --- a/src/weather_model_graphs/create/mesh/kinds/flat.py +++ b/src/poseidon_graphprep/create/mesh/kinds/flat.py @@ -1,12 +1,16 @@ import networkx import numpy as np -from ....networkx_utils import prepend_node_index +from ....networkx_utils import filter_nodes, get_chull_filter_func, prepend_node_index from .. import mesh as mesh_graph def create_flat_multiscale_mesh_graph( - xy, mesh_node_distance: float, level_refinement_factor: int, max_num_levels: int + xy, + mesh_node_distance: float, + level_refinement_factor: int, + max_num_levels: int, + filter_convex_hull: bool = False, ): """ Create flat mesh graph by merging the single-level mesh @@ -26,6 +30,10 @@ def create_flat_multiscale_mesh_graph( NOTE: Must be an odd integer >1 to create proper multiscale graph max_num_levels : int Maximum number of levels in the multi-scale graph + filter_convex_hull: bool + If the mesh nodes should be filtered to only those that lay within + the convex hull of xy. + Returns ------- G_tot : networkx.Graph @@ -86,10 +94,16 @@ def create_flat_multiscale_mesh_graph( G_tot.graph["dx"] = {i: g.graph["dx"] for i, g in enumerate(G_all_levels)} G_tot.graph["dy"] = {i: g.graph["dy"] for i, g in enumerate(G_all_levels)} + if filter_convex_hull: + # Filter mesh nodes + G_tot = filter_nodes(G_tot, get_chull_filter_func(xy)) + return G_tot -def create_flat_singlescale_mesh_graph(xy, mesh_node_distance: float): +def create_flat_singlescale_mesh_graph( + xy, mesh_node_distance: float, filter_convex_hull: bool = False +): """ Create flat mesh graph of single level @@ -102,6 +116,10 @@ def create_flat_singlescale_mesh_graph(xy, mesh_node_distance: float): mesh_node_distance: float Distance (in x- and y-direction) between created mesh nodes, in coordinate system of xy + filter_convex_hull: bool + If the mesh nodes should be filtered to only those that lay within + the convex hull of xy. + Returns ------- G_flat : networkx.Graph @@ -112,12 +130,10 @@ def create_flat_singlescale_mesh_graph(xy, mesh_node_distance: float): nx = int(range_x / mesh_node_distance) ny = int(range_y / mesh_node_distance) - if nx == 0 or ny == 0: - raise ValueError( - "The given `mesh_node_distance` is too large for the provided coordinates. " - f"Got mesh_node_distance={mesh_node_distance}, but the x-range is {range_x} " - f"and y-range is {range_y}. Maybe you want to decrease the `mesh_node_distance`" - " so that the mesh nodes are spaced closer together?" - ) + mesh = mesh_graph.create_single_level_2d_mesh_graph(xy=xy, nx=nx, ny=ny) + + if filter_convex_hull: + # Filter mesh nodes + mesh = filter_nodes(mesh, get_chull_filter_func(xy)) - return mesh_graph.create_single_level_2d_mesh_graph(xy=xy, nx=nx, ny=ny) + return mesh diff --git a/src/weather_model_graphs/create/mesh/kinds/hierarchical.py b/src/poseidon_graphprep/create/mesh/kinds/hierarchical.py similarity index 89% rename from src/weather_model_graphs/create/mesh/kinds/hierarchical.py rename to src/poseidon_graphprep/create/mesh/kinds/hierarchical.py index b897693..1fbbe39 100644 --- a/src/weather_model_graphs/create/mesh/kinds/hierarchical.py +++ b/src/poseidon_graphprep/create/mesh/kinds/hierarchical.py @@ -2,7 +2,7 @@ import numpy as np import scipy -from ....networkx_utils import prepend_node_index +from ....networkx_utils import filter_nodes, get_chull_filter_func, prepend_node_index from .. import mesh as mesh_graph @@ -11,6 +11,7 @@ def create_hierarchical_multiscale_mesh_graph( mesh_node_distance: float, level_refinement_factor: float, max_num_levels: int, + filter_convex_hull: bool = False, ): """ Create a hierarchical multiscale mesh graph with nearest neighbour @@ -31,6 +32,9 @@ def create_hierarchical_multiscale_mesh_graph( Refinement factor between grid points and bottom level of mesh hierarchy max_num_levels: int The number of levels in the hierarchical mesh graph. + filter_convex_hull: bool + If the mesh nodes should be filtered to only those that lay within + the convex hull of xy. Returns ------- @@ -54,8 +58,15 @@ def create_hierarchical_multiscale_mesh_graph( f"or number of grid points {xy.shape[0]}." ) - # Relabel nodes of each level with level index first + if filter_convex_hull: + # Filter mesh nodes + chull_filter_func = get_chull_filter_func(xy) + Gs_all_levels = [ + filter_nodes(level_graph, chull_filter_func) + for level_graph in Gs_all_levels + ] + # Relabel nodes of each level with level index first Gs_all_levels = [ prepend_node_index(graph, level_i) for level_i, graph in enumerate(Gs_all_levels) diff --git a/src/weather_model_graphs/create/mesh/mesh.py b/src/poseidon_graphprep/create/mesh/mesh.py similarity index 89% rename from src/weather_model_graphs/create/mesh/mesh.py rename to src/poseidon_graphprep/create/mesh/mesh.py index 75c82d6..1dce071 100644 --- a/src/weather_model_graphs/create/mesh/mesh.py +++ b/src/poseidon_graphprep/create/mesh/mesh.py @@ -120,8 +120,6 @@ def create_multirange_2d_mesh_graphs( max_mesh_levels_float = np.log(max_nodes_bottom) / np.log(level_refinement_factor) max_mesh_levels = max_mesh_levels_float.astype(int) # (2,) - nleaf = level_refinement_factor**max_mesh_levels - # leaves at the bottom in each direction, if using max_mesh_levels # As we can not instantiate different number of mesh levels in each # direction, create mesh levels corresponding to the minimum of the two @@ -130,7 +128,18 @@ def create_multirange_2d_mesh_graphs( if max_num_levels: # Limit the levels in mesh graph mesh_levels_to_create = min(mesh_levels_to_create, max_num_levels) + assert mesh_levels_to_create > 0, "Can not make any mesh levels" + # Now we have determined how many levels we are making: next to + # figure out number nodes at bottom level, using mesh_levels_to_create + # Start by computing maximum number of nodes at top level + nodes_at_top_float = max_nodes_bottom / ( + level_refinement_factor ** (mesh_levels_to_create - 1) + ) + # Above will in general be float, but round down to integer + nodes_at_top = np.floor(nodes_at_top_float) + # Compute nleaf based on this + nleaf = nodes_at_top * (level_refinement_factor ** (mesh_levels_to_create - 1)) logger.debug(f"mesh_levels: {mesh_levels_to_create}, nleaf: {nleaf}") # multi resolution tree levels diff --git a/src/weather_model_graphs/networkx_utils.py b/src/poseidon_graphprep/networkx_utils.py similarity index 75% rename from src/weather_model_graphs/networkx_utils.py rename to src/poseidon_graphprep/networkx_utils.py index df9fbf6..4af737c 100644 --- a/src/weather_model_graphs/networkx_utils.py +++ b/src/poseidon_graphprep/networkx_utils.py @@ -1,4 +1,5 @@ import networkx +from scipy.spatial import ConvexHull, Delaunay def prepend_node_index(graph, new_index): @@ -104,7 +105,7 @@ def split_graph_by_edge_attribute(graph, attr): def sort_nodes_in_graph(graph): """ Creates a new networkx.DiGraph that is a copy of input, but with nodes - sorted according to their label value + sorted according to their id Parameters ---------- @@ -169,3 +170,53 @@ def split_on_edge_attribute_existance(graph, attr): graph_without_attr = graph.edge_subgraph(edges_without_attr) return graph_with_attr, graph_without_attr + + +def get_chull_filter_func(xy): + """ + Create a filter function to check if nodes are within the + convex hull of given points. + + Parameters + ---------- + xy : np.array, shape (n_points, 2) + Array of xy-coordinates to construct the convex hull for. + + Returns + ------- + filter_func : function + A function that takes a node and its attributes, and returns True if + the node is within the convex hull, False otherwise. + """ + + # Filter mesh graph nodes to within convex hull + # Construct the convex hull + xy_hull = ConvexHull(xy) + # Create a Delaunay triangulation only using the points in the convex hull + xy_delaunay = Delaunay(xy[xy_hull.vertices]) + + def filter_func(node, attr): + return xy_delaunay.find_simplex(attr["pos"]) >= 0 + + return filter_func + + +def filter_nodes(graph, filter_func): + """ + Filter nodes in a graph based on a filter function. + + Parameters + ---------- + graph : networkx.Graph + Graph from which nodes are to be filtered. + filter_func : function + Function to determine if a node should be included in the filtered graph. + Function should return a boolean + + Returns + ------- + subgraph : networkx.Graph + Subgraph containing only the nodes that satisfy the filter function. + """ + filtered_nodes = [n for n, attr in graph.nodes(data=True) if filter_func(n, attr)] + return graph.subgraph(filtered_nodes) diff --git a/src/weather_model_graphs/save.py b/src/poseidon_graphprep/save.py similarity index 100% rename from src/weather_model_graphs/save.py rename to src/poseidon_graphprep/save.py diff --git a/src/weather_model_graphs/visualise/__init__.py b/src/poseidon_graphprep/visualise/__init__.py similarity index 100% rename from src/weather_model_graphs/visualise/__init__.py rename to src/poseidon_graphprep/visualise/__init__.py diff --git a/src/weather_model_graphs/visualise/plot_2d.py b/src/poseidon_graphprep/visualise/plot_2d.py similarity index 100% rename from src/weather_model_graphs/visualise/plot_2d.py rename to src/poseidon_graphprep/visualise/plot_2d.py diff --git a/src/weather_model_graphs/create/grid/__init__.py b/src/weather_model_graphs/create/grid/__init__.py deleted file mode 100644 index e1f2962..0000000 --- a/src/weather_model_graphs/create/grid/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .grid import create_grid_graph_nodes diff --git a/tests/test_graph_creation.py b/tests/test_graph_creation.py index b45d486..e4cb43a 100644 --- a/tests/test_graph_creation.py +++ b/tests/test_graph_creation.py @@ -7,7 +7,7 @@ import pytest import tests.utils as test_utils -import weather_model_graphs as wmg +import poseidon_graphprep as wmg def test_create_single_level_mesh_graph(): @@ -131,22 +131,13 @@ def test_create_exact_refinement(mesh_node_distance, level_refinement_factor): ) -@pytest.mark.parametrize( - "kind_and_num_mesh", - [ - ("keisler", 20**2), # 20 mesh nodes in bottom layer in each direction - ("graphcast", 9**2), # Can only fit 9 x 9 with level_refinement_factor=3 - ( - "oskarsson_hierarchical", - 9**2 + 3**2, - ), # As above, with additional 3 x 3 layer - ], -) -def test_create_irregular_grid(kind_and_num_mesh): +@pytest.mark.parametrize("kind", ["keisler", "graphcast", "oskarsson_hierarchical"]) +def test_create_irregular_grid(kind): """ - Tests that graphs can be created for irregular layouts of grid points + Tests that graphs can be created for irregular layouts of grid points. + The exact number of mesh nodes depends on mesh creation and zero-degree removal, + so we just verify the graph was created with both grid and mesh nodes. """ - kind, num_mesh = kind_and_num_mesh num_grid = 100 xy = test_utils.create_fake_irregular_coords(num_grid - 4) @@ -166,7 +157,13 @@ def test_create_irregular_grid(kind_and_num_mesh): graph = fn(coords=xy, mesh_node_distance=0.05) - assert len(graph.nodes) == num_grid + num_mesh + # Verify graph was created with both grid and mesh nodes + assert len(graph.nodes) > num_grid, "Graph should have more than just grid nodes" + assert len(graph.edges) > 0, "Graph should have edges" + + # Verify all graph components are present + graph_components = wmg.split_graph_by_edge_attribute(graph=graph, attr="component") + assert set(graph_components.keys()) == {"m2m", "m2g", "g2m"} @pytest.mark.parametrize("kind", ["graphcast", "keisler", "oskarsson_hierarchical"]) @@ -197,8 +194,8 @@ def test_create_lat_lon(kind): @pytest.mark.parametrize("kind", ["graphcast", "keisler", "oskarsson_hierarchical"]) def test_create_decode_mask(kind): """ - Tests that the decode mask for m2g works, resulting in less edges than - no filtering. + Tests that the decode mask for m2g works, resulting in fewer m2g edges than + without a decode mask. """ xy = test_utils.create_fake_irregular_coords(100) fn_name = f"create_{kind}_graph" @@ -214,8 +211,14 @@ def test_create_decode_mask(kind): coords=xy, mesh_node_distance=mesh_node_distance, decode_mask=decode_mask ) - # Check that some filtering has been performed - assert len(filtered_graph.edges) < len(unfiltered_graph.edges) + # Check that m2g edges (mesh-to-grid / decoding edges) are filtered + unfiltered_m2g_edges = [ + e for e in unfiltered_graph.edges(data=True) if e[2].get("component") == "m2g" + ] + filtered_m2g_edges = [ + e for e in filtered_graph.edges(data=True) if e[2].get("component") == "m2g" + ] + assert len(filtered_m2g_edges) < len(unfiltered_m2g_edges) @pytest.mark.parametrize("kind", ["graphcast", "oskarsson_hierarchical"]) diff --git a/tests/test_graph_decode_gridpoints_mask.py b/tests/test_graph_decode_gridpoints_mask.py index cfcc3bb..b646bfd 100644 --- a/tests/test_graph_decode_gridpoints_mask.py +++ b/tests/test_graph_decode_gridpoints_mask.py @@ -4,7 +4,7 @@ import numpy as np import torch -import weather_model_graphs as wmg +import poseidon_graphprep as wmg from tests import utils as test_utils diff --git a/tests/test_graph_plots.py b/tests/test_graph_plots.py index c5b1e10..6e4eb31 100644 --- a/tests/test_graph_plots.py +++ b/tests/test_graph_plots.py @@ -5,7 +5,7 @@ import pytest import tests.utils as test_utils -import weather_model_graphs as wmg +import poseidon_graphprep as wmg def test_plot(): diff --git a/tests/test_graph_utils.py b/tests/test_graph_utils.py index 39e04b0..8027772 100644 --- a/tests/test_graph_utils.py +++ b/tests/test_graph_utils.py @@ -1,5 +1,5 @@ import tests.utils as test_utils -import weather_model_graphs as wmg +import poseidon_graphprep as wmg def test_graph_splitting(): diff --git a/tests/test_save.py b/tests/test_save.py index f889ab9..4cd7540 100644 --- a/tests/test_save.py +++ b/tests/test_save.py @@ -4,15 +4,15 @@ from loguru import logger import tests.utils as test_utils -import weather_model_graphs as wmg -from weather_model_graphs.save import HAS_PYG +import poseidon_graphprep as wmg +from poseidon_graphprep.save import HAS_PYG @pytest.mark.parametrize("list_from_attribute", [None, "level"]) def test_save_to_pyg(list_from_attribute): if not HAS_PYG: logger.warning( - "Skipping test_save_to_pyg because weather-model-graphs[pytorch] is not installed." + "Skipping test_save_to_pyg because poseidon-graphprep[pytorch] is not installed." ) return