Изменить размер картинки

from math import * 
import numpy as np 
import random as rd
import matplotlib.pyplot as plt
from PIL import Image
 
 
def energy_group_num(E):
    """Возвращает массив номеров К энергетической группы, в соответствии энергии частицы.
   Е - массив с энергиями частиц (МэВ). """
    n = len(E)
    K = [0] * n
    for i in range(n):
        if 0.1 <= E[i] <= 0.15:
            K[i] = 0
        elif 0.15 < E[i] <= 0.2:
            K[i] = 1
        elif 0.2 < E[i] <= 0.3:
            K[i] = 2
        elif 0.3 < E[i] <= 0.4:
            K[i] = 3
        elif 0.4 < E[i] <= 0.5:
            K[i] = 4
        elif 0.5 < E[i] <= 0.6:
            K[i] = 5
        elif 0.6 < E[i] <= 0.8:
            K[i] = 6
        elif 0.8 < E[i] <= 1:
            K[i] = 7
        elif 1 < E[i] <= 1.5:
            K[i] = 8
        elif 1.5 < E[i] <= 2:
            K[i] = 9
        elif 2 < E[i] <= 2.1:
            K[i] = 10
        elif 2.1 < E[i] <= 2.2:
            K[i] = 11
        elif 2.2 < E[i] <= 2.3:
            K[i] = 12
        elif 2.3 < E[i] <= 2.4:
            K[i] = 13
        elif 2.4 < E[i] <= 2.5:
            K[i] = 14
        elif 2.5 < E[i] <= 2.6:
            K[i] = 15
        elif 2.6 < E[i] <= 2.7:
            K[i] = 16
        elif 2.7 < E[i] <= 2.8:
            K[i] = 17
        elif 2.8 < E[i] <= 2.9:
            K[i] = 18
        elif 2.9 < E[i] <= 3:
            K[i] = 19
        else:
            print('Энергия частица вышла за пределы. E =', E[i])
    return K
 
 
def generate_birthpoitns(n):
 
    """Разыгрывает координаты точек рождения (x0, y0, z0) для n частиц 
    на источнике в виде прямоугольника a=40 cm, b=10 cm
    Также возвращает массив начальных энергий E. 
    """
    x0 = [0] * n
    y0 = [0] * n 
    z0 = [0] * n
    G1 = [rd.uniform(0, 1) for i in range(n)]
    G2 = [rd.uniform(0, 1) for i in range(n)]
    a = 40 # длина прямоугольника (см)
    b = 10 # ширина прямоугольника (см)
 
    for i in range(n):
        r = a * G1[i]
        phi = b * G2[i]
        x0[i] = r 
        y0[i] = phi
 
    E = [2.5] * n
 
    return x0, y0, z0, E
 
 
def generate_1stInteraction_pts(x0, y0, z0, E): 
    """Вычисляет координаты первого взаимодействия частиц. 
    х0, у0, z0 - списки координат частиц.
    E - энергия частиц. """
    # Число частиц
    n = len(x0)
    # Номера энерг-х групп для частиц
    K = energy_group_num(E)
    # rho - плотность алюминия, г/см^3
 
    # Координаты точек первого взаимодействия:
    x1 = [0] * n 
    y1 = [0] * n
    z1 = [0] * n
    omega1 = [[0, 0, 0] for i in range(n)]
 
    G1 = [rd.uniform(0, 1) for i in range(n)]
    G2 = [rd.uniform(0, 1) for i in range(n)]
    G3 = [rd.uniform(0, 1) for i in range(n)]
 
    for i in range(n):
        L = -log(G1[i]) / SIGMA[K[i]] # Длина своб. пробега
        psi = 2 * pi * G2[i]
        theta = pi / 2 * G3[i]
        x1[i] = x0[i] + L * cos(theta) * cos(psi)
        y1[i] = y0[i] + L * cos(theta) * sin(psi)
        z1[i] = z0[i] + L * sin(theta)
        omega1[i][0] = (x1[i] - x0[i]) / L
        omega1[i][1] = (y1[i] - y0[i]) / L
        omega1[i][2] = (z1[i] - z0[i]) / L
 
    return x1, y1, z1, omega1
 
 
