pyg-nightly 2.6.0.dev20240511__py3-none-any.whl → 2.7.0.dev20250114__py3-none-any.whl

torch_geometric/datasets/opf.py

@@ -0,0 +1,239 @@
+import json
+import os
+import os.path as osp
+from typing import Callable, Dict, List, Literal, Optional
+
+import torch
+import tqdm
+from torch import Tensor
+
+from torch_geometric.data import (
+    HeteroData,
+    InMemoryDataset,
+    download_url,
+    extract_tar,
+)
+
+
+class OPFDataset(InMemoryDataset):
+    r"""The heterogeneous OPF data from the `"Large-scale Datasets for AC
+    Optimal Power Flow with Topological Perturbations"
+    <https://arxiv.org/abs/2406.07234>`_ paper.
+
+    :class:`OPFDataset` is a large-scale dataset of solved optimal power flow
+    problems, derived from the
+    `pglib-opf <https://github.com/power-grid-lib/pglib-opf>`_ dataset.
+
+    The physical topology of the grid is represented by the :obj:`"bus"` node
+    type, and the connecting AC lines and transformers. Additionally,
+    :obj:`"generator"`, :obj:`"load"`, and :obj:`"shunt"` nodes are connected
+    to :obj:`"bus"` nodes using a dedicated edge type each, *e.g.*,
+    :obj:`"generator_link"`.
+
+    Edge direction corresponds to the properties of the line, *e.g.*,
+    :obj:`b_fr` is the line charging susceptance at the :obj:`from`
+    (source/sender) bus.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        split (str, optional): If :obj:`"train"`, loads the training dataset.
+            If :obj:`"val"`, loads the validation dataset.
+            If :obj:`"test"`, loads the test dataset. (default: :obj:`"train"`)
+        case_name (str, optional): The name of the original pglib-opf case.
+            (default: :obj:`"pglib_opf_case14_ieee"`)
+        num_groups (int, optional): The dataset is divided into 20 groups with
+            each group containing 15,000 samples.
+            For large networks, this amount of data can be overwhelming.
+            The :obj:`num_groups` parameters controls the amount of data being
+            downloaded. Allowed values are :obj:`[1, 20]`.
+            (default: :obj:`20`)
+        topological_perturbations (bool, optional): Whether to use the dataset
+            with added topological perturbations. (default: :obj:`False`)
+        transform (callable, optional): A function/transform that takes in
+            a :obj:`torch_geometric.data.HeteroData` object and returns a
+            transformed version. The data object will be transformed before
+            every access. (default: :obj:`None`)
+        pre_transform (callable, optional): A function/transform that takes
+            in a :obj:`torch_geometric.data.HeteroData` object and returns
+            a transformed version. The data object will be transformed before
+            being saved to disk. (default: :obj:`None`)
+        pre_filter (callable, optional): A function that takes in a
+            :obj:`torch_geometric.data.HeteroData` object and returns a boolean
+            value, indicating whether the data object should be included in the
+            final dataset. (default: :obj:`None`)
+        force_reload (bool, optional): Whether to re-process the dataset.
+            (default: :obj:`False`)
+    """
+    url = 'https://storage.googleapis.com/gridopt-dataset'
+
+    def __init__(
+        self,
+        root: str,
+        split: Literal['train', 'val', 'test'] = 'train',
+        case_name: Literal[
+            'pglib_opf_case14_ieee',
+            'pglib_opf_case30_ieee',
+            'pglib_opf_case57_ieee',
+            'pglib_opf_case118_ieee',
+            'pglib_opf_case500_goc',
+            'pglib_opf_case2000_goc',
+            'pglib_opf_case6470_rte',
+            'pglib_opf_case4661_sdet'
+            'pglib_opf_case10000_goc',
+            'pglib_opf_case13659_pegase',
+        ] = 'pglib_opf_case14_ieee',
+        num_groups: int = 20,
+        topological_perturbations: bool = False,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+        force_reload: bool = False,
+    ) -> None:
+
+        self.split = split
+        self.case_name = case_name
+        self.num_groups = num_groups
+        self.topological_perturbations = topological_perturbations
+
+        self._release = 'dataset_release_1'
+        if topological_perturbations:
+            self._release += '_nminusone'
+
+        super().__init__(root, transform, pre_transform, pre_filter,
+                         force_reload=force_reload)
+
+        idx = self.processed_file_names.index(f'{split}.pt')
+        self.load(self.processed_paths[idx])
+
+    @property
+    def raw_dir(self) -> str:
+        return osp.join(self.root, self._release, self.case_name, 'raw')
+
+    @property
+    def processed_dir(self) -> str:
+        return osp.join(self.root, self._release, self.case_name,
+                        f'processed_{self.num_groups}')
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return [f'{self.case_name}_{i}.tar.gz' for i in range(self.num_groups)]
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        return ['train.pt', 'val.pt', 'test.pt']
+
+    def download(self) -> None:
+        for name in self.raw_file_names:
+            url = f'{self.url}/{self._release}/{name}'
+            path = download_url(url, self.raw_dir)
+            extract_tar(path, self.raw_dir)
+
+    def process(self) -> None:
+        train_data_list = []
+        val_data_list = []
+        test_data_list = []
+
+        for group in tqdm.tqdm(range(self.num_groups)):
+            tmp_dir = osp.join(
+                self.raw_dir,
+                'gridopt-dataset-tmp',
+                self._release,
+                self.case_name,
+                f'group_{group}',
+            )
+
+            for name in os.listdir(tmp_dir):
+                with open(osp.join(tmp_dir, name)) as f:
+                    obj = json.load(f)
+
+                grid = obj['grid']
+                solution = obj['solution']
+                metadata = obj['metadata']
+
+                # Graph-level properties:
+                data = HeteroData()
+                data.x = torch.tensor(grid['context']).view(-1)
+
+                data.objective = torch.tensor(metadata['objective'])
+
+                # Nodes (only some have a target):
+                data['bus'].x = torch.tensor(grid['nodes']['bus'])
+                data['bus'].y = torch.tensor(solution['nodes']['bus'])
+
+                data['generator'].x = torch.tensor(grid['nodes']['generator'])
+                data['generator'].y = torch.tensor(
+                    solution['nodes']['generator'])
+
+                data['load'].x = torch.tensor(grid['nodes']['load'])
+
+                data['shunt'].x = torch.tensor(grid['nodes']['shunt'])
+
+                # Edges (only ac lines and transformers have features):
+                data['bus', 'ac_line', 'bus'].edge_index = (  #
+                    extract_edge_index(obj, 'ac_line'))
+                data['bus', 'ac_line', 'bus'].edge_attr = torch.tensor(
+                    grid['edges']['ac_line']['features'])
+                data['bus', 'ac_line', 'bus'].edge_label = torch.tensor(
+                    solution['edges']['ac_line']['features'])
+
+                data['bus', 'transformer', 'bus'].edge_index = (  #
+                    extract_edge_index(obj, 'transformer'))
+                data['bus', 'transformer', 'bus'].edge_attr = torch.tensor(
+                    grid['edges']['transformer']['features'])
+                data['bus', 'transformer', 'bus'].edge_label = torch.tensor(
+                    solution['edges']['transformer']['features'])
+
+                data['generator', 'generator_link', 'bus'].edge_index = (  #
+                    extract_edge_index(obj, 'generator_link'))
+                data['bus', 'generator_link', 'generator'].edge_index = (  #
+                    extract_edge_index_rev(obj, 'generator_link'))
+
+                data['load', 'load_link', 'bus'].edge_index = (  #
+                    extract_edge_index(obj, 'load_link'))
+                data['bus', 'load_link', 'load'].edge_index = (  #
+                    extract_edge_index_rev(obj, 'load_link'))
+
+                data['shunt', 'shunt_link', 'bus'].edge_index = (  #
+                    extract_edge_index(obj, 'shunt_link'))
+                data['bus', 'shunt_link', 'shunt'].edge_index = (  #
+                    extract_edge_index_rev(obj, 'shunt_link'))
+
+                if self.pre_filter is not None and not self.pre_filter(data):
+                    continue
+
+                if self.pre_transform is not None:
+                    data = self.pre_transform(data)
+
+                i = int(name.split('.')[0].split('_')[1])
+                train_limit = int(15_000 * self.num_groups * 0.9)
+                val_limit = train_limit + int(15_000 * self.num_groups * 0.05)
+                if i < train_limit:
+                    train_data_list.append(data)
+                elif i < val_limit:
+                    val_data_list.append(data)
+                else:
+                    test_data_list.append(data)
+
+        self.save(train_data_list, self.processed_paths[0])
+        self.save(val_data_list, self.processed_paths[1])
+        self.save(test_data_list, self.processed_paths[2])
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({len(self)}, '
+                f'split={self.split}, '
+                f'case_name={self.case_name}, '
+                f'topological_perturbations={self.topological_perturbations})')
+
+
+def extract_edge_index(obj: Dict, edge_name: str) -> Tensor:
+    return torch.tensor([
+        obj['grid']['edges'][edge_name]['senders'],
+        obj['grid']['edges'][edge_name]['receivers'],
+    ])
+
+
+def extract_edge_index_rev(obj: Dict, edge_name: str) -> Tensor:
+    return torch.tensor([
+        obj['grid']['edges'][edge_name]['receivers'],
+        obj['grid']['edges'][edge_name]['senders'],
+    ])

