tsscpp_probabilities.sage

############################################################
####        Rho and Generalized Rho
############################################################


def K_restriction(K, S):
    """
    K_restriction(K, S)

    Produces the restricted K matrix on the set of points S

    Inputs:
    K: matrix K to restrict
    S: a list of integers (need not be ordered)
    """

    S = [2*(s-1) for s in S]
    S_shift = [s+1 for s in S]
    S.extend(S_shift)
    S.sort()

    return(K[S, S])

def rho(S, K, data_type):
    """
    rho(S, K, data_type)

    Computes the correlation function rho of a matrix
    K, with data type either numeric ('n') or symbolic ('s')

    Inputs:
    S: a list of integers (need not be ordered)
    K: matrix K whose Pfaffian yields rho
    data_type: numeric or symbolic
    """

    if data_type == 'n':
        n = (K.nrows() + 4)/4
        if len(S) > n-1:
            return(0)
        K_restrict_S = K_restriction(K, S)
        prob = sqrt(K_restrict_S.det())
        return(prob)

    elif data_type == 's':
        K_restrict_S = K_restriction(K, S)
        prob = K_restrict_S.pfaffian()
        num = prob.numerator()
        den = prob.denominator()
        return(num/den)

    else:
        print("Error: rho must be given data_type 'n' or 's'")
        return()

def rho_generalized(S, T, K, data_type):
    """
    rho_generalized(S, T, K, data_type)

    Computes the generalized correlation function of a matrix
    K, with data type either numeric ('n') or symbolic ('s');
    that is, the probability that the point process contains
    all the points in S and none of the points in T

    Inputs:
    S: a list of integers; points to be included
    T: a list of integers; points to be excluded
    K: matrix K whose Pfaffian yields rho
    data_type: numeric or symbolic
    """

    sets_raw = Subsets(T, submultiset=True)
    sets = sets_raw.list()
    for subset in sets:
            subset.extend(S)

    # creates all subsets of T and takes union of S with
    # all such subsets

    prob = sum((-1)^(len(subset) - len(S)) * rho(subset, K, data_type) for subset in sets)
    if data_type == 'n':
        return(prob)
    elif data_type == 's':
        num = prob.numerator()
        den = prob.denominator()
        return(num/den)
    else:
        print("Error: rho_generalized must be given data_type 'n' or 's'")
        return()


############################################################
####        Event Probabilities
############################################################


def first_k_vert(k, K, data_type):
    """
    Computes the probability that the first k paths are
    vertical, given an appropriate kernel K and its data
    type 'n' or 's'
    """

    S = [1 .. k]
    return(rho(S, K, data_type))

def first_k_hor(k, K, data_type):
    """
    Computes the probability that the first k paths are
    horizontal, given an appropriate matrix K and its data
    type 'n' or 's'
    """

    S = []
    T = [2*ii - 1 for ii in [1 .. k]]
    return(rho_generalized(S, T, K, data_type))

def set_minus(T, S):
    """
    Returns T setminus S, where T and S are lists
    """

    T_minus_S = [t for t in T if t not in S]
    return(T_minus_S)

def prob_path_k_hits_x(k, x, K, data_type):
    """
    Computes the probability that the kth path hits
    the endpoint with x-coordinate x, given an appropriate
    correlation matrix K with data type 'n' or 's'
    """

    if k == x:
        return(first_k_vert(k, K, data_type))

    lower_indices = [1 .. x-1]
    hitting_sets = Subsets(lower_indices, k-1, submultiset=True)
    hitting_sets = hitting_sets.list()
    for S in hitting_sets:
        S.append(x)

    prob = sum(rho_generalized(S, set_minus([1 .. x], S), K, data_type) for S in hitting_sets)
    if data_type == 'n':
        return(prob)
    elif data_type == 's':
        num = prob.numerator()
        den = prob.denominator()
        return(num/den)
    else:
        print("Error: prob_path_k_hits_x must be given data_type 'n' or 's'")

def prob_dist_path_k(k, K, data_type):
    """
    Computes the probability distribution of the endpoint
    of the k-th path, given an appropriate correlation
    matrix K and its data type 'n' or 's'
    """

    dist = []
    for x in [k .. 2*k]:
        dist.append(prob_path_k_hits_x(k, x, K, data_type))
    return(dist)