Compare my work to original repo

data/complex.py

@@ -58,7 +58,8 @@ class Cochain(object):
     """
     def __init__(self, dim: int, x: Tensor = None, upper_index: Adj = None, lower_index: Adj = None,
                  shared_boundaries: Tensor = None, shared_coboundaries: Tensor = None, mapping: Tensor = None,
-                 boundary_index: Adj = None, upper_orient=None, lower_orient=None, y=None, **kwargs):
+                 boundary_index: Adj = None, coboundary_index: Adj = None,
+                 upper_orient=None, lower_orient=None, y=None, **kwargs):
         if dim == 0:
             assert lower_index is None
             assert shared_boundaries is None
@@ -72,6 +73,7 @@ def __init__(self, dim: int, x: Tensor = None, upper_index: Adj = None, lower_in
         self.upper_index = upper_index
         self.lower_index = lower_index
         self.boundary_index = boundary_index
+        self.coboundary_index = coboundary_index
         self.y = y
         self.shared_boundaries = shared_boundaries
         self.shared_coboundaries = shared_coboundaries
@@ -138,7 +140,7 @@ def __cat_dim__(self, key, value):
         :obj:`key` will get concatenated when creating batches.
         """
         if key in ['upper_index', 'lower_index', 'shared_boundaries',
-                   'shared_coboundaries', 'boundary_index']:
+                   'shared_coboundaries', 'boundary_index', 'coboundary_index']:
             return -1
         # by default, concatenate sparse matrices diagonally.
         elif isinstance(value, SparseTensor):
@@ -160,7 +162,14 @@ def __inc__(self, key, value):
         elif key == 'boundary_index':
             boundary_inc = self.num_cells_down if self.num_cells_down is not None else 0
             cell_inc = self.num_cells if self.num_cells is not None else 0
-            inc = [[boundary_inc], [cell_inc]]
+            inc = [[boundary_inc], [cell_inc]] # The first number is used to increase value 
+                                               # of numbers on first row of boundary_index
+                                               # and the second number is used for the 
+                                               # second row.
+        elif key == 'coboundary_index':
+            coboundary_inc = self.num_cells_up if self.num_cells_up is not None else 0
+            cell_inc = self.num_cells if self.num_cells is not None else 0
+            inc = [[coboundary_inc], [cell_inc]]
         else:
             inc = 0
         if inc is None:
@@ -550,7 +559,8 @@ def get_cochain_params(self,
                            max_dim : int=2,
                            include_top_features=True,
                            include_down_features=True,
-                           include_boundary_features=True) -> CochainMessagePassingParams:
+                           include_boundary_features=True,
+                           include_coboundary_features=True) -> CochainMessagePassingParams:
         """
         Conveniently constructs all necessary input parameters to perform higher-dim
         message passing on the cochain of specified `dim`.
@@ -562,6 +572,7 @@ def get_cochain_params(self,
             include_top_features: Whether to include the top features from level max_dim+1.
             include_down_features: Include the features for down adjacency
             include_boundary_features: Include the features for the boundary
+            include_coboundary_features: Include the features for the coboundary
         Returns:
             An object of type CochainMessagePassingParams
         """
@@ -592,10 +603,17 @@ def get_cochain_params(self,
                 boundary_index = cells.boundary_index
                 if dim > 0 and self.cochains[dim - 1].x is not None:
                     boundary_features = self.cochains[dim - 1].x
-
+            # Add coboundary features 
+            # REVIEW: Is this correctly written?
+            coboundary_index, coboundary_features = None, None
+            if include_coboundary_features and cells.coboundary_index is not None:
+                coboundary_index = cells.coboundary_index
+                if self.cochains[dim + 1].x is not None: 
+                    coboundary_features = self.cochains[dim + 1].x
             inputs = CochainMessagePassingParams(x, upper_index, lower_index,
                                                up_attr=upper_features, down_attr=lower_features,
-                                               boundary_attr=boundary_features, boundary_index=boundary_index)
+                                               boundary_attr=boundary_features, boundary_index=boundary_index,
+                                               coboundary_attr=coboundary_features, coboundary_index=coboundary_index)
         else:
             raise NotImplementedError(
                 'Dim {} is not present in the complex or not yet supported.'.format(dim))