def solution_area_pts(x1, y1, z1, omega1, E1):
    """Возвращает координаты частиц x, y, z, 
    которые остались в области решения задачи, и их энергию E.
    """
    x = []
    y = []
 
    z = []
    omega = []
    E = []
    n0 = len(x1)
 
    for i in range(n0):
        if abs(x1[i]) <= 25 and abs(y1[i]) <= 40 and 0 <= z1[i] <= 20: 
            x.append(x1[i])
            y.append(y1[i])
            z.append(z1[i])
            omega.append(omega1[i])
            E.append(E1[i])
    n_out = n0 - len(x)
 
    print(' Из области вылетело', n_out, 'частиц.')
    
    return x, y, z, omega, E
 
def generate_interaction_type(x1, y1, z1, omega1, E1):
    """Розыгрыш типа взаимодействия.
    Возвращает координаты частиц x, y, z, для которых произошло 
    комптоновское рассеяние, их энергию E;
    число частиц n_ph, участвующих в фотоэффекте; 
    число частиц n_pair, участ-их в образовании электронно-позитронной пары.
    """
    # Номера энергетических групп для частиц
    K = energy_group_num(E1)
 
    n0 = len(x1)
    G = [rd.uniform(0, 1) for i in range(n0)]
    x = []
    y = []
    z = []
    omega = []
    E = []
    n_ph = 0
    n_pair = 0
 
    for i in range(n0):
        if sig_comp[K[i]] / SIGMA[K[i]] >= G[i]: # Произошло компт-ое рассеяние
            x.append(x1[i])
            y.append(y1[i])
            z.append(z1[i])
            omega.append(omega1[i])
            E.append(E1[i])
        elif (sig_comp[K[i]] + sig_ph[K[i]]) / SIGMA[K[i]] >= G[i]:
            # Фотоэффект
            n_ph += 1
        else: # Электронно-позитронная пара
            n_pair += 1
 
    print(' В фотоэффекте участвовало', n_ph, 'частиц.')
    print(' В образовании электронно-позитронной пары участвовало', n_pair,
    'частиц.')
 
    return x, y, z, omega, E
 
 
def scattered_photon_params(E0):
    """Рассчитывает энергию Е и косинус угла cos_theta
    рассеяния фотона.
    E0 - список энергий частиц. """
    # Число частиц
    n = len(E0)
    # mc2 - энергия покоя электрона, МэВ.
    E = [0] * n
    cos_theta = [0] * n
 
    for i in range(n): 
        alpha1 = E0[i] / mc2
        flag = True
        j = 0
        while flag and j < 10:
            G1 = rd.uniform(0, 1)
            G2 = rd.uniform(0, 1)
            x = alpha1 * (1 + 2 * alpha1 * G1) / (1 + 2 * alpha1)
            p = x / alpha1 + alpha1 / x + (1 / alpha1 - 1 / x) * (2 +
            1 / alpha1 - 1/ x)
 
            if G2 * (1 + 2 * alpha1 + 1 / (1 + 2 * alpha1)) < p:
                E[i] = x * mc2
                cos_theta[i] = 1 - 1 / x + 1 / alpha1 
                flag = False
            if j > 0:
                print('scattered_photon_params ' + 'Число итераций при расчете параметров в while', j)
 
 
    return E, cos_theta
 
 
def expulsion_small_energy(x1, y1, z1, omega1, E1, cos_theta1, W1):
    """Оставляет только частицы с энергией E > 0.1 МэВ. """
    n = len(E1)
 
    x = []
    y = []
    z = []
    omega = []
    E = []
    cos_theta = []
    W = []
 
    for i in range(n):
        if E1[i] > 0.1:
            x.append(x1[i])
            y.append(y1[i])
            z.append(z1[i])
            omega.append(omega1[i])
            E.append(E1[i]) 
            cos_theta.append(cos_theta1[i]) 
            W.append(W1[i])
 
    print(' Убрано частиц с E <= 0.1 МэВ:', n - len(E))
 
    return x, y, z, omega, E, cos_theta, W
 
 
