|
| 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)) = }") |
![pre-commit-ci[bot]](https://avatars.githubusercontent.com/in/68672?v=4&size=40){kind=avatar}
0 commit comments