torch_geometric/datasets/ose_gvcs.py

@@ -97,7 +97,7 @@ class OSE_GVCS(InMemoryDataset):
         edges = defaultdict(list)
 
         for path in self.raw_paths:
-            with open(path, 'r') as f:
+            with open(path) as f:
                 product = json.load(f)
             categories.append(self.categories.index(product['category']))
             for interaction in product['ecology']:

torch_geometric/datasets/pascal_pf.py

@@ -66,7 +66,7 @@ class PascalPF(InMemoryDataset):
         super().__init__(root, transform, pre_transform, pre_filter,
                          force_reload=force_reload)
         self.load(self.processed_paths[0])
-        self.pairs = torch.load(self.processed_paths[1])
+        self.pairs = fs.torch_load(self.processed_paths[1])
 
     @property
     def raw_file_names(self) -> List[str]:

torch_geometric/datasets/pcpnet_dataset.py

@@ -121,7 +121,7 @@ class PCPNetDataset(InMemoryDataset):
 
     def process(self) -> None:
         path_file = self.raw_paths
-        with open(path_file[0], "r") as f:
+        with open(path_file[0]) as f:
             filenames = f.read().split('\n')[:-1]
         data_list = []
         for filename in filenames:

torch_geometric/datasets/pcqm4m.py

