TheAlgorithms-Python

1
"""
2
Calculate joint probability distribution
3
https://en.wikipedia.org/wiki/Joint_probability_distribution
4
"""
5

6

7
def joint_probability_distribution(
8
    x_values: list[int],
9
    y_values: list[int],
10
    x_probabilities: list[float],
11
    y_probabilities: list[float],
12
) -> dict:
13
    """
14
    >>> joint_distribution =  joint_probability_distribution(
15
    ...     [1, 2], [-2, 5, 8], [0.7, 0.3], [0.3, 0.5, 0.2]
16
    ... )
17
    >>> from math import isclose
18
    >>> isclose(joint_distribution.pop((1, 8)), 0.14)
19
    True
20
    >>> joint_distribution
21
    {(1, -2): 0.21, (1, 5): 0.35, (2, -2): 0.09, (2, 5): 0.15, (2, 8): 0.06}
22
    """
23
    return {
24
        (x, y): x_prob * y_prob
25
        for x, x_prob in zip(x_values, x_probabilities)
26
        for y, y_prob in zip(y_values, y_probabilities)
27
    }
28

29

30
# Function to calculate the expectation (mean)
31
def expectation(values: list, probabilities: list) -> float:
32
    """
33
    >>> from math import isclose
34
    >>> isclose(expectation([1, 2], [0.7, 0.3]), 1.3)
35
    True
36
    """
37
    return sum(x * p for x, p in zip(values, probabilities))
38

39

40
# Function to calculate the variance
41
def variance(values: list[int], probabilities: list[float]) -> float:
42
    """
43
    >>> from math import isclose
44
    >>> isclose(variance([1,2],[0.7,0.3]), 0.21)
45
    True
46
    """
47
    mean = expectation(values, probabilities)
48
    return sum((x - mean) ** 2 * p for x, p in zip(values, probabilities))
49

50

51
# Function to calculate the covariance
52
def covariance(
53
    x_values: list[int],
54
    y_values: list[int],
55
    x_probabilities: list[float],
56
    y_probabilities: list[float],
57
) -> float:
58
    """
59
    >>> covariance([1, 2], [-2, 5, 8], [0.7, 0.3], [0.3, 0.5, 0.2])
60
    -2.7755575615628914e-17
61
    """
62
    mean_x = expectation(x_values, x_probabilities)
63
    mean_y = expectation(y_values, y_probabilities)
64
    return sum(
65
        (x - mean_x) * (y - mean_y) * px * py
66
        for x, px in zip(x_values, x_probabilities)
67
        for y, py in zip(y_values, y_probabilities)
68
    )
69

70

71
# Function to calculate the standard deviation
72
def standard_deviation(variance: float) -> float:
73
    """
74
    >>> standard_deviation(0.21)
75
    0.458257569495584
76
    """
77
    return variance**0.5
78

79

80
if __name__ == "__main__":
81
    from doctest import testmod
82

83
    testmod()
84
    # Input values for X and Y
85
    x_vals = input("Enter values of X separated by spaces: ").split()
86
    y_vals = input("Enter values of Y separated by spaces: ").split()
87

88
    # Convert input values to integers
89
    x_values = [int(x) for x in x_vals]
90
    y_values = [int(y) for y in y_vals]
91

92
    # Input probabilities for X and Y
93
    x_probs = input("Enter probabilities for X separated by spaces: ").split()
94
    y_probs = input("Enter probabilities for Y separated by spaces: ").split()
95
    assert len(x_values) == len(x_probs)
96
    assert len(y_values) == len(y_probs)
97

98
    # Convert input probabilities to floats
99
    x_probabilities = [float(p) for p in x_probs]
100
    y_probabilities = [float(p) for p in y_probs]
101

102
    # Calculate the joint probability distribution
103
    jpd = joint_probability_distribution(
104
        x_values, y_values, x_probabilities, y_probabilities
105
    )
106

107
    # Print the joint probability distribution
108
    print(
109
        "\n".join(
110
            f"P(X={x}, Y={y}) = {probability}" for (x, y), probability in jpd.items()
111
        )
112
    )
113
    mean_xy = expectation(
114
        [x * y for x in x_values for y in y_values],
115
        [px * py for px in x_probabilities for py in y_probabilities],
116
    )
117
    print(f"x mean: {expectation(x_values, x_probabilities) = }")
118
    print(f"y mean: {expectation(y_values, y_probabilities) = }")
119
    print(f"xy mean: {mean_xy}")
120
    print(f"x: {variance(x_values, x_probabilities) = }")
121
    print(f"y: {variance(y_values, y_probabilities) = }")
122
    print(f"{covariance(x_values, y_values, x_probabilities, y_probabilities) = }")
123
    print(f"x: {standard_deviation(variance(x_values, x_probabilities)) = }")
124
    print(f"y: {standard_deviation(variance(y_values, y_probabilities)) = }")
125