def statistical_weight(E, W):
    """Пересчитывает статистический вес W частицы. """
    n = len(E)
 
    # Номера энергетических групп для частиц
    K = energy_group_num(E)
 
    for i in range(n):
        W[i] = W[i] * sig_comp[K[i]] / SIGMA[K[i]]
 
    return W
 
 
def expulsion_small_weight(x1, y1, z1, omega1, E1, cos_theta1, W1):
    """Оставляет только частицы со стат. весом W > 1e-11. """
    n = len(W1)
 
    x = []
    y = []
    z = []
    omega = []
    E = []
    cos_theta = []
    W = []
 
    for i in range(n):
        if W1[i] > 1e-11:
            x.append(x1[i])
            y.append(y1[i])
            z.append(z1[i]) 
            omega.append(omega1[i])
            E.append(E1[i]) 
            cos_theta.append(cos_theta1[i])
            W.append(W1[i])
 
    print(' Убрано частиц с W <= 1e-11:', n - len(W))
 
    return x, y, z, omega, E, cos_theta, W
 
 
def eta_calculation(alpha, u, w, k, xi, yi, zi, rd: tuple):
    """Вычисляет один вклад в Eta"""
 
    sig_c_up = Z * r_e ** 2 / 2 / (1 + alpha * (1 - u)) ** 2 * (1 + u ** 2 + alpha  ** 2 * (1 - u) ** 2 / ( 1 +  alpha * (1 - u)))
    sig_c_down = 2 * pi * Z * r_e ** 2 * ((1 + alpha) / alpha ** 2 * log(1 + 2 * alpha) / alpha)# + 3 * alpha / (1 + 2 * alpha) ** 2)
    W_x_sigc = w * sig_c_up / sig_c_down 
    SigmaRho = SIGMA[k]
 
    #(2 * (1 + alpha) / (1 + 2 * alpha) - log(1 + 2 * alpha) / 2 / alpha - (1 +
 
 
    delta_r2 = ((xi - rd[0]) ** 2 + (yi - rd[1]) ** 2 + (zi - rd[2]) ** 2)
    Eta = (W_x_sigc * exp(-SigmaRho * sqrt(delta_r2)) / delta_r2)
    return Eta
 
 
def detector_readings(x, y, z, E, cos_theta, W):
    """Вычисляет виртуальные вклады в показание детектора. """
 
    # Номера энергетических групп для частиц
    K = energy_group_num(E)
    # Число частиц
    n = len(x)
 
    Eta1 = [0] * 20
    Eta2 = [0] * 20
 
    for i in range(n):
        alpha = E[i] / mc2
        u = cos_theta[i]
 
        Eta1[K[i]] += eta_calculation(alpha, u, W[i], K[i], x[i],        y[i], z[i], rd1)
        Eta2[K[i]] += eta_calculation(alpha, u, W[i], K[i], x[i],        y[i], z[i], rd2)
 
    return Eta1, Eta2
 
 
def generate_newInteraction_pts(x1, y1, z1, omega1, E, cos_theta):
    """Вычисляет координаты следующего взаимодействия частиц. """
    n = len(x1)
 
    omega = [[0, 0, 0] for i in range(n)]
    G1 = [rd.uniform(0, 1) for i in range(n)]
    G2 = [rd.uniform(0, 1) for i in range(n)] # Номера энерг-х групп для частиц
    K = energy_group_num(E)
    # rho - плотность алюминия, г/см^3
    # Координаты точек следующего взаимодействия:
    x0 = [0] * n
    y0 = [0] * n
    z0 = [0] * n
 
    for i in range(n):
        u = cos_theta[i]
        psi = 2 * pi * G1[i]
        omega[i][2] = omega1[i][2] * u + sqrt((1 - u ** 2) * (1 - omega1[i][2] ** 2)) * cos(psi)
        omega[i][1] = (omega1[i][1] * (u - omega1[i][2] * omega[i][2]) + omega1[i][0] * sin(psi) * sqrt((1 - u ** 2)* (1 - omega1[i][2] ** 2))) / (1 - omega1[i][2] ** 2)
        omega[i][0] = (omega1[i][0] * (u - omega1[i][2] * omega[i][2]) - omega1[i][1] * sin(psi) * sqrt((1 - u ** 2) *(1 - omega1[i][2] ** 2))) / (1 - omega1[i][2] ** 2)
 
        L = -log(G2[i]) / SIGMA[K[i]] # Длина своб. пробега
 
        x0[i] = x1[i] + L * omega[i][0]
        y0[i] = y1[i] + L * omega[i][1]
        z0[i] = z1[i] + L * omega[i][2]
 
    return x0, y0, z0, omega
 
 