@@ -7,6 +7,7 @@ from tqdm import tqdm
 
 from torch_geometric.data import Data, OnDiskDataset, download_url, extract_zip
 from torch_geometric.data.data import BaseData
+from torch_geometric.io import fs
 from torch_geometric.utils import from_smiles as _from_smiles
 
 
@@ -72,7 +73,7 @@ class PCQM4Mv2(OnDiskDataset):
         self.from_smiles = from_smiles or _from_smiles
         super().__init__(root, transform, backend=backend, schema=schema)
 
-        split_idx = torch.load(self.raw_paths[1])
+        split_idx = fs.torch_load(self.raw_paths[1])
         self._indices = split_idx[self.split_mapping[split]].tolist()
 
     @property

torch_geometric/datasets/ppi.py

@@ -106,7 +106,7 @@ class PPI(InMemoryDataset):
 
         for s, split in enumerate(['train', 'valid', 'test']):
             path = osp.join(self.raw_dir, f'{split}_graph.json')
-            with open(path, 'r') as f:
+            with open(path) as f:
                 G = nx.DiGraph(json_graph.node_link_graph(json.load(f)))
 
             x = np.load(osp.join(self.raw_dir, f'{split}_feats.npy'))

torch_geometric/datasets/qm9.py

@@ -13,6 +13,7 @@ from torch_geometric.data import (
     download_url,
     extract_zip,
 )
+from torch_geometric.io import fs
 from torch_geometric.utils import one_hot, scatter
 
 HAR2EV = 27.211386246
@@ -212,7 +213,7 @@ class QM9(InMemoryDataset):
                    "install 'rdkit' to alternatively process the raw data."),
                   file=sys.stderr)
 
-            data_list = torch.load(self.raw_paths[0])
+            data_list = fs.torch_load(self.raw_paths[0])
             data_list = [Data(**data_dict) for data_dict in data_list]
 
             if self.pre_filter is not None:
@@ -227,14 +228,14 @@ class QM9(InMemoryDataset):
         types = {'H': 0, 'C': 1, 'N': 2, 'O': 3, 'F': 4}
         bonds = {BT.SINGLE: 0, BT.DOUBLE: 1, BT.TRIPLE: 2, BT.AROMATIC: 3}
 
-        with open(self.raw_paths[1], 'r') as f:
+        with open(self.raw_paths[1]) as f:
             target = [[float(x) for x in line.split(',')[1:20]]
                       for line in f.read().split('\n')[1:-1]]
             y = torch.tensor(target, dtype=torch.float)
             y = torch.cat([y[:, 3:], y[:, :3]], dim=-1)
             y = y * conversion.view(1, -1)
 
