numerics.py 文件源码

def vech_kh(X, stack_cols=True, conserve_norm=False):
    assert X.shape[0] == X.shape[1]
    # Scale off-diagonal indexes if norm has to be preserved
    d = X.shape[0]
    if conserve_norm:
        # Scale off-diagonal
        tmp = np.copy(X)
        triu_scale_idx = np.triu_indices(d, 1)
        tmp[triu_scale_idx] = np.sqrt(2) * tmp[triu_scale_idx]
    else:
        tmp = X
    triu_idx_r = []
    triu_idx_c = []
    # Find appropriate indexes
    if stack_cols:
        for c in range(0, d):
            for r in range(0, c+1):
                triu_idx_r.append(r)
                triu_idx_c.append(c)
    else:
        for r in range(0, d):
            for c in range(r, d):
                triu_idx_r.append(r)
                triu_idx_c.append(c)
    # Extract and return upper triangular
    triu_idx = (triu_idx_r, triu_idx_c)
    return tmp[triu_idx]