def plot_diagram(phi, det_title: str):
    """Выводит диаграмму плотности потока чатиц
    для различных энергетических групп.
 
    """
    labels = ['0.1  ', ' 0.15', '   0.2', ' 0.3', ' 0.4', ' 0.5', '    0.6', '0.8', '1', '1.5', '2',
    '2.1', '2.2', '2.3', '2.4', '2.5', '2.6', '2.7', '2.8', '2.9', '3,    Мэв']
    phi.append(0)
    x = [0.1, 0.15, 0.2, 0.3, 0.4, 0.5, 0.6, 0.8, 1, 1.5, 2, 2.1, 2.2, 2.3, 2.4,
    2.5, 2.6, 2.7, 2.8, 2.9, 3]
    width = [0.05, 0.05, 0.1, 0.1, 0.1, 0.1, 0.2, 0.2, 0.5, 0.5, 0.1, 0.1, 0.1,
    0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0]
    fig, ax = plt.subplots(figsize=(8, 4))
    rect = ax.bar(x, phi, width, align='edge', label='Плотность потока частиц для разных энергетических групп',
    edgecolor='black')
 
    ax.set_ylabel(r'$\varphi, см^{-2}$') 
    ax.set_title(det_title)
    ax.set_xticks(x)
    ax.set_xticklabels(labels)
    ax.legend()
    fig.tight_layout()
 
 
def trajectories(x0, y0, z0, x1, y1, z1):
    """Рисует траекторию частицы между взаимодйствиями"""
    n0 = len(x1)
    for i in range(n0):
        ax.scatter(x1[i], y1[i], z1[i], c='g', s=1) 
        ax.plot([x0[i], x1[i]], [y0[i], y1[i]], [z0[i], z1[i]])
 
 
# Кортежи значений сечений взаимодействий фотонов # различных энергий для алюминия, [см^2/г].
# Значение элемента кортежа с индексом i соответствует 
# энергетической группе с номером i (K[j] = i).
sig_comp = (0.129, 0.118, 0.102, 0.092, 0.0893, 0.0777,
0.0681, 0.0614, 0.0498, 0.0424, 0.0399, 0.0389,
0.038, 0.0372, 0.0365, 0.0358, 0.0353, 0.0347, 0.0343, 0.0335)
 
