lab3 / main1

"""
Скрипт использует файл dataset.json.
njit нужен только для ускорения python.
Если возникли проблемы с установкой numba,
закомментировать строки 11, 19, 23, 27, 31.
Установить numba: pip3 install numba.
"""

import numpy as np
from matplotlib import pyplot as plt
from numba import njit
import time
import json


def is_line_closed(x_n, y_n, x_n_p, y_n_p, electrode_f):
    return ((y_n >= electrode_f(x_n) and y_n_p <= electrode_f(x_n_p)) or \
            (y_n <= electrode_f(x_n) and y_n_p >= electrode_f(x_n_p))) and \
            neg_x_range[0] <= x_n <= neg_x_range[1]


@njit
def get_r(x1, y1, x2, y2): return ((x1 - x2) ** 2 + (y1 - y2) ** 2) ** .5


@njit
def get_a(x1, y1, x2, y2): return np.arctan(abs(y1 - y2) / abs(x1 - x2)) if abs(x1 - x2) > 1e-20 else np.pi / 2


@njit
def get_e(r): return 1 / r ** 2 if r > 1e-10 else 0


@njit
def get_next_point(x_n, y_n):
    pos_e_x, pos_e_y, neg_e_x, neg_e_y = [0] * 4

    for pos_x, pos_y, neg_x, neg_y in zip(pos_xs, pos_ys, neg_xs, neg_ys):
        pos_r = get_r(pos_x, pos_y, x_n, y_n)
        neg_r = get_r(neg_x, neg_y, x_n, y_n)

        pos_a = get_a(x_n, y_n, pos_x, pos_y)
        neg_a = get_a(x_n, y_n, neg_x, neg_y)

        pos_e = get_e(pos_r)
        neg_e = get_e(neg_r)

        pos_e_x += pos_e * np.cos(pos_a) if x_n > pos_x else - pos_e * np.cos(pos_a)
        pos_e_y += pos_e * np.sin(pos_a) if y_n > pos_y else - pos_e * np.sin(pos_a)
        neg_e_x += neg_e * np.cos(neg_a) if x_n < neg_x else - neg_e * np.cos(neg_a)
        neg_e_y += neg_e * np.sin(neg_a) if y_n < neg_y else - neg_e * np.sin(neg_a)

    return x_n + (pos_e_x + neg_e_x) * scale, y_n + (pos_e_y + neg_e_y) * scale


def get_line(x_n, y_n):
    xs, ys, x_n_p, y_n_p = [], [], x_n, y_n

    i = 0
    while i < max_steps:
        x_n_p, y_n_p = x_n, y_n
        x_n, y_n = get_next_point(x_n, y_n)

        if is_line_closed(x_n,y_n, x_n_p, y_n_p, neg_f):
            break

        xs += [x_n]
        ys += [y_n]

        i += 1

    return xs, ys


if __name__ == '__main__':
    time_start = time.time()

    # извлечение данных из файла
    with open('dataset1.json', 'r', encoding='utf-8') as f:
        data = json.load(f)

    # данные
    exec(data['импортировать_библиотеку_для_функций'])
    scale = eval(data['масштаб'])
    n = eval(data['кол-во_точек_на_электроде'])
    lines_count = eval(data['кол-во_линий'])
    electrode_eps = eval(data['погрешность_границ_электрода'])
    pos_f = eval(data['функция_положительного_электрода'])
    neg_f = eval(data['функция_отрицательного_электрода'])
    pos_x_range = eval(data['границы_x_положительного_электрода'])
    neg_x_range = eval(data['границы_x_отрицательного_электрода'])
    max_steps = eval(data['максиальное_кол_во_шагов'])

    pos_xs = np.linspace(pos_x_range[0], pos_x_range[1], n, dtype=np.float64)
    pos_ys = np.array(list(map(pos_f, pos_xs)), dtype=np.float64)

    neg_xs = np.linspace(neg_x_range[0], neg_x_range[1], n, dtype=np.float64)
    neg_ys = np.array(list(map(neg_f, neg_xs)), dtype=np.float64)

    lines_starts = np.linspace(pos_x_range[0], pos_x_range[1], lines_count // 2)

    print(f'\rprogress is 0/{lines_count} lines (0%)', end='')
    for i in range(lines_count):
        xs, ys = [], []

        if i < lines_count // 2:
            xs, ys = get_line(lines_starts[i], pos_f(lines_starts[i]) - electrode_eps)
        else:
            xs, ys = get_line(lines_starts[i - lines_count // 2], pos_f(lines_starts[i - lines_count // 2]) + electrode_eps)

        plt.plot(xs, ys, color='black', linestyle='--')
        print(f'\rprogress is {i + 1}/{lines_count} lines ({(i + 1) / lines_count * 1e2: .1f}%)', end='')

    plt.plot(pos_xs, pos_ys, color='red', label='pos')
    plt.plot(neg_xs, neg_ys, color='blue', label='neg')
    plt.legend(loc='upper right')
    plt.savefig('img.png')

    print(f'\ntime elapsed: {time.time() - time_start: .0f}')