pyg-nightly 2.6.0.dev20240318__py3-none-any.whl → 2.7.0.dev20250115__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.py

@@ -192,19 +192,19 @@ class PascalVOCKeypoints(InMemoryDataset):
 
             child = obj.getElementsByTagName('xmin')[0].firstChild
             assert child is not None
-            xmin: float = float(child.data)  # type: ignore
+            xmin = int(child.data)  # type: ignore
 
             child = obj.getElementsByTagName('xmax')[0].firstChild
             assert child is not None
-            xmax = float(child.data)  # type: ignore
+            xmax = int(child.data)  # type: ignore
 
             child = obj.getElementsByTagName('ymin')[0].firstChild
             assert child is not None
-            ymin = float(child.data)  # type: ignore
+            ymin = int(child.data)  # type: ignore
 
             child = obj.getElementsByTagName('ymax')[0].firstChild
             assert child is not None
-            ymax = float(child.data)  # type: ignore
+            ymax = int(child.data)  # type: ignore
 
             box = (xmin, ymin, xmax, ymax)
 
@@ -227,10 +227,12 @@ class PascalVOCKeypoints(InMemoryDataset):
 
             # Add a small offset to the bounding because some keypoints lay
             # outside the bounding box intervals.
-            box = (min(float(pos[:, 0].min().floor()), box[0]) - 16,
-                   min(float(pos[:, 1].min().floor()), box[1]) - 16,
-                   max(float(pos[:, 0].max().ceil()), box[2]) + 16,
-                   max(float(pos[:, 1].max().ceil()), box[3]) + 16)
+            box = (
+                min(int(pos[:, 0].min().floor()), box[0]) - 16,
+                min(int(pos[:, 1].min().floor()), box[1]) - 16,
+                max(int(pos[:, 0].max().ceil()), box[2]) + 16,
+                max(int(pos[:, 1].max().ceil()), box[3]) + 16,
+            )
 
             # Rescale keypoints.
             pos[:, 0] = (pos[:, 0] - box[0]) * 256.0 / (box[2] - box[0])
@@ -239,7 +241,7 @@ class PascalVOCKeypoints(InMemoryDataset):
             path = osp.join(image_path, f'{filename}.jpg')
             with open(path, 'rb') as f:
                 img = Image.open(f).convert('RGB').crop(box)
-                img = img.resize((256, 256), resample=Image.BICUBIC)
+                img = img.resize((256, 256), resample=Image.Resampling.BICUBIC)
 
             img = transform(img)
 

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,7 +7,8 @@ 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.utils import from_smiles
+from torch_geometric.io import fs
+from torch_geometric.utils import from_smiles as _from_smiles
 
 
 class PCQM4Mv2(OnDiskDataset):
@@ -36,6 +37,10 @@ class PCQM4Mv2(OnDiskDataset):
             (default: :obj:`None`)
         backend (str): The :class:`Database` backend to use.
             (default: :obj:`"sqlite"`)
+        from_smiles (callable, optional): A custom function that takes a SMILES
+            string and outputs a :obj:`~torch_geometric.data.Data` object.
+            If not set, defaults to :meth:`~torch_geometric.utils.from_smiles`.
+            (default: :obj:`None`)
     """
     url = ('https://dgl-data.s3-accelerate.amazonaws.com/dataset/OGB-LSC/'
            'pcqm4m-v2.zip')
@@ -53,6 +58,7 @@ class PCQM4Mv2(OnDiskDataset):
         split: str = 'train',
         transform: Optional[Callable] = None,
         backend: str = 'sqlite',
+        from_smiles: Optional[Callable] = None,
     ) -> None:
         assert split in ['train', 'val', 'test', 'holdout']
 
@@ -64,9 +70,10 @@ class PCQM4Mv2(OnDiskDataset):
             'y': float,
         }
 
+        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
@@ -89,7 +96,7 @@ class PCQM4Mv2(OnDiskDataset):
         data_list: List[Data] = []
         iterator = enumerate(zip(df['smiles'], df['homolumogap']))
         for i, (smiles, y) in tqdm(iterator, total=len(df)):
-            data = from_smiles(smiles)
+            data = self.from_smiles(smiles)
             data.y = y
 
             data_list.append(data)

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
@@ -198,21 +199,21 @@ class QM9(InMemoryDataset):
 
     def process(self) -> None:
         try:
-            import rdkit
             from rdkit import Chem, RDLogger
             from rdkit.Chem.rdchem import BondType as BT
             from rdkit.Chem.rdchem import HybridizationType
-            RDLogger.DisableLog('rdApp.*')
+            RDLogger.DisableLog('rdApp.*')  # type: ignore
+            WITH_RDKIT = True
 
         except ImportError:
-            rdkit = None
+            WITH_RDKIT = False
 
-        if rdkit is None:
+        if not WITH_RDKIT:
             print(("Using a pre-processed version of the dataset. Please "
                    "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/rcdd.py

@@ -85,13 +85,13 @@ class RCDD(InMemoryDataset):
         mapping = torch.empty(len(node_df), dtype=torch.long)
         for node_type in node_df['node_type'].unique():
             mask = node_df['node_type'] == node_type
-            mask = torch.from_numpy(mask.values)
-            num_nodes = int(mask.sum())
-            mapping[mask] = torch.arange(num_nodes)
+            node_id = torch.from_numpy(node_df['node_id'][mask].values)
+            num_nodes = mask.sum()
+            mapping[node_id] = torch.arange(num_nodes)
             data[node_type].num_nodes = num_nodes
             x = np.vstack([
                 np.asarray(f.split(':'), dtype=np.float32)
-                for f in node_df['node_feat'][mask.numpy()]
+                for f in node_df['node_feat'][mask]
             ])
             data[node_type].x = torch.from_numpy(x)
 

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

@@ -3,6 +3,7 @@ import os.path as osp
 from typing import Callable, List, Optional
 
 import torch
+from torch import Tensor
 
 from torch_geometric.data import (
     Data,
@@ -85,13 +86,14 @@ 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, ys = [], []
+        xs: List[Tensor] = []
+        ys: List[Tensor] = []
         for filename in filenames:
             h5 = h5py.File(osp.join(self.raw_dir, filename))
             xs += torch.from_numpy(h5['data'][:]).unbind(0)

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)
 

torch_geometric/datasets/snap_dataset.py

@@ -22,6 +22,9 @@ class EgoData(Data):
 
 def read_ego(files: List[str], name: str) -> List[EgoData]:
     import pandas as pd
+    import tqdm
+
+    files = sorted(files)
 
     all_featnames = []
     files = [
@@ -38,7 +41,7 @@ def read_ego(files: List[str], name: str) -> List[EgoData]:
     all_featnames_dict = {key: i for i, key in enumerate(all_featnames)}
 
     data_list = []
-    for i in range(0, len(files), 5):
+    for i in tqdm.tqdm(range(0, len(files), 5)):
         circles_file = files[i]
         edges_file = files[i + 1]
         egofeat_file = files[i + 2]
@@ -65,6 +68,9 @@ def read_ego(files: List[str], name: str) -> List[EgoData]:
             x_all[:, torch.tensor(indices)] = x
             x = x_all
 
+            if x.size(1) > 100_000:
+                x = x.to_sparse_csr()
+
         idx = pd.read_csv(feat_file, sep=' ', header=None, dtype=str,
                           usecols=[0]).squeeze()