sig_ph = (0.00469, 0.00193, 0.000551, 0.000228, 0.000114,
0.0000672, 0.0000312, 0.0000109, 0.000004, 0.0000001,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
 
sig_pair = (0, 0, 0, 0, 0, 0, 0, 0,
0.000167, 0.00067, 0.00079, 0.00091, 0.00104,
0.00116, 0.00129, 0.00142, 0.00154, 0.00167, 0.00179, 0.00191)
 
SIGMA = [0] * 20
 
for i in range(len(SIGMA)):
    SIGMA[i] = (sig_comp[i] + sig_ph[i] + sig_pair[i]) * 2.7
 
# Число частиц
n = 100
# Энергия покоя электрона, МэВ
mc2 = 0.51099
# Атомный номер свинца
Z = 14
# Плотность свинца, г/см^3
rho = 11.3
# Kлассический радиус электрона
r_e = 2.8e-13
 
# Координаты детекторов (x,y,z), см:
rd1 = (0, 0, 75)
rd2 = (0, 0, 40)
 
W = [1] * n # Начальный стат. вес частиц
# размер перепишется в проверке на E > Emin # в expulsion_small_energy
 
x0, y0, z0, E = generate_birthpoitns(n)
print('Рождение', len(x0), 'частиц.')
x1, y1, z1, omega1 = generate_1stInteraction_pts(x0, y0, z0, E)
 
# Рисуем область
fig = plt.figure()
ax = plt.axes(projection='3d')
 
ax.set_xlim(-5, 65)
ax.set_ylim(-45, 45)
ax.set_zlim(0, 80)
 
# Рисуем начальные траектории
trajectories(x0, y0, z0, x1, y1, z1)
 
# Цилиндр
h = np.pi / 100
R_cil=30
H=20
t = np.arange(0, np.pi + h, h)
x = R_cil+R_cil * np.cos(t)
y = R_cil * np.sin(t)
z = np.arange(0, H, 0.05)
 
ax.plot(x, y, 0, color='k')
ax.plot([-R_cil, R_cil], [0, 0], color='k')
 
 
i = np.linspace(-25, 25, 100)
j = np.linspace(-10, 10, 100)
X0, Y0 = np.meshgrid(i, j)
p = 1 + (X0 * Y0) * 0
 
ax.plot_surface(X0, Y0, 0 * p, color='r', alpha=0.1)
ax.plot_surface(X0, Y0, H * p, color='r', alpha=0.1)
 
for i in z:
    ax.plot(x, y, i, color='r', alpha=0.1)
 
# Контур источника
a = 40
b = 10
x = np.arange(0, a, 0.05)
y = np.arange(0, b, 0.0125)
 
ax.plot(x, y, 0, color='k')
ax.plot([-a/2, a/2], [-b/2, b/2], color='k')
 
# Рисуем детекторы
ax.scatter(0, 0, 75, c='k')
ax.scatter(0, 0, 40, c='k')
ax.text(0, 2, 77, "Детектор\n(0, 0, 75)")
ax.text(0, 2, 40, "Детектор (0, 0, 40)")
 
# Показания детекторов с учетом всех взаимодействий
 
totEta1 = [0] * 20
totEta2 = [0] * 20
 
j = 0
while j < 100:
    x1, y1, z1, omega1, E = solution_area_pts(x1, y1, z1, omega1, E)
    if len(x1) < 1:
        break
    print(' В области осталось', len(x1), 'частиц.')
    print(j + 1, 'взаимодействие:')
    x1, y1, z1, omega1, E = generate_interaction_type(x1, y1, z1, omega1, E)
    print(' Рассеялось', len(x1), 'частиц.')
    if len(x1) < 1:
        break
    E, cos_theta = scattered_photon_params(E)
    x1, y1, z1, omega1, E, cos_theta, W = expulsion_small_energy(x1, y1, z1,   omega1, E,cos_theta, W)
    if len(x1) < 1:
        break
    W1 = statistical_weight(E, W)
    x1, y1, z1, omega1, E, cos_theta, W1 = expulsion_small_weight(x1, y1, z1,   omega1, E,    cos_theta, W1)
    if len(x1) < 1:
        break
    (Eta1, Eta2) = detector_readings(x1, y1, z1, E, cos_theta, W1)
    for i in range(20):
        totEta1[i] += Eta1[i]
        totEta2[i] += Eta2[i]
 
    x0, y0, z0, omega1 = generate_newInteraction_pts(x1, y1, z1, omega1, E, cos_theta)
    # Рисуем траектории
    trajectories(x1, y1, z1, x0, y0, z0)
    x1, y1, z1 = x0, y0, z0
 
    j += 1
print('Число взаимодействий:', j + 1)
 
ax.set_xlabel(r'$x,$ см')
ax.set_ylabel(r'$y,$ см')
ax.set_zlabel(r'$z,$ см')
ax.set_title('Траектории ' + str(n) + ' частиц')
 
 
 
# Плотность потока
phi1 = [i / n for i in totEta1]
phi2 = [i / n for i in totEta2]
 
# Рисуем диаграммы
 
plot_diagram(phi1, 'Детектор №1 (0, 0, 75)')
plot_diagram(phi2, 'Детектор №2 (0, 0, 40)')
 
plt.show()