Лаба 2 / 4_var

import warnings
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import erlang

warnings.filterwarnings('ignore')

'''
Эрланговский, k = 2 (порядок), lamdba = 2 (параметр)

Значит будет производится суммирование 2 случайных величин (распределённых экспоненциально),
с параметром lambda = 2

Экспоненциальная случ. величина создаётся по следующему закону:
-1/lambda * ln(R), где R - равномерная случ. величина

Необходимые компоненты программы:
1) Функция построения графика модулируемого закона (done)
2) Функция создания эрланговского распределения из равномерного (done)
3) Функция построения гистограммы (с возможностью подсчета мат ожидания и дисперсии) (done)
4) Табличное представление результатов моделирования (done)

'''

def create_erlang(scale, param, N: int = ...):
    uni_nums = np.random.uniform(low=0.0, high=1.0, size=scale*N)
    expon_nums = (-1/param) * np.log(uni_nums)
    erlang_nums = np.sum(expon_nums.reshape(-1, scale), axis=1)
    return erlang_nums

def plot_modulated(fig: int = ...):
    global lmbd
    x = np.linspace(0, 5, 1000)
    y = erlang.pdf(x, k, scale=1/lmbd)
    plt.plot(x, y, label=f'Erlang({k}, {lmbd})')
    plt.xlabel('Значения')
    plt.ylabel('Плотность вероятности')
    plt.title('График Эрланговского распределения')
    plt.legend()
    plt.grid()

def plot_mean_var(sample, fig: int, is_variance: bool = False):
    sizes = [i for i in range(1, len(sample)+1)]
    global k
    global lmbd

    if is_variance:
        num = k*(1/lmbd)**2
        name = "Дисперсия"
    else:
        num = k*(1/lmbd)
        name = 'Мат. ожидания'

    plt.figure(fig)
    plt.title(f'График зависимости {name} от количества')
    plt.plot(sizes, sample, label='Эксперимент')
    plt.plot(sizes, [num]*len(sizes), label='Истина')
    plt.grid()
    plt.legend()
    
def calculate_samples(scale, param, fig: int = ..., calculate_mean_var: bool = False, get_dist: bool = False, get_hist: bool = False):
    samples, exit = [], []
    sample_min, sample_max, delta, parts = 0, 0, 0, 10
    N_i = [0] * parts
    if calculate_mean_var:
        mean, var = [], []

    while min(N_i) < 10:
        # Генерация случайного числа
        alpha = create_erlang(scale, param, 1)
        samples.append(alpha[0])
        N = len(samples)
        print(N_i)
        # Расчет на интервалы
        if alpha > sample_min and delta != 0 and alpha < sample_max:
            k = int(np.floor((alpha - sample_min)/delta))
            N_i[k] += 1
        
        # Вероятность ошибки
        p = (N - np.sum(N_i))/N

        # Перерасчеты данных
        if p >= 0.01:
            sample_min, sample_max = min(samples), max(samples)
            delta = (sample_max - sample_min)/(parts)
            borders = np.linspace(sample_min, sample_max, parts, True)
            ids = np.searchsorted(borders, samples)
            N_i = [0] * parts
            for id in ids:
                N_i[id] += 1

        # Считаем среднее и дисперсию
        if calculate_mean_var:
            mean.append(np.mean(samples))
            var.append(np.var(samples))

    # Создаём массив Y
    x, y = [], [n/(N*delta) for n in N_i]
    x.append(sample_min + delta/2)
    for i in range(parts-1):
        x.append(x[i] + delta)

    # Построение гистограммы
    if get_hist:
        x_axis = np.linspace(sample_min, sample_max, 100)
        plt.figure(fig)
        plt.bar(x, y, width=delta)
        plt.plot(x_axis, erlang.pdf(x_axis, scale, scale=1/param), color='orange')
        plt.title('Гистограмма Эрланговского распределения')
        plt.grid()

    if get_dist:
        exit.append(samples)
        exit.append(N_i)
    
    if calculate_mean_var:
        exit.append(mean)
        exit.append(var)

    return exit
    

if __name__ == "__main__":
    # Глобальные параметры
    k, lmbd = 2, 2
    N = 1000

    # График модулируемого закона
    plot_modulated(fig=1)

    # Табличное представление данных + гистограмма распределения
    smpl, dist, means, vars = calculate_samples(k, lmbd, fig=2, get_dist=True, get_hist=True, calculate_mean_var=True)
    df = pd.DataFrame(dist)
    print('\nТабличное представление данных')
    print(df.T)

    # Графики зависимостей
    plot_mean_var(means, 3)
    plot_mean_var(vars, 4, True)

    plt.show()