-        with open(self.raw_paths[2], 'r') as f:
+        with open(self.raw_paths[2]) as f:
             skip = [int(x.split()[0]) - 1 for x in f.read().split('\n')[9:-2]]
 
         suppl = Chem.SDMolSupplier(self.raw_paths[0], removeHs=False,

torch_geometric/datasets/reddit.py

@@ -3,7 +3,6 @@ import os.path as osp
 from typing import Callable, List, Optional
 
 import numpy as np
-import scipy.sparse as sp
 import torch
 
 from torch_geometric.data import (
@@ -76,6 +75,8 @@ class Reddit(InMemoryDataset):
         os.unlink(path)
 
     def process(self) -> None:
+        import scipy.sparse as sp
+
         data = np.load(osp.join(self.raw_dir, 'reddit_data.npz'))
         x = torch.from_numpy(data['feature']).to(torch.float)
         y = torch.from_numpy(data['label']).to(torch.long)

torch_geometric/datasets/reddit2.py

@@ -3,7 +3,6 @@ import os.path as osp
 from typing import Callable, List, Optional
 
 import numpy as np
-import scipy.sparse as sp
 import torch
 
 from torch_geometric.data import Data, InMemoryDataset, download_google_url
@@ -81,6 +80,8 @@ class Reddit2(InMemoryDataset):
         download_google_url(self.role_id, self.raw_dir, 'role.json')
 
     def process(self) -> None:
+        import scipy.sparse as sp
+
         f = np.load(osp.join(self.raw_dir, 'adj_full.npz'))
         adj = sp.csr_matrix((f['data'], f['indices'], f['indptr']), f['shape'])
         adj = adj.tocoo()

torch_geometric/datasets/rel_link_pred_dataset.py

@@ -89,17 +89,17 @@ class RelLinkPredDataset(InMemoryDataset):
             download_url(f'{self.urls[self.name]}/{file_name}', self.raw_dir)
 
     def process(self) -> None:
-        with open(osp.join(self.raw_dir, 'entities.dict'), 'r') as f:
+        with open(osp.join(self.raw_dir, 'entities.dict')) as f:
             lines = [row.split('\t') for row in f.read().split('\n')[:-1]]
             entities_dict = {key: int(value) for value, key in lines}
 
-        with open(osp.join(self.raw_dir, 'relations.dict'), 'r') as f:
+        with open(osp.join(self.raw_dir, 'relations.dict')) as f:
             lines = [row.split('\t') for row in f.read().split('\n')[:-1]]
             relations_dict = {key: int(value) for value, key in lines}
 
         kwargs = {}
         for split in ['train', 'valid', 'test']:
-            with open(osp.join(self.raw_dir, f'{split}.txt'), 'r') as f:
+            with open(osp.join(self.raw_dir, f'{split}.txt')) as f:
                 lines = [row.split('\t') for row in f.read().split('\n')[:-1]]
                 src = [entities_dict[row[0]] for row in lines]
                 rel = [relations_dict[row[1]] for row in lines]

torch_geometric/datasets/s3dis.py

@@ -86,10 +86,10 @@ class S3DIS(InMemoryDataset):
     def process(self) -> None:
         import h5py
 
-        with open(self.raw_paths[0], 'r') as f:
+        with open(self.raw_paths[0]) as f:
             filenames = [x.split('/')[-1] for x in f.read().split('\n')[:-1]]
 
-        with open(self.raw_paths[1], 'r') as f:
+        with open(self.raw_paths[1]) as f:
             rooms = f.read().split('\n')[:-1]
 
         xs: List[Tensor] = []

torch_geometric/datasets/shapenet.py

@@ -148,8 +148,8 @@ class ShapeNet(InMemoryDataset):
         elif split == 'trainval':
             path = self.processed_paths[3]
         else:
-            raise ValueError((f'Split {split} found, but expected either '
-                              'train, val, trainval or test'))
+            raise ValueError(f'Split {split} found, but expected either '
+                             'train, val, trainval or test')
 
         self.load(path)
 
@@ -213,7 +213,7 @@ class ShapeNet(InMemoryDataset):
         for i, split in enumerate(['train', 'val', 'test']):
             path = osp.join(self.raw_dir, 'train_test_split',
                             f'shuffled_{split}_file_list.json')
-            with open(path, 'r') as f:
+            with open(path) as f:
                 filenames = [
                     osp.sep.join(name.split('/')[1:]) + '.txt'
                     for name in json.load(f)

torch_geometric/datasets/shrec2016.py

@@ -6,7 +6,7 @@ from typing import Callable, List, Optional
 import torch
 
 from torch_geometric.data import InMemoryDataset, download_url, extract_zip
-from torch_geometric.io import read_off, read_txt_array
+from torch_geometric.io import fs, read_off, read_txt_array
 
 
 class SHREC2016(InMemoryDataset):
@@ -79,7 +79,7 @@ class SHREC2016(InMemoryDataset):
         self.cat = category.lower()
         super().__init__(root, transform, pre_transform, pre_filter,
                          force_reload=force_reload)
-        self.__ref__ = torch.load(self.processed_paths[0])
+        self.__ref__ = fs.torch_load(self.processed_paths[0])
         path = self.processed_paths[1] if train else self.processed_paths[2]
         self.load(path)