@@ -605,7 +623,8 @@ def get_all_cochain_params(self,
                                max_dim:int=2,
                                include_top_features=True,
                                include_down_features=True,
-                               include_boundary_features=True) -> List[CochainMessagePassingParams]:
+                               include_boundary_features=True,
+                               include_coboundary_features=False) -> List[CochainMessagePassingParams]:
         """Extracts the cochain parameters for message passing on the cochains up to max_dim.
 
         Args:
@@ -622,7 +641,8 @@ def get_all_cochain_params(self,
             all_params.append(self.get_cochain_params(dim, max_dim=max_dim,
                                                     include_top_features=include_top_features,
                                                     include_down_features=include_down_features,
-                                                    include_boundary_features=include_boundary_features))
+                                                    include_boundary_features=include_boundary_features,
+                                                    include_coboundary_features=include_coboundary_features))
         return all_params
 
     def get_labels(self, dim=None):

data/data_loading.py

@@ -159,10 +159,14 @@ def load_dataset(name, root=os.path.join(ROOT_DIR, 'datasets'), max_dim=2, fold=
         dataset = RingLookupDataset(os.path.join(root, name), nodes=kwargs['max_ring_size'])
     elif name == 'ZINC':
         dataset = ZincDataset(os.path.join(root, name), max_ring_size=kwargs['max_ring_size'],
-                              use_edge_features=kwargs['use_edge_features'], n_jobs=n_jobs)
+                              use_edge_features=kwargs['use_edge_features'], n_jobs=n_jobs,
+                              include_down_adj=kwargs['include_down_adj'], 
+                              include_coboundary_links=kwargs['include_coboundary_links'])
     elif name == 'ZINC-FULL':
         dataset = ZincDataset(os.path.join(root, name), subset=False, max_ring_size=kwargs['max_ring_size'],
-                              use_edge_features=kwargs['use_edge_features'], n_jobs=n_jobs)
+                              use_edge_features=kwargs['use_edge_features'], n_jobs=n_jobs,
+                              include_down_adj=kwargs['include_down_adj'], 
+                              include_coboundary_links=kwargs['include_coboundary_links'])
     elif name == 'CSL':
         dataset = CSLDataset(os.path.join(root, name), max_ring_size=kwargs['max_ring_size'],
                              fold=fold, n_jobs=n_jobs)
@@ -172,6 +176,8 @@ def load_dataset(name, root=os.path.join(ROOT_DIR, 'datasets'), max_dim=2, fold=
         official_name = 'ogbg-'+name.lower()
         dataset = OGBDataset(os.path.join(root, name), official_name, max_ring_size=kwargs['max_ring_size'],
                              use_edge_features=kwargs['use_edge_features'], simple=kwargs['simple_features'],
+                              include_down_adj=kwargs['include_down_adj'], 
+                              include_coboundary_links=kwargs['include_coboundary_links'],
                              init_method=init_method, n_jobs=n_jobs)
     elif name == 'DUMMY':
         dataset = DummyDataset(os.path.join(root, name))

data/datasets/dataset.py

@@ -135,8 +135,10 @@ def process(self):
 
     def __init__(self, root=None, transform=None, pre_transform=None,
                  pre_filter=None, max_dim: int = None, num_classes: int = None,
-                 include_down_adj=False, init_method=None, cellular: bool = False):
+                 include_down_adj=False, init_method=None, cellular: bool = False, 
+                 include_coboundary_links=False):
         self.include_down_adj = include_down_adj
+        self.include_coboundary_links = include_coboundary_links
         super(InMemoryComplexDataset, self).__init__(root, transform, pre_transform, pre_filter,
                                                      max_dim, num_classes, init_method=init_method,
                                                      cellular=cellular)

data/datasets/ogb.py

@@ -10,14 +10,17 @@ class OGBDataset(InMemoryComplexDataset):
     """This is OGB graph-property prediction. This are graph-wise classification tasks."""
 
     def __init__(self, root, name, max_ring_size, use_edge_features=False, transform=None,
-                 pre_transform=None, pre_filter=None, init_method='sum', simple=False, n_jobs=2):
+                 pre_transform=None, pre_filter=None, init_method='sum', simple=False, n_jobs=2,
+                 include_down_adj=False, include_coboundary_links=False):
         self.name = name
         self._max_ring_size = max_ring_size
         self._use_edge_features = use_edge_features
         self._simple = simple
         self._n_jobs = n_jobs
         super(OGBDataset, self).__init__(root, transform, pre_transform, pre_filter,
-                                         max_dim=2, init_method=init_method, cellular=True)
+                                         max_dim=2, init_method=init_method, cellular=True,
+                                         include_down_adj=include_down_adj,
+                                         include_coboundary_links=include_coboundary_links)
         self.data, self.slices, idx, self.num_tasks = self.load_dataset()
         self.train_ids = idx['train']
         self.val_ids = idx['valid']
@@ -40,7 +43,9 @@ def processed_dir(self):
         suffix1 = f"_{self._max_ring_size}rings" if self._cellular else ""
         suffix2 = "-E" if self._use_edge_features else ""
         suffix3 = "-S" if self._simple else ""
-        return directory + suffix1 + suffix2 + suffix3
+        suffix4 = "-DA" if self.include_down_adj else ""
+        suffix5= "-CBL" if self.include_coboundary_links else ""
+        return directory + suffix1 + suffix2 + suffix3 + suffix4 + suffix5
 
     def download(self):
         # Instantiating this will download and process the graph dataset.
@@ -73,6 +78,7 @@ def process(self):
             include_down_adj=self.include_down_adj,
             init_method=self._init_method,
             init_edges=self._use_edge_features,
+            include_coboundary_links=self.include_coboundary_links,
             init_rings=False,
             n_jobs=self._n_jobs)
 

data/datasets/test_zinc.py

@@ -11,7 +11,9 @@
 
 @pytest.mark.slow
 def test_zinc_splits_are_retained():
-    dataset1 = load_dataset("ZINC", max_ring_size=7, use_edge_features=True)
+    dataset1 = load_dataset("ZINC", max_ring_size=7, use_edge_features=True, 
+                            include_down_adj=False, 
+                            include_coboundary_links=False)
     dataset1_train = dataset1.get_split('train')
     dataset1_valid = dataset1.get_split('valid')
     dataset1_test = dataset1.get_split('test')
@@ -40,7 +42,9 @@ def test_zinc_splits_are_retained():
 @pytest.mark.slow
 def test_we_find_only_the_induced_cycles_on_zinc():
     max_ring = 7
-    dataset = load_dataset("ZINC", max_ring_size=max_ring, use_edge_features=True)
+    dataset = load_dataset("ZINC", max_ring_size=max_ring, use_edge_features=True, 
+                            include_down_adj=False, 
+                            include_coboundary_links=False)
     # Check only on validation to save time. I've also run once on the whole dataset and passes.
     dataset = dataset.get_split('valid')
 

data/datasets/tu.py

@@ -38,7 +38,8 @@ class TUDataset(InMemoryComplexDataset):
     """A dataset of complexes obtained by lifting graphs from TUDatasets."""
 
     def __init__(self, root, name, max_dim=2, num_classes=2, degree_as_tag=False, fold=0,
-                 init_method='sum', seed=0, include_down_adj=False, max_ring_size=None):
+                 init_method='sum', seed=0, include_down_adj=False, max_ring_size=None, 
+                 include_coboundary_links=False):
         self.name = name
         self.degree_as_tag = degree_as_tag
         assert max_ring_size is None or max_ring_size > 3
@@ -48,7 +49,8 @@ def __init__(self, root, name, max_dim=2, num_classes=2, degree_as_tag=False, fo
             assert max_dim == 2
 
         super(TUDataset, self).__init__(root, max_dim=max_dim, num_classes=num_classes,
-            init_method=init_method, include_down_adj=include_down_adj, cellular=cellular)
+            init_method=init_method, include_down_adj=include_down_adj, cellular=cellular,
+            include_coboundary_links=include_coboundary_links)
 
         self.data, self.slices = torch.load(self.processed_paths[0])
 
@@ -79,6 +81,7 @@ def processed_dir(self):
         directory = super(TUDataset, self).processed_dir
         suffix = f"_{self._max_ring_size}rings" if self._cellular else ""
         suffix += f"_down_adj" if self.include_down_adj else ""
+        suffix += f"_CBL" if self.include_coboundary_links else ""
         return directory + suffix
 
     @property
@@ -108,6 +111,7 @@ def process(self):
             print("Converting the dataset accounting for rings...")
             complexes, _, _ = convert_graph_dataset_with_rings(graph_list, max_ring_size=self._max_ring_size,
                                                                include_down_adj=self.include_down_adj,
+                                                               include_coboundary_links=self.include_coboundary_links,
                                                                init_method=self._init_method,
                                                                init_edges=True, init_rings=True)
         else:
@@ -116,6 +120,7 @@ def process(self):
             # What about the init_method here? Adding now, although I remember we had handled this
             complexes, _, _ = convert_graph_dataset_with_gudhi(graph_list, expansion_dim=self.max_dim,
                                                                include_down_adj=self.include_down_adj,
+                                                               include_coboundary_links=self.include_coboundary_links,
                                                                init_method=self._init_method)
 
         torch.save(self.collate(complexes, self.max_dim), self.processed_paths[0])

data/datasets/zinc.py

@@ -10,14 +10,18 @@ class ZincDataset(InMemoryComplexDataset):
     """This is ZINC from the Benchmarking GNNs paper. This is a graph regression task."""
 
     def __init__(self, root, max_ring_size, use_edge_features=False, transform=None,
-                 pre_transform=None, pre_filter=None, subset=True, n_jobs=2):
+                 pre_transform=None, pre_filter=None, subset=True, n_jobs=2,
+                 include_down_adj=False, include_coboundary_links=False):
         self.name = 'ZINC'
+        self.include_coboundary_links = include_coboundary_links
         self._max_ring_size = max_ring_size
         self._use_edge_features = use_edge_features
         self._subset = subset
         self._n_jobs = n_jobs
         super(ZincDataset, self).__init__(root, transform, pre_transform, pre_filter,
-                                          max_dim=2, cellular=True, num_classes=1)
+                                          max_dim=2, cellular=True, num_classes=1, 
+                                          include_down_adj=include_down_adj, 
+                                          include_coboundary_links=include_coboundary_links)
 
         self.data, self.slices, idx = self.load_dataset()
         self.train_ids = idx[0]
@@ -62,6 +66,7 @@ def process(self):
             train_data,
             max_ring_size=self._max_ring_size,
             include_down_adj=self.include_down_adj,
+            include_coboundary_links=self.include_coboundary_links,
             init_edges=self._use_edge_features,
             init_rings=False,
             n_jobs=self._n_jobs)
@@ -73,6 +78,7 @@ def process(self):
             val_data,
             max_ring_size=self._max_ring_size,
             include_down_adj=self.include_down_adj,
+            include_coboundary_links=self.include_coboundary_links,
             init_edges=self._use_edge_features,
             init_rings=False,
             n_jobs=self._n_jobs)
@@ -84,6 +90,7 @@ def process(self):
             test_data,
             max_ring_size=self._max_ring_size,
             include_down_adj=self.include_down_adj,
+            include_coboundary_links=self.include_coboundary_links,
             init_edges=self._use_edge_features,
             init_rings=False,
             n_jobs=self._n_jobs)
@@ -105,7 +112,9 @@ def processed_dir(self):
         suffix0 = "_full" if self._subset is False else ""
         suffix1 = f"_{self._max_ring_size}rings" if self._cellular else ""
         suffix2 = "-E" if self._use_edge_features else ""
-        return directory + suffix0 + suffix1 + suffix2
+        suffix3 = f"_down_adj" if self.include_down_adj else ""
+        suffix4 = f"_CBL" if self.include_coboundary_links else ""
+        return directory + suffix0 + suffix1 + suffix2 + suffix3 + suffix4
 
 
 def load_zinc_graph_dataset(root, subset=True):