Coverage for src/abcd_graph/models.py: 100%

1# Copyright (c) 2024 Jordan Barrett & Aleksander Wojnarowicz

2#

3# Permission is hereby granted, free of charge, to any person obtaining a copy

4# of this software and associated documentation files (the "Software"), to deal

5# in the Software without restriction, including without limitation the rights

6# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

7# copies of the Software, and to permit persons to whom the Software is

8# furnished to do so, subject to the following conditions:

9#

10# The above copyright notice and this permission notice shall be included in all

11# copies or substantial portions of the Software.

12#

13# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

14# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

15# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

16# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

17# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

18# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE

19# SOFTWARE.

20

21 __all__ = [

22 "Model",

23 "configuration_model",

24 "chung_lu",

25]

26

27 from typing import Protocol

28

29 import numpy as np

30 from numpy.typing import NDArray

31

32

33 class Model(Protocol):

34 def __call__(self, degree_sequence: dict[int, int]) -> NDArray[np.int64]: ...

35

36 @property

37 def __name__(self) -> str: ...

38

39

40 def configuration_model(degree_sequence: dict[int, int]) -> NDArray[np.int64]:

41 labels = list(degree_sequence.keys())

42 counts = list(degree_sequence.values())

43

44 vertices = np.repeat(labels, counts)

45

46 np.random.shuffle(vertices)

47

48 edges = np.array(vertices).reshape(-1, 2)

49

50 return edges

51

52

53 def normalize(degrees: list[int]) -> NDArray[np.float64]:

54 """Normalize the degree sequence."""

55 degrees_array: NDArray[np.int64] = np.array(degrees)

56 norm = degrees_array.sum()

57 result: NDArray[np.float64] = np.divide(degrees_array, norm)

58 return result

59

60

61 def chung_lu(degree_sequence: dict[int, int]) -> NDArray[np.int64]:

62 """Generate a Chung-Lu random graph based on a given degree sequence."""

63 nodes = list(degree_sequence.keys())

64 degrees = list(degree_sequence.values())

65

66 return np.random.choice(a=nodes, size=int(sum(degrees)), p=normalize(degrees)).reshape(-1, 2)