marriage_analysis

import numpy as np
data = np.loadtxt('marriage.txt').reshape((4, 5, 5))

def rake(A, epsilon=0.001):
    """Return a raked version of a given array, i.e., with row and column
    sums adjusted to unity.
    """
    B = np.empty(A.shape)
    B[:,:] = A[:,:]
    while (abs(B.sum(axis=0) - 1).max() > epsilon or
           abs(B.sum(axis=1) - 1).max() > epsilon):
        B = B / B.sum(axis=0)
        B = (B.transpose() / B.sum(axis=1)).transpose()
    return B

def frac_data(A):
    """Return the fractions of population found in the lower triangle, the
    main diagonal, and the upper triangle.
    """
    assert len(A.shape) == 2
    assert A.shape[0] == A.shape[1]
    n = A.shape[0]
    diag = np.trace(A) / A.sum()
    lower = (np.tril(A).sum() - np.trace(A)) / A.sum()
    upper = (np.triu(A).sum() - np.trace(A)) / A.sum()
    return lower, diag, upper

def split_data(A):
    """Return the mean values found in the lower triangle, the
    main diagonal, and the upper triangle.
    """
    assert len(A.shape) == 2
    assert A.shape[0] == A.shape[1]
    n = A.shape[0]
    diag = np.trace(A) / n
    lower = (np.tril(A).sum() - np.trace(A)) / ((n * n - n) / 2)
    upper = (np.triu(A).sum() - np.trace(A)) / ((n * n - n) / 2)
    return lower, diag, upper

def high_vals(A):
    """Return a version of A with all "low" values set to zero (i.e., all values with
    (i + j) < n).  N.B. the cross diagonal will not be set to zero.
    """
    B = np.empty(A.shape)
    B[:,:] = A[:,:]
    n = len(A)
    for i in range(n):
        for j in range(n-i-1):
            B[i,j] = 0
    return B

def low_vals(A):
    """Return a version of A with all "high" values set to zero (i.e., all values with
    (i + j) >= n-1).  N.B. the cross diagonal will not be set to zero.
    """
    B = np.empty(A.shape)
    B[:,:] = A[:,:]
    n = len(A)
    for i in range(n):
        for j in range(n-i, n):
            B[i,j] = 0
    return B