jit_circular.py

def _get_similar_ones(data: pandas.DataFrame, x_name, y_name, index):
    return data[(data[x_name] == data.iloc[index][x_name]) & (data[y_name] == data.iloc[index][y_name])]


def circular_layout(data: pandas.DataFrame, x_name, y_name, x_offset, y_offset, radius):
    data[x_offset] = None
    data[y_offset] = None
    for index, row in data.iterrows():
        if not pandas.isnull(data.iloc[index][x_offset]) and not pandas.isnull(data.iloc[index][y_offset]):
            continue
        similar_ones = _get_similar_ones(data, x_name, y_name, index)
        if len(similar_ones) <= 1:
            for similar_index, similar_row in similar_ones.iterrows():
                data.loc[similar_index, x_offset] = 0
                data.loc[similar_index, y_offset] = 0
            continue
        step = numpy.pi * 2 / len(similar_ones)
        current_angle = 0
        for similar_index, similar_row in similar_ones.iterrows():
            data.loc[similar_index, x_offset] = radius * numpy.cos(current_angle)
            data.loc[similar_index, y_offset] = radius * numpy.sin(current_angle)
            current_angle += step


def jitt(data: pandas.DataFrame, x_jit_name: str, y_jit_name: str, jit_radius: float):
    data[x_jit_name] = 0
    data[y_jit_name] = 0
    data[x_jit_name] = data[x_jit_name].apply(lambda _: random.uniform(-jit_radius, jit_radius))
    data[y_jit_name] = data[y_jit_name].apply(lambda _: random.uniform(-jit_radius, jit_radius))