IF YOU WOULD LIKE TO GET AN ACCOUNT, please write an email to s dot adaszewski at gmail dot com. User accounts are meant only to report issues and/or generate pull requests. This is a purpose-specific Git hosting for ADARED projects. Thank you for your understanding!
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

117 line
4.8KB

  1. #
  2. # Copyright (C) Stanislaw Adaszewski, 2020
  3. # License: GPLv3
  4. #
  5. from icosagon.trainprep import TrainValTest, \
  6. train_val_test_split_edges, \
  7. get_edges_and_degrees, \
  8. prepare_adj_mat
  9. import torch
  10. import pytest
  11. import numpy as np
  12. from itertools import chain
  13. def test_train_val_test_split_edges_01():
  14. edges = torch.randint(0, 10, (10, 2))
  15. with pytest.raises(ValueError):
  16. _ = train_val_test_split_edges(edges, TrainValTest(.5, .5, .5))
  17. with pytest.raises(ValueError):
  18. _ = train_val_test_split_edges(edges, TrainValTest(.2, .2, .2))
  19. with pytest.raises(ValueError):
  20. _ = train_val_test_split_edges(edges, None)
  21. with pytest.raises(ValueError):
  22. _ = train_val_test_split_edges(edges, (.8, .1, .1))
  23. with pytest.raises(ValueError):
  24. _ = train_val_test_split_edges(np.random.randint(0, 10, (10, 2)), TrainValTest(.8, .1, .1))
  25. with pytest.raises(ValueError):
  26. _ = train_val_test_split_edges(torch.randint(0, 10, (10, 3)), TrainValTest(.8, .1, .1))
  27. with pytest.raises(ValueError):
  28. _ = train_val_test_split_edges(torch.randint(0, 10, (10, 2, 1)), TrainValTest(.8, .1, .1))
  29. with pytest.raises(ValueError):
  30. _ = train_val_test_split_edges(None, TrainValTest(.8, .2, .2))
  31. res = train_val_test_split_edges(edges, TrainValTest(.8, .1, .1))
  32. assert res.train.shape == (8, 2) and res.val.shape == (1, 2) and \
  33. res.test.shape == (1, 2)
  34. res = train_val_test_split_edges(edges, TrainValTest(.8, .0, .2))
  35. assert res.train.shape == (8, 2) and res.val.shape == (0, 2) and \
  36. res.test.shape == (2, 2)
  37. res = train_val_test_split_edges(edges, TrainValTest(.8, .2, .0))
  38. assert res.train.shape == (8, 2) and res.val.shape == (2, 2) and \
  39. res.test.shape == (0, 2)
  40. res = train_val_test_split_edges(edges, TrainValTest(.0, .5, .5))
  41. assert res.train.shape == (0, 2) and res.val.shape == (5, 2) and \
  42. res.test.shape == (5, 2)
  43. res = train_val_test_split_edges(edges, TrainValTest(.0, .0, 1.))
  44. assert res.train.shape == (0, 2) and res.val.shape == (0, 2) and \
  45. res.test.shape == (10, 2)
  46. res = train_val_test_split_edges(edges, TrainValTest(.0, 1., .0))
  47. assert res.train.shape == (0, 2) and res.val.shape == (10, 2) and \
  48. res.test.shape == (0, 2)
  49. def test_train_val_test_split_edges_02():
  50. edges = torch.randint(0, 30, (30, 2))
  51. ratios = TrainValTest(.8, .1, .1)
  52. res = train_val_test_split_edges(edges, ratios)
  53. edges = [ tuple(a) for a in edges ]
  54. res = [ tuple(a) for a in chain(res.train, res.val, res.test) ]
  55. assert all([ a in edges for a in res ])
  56. def test_get_edges_and_degrees_01():
  57. adj_mat_dense = (torch.rand((10, 10)) > .5)
  58. adj_mat_sparse = adj_mat_dense.to_sparse()
  59. edges_dense, degrees_dense = get_edges_and_degrees(adj_mat_dense)
  60. edges_sparse, degrees_sparse = get_edges_and_degrees(adj_mat_sparse)
  61. assert torch.all(degrees_dense == degrees_sparse)
  62. edges_dense = [ tuple(a) for a in edges_dense ]
  63. edges_sparse = [ tuple(a) for a in edges_dense ]
  64. assert len(edges_dense) == len(edges_sparse)
  65. assert all([ a in edges_dense for a in edges_sparse ])
  66. assert all([ a in edges_sparse for a in edges_dense ])
  67. # assert torch.all(edges_dense == edges_sparse)
  68. def test_prepare_adj_mat_01():
  69. adj_mat = (torch.rand((10, 10)) > .5)
  70. adj_mat = adj_mat.to_sparse()
  71. ratios = TrainValTest(.8, .1, .1)
  72. _ = prepare_adj_mat(adj_mat, ratios)
  73. def test_prepare_adj_mat_02():
  74. adj_mat = (torch.rand((10, 10)) > .5)
  75. adj_mat = adj_mat.to_sparse()
  76. ratios = TrainValTest(.8, .1, .1)
  77. (adj_mat_train, edges_pos, edges_neg) = prepare_adj_mat(adj_mat, ratios)
  78. assert isinstance(adj_mat_train, torch.Tensor)
  79. assert adj_mat_train.is_sparse
  80. assert adj_mat_train.shape == adj_mat.shape
  81. assert adj_mat_train.dtype == adj_mat.dtype
  82. assert isinstance(edges_pos, TrainValTest)
  83. assert isinstance(edges_neg, TrainValTest)
  84. for a in ['train', 'val', 'test']:
  85. for b in [edges_pos, edges_neg]:
  86. edges = getattr(b, a)
  87. assert isinstance(edges, torch.Tensor)
  88. assert len(edges.shape) == 2
  89. assert edges.shape[1] == 2
  90. # def prepare_adj_mat(adj_mat: torch.Tensor,
  91. # ratios: TrainValTest) -> Tuple[TrainValTest, TrainValTest]:
  92. #
  93. # degrees = adj_mat.sum(0)
  94. # edges_pos = torch.nonzero(adj_mat)
  95. #
  96. # neg_neighbors = fixed_unigram_candidate_sampler(edges_pos[:, 1],
  97. # len(edges), degrees, 0.75)
  98. # edges_neg = torch.cat((edges_pos[:, 0], neg_neighbors.view(-1, 1)), 1)
  99. #
  100. # edges_pos = train_val_test_split_edges(edges_pos, ratios)
  101. # edges_neg = train_val_test_split_edges(edges_neg, ratios)
  102. #
  103. # return edges_pos, edges_neg