joint PMF functions

joint_pmf.py

import numpy as np

expected_value = 0

jpmf = np.array([
                [0.025, 0.015, 0.010],
                [0.050, 0.030, 0.020],
                [0.125, 0.075, 0.050],
                [0.150, 0.090, 0.060],
                [0.100, 0.060, 0.040],
                [0.050, 0.030, 0.020]
])

# find the expected value of h(x,y) given joint pmf
def joint_functional_E(h, jpmf):

    expected_value = 0
    x_range, y_range = jpmf.shape

    for x in range(x_range):
        for y in range(y_range):
            print("adding (x,y) ({0},{1})".format(x,y))
            expected_value += h(x,y) * jpmf[x][y]

    return expected_value

print("Expected number of total vehicles: {0}".format(
        joint_functional_E(lambda x,y: x+y, jpmf)))

exy = joint_functional_E(lambda x,y: x*y, jpmf)

# check independence of drv's, given joint pmf
def check_independence(jpmf):

    print("Checking independence of X and Y...")

    x_pmf, y_pmf = np.sum(jpmf, axis=1), np.sum(jpmf, axis=0)

    print(x_pmf)
    print(y_pmf)

    multiplied_probs = []

    for x in range(len(x_pmf)):
        for y in range(len(y_pmf)):
            multiplied_probs.append(x_pmf[x]*y_pmf[y])
    print(multiplied_probs)
    multiplied_probs = np.reshape(
                        np.array(multiplied_probs),
                        (-1, 3))
    print(multiplied_probs)
    print(jpmf)

    if np.allclose(multiplied_probs,
                   jpmf, 0.001, 0.001):
        print("drv's are independent!")
        return True
    else:
        print("drv's are dependent!")
        return False

check_independence(jpmf)

# generalized E(h(x))
def expected_val(h, pmf):
    print("range: {0}".format(range(len(pmf))))
    return sum([h(index)*pmf[index] for index in range(len(pmf))])

# return variances of drv's of a joint pmf
def component_variances(jpmf):
    x_pmf, y_pmf = np.sum(jpmf, axis=1), np.sum(jpmf, axis=0)
    Ex = expected_val(lambda x: x, x_pmf)
    Ey = expected_val(lambda y: y, y_pmf)

    Vx = expected_val(lambda x: x**2, x_pmf) - Ex**2
    Vy = expected_val(lambda x: x**2, y_pmf) - Ey**2

    return Vx, Vy

print(component_variances(jpmf))

# return covariance of drv's of a joint pmf
def covariance(jpmf):
    x_pmf, y_pmf = np.sum(jpmf, axis=1), np.sum(jpmf, axis=0)
    Ex = expected_val(lambda x: x, x_pmf)
    Ey = expected_val(lambda y: y, y_pmf)

    print(Ex,Ey)

    covariance = joint_functional_E(
        lambda x,y: (x - Ex)*(y-Ey),
        jpmf
        )
    print(covariance)
    return covariance

component_variances(jpmf)