Coverage for tests/test_nw_ops.py: 100%

1import sys 

2from os.path import dirname as up 

3 

4import numpy as np 

5import xarray as xr 

6from dantro._import_tools import import_module_from_path 

7from pkg_resources import resource_filename 

8 

9from utopya.yaml import load_yml 

10 

11sys.path.insert(0, up(up(up(__file__)))) 

12 

13ops = import_module_from_path( 

14 mod_path=up(up(up(__file__))), mod_str="model_plots.data_ops" 

15) 

16 

17nw_ops = import_module_from_path( 

18 mod_path=up(up(up(__file__))), mod_str="model_plots.nw_ops" 

19) 

20 

21# Load the test config 

22CFG_FILENAME = resource_filename("tests", "cfgs/data_ops.yml") 

23test_cfg = load_yml(CFG_FILENAME) 

24 

25 

26# ---------------------------------------------------------------------------------------------------------------------- 

27# ADJACENCY MATRIX OPERATIONS 

28# ---------------------------------------------------------------------------------------------------------------------- 

29def test_triangles(): 

30 """Tests the calculation of triangles on adjacency matrices.""" 

31 

32 # Generate some samples of adjacency matrices 

33 N_samples, N_vertices = 10, 10 

34 A_samples = np.ones((N_samples, N_vertices, N_vertices)) 

35 for i in range(len(A_samples)): 

36 np.fill_diagonal(A_samples[i], 0) 

37 

38 A_samples = xr.DataArray( 

39 A_samples, 

40 dims=["samples", "i", "j"], 

41 coords=dict( 

42 samples=np.arange(N_samples), 

43 i=np.arange(N_vertices), 

44 j=np.arange(N_vertices), 

45 ), 

46 ) 

47 

48 triangles = nw_ops.triangles(A_samples, input_core_dims=["j"]) 

49 assert set(triangles.coords.keys()) == {"samples", "i"} 

50 

51 # Assert number of triangles on each node for complete graph is equal to (N_vertices-1)*(N_vertices-2) 

52 # (counting directed triangles) 

53 assert all(triangles.data.flatten() == (N_vertices - 1) * (N_vertices - 2)) 

54 

55 # Bin the triangles along a dimension 

56 triangles_binned = ops.hist( 

57 triangles, dim="i", bins=100, ranges=[0, 100], use_bins_as_coords=True 

58 ) 

59 

60 # Assert that all triangles are in one bin, namely the one at (N_vertices-1)*(N_vertices-2) 

61 assert all( 

62 triangles_binned.sel( 

63 {"x": (N_vertices - 1) * (N_vertices - 2)}, method="nearest" 

64 ).data 

65 == N_vertices 

66 ) 

67 assert triangles_binned.equals( 

68 nw_ops.binned_nw_statistic( 

69 triangles, bins=100, ranges=[0, 100], sample_dim="samples" 

70 ) 

71 ) 

72 

73 

74def test_sel_matrix_indices(): 

75 """Tests selecting matrix indices""" 

76 N_samples, N_vertices = 10, 10 

77 n_indices = 5 

78 A_samples = xr.DataArray( 

79 np.random.rand(N_samples, N_vertices, N_vertices), 

80 dims=["sample", "i", "j"], 

81 coords=dict( 

82 sample=np.arange(N_samples), 

83 i=np.arange(N_vertices), 

84 j=np.arange(N_vertices), 

85 ), 

86 name="_A", 

87 ) 

88 

89 indices = xr.Dataset( 

90 data_vars=dict( 

91 i=(["idx"], np.random.randint(0, N_vertices, n_indices)), 

92 j=(["idx"], np.random.randint(0, N_vertices, n_indices)), 

93 ), 

94 coords=dict(idx=np.arange(n_indices)), 

95 ).expand_dims(dict(sample=np.arange(N_samples))) 

96 

97 selection = nw_ops.sel_matrix_indices(A_samples, indices, exclude_dim=["sample"]) 

98 assert set(selection.coords) == {"idx", "sample", "i", "j"} 

99 assert len(selection.coords["idx"]) == n_indices 

100 

101 selection = nw_ops.sel_matrix_indices( 

102 A_samples, indices, exclude_dim=["sample"], drop=True 

103 ) 

104 assert set(selection.coords) == {"idx", "sample"} 

105 

106 

107def test_largest_entry_indices(): 

108 """Tests returning the indices of the largest entries""" 

109 N_vertices = 10 

110 A = xr.DataArray( 

111 [[i + j for j in range(N_vertices)] for i in range(N_vertices)], 

112 dims=["i", "j"], 

113 coords=dict(i=np.arange(N_vertices), j=np.arange(N_vertices)), 

114 ) 

115 

116 # Get the indices for the symmetric case 

117 indices = nw_ops.largest_entry_indices(A, 4, symmetric=False) 

118 assert set(indices.data_vars) == {"i", "j", "value"} 

119 assert set(indices.coords.keys()) == {"idx"} 

120 

121 assert indices.isel({"idx": 0})["i"] == N_vertices - 1 

122 assert indices.isel({"idx": 0})["j"] == N_vertices - 1 

123 assert indices.isel({"idx": 0})["value"] == 2 * (N_vertices - 1) 

124 

125 # Repeat for asymmetric case 

126 indices = nw_ops.largest_entry_indices(A, 4, symmetric=True) 

127 assert indices.isel({"idx": -1})["i"] == N_vertices - 3 

128 assert indices.isel({"idx": -1})["value"] == 16 

129 

130 # Select the entries from A and check they are equal to values obtained 

131 assert all( 

132 indices["value"].data == nw_ops.sel_matrix_indices(A, indices[["i", "j"]]).data 

133 ) 

134 

135 

136# ---------------------------------------------------------------------------------------------------------------------- 

137# DISTRIBUTION OPERATIONS 

138# ---------------------------------------------------------------------------------------------------------------------- 

139 

140# TODO: Add tests