[Code Coverage] test_home.py (#6638)

Adding test for `torch_geometric/home.py`

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Matthias Fey <matthias.fey@tu-dortmund.de>

CHANGELOG.md

@@ -80,7 +80,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed a bug in `Data.subgraph()` and `HeteroData.subgraph()` ([#6613](https://github.com/pyg-team/pytorch_geometric/pull/6613)
 - Fixed a bug in `PNAConv` and `DegreeScalerAggregation` to correctly incorporate degree statistics of isolated nodes ([#6609](https://github.com/pyg-team/pytorch_geometric/pull/6609))
-- Improved code coverage ([#6523](https://github.com/pyg-team/pytorch_geometric/pull/6523), [#6538](https://github.com/pyg-team/pytorch_geometric/pull/6538), [#6555](https://github.com/pyg-team/pytorch_geometric/pull/6555), [#6558](https://github.com/pyg-team/pytorch_geometric/pull/6558), [#6568](https://github.com/pyg-team/pytorch_geometric/pull/6568), [#6573](https://github.com/pyg-team/pytorch_geometric/pull/6573), [#6578](https://github.com/pyg-team/pytorch_geometric/pull/6578), [#6597](https://github.com/pyg-team/pytorch_geometric/pull/6597), [#6600](https://github.com/pyg-team/pytorch_geometric/pull/6600), [#6618](https://github.com/pyg-team/pytorch_geometric/pull/6618), [#6619](https://github.com/pyg-team/pytorch_geometric/pull/6619), [#6621](https://github.com/pyg-team/pytorch_geometric/pull/6621), [#6623](https://github.com/pyg-team/pytorch_geometric/pull/6623), [#6640](https://github.com/pyg-team/pytorch_geometric/pull/6640),[#6637](https://github.com/pyg-team/pytorch_geometric/pull/6637))
+- Improved code coverage ([#6523](https://github.com/pyg-team/pytorch_geometric/pull/6523), [#6538](https://github.com/pyg-team/pytorch_geometric/pull/6538), [#6555](https://github.com/pyg-team/pytorch_geometric/pull/6555), [#6558](https://github.com/pyg-team/pytorch_geometric/pull/6558), [#6568](https://github.com/pyg-team/pytorch_geometric/pull/6568), [#6573](https://github.com/pyg-team/pytorch_geometric/pull/6573), [#6578](https://github.com/pyg-team/pytorch_geometric/pull/6578), [#6597](https://github.com/pyg-team/pytorch_geometric/pull/6597), [#6600](https://github.com/pyg-team/pytorch_geometric/pull/6600), [#6618](https://github.com/pyg-team/pytorch_geometric/pull/6618), [#6619](https://github.com/pyg-team/pytorch_geometric/pull/6619), [#6621](https://github.com/pyg-team/pytorch_geometric/pull/6621), [#6623](https://github.com/pyg-team/pytorch_geometric/pull/6623), [#6637](https://github.com/pyg-team/pytorch_geometric/pull/6637), [#6638](https://github.com/pyg-team/pytorch_geometric/pull/6638),[#6640](https://github.com/pyg-team/pytorch_geometric/pull/6640))
 - Fixed a bug in which `data.to_heterogeneous()` filtered attributs in the wrong dimension ([#6522](https://github.com/pyg-team/pytorch_geometric/pull/6522))
 - Breaking Change: Temporal sampling will now also sample nodes with an equal timestamp to the seed time (requires `pyg-lib>0.1.0`) ([#6517](https://github.com/pyg-team/pytorch_geometric/pull/6517))
 - Changed `DataLoader` workers with affinity to start at `cpu0` ([#6512](https://github.com/pyg-team/pytorch_geometric/pull/6512))

test/nn/models/test_correct_and_smooth.py

@@ -8,28 +8,37 @@ def test_correct_and_smooth():
     y_soft = torch.tensor([0.1, 0.5, 0.4]).repeat(6, 1)
     y_true = torch.tensor([1, 0, 0, 2, 1, 1])
     edge_index = torch.tensor([[0, 1, 1, 2, 4, 5], [1, 0, 2, 1, 5, 4]])
-    adj = SparseTensor.from_edge_index(edge_index, sparse_sizes=(6, 6))
+    adj = SparseTensor.from_edge_index(edge_index, sparse_sizes=(6, 6)).t()
     mask = torch.randint(0, 2, (6, ), dtype=torch.bool)
 
-    model = CorrectAndSmooth(num_correction_layers=2, correction_alpha=0.5,
-                             num_smoothing_layers=2, smoothing_alpha=0.5)
+    model = CorrectAndSmooth(
+        num_correction_layers=2,
+        correction_alpha=0.5,
+        num_smoothing_layers=2,
+        smoothing_alpha=0.5,
+    )
     assert str(model) == ('CorrectAndSmooth(\n'
-                          '    correct: num_layers=2, alpha=0.5\n'
-                          '    smooth:  num_layers=2, alpha=0.5\n'
-                          '    autoscale=True, scale=1.0\n'
+                          '  correct: num_layers=2, alpha=0.5\n'
+                          '  smooth:  num_layers=2, alpha=0.5\n'
+                          '  autoscale=True, scale=1.0\n'
                           ')')
 
-    correct_out = model.correct(y_soft, y_true[mask], mask, edge_index)
-    assert correct_out.size() == (6, 3)
-    assert torch.allclose(correct_out,
-                          model.correct(y_soft, y_true[mask], mask, adj))
-    smooth_out = model.smooth(y_soft, y_true[mask], mask, edge_index)
-    assert smooth_out.size() == (6, 3)
-    assert torch.allclose(smooth_out,
-                          model.smooth(y_soft, y_true[mask], mask, adj))
+    out = model.correct(y_soft, y_true[mask], mask, edge_index)
+    assert out.size() == (6, 3)
+    assert torch.allclose(out, model.correct(y_soft, y_true[mask], mask, adj))
+
+    out = model.smooth(y_soft, y_true[mask], mask, edge_index)
+    assert out.size() == (6, 3)
+    assert torch.allclose(out, model.smooth(y_soft, y_true[mask], mask, adj))
 
     # Test without autoscale:
-    model = CorrectAndSmooth(num_correction_layers=2, correction_alpha=0.5,
-                             num_smoothing_layers=2, smoothing_alpha=0.5,
-                             autoscale=False)
-    model.correct(y_soft, y_true[mask], mask, edge_index)
+    model = CorrectAndSmooth(
+        num_correction_layers=2,
+        correction_alpha=0.5,
+        num_smoothing_layers=2,
+        smoothing_alpha=0.5,
+        autoscale=False,
+    )
+    out = model.correct(y_soft, y_true[mask], mask, edge_index)
+    assert out.size() == (6, 3)
+    assert torch.allclose(out, model.correct(y_soft, y_true[mask], mask, adj))

test/test_home.py

@@ -0,0 +1,19 @@
+import os
+import os.path as osp
+
+from torch_geometric import get_home_dir, set_home_dir
+from torch_geometric.home import DEFAULT_CACHE_DIR
+
+
+def test_home():
+    os.environ.pop('PYG_HOME', None)
+    home_dir = osp.expanduser(DEFAULT_CACHE_DIR)
+    assert get_home_dir() == home_dir
+
+    home_dir = '/tmp/test_pyg1'
+    os.environ['PYG_HOME'] = home_dir
+    assert get_home_dir() == home_dir
+
+    home_dir = '/tmp/test_pyg2'
+    set_home_dir(home_dir)
+    assert get_home_dir() == home_dir

test/transforms/test_remove_training_classes.py

@@ -6,13 +6,14 @@
 
 def test_remove_training_classes():
     y = torch.tensor([1, 0, 0, 2, 1, 3])
+    train_mask = torch.tensor([False, False, True, True, True, True])
 
-    data = Data(y=y)
-    data.train_mask = torch.tensor([0, 0, 1, 1, 1, 1]).to(torch.bool)
+    data = Data(y=y, train_mask=train_mask)
 
-    classes = [0, 1]
-    transform = RemoveTrainingClasses(classes)
-    assert str(transform) == f'RemoveTrainingClasses({classes})'
+    transform = RemoveTrainingClasses(classes=[0, 1])
+    assert str(transform) == 'RemoveTrainingClasses([0, 1])'
 
     data = transform(data)
-    assert data.y[data.train_mask].unique().tolist() == [2, 3]
+    assert len(data) == 2
+    assert torch.equal(data.y, y)
+    assert data.train_mask.tolist() == [False, False, False, True, False, True]

torch_geometric/nn/models/correct_and_smooth.py

@@ -137,7 +137,7 @@ def __repr__(self):
         L1, alpha1 = self.prop1.num_layers, self.prop1.alpha
         L2, alpha2 = self.prop2.num_layers, self.prop2.alpha
         return (f'{self.__class__.__name__}(\n'
-                f'    correct: num_layers={L1}, alpha={alpha1}\n'
-                f'    smooth:  num_layers={L2}, alpha={alpha2}\n'
-                f'    autoscale={self.autoscale}, scale={self.scale}\n'
+                f'  correct: num_layers={L1}, alpha={alpha1}\n'
+                f'  smooth:  num_layers={L2}, alpha={alpha2}\n'
+                f'  autoscale={self.autoscale}, scale={self.scale}\n'
                 ')')