Inclusion of the AirfRANS dataset

python/dgl/data/__init__.py

@@ -8,6 +8,7 @@
 from .actor import ActorDataset
 from .movielens import MovieLensDataset
 from .adapter import *
+from .airfrans import AirfRANSDataset
 from .bitcoinotc import BitcoinOTC, BitcoinOTCDataset
 from .citation_graph import (
     CitationGraphDataset,

python/dgl/data/airfrans.py

@@ -0,0 +1,254 @@
+""" PPIDataset for inductive learning. """
+import json
+import os
+
+import numpy as np
+
+from dgl import graph
+from .. import backend as F
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_graphs, load_info, save_graphs, save_info
+
+class AirfRANSDataset(DGLBuiltinDataset):
+    r"""The AirfRANS dataset from the "AirfRANS: High Fidelity Computational
+    Fluid Dynamics Dataset for Approximating Reynolds-Averaged Navier-Stokes
+    Solutions" paper, consisting of 1,000
+    simulations of steady-state aerodynamics over 2D airfoils in a subsonic
+    flight regime.
+
+    The different tasks (:obj:`"full"`, :obj:`"scarce"`, :obj:`"reynolds"`,
+    :obj:`"aoa"`) define the utilized training and test splits.
+
+    Each simulation is given as a point cloud defined as the nodes of the
+    simulation mesh. Each point of a point cloud is described via 5
+    features: the inlet velocity (two components in meter per second), the
+    distance to the airfoil (one component in meter), and the normals (two
+    components in meter, set to :obj:`0` if the point is not on the airfoil).
+    Each point is given a target of 4 components for the underyling regression
+    task: the velocity (two components in meter per second), the pressure
+    divided by the specific mass (one component in meter squared per second
+    squared), the turbulent kinematic viscosity (one component in meter squared
+    per second).
+    Finaly, a boolean is attached to each point to inform if this point lies on
+    the airfoil or not.
+
+    Reference: 
+    `NeurIPS Paper<https://arxiv.org/abs/2212.07564>`_
+
+    A library for manipulating simulations of the dataset is available `here
+    <https://airfrans.readthedocs.io/en/latest/index.html>`_.
+
+    The dataset is released under the `ODbL v1.0 License
+    <https://opendatacommons.org/licenses/odbl/1-0/>`_.
+
+    Statistics:
+
+    .. list-table::
+        :widths: 10 10 10 10 10
+        :header-rows: 1
+
+        * - #graphs
+          - #nodes
+          - #edges
+          - #features
+          - #labels
+        * - 1,000
+          - ~180,000
+          - 0
+          - 5
+          - 4
+
+    Notes
+    -----
+        Data objects contain no edge indices to be agnostic to the simulation
+        mesh. You are free to build a graph upon it.
+
+    Parameters
+    ----------
+    mode : str
+        Must be one of ('train', 'test').
+        Default: 'train'
+    task : str
+        The task to study that defines the train and test splits ('full', 'scarce', 'reynolds', 'aoa').
+        Default: 'full'
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False
+    verbose : bool
+        Whether to print out progress information.
+        Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_features : int
+        Number of features for each node
+    num_labels : int
+        Number of labels for each node
+    positions : Tensor
+        Node positions
+    labels : Tensor
+        Node labels
+    features : Tensor
+        Node features
+    surfaces : Tensor
+        Boolean attached to each node to specify if it lies on the surface of the airfoil
+
+    Examples
+    --------
+    >>> dataset = AirfRANSDataset(mode='test', task='scarce')
+    >>> graph_names = dataset.graph_names
+    >>> for g in dataset:
+    ....    name = g.name
+    ....    pos = g.ndata['pos']
+    ....    feat = g.ndata['feat']
+    ....    label = g.ndata['label']
+    ....    surf = g.ndata['surf']
+    ....    # your code here
+    >>>
+    """
+
+    def __init__(
+        self,
+        mode="train",
+        task="full",
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        assert mode in ["train", "test"]
+        assert task in ["full", "scarce", "reynolds", "aoa"]
+        self.mode = mode
+        self.task = task
+        _url = _get_dgl_url("dataset/airfrans.zip")
+        super(AirfRANSDataset, self).__init__(
+            name="airfrans",
+            url=_url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        position_file = os.path.join(
+            self.save_path, "airfrans_positions.npy"
+        )
+        label_file = os.path.join(
+            self.save_path, "airfrans_labels.npy"
+        )
+        feat_file = os.path.join(
+            self.save_path, "airfrans_feats.npy"
+        )
+        surf_file = os.path.join(
+            self.save_path, "airfrans_surfaces.npy"
+        )
+        graph_name_file = os.path.join(
+            self.save_path, "manifest.json"
+        )
+
+        self._positions = np.load(position_file, allow_pickle = True)
+        self._labels = np.load(label_file, allow_pickle = True)
+        self._feats = np.load(feat_file, allow_pickle = True)
+        self._surfaces = np.load(surf_file, allow_pickle = True)
+
+        with open(graph_name_file, 'r') as f:
+            self._graph_names = json.load(f)
+
+        total = self._graph_names['full_train'] + self._graph_names['full_test']
+        partial = set(self._graph_names[f'{self.task}_{self.mode}'])
+        self.graphs = []
+        for k, s in enumerate(total):
+            if s in partial:
+                g = graph(([], []), num_nodes=self._positions[k].shape[0])
+                g.ndata["pos"] = F.tensor(
+                    self._positions[k], dtype=F.data_type_dict["float32"]
+                )
+                g.ndata["feat"] = F.tensor(
+                    self._feats[k], dtype=F.data_type_dict["float32"]
+                )
+                g.ndata["label"] = F.tensor(
+                    self._labels[k], dtype=F.data_type_dict["float32"]
+                )
+                g.ndata["surf"] = F.tensor(
+                    self._surfaces[k], dtype=F.data_type_dict["float32"]
+                )
+
+                self.graphs.append(g)
+
+    @property
+    def graph_list_path(self):
+        return os.path.join(
+            self.save_path, "{}_dgl_graph_list.bin".format(self.mode)
+        )
+
+    @property
+    def info_path(self):
+        return os.path.join(self.save_path, "{}_info.pkl".format(self.mode))
+
+    def has_cache(self):
+        return (
+            os.path.exists(self.graph_list_path)
+            and os.path.exists(self.g_path)
+            and os.path.exists(self.info_path)
+        )
+
+    def save(self):
+        save_graphs(self.graph_list_path, self.graphs)
+        save_info(
+            self.info_path, {"positions": self._positions, "labels": self._labels, "feats": self._feats, "surfaces": self._surfaces}
+        )
+
+    def load(self):
+        self.graphs = load_graphs(self.graph_list_path)
+        info = load_info(self.info_path)
+        self._positions = info["positions"]
+        self._labels = info["labels"]
+        self._feats = info["feats"]
+        self._surfaces = info["surfaces"]
+
+    @property
+    def num_features(self):
+        return 5
+
+    @property
+    def num_labels(self):
+        return 4
+
+    @property
+    def graph_names(self):
+        return self._graph_names
+
+    def __len__(self):
+        """Return number of samples in this dataset."""
+        return len(self.graphs)
+
+    def __getitem__(self, item):
+        """Get the item^th sample.
+
+        Parameters
+        ---------
+        item : int
+            The sample index.
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+            graph structure, node features and node labels.
+
+            - ``ndata['pos']``: node positions
+            - ``ndata['feat']``: node features
+            - ``ndata['label']``: node labels
+            - ``ndata['surf']``: node surfaces boolean (``True`` if the node lies on the surface of the airfoil)
+        """
+        if self._transform is None:
+            return self.graphs[item]
+        else:
+            return self._transform(self.graphs[item])